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": [], }