dify/dify-agent/tests/local/dify_agent/layers/knowledge/test_layer.py

import asyncio
import json

import httpx
import pytest
from pydantic_ai import Tool

from agenton.compositor import Compositor, LayerNode, LayerProvider
from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig
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 DifyKnowledgeBaseLayerConfig
from dify_agent.layers.knowledge.layer import (
    BLANK_QUERY_OBSERVATION,
    DifyKnowledgeBaseLayer,
    NO_RESULTS_OBSERVATION,
    TEMPORARY_UNAVAILABLE_OBSERVATION,
)


def _execution_context_config(**overrides: object) -> DifyExecutionContextLayerConfig:
    payload: dict[str, object] = {
        "tenant_id": "tenant-1",
        "user_id": "user-1",
        "user_from": "account",
        "app_id": "app-1",
        "agent_mode": "agent_app",
        "invoke_from": "web-app",
    }
    payload.update(overrides)
    return DifyExecutionContextLayerConfig.model_validate(payload)


def _knowledge_config(**overrides: object) -> DifyKnowledgeBaseLayerConfig:
    set_payload: dict[str, object] = {
        "id": "support",
        "name": "Support KB",
        "datasets": [{"id": "dataset-1"}],
        "query": {"mode": "generated_query"},
        "retrieval": {"mode": "multiple", "top_k": 4},
    }
    for key in ("id", "name", "description", "datasets", "query", "retrieval", "metadata_filtering"):
        if key in overrides:
            set_payload[key] = overrides.pop(key)
    if "dataset_ids" in overrides:
        dataset_ids = overrides.pop("dataset_ids")
        assert isinstance(dataset_ids, list)
        set_payload["datasets"] = [{"id": dataset_id} for dataset_id in dataset_ids]
    payload: dict[str, object] = {
        "sets": [set_payload],
    }
    payload.update(overrides)
    return DifyKnowledgeBaseLayerConfig.model_validate(payload)


def _execution_context_provider() -> LayerProvider[DifyExecutionContextLayer]:
    return LayerProvider.from_factory(
        layer_type=DifyExecutionContextLayer,
        create=lambda config: DifyExecutionContextLayer.from_config_with_settings(
            DifyExecutionContextLayerConfig.model_validate(config),
            daemon_url="http://plugin-daemon",
            daemon_api_key="daemon-secret",
        ),
    )


def _knowledge_provider() -> LayerProvider[DifyKnowledgeBaseLayer]:
    return LayerProvider.from_factory(
        layer_type=DifyKnowledgeBaseLayer,
        create=lambda config: DifyKnowledgeBaseLayer.from_config_with_settings(
            DifyKnowledgeBaseLayerConfig.model_validate(config),
            inner_api_url="http://dify-api",
            inner_api_key="inner-secret",
        ),
    )


def test_knowledge_layer_exposes_one_set_scoped_tool_definition() -> None:
    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient() as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                tool = (await knowledge_layer.get_tools(http_client=http_client))[0]
                tool_def = await tool.prepare_tool_def(None)  # pyright: ignore[reportArgumentType]
                assert isinstance(tool, Tool)
                assert tool.name == "knowledge_base_search"
                assert "Pick one configured set_name" in tool.description
                assert tool_def is not None
                assert "Pick one configured set_name" in tool_def.description
                assert tool_def.parameters_json_schema == {
                    "type": "object",
                    "properties": {
                        "set_name": {
                            "type": "string",
                            "enum": ["Support KB"],
                            "description": "Knowledge set to search.",
                        },
                        "query": {
                            "type": "string",
                            "description": "Search query for the selected knowledge set.",
                        },
                    },
                    "required": ["set_name", "query"],
                    "additionalProperties": False,
                }

    asyncio.run(scenario())


def test_knowledge_layer_rejects_blank_query_locally() -> None:
    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient() as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                tool = (await knowledge_layer.get_tools(http_client=http_client))[0]
                result = await tool.function_schema.call(  # pyright: ignore[reportArgumentType]
                    {"set_name": "Support KB", "query": "   "}, None
                )
                assert result == BLANK_QUERY_OBSERVATION

    asyncio.run(scenario())


def test_knowledge_layer_exposes_no_tool_when_all_sets_are_user_query(monkeypatch: pytest.MonkeyPatch) -> None:
    async def fake_retrieve(self: DifyKnowledgeBaseClient, **_kwargs: object) -> DifyKnowledgeRetrieveResponse:
        del self
        return DifyKnowledgeRetrieveResponse.model_validate({"results": [], "usage": {}})

    monkeypatch.setattr(DifyKnowledgeBaseClient, "retrieve", fake_retrieve)

    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient() as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(query={"mode": "user_query", "value": "release notes"}),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                assert await knowledge_layer.get_tools(http_client=http_client) == []

    asyncio.run(scenario())


