dify/api/tasks/refresh_billing_vector_space_task.py

import logging

from celery import shared_task
from opentelemetry import metrics

from configs import dify_config
from services.billing_service import BillingService

logger = logging.getLogger(__name__)

_MAX_RETRIES = 3
_RETRY_DELAY_SECONDS = 30
_refresh_counter = metrics.get_meter(__name__).create_counter(
    "billing.vector_space_cache_refresh.count",
    description="Number of billing vector-space cache refresh outcomes",
)


@shared_task(queue="dataset", bind=True, max_retries=_MAX_RETRIES, default_retry_delay=_RETRY_DELAY_SECONDS)
def refresh_billing_vector_space_task(self, tenant_id: str) -> None:
    """Refresh billing vector-space usage after vector cleanup has completed."""
    if not dify_config.BILLING_ENABLED:
        return

    try:
        BillingService.invalidate_vector_space_cache(tenant_id)
    except Exception as exc:
        if self.request.retries >= _MAX_RETRIES:
            _refresh_counter.add(1, {"outcome": "exhausted"})
            logger.exception(
                "Billing vector-space cache refresh retry budget exhausted, tenant_id=%s",
                tenant_id,
            )
            raise

        _refresh_counter.add(1, {"outcome": "retry"})
        logger.warning(
            "Billing vector-space cache refresh failed, scheduling retry %d/%d, tenant_id=%s",
            self.request.retries + 1,
            _MAX_RETRIES,
            tenant_id,
            exc_info=True,
        )
        raise self.retry(exc=exc, countdown=_RETRY_DELAY_SECONDS * (2**self.request.retries))

    _refresh_counter.add(1, {"outcome": "success"})
    logger.info("Billing vector-space cache refreshed, tenant_id=%s", tenant_id)


def schedule_billing_vector_space_refresh(tenant_id: str) -> None:
    """Dispatch a best-effort billing refresh without changing cleanup status."""
    if not dify_config.BILLING_ENABLED:
        return

    try:
        refresh_billing_vector_space_task.delay(tenant_id)
    except Exception:
        _refresh_counter.add(1, {"outcome": "dispatch_failure"})
        logger.exception("Failed to dispatch billing vector-space cache refresh, tenant_id=%s", tenant_id)