knowledge-adapter-service/tests/test_retrieval.py

import hashlib
from collections.abc import AsyncIterator
from uuid import UUID

import httpx
import pytest

from knowledge_adapter.app import create_app
from knowledge_adapter.asset_content_processor import ContentProcessingError
from knowledge_adapter.auth import DifyKeyAuthenticator
from knowledge_adapter.service import RetrievalService
from knowledge_adapter.weknora_client import WeKnoraBadResponseError, WeKnoraUnavailableError

DIFY_KEY = "d" * 32
ALLOWED_KB_ID = "3904b40e-278b-4273-a6fa-db6cbc975587"


@pytest.fixture
async def client() -> AsyncIterator[httpx.AsyncClient]:
    authenticator = DifyKeyAuthenticator(
        current_hash=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
    )
    transport = httpx.ASGITransport(app=create_app(authenticator=authenticator))
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as value:
        yield value


def auth_headers() -> dict[str, str]:
    return {"Authorization": f"Bearer {DIFY_KEY}"}


def retrieval_payload(**overrides: object) -> dict[str, object]:
    payload: dict[str, object] = {
        "knowledge_id": ALLOWED_KB_ID,
        "query": "test query",
        "retrieval_setting": {"top_k": 3, "score_threshold": 0.5},
    }
    payload.update(overrides)
    return payload


class RecordingSearchClient:
    def __init__(self) -> None:
        self.calls: list[tuple[str, UUID]] = []

    async def search(self, *, query: str, knowledge_base_id: UUID) -> list[dict[str, object]]:
        self.calls.append((query, knowledge_base_id))
        return []


class UnavailableSearchClient:
    async def search(self, *, query: str, knowledge_base_id: UUID) -> list[dict[str, object]]:
        raise WeKnoraUnavailableError("WeKnora unavailable")


class BadResponseSearchClient:
    async def search(self, *, query: str, knowledge_base_id: UUID) -> list[dict[str, object]]:
        raise WeKnoraBadResponseError("invalid WeKnora response")


class OneChunkSearchClient:
    async def search(self, *, query: str, knowledge_base_id: UUID) -> list[dict[str, object]]:
        return [
            {
                "id": "09ab721f-7d05-45de-ae0f-dac303276f6d",
                "content": "safe text",
                "score": 0.9,
                "knowledge_base_id": str(knowledge_base_id),
                "knowledge_id": "6f77a258-27ff-4278-a667-8fe242ff8e10",
                "knowledge_title": "Manual",
                "knowledge_filename": "manual.pdf",
                "chunk_index": 0,
                "chunk_type": "text",
                "image_info": "",
            }
        ]


class FailingContentProcessor:
    async def process(self, chunk: object) -> str:
        raise ContentProcessingError("content processing failed")


async def test_valid_request_reaches_retrieval_service() -> None:
    search_client = RecordingSearchClient()
    service = RetrievalService(
        search_client=search_client,
        allowed_kb_ids=frozenset({UUID(ALLOWED_KB_ID)}),
    )
    authenticator = DifyKeyAuthenticator(
        current_hash=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
    )
    transport = httpx.ASGITransport(
        app=create_app(authenticator=authenticator, retrieval_service=service)
    )

    async with httpx.AsyncClient(transport=transport, base_url="http://test") as app_client:
        response = await app_client.post(
            "/retrieval",
            headers=auth_headers(),
            json=retrieval_payload(query="  normalized query  "),
        )

    assert response.status_code == 200
    assert response.json() == {"records": []}
    assert search_client.calls == [("normalized query", UUID(ALLOWED_KB_ID))]


async def test_disallowed_kb_returns_400_without_calling_upstream() -> None:
    search_client = RecordingSearchClient()
    service = RetrievalService(
        search_client=search_client,
        allowed_kb_ids=frozenset({UUID(ALLOWED_KB_ID)}),
    )
    authenticator = DifyKeyAuthenticator(
        current_hash=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
    )
    transport = httpx.ASGITransport(
        app=create_app(authenticator=authenticator, retrieval_service=service),
        raise_app_exceptions=False,
    )

    async with httpx.AsyncClient(transport=transport, base_url="http://test") as app_client:
        response = await app_client.post(
            "/retrieval",
            headers=auth_headers(),
            json=retrieval_payload(knowledge_id="7fa6e8d2-05d8-4e74-b98c-d8591b680db9"),
        )

    assert response.status_code == 400
    assert response.json() == {"detail": "invalid request"}
    assert search_client.calls == []


