dify/api/services/knowledge_fs/app_execution_capability.py

"""Single production entry point for Agent and Workflow KnowledgeFS capabilities."""

from __future__ import annotations

from typing import Literal, Protocol, cast

from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator

from core.app.entities.app_invoke_entities import DifyRunContext
from models.knowledge_fs import KnowledgeFSAppSpaceJoinType
from services.knowledge_fs.app_admission_service import KnowledgeFSAppAdmissionService
from services.knowledge_fs.capability_broker import KnowledgeFSCapabilityBroker, KnowledgeFSIssuedProductCapability
from services.knowledge_fs.product_dto import (
    KnowledgeFSResearchTaskCreatePayload,
    KnowledgeFSResearchTaskResponse,
    KnowledgeFSRetrievalTestPayload,
    KnowledgeFSRetrievalTestResponse,
    KnowledgeFSWorkflowFailedRetrievalCapturePayload,
    KnowledgeFSWorkflowFailedRetrievalCaptureResponse,
)
from services.knowledge_fs.product_operations import KNOWLEDGE_FS_PRODUCT_OPERATIONS, is_product_operation_ready
from services.knowledge_fs.product_remote import (
    KnowledgeFSOperationUnavailableError,
    KnowledgeFSProductRemotePort,
    KnowledgeFSRemoteJSONRequest,
)


class KnowledgeResourceRef(BaseModel):
    """A typed app configuration reference to one Dify-owned KnowledgeFS control-space."""

    kind: Literal["knowledge_fs"]
    control_space_id: str = Field(min_length=1, max_length=1_000)

    model_config = ConfigDict(extra="forbid", frozen=True)

    @field_validator("control_space_id")
    @classmethod
    def normalize_control_space_id(cls, value: str) -> str:
        normalized = value.strip()
        if not normalized:
            raise ValueError("KnowledgeFS control-space reference is required")
        return normalized


class KnowledgeFSWorkflowFailedRetrievalCaptureCapability(Protocol):
    def capture_workflow_failed_retrieval(
        self,
        *,
        tenant_id: str,
        app_id: str,
        resource: KnowledgeResourceRef,
        payload: KnowledgeFSWorkflowFailedRetrievalCapturePayload,
    ) -> KnowledgeFSWorkflowFailedRetrievalCaptureResponse: ...