def test_knowledge_layer_fetches_user_query_sets_on_context_entry(monkeypatch: pytest.MonkeyPatch) -> None:
    seen_requests: list[dict[str, object]] = []

    async def fake_retrieve(self: DifyKnowledgeBaseClient, **kwargs: object) -> DifyKnowledgeRetrieveResponse:
        del self
        seen_requests.append(kwargs)
        return DifyKnowledgeRetrieveResponse.model_validate(
            {
                "results": [
                    {
                        "metadata": {
                            "_source": "knowledge",
                            "dataset_name": "Docs",
                            "document_name": "Release.md",
                            "score": 0.8,
                        },
                        "title": "Release",
                        "files": [],
                        "content": "Version notes",
                        "summary": None,
                    }
                ],
                "usage": {},
            }
        )

    monkeypatch.setattr(DifyKnowledgeBaseClient, "retrieve", fake_retrieve)

    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with compositor.enter(
            configs={
                "execution_context": _execution_context_config(),
                "knowledge": _knowledge_config(query={"mode": "user_query", "value": "release notes"}),
            }
        ) as run:
            knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
            assert len(seen_requests) == 1
            assert seen_requests[0]["query"] == "release notes"
            assert seen_requests[0]["dataset_ids"] == ["dataset-1"]
            assert knowledge_layer.runtime_state.eager_config_fingerprint
            assert knowledge_layer.runtime_state.eager_results[0].status == "success"
            assert knowledge_layer.user_prompts == [
                "Knowledge retrieval results:\n\n"
                "Set: Support KB\n"
                "Query: release notes\n"
                "Results:\n"
                "1. Title: Release\n"
                "   Dataset: Docs\n"
                "   Document: Release.md\n"
                "   Score: 0.8\n"
                "   Content: Version notes"
            ]
            await knowledge_layer.on_context_resume()
            assert len(seen_requests) == 1

    asyncio.run(scenario())


@pytest.mark.parametrize(
    ("field_name", "field_value"),
    [
        ("user_id", None),
        ("user_from", None),
        ("app_id", None),
    ],
)
def test_knowledge_layer_fails_fast_when_execution_context_is_missing_required_fields(
    field_name: str,
    field_value: object,
) -> None:
    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient() as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(),
                }
            ) as run:
                execution_context_layer = run.get_layer("execution_context", DifyExecutionContextLayer)
                setattr(execution_context_layer.config, field_name, field_value)
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                with pytest.raises(ValueError, match=field_name):
                    _ = await knowledge_layer.get_tools(http_client=http_client)

    asyncio.run(scenario())


def test_knowledge_layer_formats_results_and_truncates_observation() -> None:
    def handler(_request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={
                "results": [
                    {
                        "metadata": {
                            "_source": "knowledge",
                            "dataset_name": "Docs",
                            "document_name": "Guide.md",
                            "score": 0.9,
                        },
                        "title": "Guide",
                        "files": [],
                        "content": "ABCDEFGHIJKL",
                        "summary": None,
                    }
                ],
                "usage": {},
            },
        )

    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(max_result_content_chars=8, max_observation_chars=160),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                tool = (await knowledge_layer.get_tools(http_client=http_client))[0]
                result = await tool.function_schema.call(  # pyright: ignore[reportArgumentType]
                    {"set_name": "Support KB", "query": "reset"}, None
                )
                assert result.startswith("Knowledge base search results:\n1. Title: Guide")
                assert "Dataset: Docs" in result
                assert "Document: Guide.md" in result
                assert "Score: 0.9" in result
                assert "Content: ABCDE..." in result
                assert len(result) <= 160

    asyncio.run(scenario())


def test_knowledge_layer_returns_no_results_observation() -> None:
    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient(
            transport=httpx.MockTransport(lambda _request: httpx.Response(200, json={"results": [], "usage": {}}))
        ) as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                tool = (await knowledge_layer.get_tools(http_client=http_client))[0]
                result = await tool.function_schema.call(  # pyright: ignore[reportArgumentType]
                    {"set_name": "Support KB", "query": "reset"}, None
                )
                assert result == NO_RESULTS_OBSERVATION

    asyncio.run(scenario())


