dify/dify-agent/src/dify_agent/layers/knowledge/layer.py

"""Dify knowledge-base layer exposing set-aware retrieval.

The layer depends on ``DifyExecutionContextLayer`` for tenant/app/user/invoke
identity. Generated-query sets become one stable model-visible
``knowledge_base_search(set_name, query)`` tool, while user-query sets are
retrieved eagerly during context entry and exposed as additional user prompt
content. Eager observations are persisted only as JSON-safe runtime state so
Agenton session snapshots can resume without repeating unchanged retrievals.
"""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import json
import logging
from typing import ClassVar, cast

import httpx
from pydantic_ai import RunContext, Tool
from pydantic_ai.tools import ToolDefinition
from typing_extensions import Self, override

from agenton.layers import LayerDeps, PlainLayer
from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer
from dify_agent.layers.knowledge.client import (
    DifyKnowledgeBaseClient,
    DifyKnowledgeBaseClientError,
    DifyKnowledgeRetrieveResponse,
)
from dify_agent.layers.knowledge.configs import (
    DIFY_KNOWLEDGE_BASE_LAYER_TYPE_ID,
    DifyKnowledgeBaseLayerConfig,
    DifyKnowledgeEagerResult,
    DifyKnowledgeRuntimeState,
    DifyKnowledgeSetConfig,
)

logger = logging.getLogger(__name__)

# Fixed model-visible tool identity. These stay module-private on purpose so the
# public DTO cannot grow a parallel naming contract that diverges from the
# runtime knowledge-search surface.
_KNOWLEDGE_BASE_TOOL_NAME = "knowledge_base_search"
_KNOWLEDGE_BASE_TOOL_DESCRIPTION = (
    "Search a configured knowledge set. Pick one configured set_name and provide a focused search query."
)
BLANK_QUERY_OBSERVATION = "knowledge base search requires a non-empty query"
NO_RESULTS_OBSERVATION = "No relevant knowledge base results were found."
TEMPORARY_UNAVAILABLE_OBSERVATION = (
    "Knowledge base search is temporarily unavailable. Please continue without it if possible."
)


class DifyKnowledgeBaseDeps(LayerDeps):
    """Dependencies required by ``DifyKnowledgeBaseLayer``."""

    execution_context: DifyExecutionContextLayer  # pyright: ignore[reportUninitializedInstanceVariable]


