Make the governed RAG evidence path executable end to end
Some checks failed
verify / verify (push) Has been cancelled
Some checks failed
verify / verify (push) Has been cancelled
Separate local parsing from model indexing, bind review decisions to immutable manifests, persist vectors behind active profiles, and expose retrieval, chat, evaluation, and document workflows through the React workbench. Constraint: Live Bailian authentication currently fails for all three configured capabilities Rejected: Direct upload-to-embedding flow | bypasses local review and manifest binding Confidence: high Scope-risk: broad Directive: Keep private-data deployment blocked until authentication, RBAC, and separate database roles land Tested: make verify; fresh and replay Docker document smoke; worker recovery smoke; frozen synthetic evaluation; migration 0003-0004 roundtrip Not-tested: Successful live Bailian calls, OCR, real multi-user authorization
This commit is contained in:
519
backend/app/services/retrieval.py
Normal file
519
backend/app/services/retrieval.py
Normal file
@@ -0,0 +1,519 @@
|
||||
"""Formal two-stage retrieval use case with server-owned authorization scope."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
from app.core.problems import ApiProblem
|
||||
from app.persistence.retrieval import (
|
||||
ActiveEmbeddingProfile,
|
||||
RetrievalCandidate,
|
||||
RetrievalPersistenceError,
|
||||
RetrievalRepository,
|
||||
)
|
||||
from app.ports.model_providers import (
|
||||
EmbeddingProvider,
|
||||
ModelProviderError,
|
||||
RankedItem,
|
||||
Reranker,
|
||||
RerankResult,
|
||||
)
|
||||
|
||||
QUERY_MAX_LENGTH = 500
|
||||
VECTOR_TOP_K_DEFAULT = 50
|
||||
VECTOR_TOP_K_MAX = 50
|
||||
RERANK_TOP_N_DEFAULT = 10
|
||||
RERANK_TOP_N_MAX = 10
|
||||
RERANK_TEXT_MAX_BYTES = 4_000
|
||||
RERANK_REQUEST_MAX_BYTES = 120_000
|
||||
SNIPPET_MAX_LENGTH = 1_200
|
||||
SOURCE_NAME_MAX_LENGTH = 240
|
||||
RERANK_INSTRUCT = (
|
||||
"Given a geological exploration question, rank passages that directly support "
|
||||
"an evidence-grounded answer."
|
||||
)
|
||||
_SPACE_PATTERN = re.compile(r"\s+")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RetrievalGrant:
|
||||
"""A server-resolved knowledge-base grant; never built from request scope fields."""
|
||||
|
||||
knowledge_base_id: uuid.UUID
|
||||
access_scope_ids: tuple[uuid.UUID, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RetrievalActor:
|
||||
"""Authenticated identity projection consumed by the retrieval service."""
|
||||
|
||||
subject: str
|
||||
grants: tuple[RetrievalGrant, ...]
|
||||
|
||||
def scopes_for(self, knowledge_base_id: uuid.UUID) -> tuple[uuid.UUID, ...]:
|
||||
scopes: list[uuid.UUID] = []
|
||||
for grant in self.grants:
|
||||
if grant.knowledge_base_id == knowledge_base_id:
|
||||
scopes.extend(grant.access_scope_ids)
|
||||
return tuple(dict.fromkeys(scopes))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class EffectiveRetrievalParameters:
|
||||
vector_top_k: int
|
||||
rerank_top_n: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RetrievalTimings:
|
||||
embedding_ms: float
|
||||
database_ms: float
|
||||
rerank_ms: float
|
||||
total_ms: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RetrievalHit:
|
||||
rank: int
|
||||
vector_rank: int
|
||||
citation_id: uuid.UUID
|
||||
document_id: uuid.UUID
|
||||
source_name: str
|
||||
snippet: str
|
||||
section_path: tuple[str, ...]
|
||||
page_start: int | None
|
||||
page_end: int | None
|
||||
page_label: str
|
||||
vector_score: float
|
||||
rerank_score: float | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RetrievalResult:
|
||||
status: Literal["ok", "empty"]
|
||||
knowledge_base_id: uuid.UUID
|
||||
access_scope_count: int
|
||||
profile: ActiveEmbeddingProfile
|
||||
parameters: EffectiveRetrievalParameters
|
||||
rerank_status: Literal["applied", "degraded", "skipped_empty"]
|
||||
degradation_reason: Literal["rerank_unavailable"] | None
|
||||
embedding_request_id: str | None
|
||||
rerank_request_id: str | None
|
||||
embedding_model: str
|
||||
rerank_model: str | None
|
||||
timings: RetrievalTimings
|
||||
results: tuple[RetrievalHit, ...]
|
||||
|
||||
|
||||
class RetrievalService:
|
||||
"""Coordinate query embedding, authorized vector search, and bounded reranking."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
repository: RetrievalRepository,
|
||||
embedding_provider: EmbeddingProvider,
|
||||
reranker: Reranker,
|
||||
synthetic_embedding_provider: EmbeddingProvider | None = None,
|
||||
synthetic_reranker: Reranker | None = None,
|
||||
) -> None:
|
||||
self._repository = repository
|
||||
self._embedding_provider = embedding_provider
|
||||
self._reranker = reranker
|
||||
self._synthetic_embedding_provider = synthetic_embedding_provider
|
||||
self._synthetic_reranker = synthetic_reranker
|
||||
|
||||
async def search(
|
||||
self,
|
||||
*,
|
||||
actor: RetrievalActor,
|
||||
knowledge_base_id: uuid.UUID,
|
||||
query: str,
|
||||
vector_top_k: int = VECTOR_TOP_K_DEFAULT,
|
||||
rerank_top_n: int = RERANK_TOP_N_DEFAULT,
|
||||
) -> RetrievalResult:
|
||||
started = time.perf_counter()
|
||||
normalized_query = _normalize_query(query)
|
||||
parameters = _effective_parameters(vector_top_k, rerank_top_n)
|
||||
allowed_scope_ids = actor.scopes_for(knowledge_base_id)
|
||||
if not allowed_scope_ids:
|
||||
raise ApiProblem(
|
||||
status=403,
|
||||
code="RETRIEVAL_SCOPE_FORBIDDEN",
|
||||
title="Knowledge base access denied",
|
||||
detail="The current identity cannot search this knowledge base.",
|
||||
)
|
||||
|
||||
database_started = time.perf_counter()
|
||||
try:
|
||||
profile = await asyncio.to_thread(
|
||||
self._repository.resolve_active_profile,
|
||||
knowledge_base_id,
|
||||
allowed_scope_ids=allowed_scope_ids,
|
||||
)
|
||||
except RetrievalPersistenceError as exc:
|
||||
raise _storage_unavailable() from exc
|
||||
profile_database_ms = (time.perf_counter() - database_started) * 1_000
|
||||
if profile is None:
|
||||
raise ApiProblem(
|
||||
status=409,
|
||||
code="KNOWLEDGE_BASE_NOT_SEARCHABLE",
|
||||
title="Knowledge base is not searchable",
|
||||
detail="No enabled active embedding profile is available for this knowledge base.",
|
||||
)
|
||||
|
||||
embedding_provider, reranker = self._providers(profile)
|
||||
try:
|
||||
embedding = await embedding_provider.embed_query(normalized_query)
|
||||
except ModelProviderError as exc:
|
||||
raise _embedding_problem(exc) from exc
|
||||
query_vector = _validated_query_vector(embedding.vectors, profile)
|
||||
if embedding.model != profile.model:
|
||||
raise ApiProblem(
|
||||
status=502,
|
||||
code="EMBEDDING_PROFILE_MISMATCH",
|
||||
title="Embedding response did not match the active profile",
|
||||
detail="The query embedding model did not match the knowledge base profile.",
|
||||
)
|
||||
|
||||
search_started = time.perf_counter()
|
||||
try:
|
||||
candidates = await asyncio.to_thread(
|
||||
self._repository.search_candidates,
|
||||
knowledge_base_id,
|
||||
allowed_scope_ids=allowed_scope_ids,
|
||||
profile_hash=profile.profile_hash,
|
||||
query_vector=query_vector,
|
||||
limit=parameters.vector_top_k,
|
||||
)
|
||||
except RetrievalPersistenceError as exc:
|
||||
raise _storage_unavailable() from exc
|
||||
database_ms = profile_database_ms + (time.perf_counter() - search_started) * 1_000
|
||||
|
||||
candidates = _unique_candidates(candidates)
|
||||
if not candidates:
|
||||
total_ms = (time.perf_counter() - started) * 1_000
|
||||
return RetrievalResult(
|
||||
status="empty",
|
||||
knowledge_base_id=knowledge_base_id,
|
||||
access_scope_count=len(allowed_scope_ids),
|
||||
profile=profile,
|
||||
parameters=parameters,
|
||||
rerank_status="skipped_empty",
|
||||
degradation_reason=None,
|
||||
embedding_request_id=embedding.request_id,
|
||||
rerank_request_id=None,
|
||||
embedding_model=embedding.model,
|
||||
rerank_model=None,
|
||||
timings=RetrievalTimings(
|
||||
embedding_ms=_safe_elapsed(embedding.elapsed_ms),
|
||||
database_ms=max(0.0, database_ms),
|
||||
rerank_ms=0.0,
|
||||
total_ms=max(0.0, total_ms),
|
||||
),
|
||||
results=(),
|
||||
)
|
||||
|
||||
effective_top_n = min(parameters.rerank_top_n, len(candidates))
|
||||
documents = _bounded_rerank_documents(normalized_query, candidates)
|
||||
rerank_result: RerankResult | None = None
|
||||
try:
|
||||
attempted = await reranker.rerank(
|
||||
normalized_query,
|
||||
documents,
|
||||
top_n=effective_top_n,
|
||||
instruct=RERANK_INSTRUCT,
|
||||
)
|
||||
if _valid_rerank(attempted, documents, effective_top_n):
|
||||
rerank_result = attempted
|
||||
except ModelProviderError:
|
||||
pass
|
||||
|
||||
selected: tuple[tuple[int, float | None], ...]
|
||||
if rerank_result is None:
|
||||
selected = tuple((index, None) for index in range(effective_top_n))
|
||||
rerank_status: Literal["applied", "degraded"] = "degraded"
|
||||
degradation_reason: Literal["rerank_unavailable"] | None = "rerank_unavailable"
|
||||
rerank_request_id = None
|
||||
rerank_model = None
|
||||
rerank_ms = 0.0
|
||||
else:
|
||||
selected = tuple((item.index, item.relevance_score) for item in rerank_result.items)
|
||||
rerank_status = "applied"
|
||||
degradation_reason = None
|
||||
rerank_request_id = rerank_result.request_id
|
||||
rerank_model = rerank_result.model
|
||||
rerank_ms = _safe_elapsed(rerank_result.elapsed_ms)
|
||||
|
||||
hits = tuple(
|
||||
_hit(
|
||||
candidate=candidates[candidate_index],
|
||||
rank=rank,
|
||||
vector_rank=candidate_index + 1,
|
||||
rerank_score=rerank_score,
|
||||
)
|
||||
for rank, (candidate_index, rerank_score) in enumerate(selected, start=1)
|
||||
)
|
||||
total_ms = (time.perf_counter() - started) * 1_000
|
||||
return RetrievalResult(
|
||||
status="ok",
|
||||
knowledge_base_id=knowledge_base_id,
|
||||
access_scope_count=len(allowed_scope_ids),
|
||||
profile=profile,
|
||||
parameters=parameters,
|
||||
rerank_status=rerank_status,
|
||||
degradation_reason=degradation_reason,
|
||||
embedding_request_id=embedding.request_id,
|
||||
rerank_request_id=rerank_request_id,
|
||||
embedding_model=embedding.model,
|
||||
rerank_model=rerank_model,
|
||||
timings=RetrievalTimings(
|
||||
embedding_ms=_safe_elapsed(embedding.elapsed_ms),
|
||||
database_ms=max(0.0, database_ms),
|
||||
rerank_ms=rerank_ms,
|
||||
total_ms=max(0.0, total_ms),
|
||||
),
|
||||
results=hits,
|
||||
)
|
||||
|
||||
def _providers(
|
||||
self,
|
||||
profile: ActiveEmbeddingProfile,
|
||||
) -> tuple[EmbeddingProvider, Reranker]:
|
||||
if not profile.synthetic:
|
||||
return self._embedding_provider, self._reranker
|
||||
if self._synthetic_embedding_provider is None or self._synthetic_reranker is None:
|
||||
raise ApiProblem(
|
||||
status=503,
|
||||
code="SYNTHETIC_PROVIDER_UNAVAILABLE",
|
||||
title="Synthetic retrieval provider unavailable",
|
||||
detail="The active synthetic profile has no matching local provider.",
|
||||
)
|
||||
return self._synthetic_embedding_provider, self._synthetic_reranker
|
||||
|
||||
|
||||
def _normalize_query(value: str) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise ApiProblem(
|
||||
status=400,
|
||||
code="INVALID_RETRIEVAL_QUERY",
|
||||
title="Invalid retrieval query",
|
||||
detail="The query must be non-empty text.",
|
||||
)
|
||||
normalized = _SPACE_PATTERN.sub(" ", value).strip()
|
||||
if not normalized or len(normalized) > QUERY_MAX_LENGTH:
|
||||
raise ApiProblem(
|
||||
status=400,
|
||||
code="INVALID_RETRIEVAL_QUERY",
|
||||
title="Invalid retrieval query",
|
||||
detail=f"The query must contain between 1 and {QUERY_MAX_LENGTH} characters.",
|
||||
)
|
||||
return normalized
|
||||
|
||||
|
||||
def _effective_parameters(vector_top_k: int, rerank_top_n: int) -> EffectiveRetrievalParameters:
|
||||
for value in (vector_top_k, rerank_top_n):
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
|
||||
raise ApiProblem(
|
||||
status=400,
|
||||
code="INVALID_RETRIEVAL_PARAMETERS",
|
||||
title="Invalid retrieval parameters",
|
||||
detail="Retrieval limits must be positive integers.",
|
||||
)
|
||||
bounded_vector_top_k = min(vector_top_k, VECTOR_TOP_K_MAX)
|
||||
bounded_rerank_top_n = min(rerank_top_n, RERANK_TOP_N_MAX, bounded_vector_top_k)
|
||||
return EffectiveRetrievalParameters(
|
||||
vector_top_k=bounded_vector_top_k,
|
||||
rerank_top_n=bounded_rerank_top_n,
|
||||
)
|
||||
|
||||
|
||||
def _validated_query_vector(
|
||||
vectors: tuple[tuple[float, ...], ...],
|
||||
profile: ActiveEmbeddingProfile,
|
||||
) -> tuple[float, ...]:
|
||||
if len(vectors) != 1 or len(vectors[0]) != profile.dimension:
|
||||
raise ApiProblem(
|
||||
status=502,
|
||||
code="INVALID_EMBEDDING_RESPONSE",
|
||||
title="Invalid embedding response",
|
||||
detail="The embedding provider returned an unexpected vector shape.",
|
||||
)
|
||||
vector = vectors[0]
|
||||
if any(
|
||||
isinstance(value, bool)
|
||||
or not isinstance(value, (int, float))
|
||||
or not math.isfinite(float(value))
|
||||
for value in vector
|
||||
):
|
||||
raise ApiProblem(
|
||||
status=502,
|
||||
code="INVALID_EMBEDDING_RESPONSE",
|
||||
title="Invalid embedding response",
|
||||
detail="The embedding provider returned an invalid vector.",
|
||||
)
|
||||
normalized = tuple(float(value) for value in vector)
|
||||
if math.hypot(*normalized) <= 0:
|
||||
raise ApiProblem(
|
||||
status=502,
|
||||
code="INVALID_EMBEDDING_RESPONSE",
|
||||
title="Invalid embedding response",
|
||||
detail="The embedding provider returned a zero vector.",
|
||||
)
|
||||
return normalized
|
||||
|
||||
|
||||
def _unique_candidates(candidates: list[RetrievalCandidate]) -> list[RetrievalCandidate]:
|
||||
seen: set[uuid.UUID] = set()
|
||||
unique: list[RetrievalCandidate] = []
|
||||
for candidate in candidates:
|
||||
if candidate.citation_id not in seen:
|
||||
seen.add(candidate.citation_id)
|
||||
unique.append(candidate)
|
||||
return unique
|
||||
|
||||
|
||||
def _bounded_rerank_documents(
|
||||
query: str,
|
||||
candidates: list[RetrievalCandidate],
|
||||
) -> tuple[str, ...]:
|
||||
query_bytes = len(query.encode("utf-8"))
|
||||
available = RERANK_REQUEST_MAX_BYTES - query_bytes * len(candidates)
|
||||
per_document = min(RERANK_TEXT_MAX_BYTES, available // len(candidates))
|
||||
if per_document < 1:
|
||||
# The public query and candidate limits make this unreachable. Keep a
|
||||
# fail-closed guard so future limit changes cannot exceed provider bounds.
|
||||
raise ApiProblem(
|
||||
status=400,
|
||||
code="RERANK_BUDGET_EXCEEDED",
|
||||
title="Rerank request is too large",
|
||||
detail="The effective retrieval request exceeds the rerank input budget.",
|
||||
)
|
||||
return tuple(_truncate_utf8(_rerank_text(candidate), per_document) for candidate in candidates)
|
||||
|
||||
|
||||
def _rerank_text(candidate: RetrievalCandidate) -> str:
|
||||
section = " > ".join(candidate.section_path) if candidate.section_path else "章节未知"
|
||||
page = _page_label(candidate.page_start, candidate.page_end)
|
||||
return f"章节:{section}\n页码:{page}\n{candidate.cloud_text}"
|
||||
|
||||
|
||||
def _truncate_utf8(value: str, maximum_bytes: int) -> str:
|
||||
encoded = value.encode("utf-8")
|
||||
if len(encoded) <= maximum_bytes:
|
||||
return value
|
||||
truncated = encoded[:maximum_bytes].decode("utf-8", errors="ignore").rstrip()
|
||||
return truncated or value[0]
|
||||
|
||||
|
||||
def _valid_rerank(
|
||||
result: RerankResult,
|
||||
documents: tuple[str, ...],
|
||||
expected: int,
|
||||
) -> bool:
|
||||
if len(result.items) != expected or not result.model.strip():
|
||||
return False
|
||||
seen: set[int] = set()
|
||||
previous = math.inf
|
||||
for item in result.items:
|
||||
if not _valid_ranked_item(item, documents, seen, previous):
|
||||
return False
|
||||
seen.add(item.index)
|
||||
previous = item.relevance_score
|
||||
return True
|
||||
|
||||
|
||||
def _valid_ranked_item(
|
||||
item: RankedItem,
|
||||
documents: tuple[str, ...],
|
||||
seen: set[int],
|
||||
previous: float,
|
||||
) -> bool:
|
||||
return (
|
||||
not isinstance(item.index, bool)
|
||||
and isinstance(item.index, int)
|
||||
and 0 <= item.index < len(documents)
|
||||
and item.index not in seen
|
||||
and item.document == documents[item.index]
|
||||
and isinstance(item.relevance_score, (int, float))
|
||||
and not isinstance(item.relevance_score, bool)
|
||||
and math.isfinite(float(item.relevance_score))
|
||||
and 0.0 <= item.relevance_score <= 1.0
|
||||
and item.relevance_score <= previous
|
||||
)
|
||||
|
||||
|
||||
def _hit(
|
||||
*,
|
||||
candidate: RetrievalCandidate,
|
||||
rank: int,
|
||||
vector_rank: int,
|
||||
rerank_score: float | None,
|
||||
) -> RetrievalHit:
|
||||
return RetrievalHit(
|
||||
rank=rank,
|
||||
vector_rank=vector_rank,
|
||||
citation_id=candidate.citation_id,
|
||||
document_id=candidate.document_id,
|
||||
source_name=_bounded_text(candidate.source_name, SOURCE_NAME_MAX_LENGTH),
|
||||
snippet=_bounded_text(candidate.cloud_text, SNIPPET_MAX_LENGTH),
|
||||
section_path=candidate.section_path,
|
||||
page_start=candidate.page_start,
|
||||
page_end=candidate.page_end,
|
||||
page_label=_page_label(candidate.page_start, candidate.page_end),
|
||||
vector_score=round(max(-1.0, min(1.0, candidate.vector_score)), 6),
|
||||
rerank_score=round(rerank_score, 6) if rerank_score is not None else None,
|
||||
)
|
||||
|
||||
|
||||
def _bounded_text(value: str, maximum: int) -> str:
|
||||
normalized = _SPACE_PATTERN.sub(" ", value).strip()
|
||||
if len(normalized) <= maximum:
|
||||
return normalized
|
||||
return f"{normalized[: maximum - 1]}…"
|
||||
|
||||
|
||||
def _safe_elapsed(value: float) -> float:
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
return 0.0
|
||||
elapsed = float(value)
|
||||
return elapsed if math.isfinite(elapsed) and elapsed >= 0 else 0.0
|
||||
|
||||
|
||||
def _page_label(page_start: int | None, page_end: int | None) -> str:
|
||||
if page_start is None or page_end is None:
|
||||
return "页码未知"
|
||||
if page_start == page_end:
|
||||
return f"第 {page_start} 页"
|
||||
return f"第 {page_start}-{page_end} 页"
|
||||
|
||||
|
||||
def _storage_unavailable() -> ApiProblem:
|
||||
return ApiProblem(
|
||||
status=503,
|
||||
code="RETRIEVAL_STORAGE_UNAVAILABLE",
|
||||
title="Retrieval storage unavailable",
|
||||
detail="The retrieval index is temporarily unavailable.",
|
||||
)
|
||||
|
||||
|
||||
def _embedding_problem(exc: ModelProviderError) -> ApiProblem:
|
||||
if exc.kind.value == "invalid_response":
|
||||
return ApiProblem(
|
||||
status=502,
|
||||
code="INVALID_EMBEDDING_RESPONSE",
|
||||
title="Invalid embedding response",
|
||||
detail="The embedding provider returned an invalid response.",
|
||||
)
|
||||
return ApiProblem(
|
||||
status=503,
|
||||
code="EMBEDDING_UNAVAILABLE",
|
||||
title="Embedding service unavailable",
|
||||
detail="The query embedding service is temporarily unavailable.",
|
||||
)
|
||||
Reference in New Issue
Block a user