class KnowledgeFSAppExecutionCapabilityService:
    def __init__(
        self,
        *,
        admission: KnowledgeFSAppAdmissionService,
        broker: KnowledgeFSCapabilityBroker,
        remote: KnowledgeFSProductRemotePort,
    ) -> None:
        self._admission = admission
        self._broker = broker
        self._remote = remote

    def issue(
        self,
        *,
        tenant_id: str,
        app_id: str,
        control_space_id: str,
        caller_kind: KnowledgeFSAppSpaceJoinType,
        operation_id: str,
        resource_id: str | None = None,
        trace_id: str | None = None,
    ) -> KnowledgeFSIssuedProductCapability:
        profile = self._admission.admit(
            tenant_id=tenant_id,
            app_id=app_id,
            control_space_id=control_space_id,
            caller_kind=caller_kind,
            operation_id=operation_id,
        )
        return self._broker.issue_app(
            profile=profile,
            operation_id=operation_id,
            resource_id=resource_id,
            trace_id=trace_id,
        )

    def create_research_task(
        self,
        *,
        run_context: DifyRunContext,
        caller_kind: KnowledgeFSAppSpaceJoinType,
        resource: KnowledgeResourceRef,
        payload: KnowledgeFSResearchTaskCreatePayload,
    ) -> KnowledgeFSResearchTaskResponse:
        """Create one Research task through app admission and a bounded product operation."""

        operation_id = "createResearchTask"
        operation = KNOWLEDGE_FS_PRODUCT_OPERATIONS[operation_id]
        if (
            not is_product_operation_ready(operation_id)
            or operation.transport != "json"
            or operation.kfs_path is None
            or "{" in operation.kfs_path
        ):
            raise KnowledgeFSOperationUnavailableError("KnowledgeFS app Research task creation is unavailable")
        issued = self.issue(
            tenant_id=run_context.tenant_id,
            app_id=run_context.app_id,
            control_space_id=resource.control_space_id,
            caller_kind=caller_kind,
            operation_id=operation_id,
            trace_id=run_context.trace_session_id,
        )
        remote_payload = cast(
            dict[str, JsonValue],
            payload.model_dump(mode="json", exclude_none=True, by_alias=True),
        )
        remote_payload["knowledgeSpaceId"] = issued.knowledge_space_id
        raw = self._remote.execute_json(
            KnowledgeFSRemoteJSONRequest(
                operation_id=operation_id,
                method=operation.method,
                path=operation.kfs_path,
                namespace_id=run_context.tenant_id,
                knowledge_space_id=issued.knowledge_space_id,
                capability_token=issued.token,
                trace_id=issued.trace_id,
                payload=remote_payload,
            )
        )
        return KnowledgeFSResearchTaskResponse.model_validate(raw)

    def run_retrieval(
        self,
        *,
        run_context: DifyRunContext,
        caller_kind: KnowledgeFSAppSpaceJoinType,
        resource: KnowledgeResourceRef,
        payload: KnowledgeFSRetrievalTestPayload,
    ) -> KnowledgeFSRetrievalTestResponse:
        """Retrieve bounded evidence without invoking KnowledgeFS answer generation."""

        operation_id = "retrieveEvidence"
        operation = KNOWLEDGE_FS_PRODUCT_OPERATIONS[operation_id]
        if (
            not is_product_operation_ready(operation_id)
            or operation.transport != "json"
            or operation.kfs_path != "/knowledge-spaces/{id}/retrieval-tests"
        ):
            raise KnowledgeFSOperationUnavailableError("KnowledgeFS app evidence retrieval is unavailable")
        issued = self.issue(
            tenant_id=run_context.tenant_id,
            app_id=run_context.app_id,
            control_space_id=resource.control_space_id,
            caller_kind=caller_kind,
            operation_id=operation_id,
            trace_id=run_context.trace_session_id,
        )
        remote_payload = cast(
            dict[str, JsonValue],
            payload.model_dump(mode="json", exclude_none=True, by_alias=True),
        )
        raw = self._remote.execute_json(
            KnowledgeFSRemoteJSONRequest(
                operation_id=operation_id,
                method=operation.method,
                path=operation.kfs_path.replace("{id}", issued.knowledge_space_id),
                namespace_id=run_context.tenant_id,
                knowledge_space_id=issued.knowledge_space_id,
                capability_token=issued.token,
                trace_id=issued.trace_id,
                payload=remote_payload,
            )
        )
        return KnowledgeFSRetrievalTestResponse.model_validate(raw)

    def capture_workflow_failed_retrieval(
        self,
        *,
        tenant_id: str,
        app_id: str,
        resource: KnowledgeResourceRef,
        payload: KnowledgeFSWorkflowFailedRetrievalCapturePayload,
    ) -> KnowledgeFSWorkflowFailedRetrievalCaptureResponse:
        """Capture and classify one empty Workflow retrieval outside the node hot path."""

        operation_id = "captureWorkflowFailedRetrieval"
        operation = KNOWLEDGE_FS_PRODUCT_OPERATIONS[operation_id]
        expected_path = "/knowledge-spaces/{id}/failed-queries/workflow-retrieval-misses"
        if (
            not is_product_operation_ready(operation_id)
            or operation.transport != "json"
            or operation.kfs_path != expected_path
        ):
            raise KnowledgeFSOperationUnavailableError("KnowledgeFS Workflow failed-retrieval capture is unavailable")
        # Use a fresh transport trace for each task attempt. ``event_id`` remains the durable,
        # retry-safe business idempotency key owned by KnowledgeFS.
        issued = self.issue(
            tenant_id=tenant_id,
            app_id=app_id,
            control_space_id=resource.control_space_id,
            caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW,
            operation_id=operation_id,
        )
        remote_payload = cast(
            dict[str, JsonValue],
            payload.model_dump(mode="json", exclude_none=True, by_alias=True),
        )
        raw = self._remote.execute_json(
            KnowledgeFSRemoteJSONRequest(
                operation_id=operation_id,
                method=operation.method,
                path=expected_path.replace("{id}", issued.knowledge_space_id),
                namespace_id=tenant_id,
                knowledge_space_id=issued.knowledge_space_id,
                capability_token=issued.token,
                trace_id=issued.trace_id,
                payload=remote_payload,
            )
        )
        return KnowledgeFSWorkflowFailedRetrievalCaptureResponse.model_validate(raw)


__all__ = [
    "KnowledgeFSAppExecutionCapabilityService",
    "KnowledgeFSWorkflowFailedRetrievalCaptureCapability",
    "KnowledgeResourceRef",
]