dify/api/tasks/knowledge_fs_failed_retrieval_tasks.py

"""Asynchronously classify empty KnowledgeFS Workflow retrievals for quality follow-up."""

from __future__ import annotations

import logging
import uuid
from typing import Literal

from celery import shared_task
from pydantic import ValidationError

from core.db.session_factory import session_factory
from services.knowledge_fs.app_admission_service import KnowledgeFSAppAdmissionError
from services.knowledge_fs.app_execution_capability import KnowledgeResourceRef
from services.knowledge_fs.product_dto import KnowledgeFSWorkflowFailedRetrievalCapturePayload
from services.knowledge_fs.product_remote import (
    KnowledgeFSOperationUnavailableError,
    KnowledgeFSProductRequestRejectedError,
    KnowledgeFSProductResourceNotFoundError,
)
from services.knowledge_fs.runtime import get_knowledge_fs_runtime

logger = logging.getLogger(__name__)

_MAX_RETRIES = 3
_RETRY_DELAY_SECONDS = 30


@shared_task(queue="dataset", bind=True, max_retries=_MAX_RETRIES, default_retry_delay=_RETRY_DELAY_SECONDS)
def capture_workflow_failed_retrieval_task(
    self,
    *,
    event_id: str,
    tenant_id: str,
    app_id: str,
    control_space_id: str,
    query: str,
    mode: Literal["deep", "fast", "research"],
    retrieval_trace_id: str,
) -> None:
    """Re-authorize a Workflow app and ask KnowledgeFS to classify one empty retrieval.

    ``event_id`` is the stable business idempotency key across Celery retries. Capability transport
    traces are minted independently on every attempt so uncertain delivery can be retried safely.
    """

    context = {
        "app_id": app_id,
        "control_space_id": control_space_id,
        "event_id": event_id,
        "tenant_id": tenant_id,
    }
    try:
        payload = KnowledgeFSWorkflowFailedRetrievalCapturePayload.model_validate(
            {
                "eventId": event_id,
                "query": query,
                "mode": mode,
                "retrievalTraceId": retrieval_trace_id,
            }
        )
        capability = get_knowledge_fs_runtime(session_factory.get_session_maker()).app_capabilities
        result = capability.capture_workflow_failed_retrieval(
            tenant_id=tenant_id,
            app_id=app_id,
            resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id=control_space_id),
            payload=payload,
        )
    except (
        KnowledgeFSAppAdmissionError,
        KnowledgeFSOperationUnavailableError,
        KnowledgeFSProductResourceNotFoundError,
        ValidationError,
    ):
        logger.exception("KnowledgeFS Workflow failed-retrieval capture was rejected", extra=context)
        return
    except KnowledgeFSProductRequestRejectedError as exc:
        if exc.status_code != 429:
            logger.exception("KnowledgeFS Workflow failed-retrieval capture was rejected", extra=context)
            return
        _retry_capture(self, exc=exc, context=context)
    except Exception as exc:
        _retry_capture(self, exc=exc, context=context)
    else:
        logger.info(
            "KnowledgeFS Workflow failed-retrieval capture completed",
            extra={
                **context,
                "bad_case_id": str(result.bad_case_id) if result.bad_case_id else None,
                "failed_query_id": str(result.failed_query_id),
                "verdict": result.verdict,
            },
        )


def enqueue_workflow_failed_retrieval_capture(
    *,
    tenant_id: str,
    app_id: str,
    control_space_id: str,
    query: str,
    mode: Literal["deep", "fast", "research"],
    retrieval_trace_id: str,
    event_id: str | None = None,
) -> None:
    """Best-effort dispatch that never changes the completed Workflow node outcome."""

    capture_event_id = event_id or str(uuid.uuid4())
    try:
        payload = KnowledgeFSWorkflowFailedRetrievalCapturePayload.model_validate(
            {
                "eventId": capture_event_id,
                "query": query,
                "mode": mode,
                "retrievalTraceId": retrieval_trace_id,
            }
        )
        capture_workflow_failed_retrieval_task.delay(
            event_id=str(payload.event_id),
            tenant_id=tenant_id,
            app_id=app_id,
            control_space_id=control_space_id,
            query=payload.query,
            mode=payload.mode,
            retrieval_trace_id=payload.retrieval_trace_id,
        )
    except Exception:
        logger.exception(
            "Failed to enqueue KnowledgeFS Workflow failed-retrieval capture",
            extra={
                "app_id": app_id,
                "control_space_id": control_space_id,
                "event_id": capture_event_id,
                "tenant_id": tenant_id,
            },
        )


def _retry_capture(self, *, exc: Exception, context: dict[str, str]) -> None:
    if self.request.retries >= _MAX_RETRIES:
        logger.exception(
            "KnowledgeFS Workflow failed-retrieval capture retry budget exhausted",
            extra=context,
        )
        raise exc
    logger.warning(
        "KnowledgeFS Workflow failed-retrieval capture failed; scheduling retry %d/%d",
        self.request.retries + 1,
        _MAX_RETRIES,
        extra=context,
        exc_info=True,
    )
    raise self.retry(exc=exc, countdown=_RETRY_DELAY_SECONDS * (2**self.request.retries))


__all__ = [
    "capture_workflow_failed_retrieval_task",
    "enqueue_workflow_failed_retrieval_capture",
]