from __future__ import annotations import uuid from collections.abc import Sequence from dataclasses import dataclass, field import pytest from app.core.problems import ApiProblem from app.persistence.retrieval import ( CANDIDATE_SEARCH_SQL, ActiveEmbeddingProfile, RetrievalCandidate, RetrievalPersistenceError, ) from app.ports.model_providers import ( EmbeddingResult, ModelProviderError, ProviderErrorKind, ProviderUsage, RankedItem, RerankResult, ) from app.services.retrieval import RetrievalActor, RetrievalGrant, RetrievalService KNOWLEDGE_BASE_ID = uuid.UUID("10000000-0000-0000-0000-000000000001") OTHER_KNOWLEDGE_BASE_ID = uuid.UUID("10000000-0000-0000-0000-000000000002") SCOPE_ID = uuid.UUID("20000000-0000-0000-0000-000000000001") PROFILE = ActiveEmbeddingProfile( profile_hash="a" * 64, model="text-embedding-v4", dimension=1024, ) QUERY_VECTOR = (1.0,) + (0.0,) * 1023 def _candidate(index: int, *, score: float) -> RetrievalCandidate: return RetrievalCandidate( citation_id=uuid.UUID(f"30000000-0000-0000-0000-{index + 1:012d}"), document_id=uuid.UUID(f"40000000-0000-0000-0000-{index + 1:012d}"), source_name=f"报告-{index + 1}.pdf", cloud_text=f"第 {index + 1} 条已批准的斑岩铜矿证据。", section_path=("区域地质", f"矿化特征 {index + 1}"), page_start=index + 2, page_end=index + 2, vector_score=score, ) @dataclass class StubRepository: profile: ActiveEmbeddingProfile | None = PROFILE candidates: list[RetrievalCandidate] = field(default_factory=list) failure: bool = False profile_calls: list[tuple[uuid.UUID, tuple[uuid.UUID, ...]]] = field(default_factory=list) search_calls: list[tuple[uuid.UUID, tuple[uuid.UUID, ...], str, tuple[float, ...], int]] = ( field(default_factory=list) ) def resolve_active_profile( self, knowledge_base_id: uuid.UUID, *, allowed_scope_ids: Sequence[uuid.UUID], ) -> ActiveEmbeddingProfile | None: self.profile_calls.append((knowledge_base_id, tuple(allowed_scope_ids))) if self.failure: raise RetrievalPersistenceError return self.profile def search_candidates( self, knowledge_base_id: uuid.UUID, *, allowed_scope_ids: Sequence[uuid.UUID], profile_hash: str, query_vector: tuple[float, ...], limit: int, ) -> list[RetrievalCandidate]: self.search_calls.append( (knowledge_base_id, tuple(allowed_scope_ids), profile_hash, query_vector, limit) ) if self.failure: raise RetrievalPersistenceError return self.candidates[:limit] @dataclass class StubEmbeddingProvider: result: EmbeddingResult = EmbeddingResult( vectors=(QUERY_VECTOR,), model="text-embedding-v4", request_id="embed-request", usage=ProviderUsage(input_tokens=3, total_tokens=3), elapsed_ms=4.0, ) failure: ModelProviderError | None = None queries: list[str] = field(default_factory=list) async def embed_query(self, text: str) -> EmbeddingResult: self.queries.append(text) if self.failure is not None: raise self.failure return self.result async def embed_documents(self, texts: Sequence[str]) -> EmbeddingResult: del texts return self.result @dataclass class StubReranker: indices: tuple[int, ...] = (0,) failure: ModelProviderError | None = None calls: list[tuple[str, tuple[str, ...], int, str | None]] = field(default_factory=list) async def rerank( self, query: str, documents: Sequence[str], *, top_n: int, instruct: str | None = None, ) -> RerankResult: self.calls.append((query, tuple(documents), top_n, instruct)) if self.failure is not None: raise self.failure items = tuple( RankedItem( index=index, relevance_score=round(0.95 - rank * 0.1, 2), document=documents[index], ) for rank, index in enumerate(self.indices[:top_n]) ) return RerankResult( items=items, model="qwen3-rerank", request_id="rerank-request", usage=ProviderUsage(input_tokens=12, total_tokens=12), elapsed_ms=8.0, ) def _actor(*, knowledge_base_id: uuid.UUID = KNOWLEDGE_BASE_ID) -> RetrievalActor: return RetrievalActor( subject="synthetic-test-actor", grants=( RetrievalGrant( knowledge_base_id=knowledge_base_id, access_scope_ids=(SCOPE_ID,), ), ), ) @pytest.mark.asyncio async def test_service_derives_scope_clamps_parameters_and_maps_rerank() -> None: repository = StubRepository(candidates=[_candidate(0, score=0.80), _candidate(1, score=0.75)]) embedder = StubEmbeddingProvider() reranker = StubReranker(indices=(1, 0)) service = RetrievalService( repository=repository, embedding_provider=embedder, reranker=reranker, ) result = await service.search( actor=_actor(), knowledge_base_id=KNOWLEDGE_BASE_ID, query=" 斑岩铜矿\n成矿 ", vector_top_k=9_999, rerank_top_n=9_999, ) assert embedder.queries == ["斑岩铜矿 成矿"] assert repository.profile_calls == [(KNOWLEDGE_BASE_ID, (SCOPE_ID,))] assert repository.search_calls[0][:3] == (KNOWLEDGE_BASE_ID, (SCOPE_ID,), "a" * 64) assert repository.search_calls[0][4] == 50 assert result.parameters.vector_top_k == 50 assert result.parameters.rerank_top_n == 10 assert result.rerank_status == "applied" assert result.rerank_request_id == "rerank-request" assert [hit.vector_rank for hit in result.results] == [2, 1] assert [hit.rerank_score for hit in result.results] == [0.95, 0.85] assert result.results[0].citation_id == _candidate(1, score=0.75).citation_id assert result.results[0].section_path == ("区域地质", "矿化特征 2") assert result.results[0].page_label == "第 3 页" query, documents, top_n, instruct = reranker.calls[0] assert query == "斑岩铜矿 成矿" assert top_n == 2 assert instruct is not None assert all(len(document.encode("utf-8")) <= 4_000 for document in documents) assert ( len(query.encode("utf-8")) * len(documents) + sum(len(document.encode("utf-8")) for document in documents) <= 120_000 ) @pytest.mark.asyncio async def test_synthetic_active_profile_uses_only_explicit_local_providers() -> None: synthetic_profile = ActiveEmbeddingProfile( profile_hash="c" * 64, model="fake-feature-hash-v1", dimension=1024, synthetic=True, ) real_embedder = StubEmbeddingProvider() real_reranker = StubReranker() synthetic_embedder = StubEmbeddingProvider( result=EmbeddingResult( vectors=(QUERY_VECTOR,), model="fake-feature-hash-v1", request_id=None, usage=ProviderUsage(input_tokens=2, total_tokens=2), elapsed_ms=1, ) ) synthetic_reranker = StubReranker() service = RetrievalService( repository=StubRepository( profile=synthetic_profile, candidates=[_candidate(0, score=0.8)], ), embedding_provider=real_embedder, reranker=real_reranker, synthetic_embedding_provider=synthetic_embedder, synthetic_reranker=synthetic_reranker, ) result = await service.search( actor=_actor(), knowledge_base_id=KNOWLEDGE_BASE_ID, query="离线铜矿证据", ) assert result.status == "ok" assert result.profile.synthetic is True assert result.embedding_model == "fake-feature-hash-v1" assert real_embedder.queries == [] assert real_reranker.calls == [] assert synthetic_embedder.queries == ["离线铜矿证据"] assert len(synthetic_reranker.calls) == 1 @pytest.mark.asyncio async def test_unauthorized_knowledge_base_is_rejected_before_database_or_models() -> None: repository = StubRepository() embedder = StubEmbeddingProvider() reranker = StubReranker() service = RetrievalService( repository=repository, embedding_provider=embedder, reranker=reranker, ) with pytest.raises(ApiProblem) as caught: await service.search( actor=_actor(knowledge_base_id=OTHER_KNOWLEDGE_BASE_ID), knowledge_base_id=KNOWLEDGE_BASE_ID, query="铜矿", ) assert caught.value.status == 403 assert caught.value.code == "RETRIEVAL_SCOPE_FORBIDDEN" assert repository.profile_calls == [] assert embedder.queries == [] @pytest.mark.asyncio async def test_missing_active_profile_is_a_stable_problem() -> None: service = RetrievalService( repository=StubRepository(profile=None), embedding_provider=StubEmbeddingProvider(), reranker=StubReranker(), ) with pytest.raises(ApiProblem) as caught: await service.search( actor=_actor(), knowledge_base_id=KNOWLEDGE_BASE_ID, query="金矿", ) assert caught.value.status == 409 assert caught.value.code == "KNOWLEDGE_BASE_NOT_SEARCHABLE" @pytest.mark.asyncio @pytest.mark.parametrize( ("embedding", "expected_code"), [ ( EmbeddingResult( vectors=((1.0, 0.0),), model="text-embedding-v4", request_id=None, usage=ProviderUsage(), elapsed_ms=1, ), "INVALID_EMBEDDING_RESPONSE", ), ( EmbeddingResult( vectors=(QUERY_VECTOR,), model="another-model", request_id=None, usage=ProviderUsage(), elapsed_ms=1, ), "EMBEDDING_PROFILE_MISMATCH", ), ], ) async def test_embedding_must_match_active_profile( embedding: EmbeddingResult, expected_code: str, ) -> None: service = RetrievalService( repository=StubRepository(), embedding_provider=StubEmbeddingProvider(result=embedding), reranker=StubReranker(), ) with pytest.raises(ApiProblem) as caught: await service.search( actor=_actor(), knowledge_base_id=KNOWLEDGE_BASE_ID, query="铜矿", ) assert caught.value.status == 502 assert caught.value.code == expected_code @pytest.mark.asyncio async def test_rerank_provider_failure_degrades_to_vector_order() -> None: failure = ModelProviderError( operation="rerank.create", kind=ProviderErrorKind.RATE_LIMITED, provider_code="private-provider-code", retryable=True, ) repository = StubRepository(candidates=[_candidate(0, score=0.9), _candidate(1, score=0.8)]) service = RetrievalService( repository=repository, embedding_provider=StubEmbeddingProvider(), reranker=StubReranker(indices=(1, 0), failure=failure), ) result = await service.search( actor=_actor(), knowledge_base_id=KNOWLEDGE_BASE_ID, query="铜矿", rerank_top_n=2, ) assert result.status == "ok" assert result.rerank_status == "degraded" assert result.degradation_reason == "rerank_unavailable" assert result.rerank_request_id is None assert [hit.vector_rank for hit in result.results] == [1, 2] assert [hit.rerank_score for hit in result.results] == [None, None] assert "private-provider-code" not in repr(result) @pytest.mark.asyncio async def test_empty_candidates_skip_rerank_but_keep_profile_and_trace_metadata() -> None: reranker = StubReranker() service = RetrievalService( repository=StubRepository(candidates=[]), embedding_provider=StubEmbeddingProvider(), reranker=reranker, ) result = await service.search( actor=_actor(), knowledge_base_id=KNOWLEDGE_BASE_ID, query="无结果", ) assert result.status == "empty" assert result.rerank_status == "skipped_empty" assert result.profile == PROFILE assert result.results == () assert reranker.calls == [] @pytest.mark.asyncio async def test_storage_failure_is_sanitized_as_problem() -> None: service = RetrievalService( repository=StubRepository(failure=True), embedding_provider=StubEmbeddingProvider(), reranker=StubReranker(), ) with pytest.raises(ApiProblem) as caught: await service.search( actor=_actor(), knowledge_base_id=KNOWLEDGE_BASE_ID, query="铜矿", ) assert caught.value.status == 503 assert caught.value.code == "RETRIEVAL_STORAGE_UNAVAILABLE" assert "password" not in caught.value.detail.lower() def test_candidate_sql_enforces_acl_lifecycle_and_active_profile_before_limit() -> None: sql = " ".join(CANDIDATE_SEARCH_SQL.lower().split()) required_predicates = ( "chunk.knowledge_base_id = %s", "chunk.access_scope_id = any(%s::uuid[])", "knowledge_base.active_embedding_profile_hash = %s", "chunk.embedding_profile_hash = knowledge_base.active_embedding_profile_hash", "profile.enabled is true", "chunk.searchable is true", "chunk.index_status = 'ready'", "chunk.approval_status = 'cloud_approved'", "document.active_version_id = chunk.document_version_id", "document_version.review_state = 'cloud_approved'", "document_version.embedding_profile_hash = knowledge_base.active_embedding_profile_hash", ) for predicate in required_predicates: assert predicate in sql assert sql.index("chunk.access_scope_id = any(%s::uuid[])") < sql.index("limit %s")