def test_knowledge_layer_converts_retryable_failures_into_observation() -> None:
    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient(
            transport=httpx.MockTransport(
                lambda _request: httpx.Response(429, json={"code": "knowledge_rate_limited", "message": "slow down"})
            )
        ) as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                tool = (await knowledge_layer.get_tools(http_client=http_client))[0]
                result = await tool.function_schema.call(  # pyright: ignore[reportArgumentType]
                    {"set_name": "Support KB", "query": "reset"}, None
                )
                assert result == TEMPORARY_UNAVAILABLE_OBSERVATION

    asyncio.run(scenario())


@pytest.mark.parametrize(
    "transport_error",
    [
        lambda request: httpx.ReadTimeout("timed out", request=request),
        lambda request: httpx.ConnectError("connection failed", request=request),
    ],
)
def test_knowledge_layer_converts_retryable_transport_failures_into_observation(transport_error) -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        raise transport_error(request)

    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                tool = (await knowledge_layer.get_tools(http_client=http_client))[0]
                result = await tool.function_schema.call(  # pyright: ignore[reportArgumentType]
                    {"set_name": "Support KB", "query": "reset"}, None
                )
                assert result == TEMPORARY_UNAVAILABLE_OBSERVATION

    asyncio.run(scenario())


def test_knowledge_layer_raises_non_retryable_client_errors() -> None:
    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient(
            transport=httpx.MockTransport(
                lambda _request: httpx.Response(403, json={"code": "dataset_tenant_mismatch", "message": "forbidden"})
            )
        ) as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                tool = (await knowledge_layer.get_tools(http_client=http_client))[0]
                with pytest.raises(DifyKnowledgeBaseClientError) as exc_info:
                    await tool.function_schema.call(  # pyright: ignore[reportArgumentType]
                        {"set_name": "Support KB", "query": "reset"}, None
                    )
                assert exc_info.value.status_code == 403

    asyncio.run(scenario())


def test_knowledge_layer_raises_for_malformed_success_responses() -> None:
    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient(
            transport=httpx.MockTransport(lambda _request: httpx.Response(200, json={"bad": []}))
        ) as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                tool = (await knowledge_layer.get_tools(http_client=http_client))[0]
                with pytest.raises(DifyKnowledgeBaseClientError) as exc_info:
                    await tool.function_schema.call(  # pyright: ignore[reportArgumentType]
                        {"set_name": "Support KB", "query": "reset"}, None
                    )
                assert exc_info.value.error_code == "invalid_response"
                assert exc_info.value.retryable is False

    asyncio.run(scenario())


def test_knowledge_layer_sends_execution_context_and_static_config_to_inner_api() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        payload = json.loads(request.content.decode("utf-8"))
        assert request.headers["X-Inner-Api-Key"] == "inner-secret"
        assert payload["caller"] == {
            "tenant_id": "tenant-1",
            "user_id": "user-1",
            "app_id": "app-1",
            "user_from": "account",
            "invoke_from": "web-app",
        }
        assert payload["dataset_ids"] == ["dataset-1", "dataset-2"]
        assert payload["query"] == "reset"
        assert payload["retrieval"]["top_k"] == 2
        assert payload["metadata_filtering"] == {
            "mode": "manual",
            "conditions": {
                "logical_operator": "and",
                "conditions": [
                    {
                        "name": "category",
                        "comparison_operator": "contains",
                        "value": "auth",
                    }
                ],
            },
        }
        return httpx.Response(200, json={"results": [], "usage": {}})

    async def scenario() -> None:
        compositor = Compositor(
            [
                LayerNode("execution_context", _execution_context_provider()),
                LayerNode("knowledge", _knowledge_provider(), deps={"execution_context": "execution_context"}),
            ]
        )
        async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
            async with compositor.enter(
                configs={
                    "execution_context": _execution_context_config(),
                    "knowledge": _knowledge_config(
                        dataset_ids=["dataset-1", "dataset-2"],
                        retrieval={"mode": "multiple", "top_k": 2},
                        metadata_filtering={
                            "mode": "manual",
                            "conditions": {
                                "logical_operator": "and",
                                "conditions": [
                                    {
                                        "name": "category",
                                        "comparison_operator": "contains",
                                        "value": "auth",
                                    }
                                ],
                            },
                        },
                    ),
                }
            ) as run:
                knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer)
                tool = (await knowledge_layer.get_tools(http_client=http_client))[0]
                result = await tool.function_schema.call(  # pyright: ignore[reportArgumentType]
                    {"set_name": "Support KB", "query": "reset"}, None
                )
                assert result == NO_RESULTS_OBSERVATION

    asyncio.run(scenario())