async def test_weknora_timeout_returns_sanitized_503() -> None:
    service = RetrievalService(
        search_client=UnavailableSearchClient(),
        allowed_kb_ids=frozenset({UUID(ALLOWED_KB_ID)}),
    )
    authenticator = DifyKeyAuthenticator(
        current_hash=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
    )
    transport = httpx.ASGITransport(
        app=create_app(authenticator=authenticator, retrieval_service=service),
        raise_app_exceptions=False,
    )

    async with httpx.AsyncClient(transport=transport, base_url="http://test") as app_client:
        response = await app_client.post(
            "/retrieval",
            headers=auth_headers(),
            json=retrieval_payload(),
        )

    assert response.status_code == 503
    assert response.json() == {"detail": "upstream unavailable"}


async def test_weknora_bad_response_returns_sanitized_502() -> None:
    service = RetrievalService(
        search_client=BadResponseSearchClient(),
        allowed_kb_ids=frozenset({UUID(ALLOWED_KB_ID)}),
    )
    authenticator = DifyKeyAuthenticator(
        current_hash=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
    )
    transport = httpx.ASGITransport(
        app=create_app(authenticator=authenticator, retrieval_service=service),
        raise_app_exceptions=False,
    )

    async with httpx.AsyncClient(transport=transport, base_url="http://test") as app_client:
        response = await app_client.post(
            "/retrieval",
            headers=auth_headers(),
            json=retrieval_payload(),
        )

    assert response.status_code == 502
    assert response.json() == {"detail": "invalid upstream response"}


async def test_content_processing_failure_returns_sanitized_502() -> None:
    service = RetrievalService(
        search_client=OneChunkSearchClient(),
        allowed_kb_ids=frozenset({UUID(ALLOWED_KB_ID)}),
        content_processor=FailingContentProcessor(),
    )
    authenticator = DifyKeyAuthenticator(
        current_hash=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
    )
    transport = httpx.ASGITransport(
        app=create_app(authenticator=authenticator, retrieval_service=service),
        raise_app_exceptions=False,
    )

    async with httpx.AsyncClient(transport=transport, base_url="http://test") as app_client:
        response = await app_client.post(
            "/retrieval",
            headers=auth_headers(),
            json=retrieval_payload(),
        )

    assert response.status_code == 502
    assert response.json() == {"detail": "content processing failed"}


async def test_empty_query_returns_400_without_echoing_query(client: httpx.AsyncClient) -> None:
    response = await client.post(
        "/retrieval",
        headers=auth_headers(),
        json=retrieval_payload(query="   "),
    )

    assert response.status_code == 400
    assert "query" not in response.text.lower()


async def test_invalid_json_returns_400(client: httpx.AsyncClient) -> None:
    response = await client.post(
        "/retrieval",
        headers={**auth_headers(), "Content-Type": "application/json"},
        content=b'{"query":',
    )

    assert response.status_code == 400
    assert response.json() == {"detail": "invalid request"}


async def test_query_over_maximum_returns_400_without_echoing_content(
    client: httpx.AsyncClient,
) -> None:
    oversized_query = "敏" * 2001

    response = await client.post(
        "/retrieval",
        headers=auth_headers(),
        json=retrieval_payload(query=oversized_query),
    )

    assert response.status_code == 400
    assert oversized_query not in response.text


async def test_user_cannot_supply_object_key_or_upstream_parameters(
    client: httpx.AsyncClient,
) -> None:
    response = await client.post(
        "/retrieval",
        headers=auth_headers(),
        json=retrieval_payload(object_key="private/path.png"),
    )

    assert response.status_code == 400


async def test_nonempty_metadata_conditions_are_rejected(client: httpx.AsyncClient) -> None:
    response = await client.post(
        "/retrieval",
        headers=auth_headers(),
        json=retrieval_payload(
            metadata_condition={
                "logical_operator": "and",
                "conditions": [{"name": "product", "value": "secret"}],
            }
        ),
    )

    assert response.status_code == 400