@dataclass(slots=True)
class DifyKnowledgeBaseLayer(
    PlainLayer[DifyKnowledgeBaseDeps, DifyKnowledgeBaseLayerConfig, DifyKnowledgeRuntimeState]
):
    """Layer that resolves set-scoped knowledge tools and eager user prompts."""

    type_id: ClassVar[str | None] = DIFY_KNOWLEDGE_BASE_LAYER_TYPE_ID

    config: DifyKnowledgeBaseLayerConfig
    inner_api_url: str
    inner_api_key: str

    @classmethod
    @override
    def from_config(cls, config: DifyKnowledgeBaseLayerConfig) -> Self:
        """Reject construction without server-injected Dify API settings."""
        del config
        raise TypeError(
            "DifyKnowledgeBaseLayer requires server-side Dify API settings and must use a provider factory."
        )

    @classmethod
    def from_config_with_settings(
        cls,
        config: DifyKnowledgeBaseLayerConfig,
        *,
        inner_api_url: str,
        inner_api_key: str,
    ) -> Self:
        """Create the layer from public config plus server-only API settings."""
        return cls(
            config=DifyKnowledgeBaseLayerConfig.model_validate(config),
            inner_api_url=inner_api_url,
            inner_api_key=inner_api_key,
        )

    async def get_tools(self, *, http_client: httpx.AsyncClient) -> list[Tool[object]]:
        """Build the unified generated-query Pydantic AI tool, when needed.

        Knowledge tools depend on execution-context identity that is optional for
        other run types but mandatory here: ``tenant_id``, ``user_id``,
        ``user_from``, ``app_id``, and ``invoke_from`` must all be present before
        any HTTP request is attempted. Tool execution then follows a strict
        observation policy:

        - unknown ``set_name`` returns a local validation observation;
        - blank ``query`` returns a local validation observation;
        - retryable client failures (timeouts, connection failures, HTTP
          ``429``/``502``) become a temporary-unavailable observation;
        - non-retryable client failures are raised so the run fails fast.
        """
        generated_sets = self._generated_query_sets()
        if not generated_sets:
            return []
        if http_client.is_closed:
            raise RuntimeError("DifyKnowledgeBaseLayer.get_tools() requires an open shared HTTP client.")

        execution_context = self.deps.execution_context.config
        caller = _build_caller_context(execution_context)
        client = DifyKnowledgeBaseClient(
            base_url=self.inner_api_url,
            api_key=self.inner_api_key,
            http_client=http_client,
        )
        set_by_name = {knowledge_set.name: knowledge_set for knowledge_set in generated_sets}

        async def knowledge_base_search(_ctx: RunContext[object], set_name: str, query: str) -> str:
            knowledge_set = set_by_name.get(set_name)
            if knowledge_set is None:
                return f"unknown knowledge set: {set_name}"
            normalized_query = query.strip()
            if not normalized_query:
                return BLANK_QUERY_OBSERVATION
            return await self._retrieve_for_set(
                client=client,
                caller=caller,
                knowledge_set=knowledge_set,
                query=normalized_query,
                retryable_observation=True,
            )

        async def prepare_tool_definition(_ctx: RunContext[object], tool_def: ToolDefinition) -> ToolDefinition:
            return ToolDefinition(
                name=tool_def.name,
                description=tool_def.description,
                parameters_json_schema=_tool_schema(generated_sets),
                strict=tool_def.strict,
                sequential=tool_def.sequential,
                metadata=tool_def.metadata,
                timeout=tool_def.timeout,
                defer_loading=tool_def.defer_loading,
                kind=tool_def.kind,
                return_schema=tool_def.return_schema,
                include_return_schema=tool_def.include_return_schema,
            )

        return [
            Tool(
                knowledge_base_search,
                takes_ctx=True,
                name=_KNOWLEDGE_BASE_TOOL_NAME,
                description=_tool_description(generated_sets),
                prepare=prepare_tool_definition,
            )
        ]

    @property
    @override
    def user_prompts(self) -> list[str]:
        """Expose eager user-query results as an additional user prompt."""
        if not self.runtime_state.eager_results:
            return []

        sections: list[str] = []
        for result in self.runtime_state.eager_results:
            sections.append(
                "\n".join(
                    [
                        f"Set: {result.set_name}",
                        f"Query: {result.query}",
                        "Results:",
                        result.observation,
                    ]
                )
            )
        return ["Knowledge retrieval results:\n\n" + "\n\n".join(sections)]

    @override
    async def on_context_create(self) -> None:
        await self._refresh_eager_results_if_needed()

    @override
    async def on_context_resume(self) -> None:
        await self._refresh_eager_results_if_needed()

    def _generated_query_sets(self) -> list[DifyKnowledgeSetConfig]:
        return [knowledge_set for knowledge_set in self.config.sets if knowledge_set.query.mode == "generated_query"]

    def _user_query_sets(self) -> list[DifyKnowledgeSetConfig]:
        return [knowledge_set for knowledge_set in self.config.sets if knowledge_set.query.mode == "user_query"]

    async def _refresh_eager_results_if_needed(self) -> None:
        user_query_sets = self._user_query_sets()
        if not user_query_sets:
            self.runtime_state.eager_config_fingerprint = None
            self.runtime_state.eager_results = []
            return

        fingerprint = _eager_config_fingerprint(user_query_sets)
        if self.runtime_state.eager_config_fingerprint == fingerprint:
            return

        caller = _build_caller_context(self.deps.execution_context.config)
        async with httpx.AsyncClient() as http_client:
            client = DifyKnowledgeBaseClient(
                base_url=self.inner_api_url,
                api_key=self.inner_api_key,
                http_client=http_client,
            )
            eager_results: list[DifyKnowledgeEagerResult] = []
            for knowledge_set in user_query_sets:
                query = (knowledge_set.query.value or "").strip()
                try:
                    response = await client.retrieve(
                        tenant_id=caller["tenant_id"],
                        user_id=caller["user_id"],
                        app_id=caller["app_id"],
                        user_from=caller["user_from"],
                        invoke_from=caller["invoke_from"],
                        dataset_ids=knowledge_set.dataset_ids,
                        query=query,
                        retrieval=knowledge_set.retrieval,
                        metadata_filtering=knowledge_set.metadata_filtering,
                    )
                except DifyKnowledgeBaseClientError as exc:
                    if exc.retryable:
                        logger.warning(
                            "eager knowledge retrieval temporarily unavailable",
                            extra={
                                "tenant_id": caller["tenant_id"],
                                "app_id": caller["app_id"],
                                "invoke_from": caller["invoke_from"],
                                "knowledge_set_id": knowledge_set.id,
                                "error_code": exc.error_code,
                                "status_code": exc.status_code,
                            },
                        )
                        eager_results.append(
                            DifyKnowledgeEagerResult(
                                set_id=knowledge_set.id,
                                set_name=knowledge_set.name,
                                query=query,
                                observation=TEMPORARY_UNAVAILABLE_OBSERVATION,
                                status="temporarily_unavailable",
                            )
                        )
                        continue
                    logger.error(
                        "eager knowledge retrieval failed",
                        extra={
                            "tenant_id": caller["tenant_id"],
                            "app_id": caller["app_id"],
                            "invoke_from": caller["invoke_from"],
                            "knowledge_set_id": knowledge_set.id,
                            "error_code": exc.error_code,
                            "status_code": exc.status_code,
                        },
                    )
                    raise

                eager_results.append(
                    DifyKnowledgeEagerResult(
                        set_id=knowledge_set.id,
                        set_name=knowledge_set.name,
                        query=query,
                        observation=_format_observation(response, self.config, include_heading=False),
                        status="success" if response.results else "empty",
                    )
                )

        self.runtime_state.eager_results = eager_results
        self.runtime_state.eager_config_fingerprint = fingerprint

    async def _retrieve_for_set(
        self,
        *,
        client: DifyKnowledgeBaseClient,
        caller: dict[str, str],
        knowledge_set: DifyKnowledgeSetConfig,
        query: str,
        retryable_observation: bool,
    ) -> str:
        try:
            response = await client.retrieve(
                tenant_id=caller["tenant_id"],
                user_id=caller["user_id"],
                app_id=caller["app_id"],
                user_from=caller["user_from"],
                invoke_from=caller["invoke_from"],
                dataset_ids=knowledge_set.dataset_ids,
                query=query,
                retrieval=knowledge_set.retrieval,
                metadata_filtering=knowledge_set.metadata_filtering,
            )
        except DifyKnowledgeBaseClientError as exc:
            if exc.retryable and retryable_observation:
                logger.warning(
                    "knowledge base search temporarily unavailable",
                    extra={
                        "tenant_id": caller["tenant_id"],
                        "app_id": caller["app_id"],
                        "invoke_from": caller["invoke_from"],
                        "knowledge_set_id": knowledge_set.id,
                        "error_code": exc.error_code,
                        "status_code": exc.status_code,
                    },
                )
                return TEMPORARY_UNAVAILABLE_OBSERVATION
            logger.error(
                "knowledge base search failed",
                extra={
                    "tenant_id": caller["tenant_id"],
                    "app_id": caller["app_id"],
                    "invoke_from": caller["invoke_from"],
                    "knowledge_set_id": knowledge_set.id,
                    "error_code": exc.error_code,
                    "status_code": exc.status_code,
                },
            )
            raise
        return _format_observation(response, self.config)


