RAG/backend/tests/unit/test_retrieval_api.py

from __future__ import annotations

import uuid
from dataclasses import dataclass, field

import httpx
import pytest
from fastapi import FastAPI

from app.api.v1.retrieval import (
    get_retrieval_actor,
    get_retrieval_service,
    router,
)
from app.core.demo_identity import (
    ACCESS_SCOPE_ID,
    BAILIAN_ACCESS_SCOPE_ID,
    BAILIAN_KNOWLEDGE_BASE_ID,
    KNOWLEDGE_BASE_ID,
)
from app.core.problems import ApiProblem, api_problem_handler
from app.core.request_context import trace_request
from app.persistence.retrieval import ActiveEmbeddingProfile
from app.services.retrieval import (
    EffectiveRetrievalParameters,
    RetrievalActor,
    RetrievalHit,
    RetrievalResult,
    RetrievalTimings,
)

TRACE_ID = "50000000-0000-0000-0000-000000000001"
CITATION_ID = uuid.UUID("60000000-0000-0000-0000-000000000001")
DOCUMENT_ID = uuid.UUID("70000000-0000-0000-0000-000000000001")


def _result() -> RetrievalResult:
    return RetrievalResult(
        status="ok",
        knowledge_base_id=KNOWLEDGE_BASE_ID,
        access_scope_count=1,
        profile=ActiveEmbeddingProfile(
            profile_hash="b" * 64,
            model="text-embedding-v4",
            dimension=1024,
        ),
        parameters=EffectiveRetrievalParameters(vector_top_k=50, rerank_top_n=10),
        rerank_status="applied",
        degradation_reason=None,
        embedding_request_id="embed-safe-id",
        rerank_request_id="rerank-safe-id",
        embedding_model="text-embedding-v4",
        rerank_model="qwen3-rerank",
        timings=RetrievalTimings(
            embedding_ms=3.0,
            database_ms=4.0,
            rerank_ms=5.0,
            total_ms=12.0,
        ),
        results=(
            RetrievalHit(
                rank=1,
                vector_rank=2,
                citation_id=CITATION_ID,
                document_id=DOCUMENT_ID,
                source_name="西岭铜矿报告.pdf",
                snippet="已批准且脱敏的铜矿证据。",
                section_path=("区域地质", "矿化特征"),
                page_start=8,
                page_end=9,
                page_label="第 8-9 页",
                vector_score=0.81,
                rerank_score=0.94,
            ),
        ),
    )


@dataclass
class StubService:
    result: RetrievalResult = field(default_factory=_result)
    problem: ApiProblem | None = None
    calls: list[tuple[RetrievalActor, uuid.UUID, str, int, int]] = field(default_factory=list)

    async def search(
        self,
        *,
        actor: RetrievalActor,
        knowledge_base_id: uuid.UUID,
        query: str,
        vector_top_k: int,
        rerank_top_n: int,
    ) -> RetrievalResult:
        self.calls.append((actor, knowledge_base_id, query, vector_top_k, rerank_top_n))
        if self.problem is not None:
            raise self.problem
        return self.result


def _app(service: StubService) -> FastAPI:
    app = FastAPI()
    app.middleware("http")(trace_request)
    app.add_exception_handler(ApiProblem, api_problem_handler)  # type: ignore[arg-type]
    app.include_router(router)
    app.dependency_overrides[get_retrieval_service] = lambda: service
    return app


def test_server_actor_grants_only_the_two_separated_synthetic_namespaces() -> None:
    actor = get_retrieval_actor()

    assert actor.scopes_for(KNOWLEDGE_BASE_ID) == (ACCESS_SCOPE_ID,)
    assert actor.scopes_for(BAILIAN_KNOWLEDGE_BASE_ID) == (BAILIAN_ACCESS_SCOPE_ID,)
    assert actor.scopes_for(uuid.uuid4()) == ()


@pytest.mark.asyncio
async def test_formal_search_exposes_trace_profile_ranks_and_stable_citation() -> None:
    service = StubService()
    async with httpx.AsyncClient(
        transport=httpx.ASGITransport(app=_app(service)),
        base_url="http://test",
    ) as client:
        response = await client.post(
            "/api/v1/retrieval/search",
            headers={"x-request-id": TRACE_ID},
            json={
                "knowledge_base_id": str(KNOWLEDGE_BASE_ID),
                "query": "  斑岩铜矿\n成矿  ",
                "vector_top_k": 999,
                "rerank_top_n": 999,
            },
        )

    assert response.status_code == 200
    payload = response.json()
    assert payload["trace_id"] == TRACE_ID
    assert payload["knowledge_base_id"] == str(KNOWLEDGE_BASE_ID)
    assert payload["access_scope_count"] == 1
    assert payload["profile"] == {
        "profile_hash": "b" * 64,
        "model": "text-embedding-v4",
        "dimension": 1024,
        "synthetic": False,
    }
    assert payload["parameters"] == {"vector_top_k": 50, "rerank_top_n": 10}
    assert payload["rerank_status"] == "applied"
    assert payload["results"][0]["citation_id"] == str(CITATION_ID)
    assert payload["results"][0]["page_label"] == "第 8-9 页"
    assert payload["results"][0]["section_path"] == ["区域地质", "矿化特征"]
    assert "access_scope_id" not in response.text
    assert "chunk_id" not in response.text

    actor, knowledge_base_id, query, vector_top_k, rerank_top_n = service.calls[0]
    assert actor.subject == "synthetic-demo-reader"
    assert knowledge_base_id == KNOWLEDGE_BASE_ID
    assert query == "斑岩铜矿 成矿"
    assert (vector_top_k, rerank_top_n) == (999, 999)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "forbidden_field",
    ["access_scope_id", "access_scope_ids", "scope", "allowed_scope_ids"],
)
async def test_request_cannot_supply_an_access_scope(forbidden_field: str) -> None:
    service = StubService()
    body = {
        "knowledge_base_id": str(KNOWLEDGE_BASE_ID),
        "query": "铜矿",
        forbidden_field: str(uuid.uuid4()),
    }
    async with httpx.AsyncClient(
        transport=httpx.ASGITransport(app=_app(service)),
        base_url="http://test",
    ) as client:
        response = await client.post("/api/v1/retrieval/search", json=body)

    assert response.status_code == 422
    assert service.calls == []


@pytest.mark.asyncio
async def test_api_problem_uses_sanitized_problem_json_and_trace_id() -> None:
    service = StubService(
        problem=ApiProblem(
            status=403,
            code="RETRIEVAL_SCOPE_FORBIDDEN",
            title="Knowledge base access denied",
            detail="The current identity cannot search this knowledge base.",
        )
    )
    async with httpx.AsyncClient(
        transport=httpx.ASGITransport(app=_app(service)),
        base_url="http://test",
    ) as client:
        response = await client.post(
            "/api/v1/retrieval/search",
            headers={"x-request-id": TRACE_ID},
            json={
                "knowledge_base_id": str(KNOWLEDGE_BASE_ID),
                "query": "铜矿",
            },
        )

    assert response.status_code == 403
    assert response.headers["content-type"].startswith("application/problem+json")
    assert response.json() == {
        "type": "https://geological-rag.local/problems/retrieval-scope-forbidden",
        "title": "Knowledge base access denied",
        "status": 403,
        "code": "RETRIEVAL_SCOPE_FORBIDDEN",
        "detail": "The current identity cannot search this knowledge base.",
        "trace_id": TRACE_ID,
        "field_errors": [],
    }