def _build_caller_context(execution_context: object) -> dict[str, str]:
    """Extract the inner-API caller identity from execution-context config.

    The public execution-context DTO keeps several fields optional for general
    runs, but knowledge retrieval requires all of ``tenant_id``, ``user_id``,
    ``user_from``, ``app_id``, and ``invoke_from``. Missing or blank values are
    rejected here so misconfigured runs fail before transport rather than being
    softened into tool observations.
    """
    tenant_id = getattr(execution_context, "tenant_id", None)
    user_id = getattr(execution_context, "user_id", None)
    user_from = getattr(execution_context, "user_from", None)
    app_id = getattr(execution_context, "app_id", None)
    invoke_from = getattr(execution_context, "invoke_from", None)

    missing_fields = [
        field_name
        for field_name, value in (
            ("tenant_id", tenant_id),
            ("user_id", user_id),
            ("user_from", user_from),
            ("app_id", app_id),
            ("invoke_from", invoke_from),
        )
        if not isinstance(value, str) or not value.strip()
    ]
    if missing_fields:
        joined_fields = ", ".join(missing_fields)
        raise ValueError(f"Dify knowledge base layer requires execution context fields: {joined_fields}")

    normalized_tenant_id = cast(str, tenant_id).strip()
    normalized_user_id = cast(str, user_id).strip()
    normalized_user_from = cast(str, user_from).strip()
    normalized_app_id = cast(str, app_id).strip()
    normalized_invoke_from = cast(str, invoke_from).strip()

    return {
        "tenant_id": normalized_tenant_id,
        "user_id": normalized_user_id,
        "user_from": normalized_user_from,
        "app_id": normalized_app_id,
        "invoke_from": normalized_invoke_from,
    }


def _tool_schema(generated_sets: list[DifyKnowledgeSetConfig]) -> dict[str, object]:
    return {
        "type": "object",
        "properties": {
            "set_name": {
                "type": "string",
                "enum": [knowledge_set.name for knowledge_set in generated_sets],
                "description": "Knowledge set to search.",
            },
            "query": {
                "type": "string",
                "description": "Search query for the selected knowledge set.",
            },
        },
        "required": ["set_name", "query"],
        "additionalProperties": False,
    }


def _tool_description(generated_sets: list[DifyKnowledgeSetConfig]) -> str:
    set_descriptions = []
    for knowledge_set in generated_sets:
        if knowledge_set.description:
            set_descriptions.append(f"{knowledge_set.name}: {knowledge_set.description}")
        else:
            set_descriptions.append(knowledge_set.name)
    return f"{_KNOWLEDGE_BASE_TOOL_DESCRIPTION} Configured sets: {', '.join(set_descriptions)}."


def _eager_config_fingerprint(user_query_sets: list[DifyKnowledgeSetConfig]) -> str:
    payload = [
        {
            "id": knowledge_set.id,
            "query": knowledge_set.query.model_dump(mode="json"),
            "dataset_ids": knowledge_set.dataset_ids,
            "retrieval": knowledge_set.retrieval.model_dump(mode="json"),
            "metadata_filtering": knowledge_set.metadata_filtering.model_dump(mode="json", by_alias=True),
        }
        for knowledge_set in user_query_sets
    ]
    serialized = json.dumps(payload, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(serialized.encode("utf-8")).hexdigest()


def _format_observation(
    response: DifyKnowledgeRetrieveResponse,
    config: DifyKnowledgeBaseLayerConfig,
    *,
    include_heading: bool = True,
) -> str:
    """Render inner-API retrieval results into the model-visible tool response.

    The formatting contract is intentionally simple and stable for the model:

    - empty ``results`` returns ``NO_RESULTS_OBSERVATION``;
    - non-empty results become a numbered list headed by
      ``"Knowledge base search results:"``;
    - each item includes title plus dataset/document/score metadata when those
      fields are present;
    - each content snippet is truncated by ``max_result_content_chars``;
    - the final observation is truncated by ``max_observation_chars``.
    """
    if not response.results:
        return NO_RESULTS_OBSERVATION

    lines = ["Knowledge base search results:"] if include_heading else []
    for index, result in enumerate(response.results, start=1):
        metadata = result.metadata
        title = result.title or metadata.document_name or "Untitled"
        lines.append(f"{index}. Title: {title}")
        if metadata.dataset_name:
            lines.append(f"   Dataset: {metadata.dataset_name}")
        if metadata.document_name:
            lines.append(f"   Document: {metadata.document_name}")
        if metadata.score is not None:
            lines.append(f"   Score: {metadata.score}")
        content = _truncate_text(result.content or result.summary or "", config.max_result_content_chars)
        if content:
            lines.append(f"   Content: {content}")
        lines.append("")

    return _truncate_text("\n".join(lines).rstrip(), config.max_observation_chars)


def _truncate_text(text: str, max_chars: int) -> str:
    if len(text) <= max_chars:
        return text
    if max_chars <= 3:
        return text[:max_chars]
    return f"{text[: max_chars - 3]}..."


__all__ = [
    "BLANK_QUERY_OBSERVATION",
    "DifyKnowledgeBaseDeps",
    "DifyKnowledgeBaseLayer",
    "NO_RESULTS_OBSERVATION",
    "TEMPORARY_UNAVAILABLE_OBSERVATION",
]