Files
RAG/backend/app/services/indexing.py
YoVinchen ecdb10c37a
Some checks failed
verify / verify (push) Has been cancelled
Make the governed RAG evidence path executable end to end
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
2026-07-13 05:58:11 +08:00

638 lines
22 KiB
Python

"""Lease-fenced document embedding orchestration.
This module deliberately owns no database transaction. Every repository call
is a complete short operation, while provider I/O happens only after that call
has returned. Persistence implementations must atomically validate the full
``JobLease`` on every mutation.
"""
from __future__ import annotations
import asyncio
import hashlib
import math
import re
import uuid
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Literal, Protocol
from app.persistence.job_queue import JobLease
from app.persistence.retrieval import ActiveEmbeddingProfile
from app.ports.model_providers import (
EmbeddingProvider,
EmbeddingResult,
ModelProviderError,
ProviderErrorKind,
ProviderUsage,
)
EMBEDDING_BATCH_SIZE = 10
EMBEDDING_DIMENSION = 1024
_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$")
type AssignmentStatus = Literal["PENDING", "EMBEDDING", "READY", "FAILED", "STALE"]
type InvocationStatus = Literal["SUCCEEDED", "FAILED", "UNKNOWN"]
type WriteSource = Literal["cache", "provider"]
class IndexingError(RuntimeError):
"""Base class for safe indexing failures."""
class InvalidIndexingPlanError(IndexingError):
"""Approved indexing inputs violated their immutable contract."""
class InvalidEmbeddingResponseError(IndexingError):
"""The provider returned an unusable or profile-incompatible embedding."""
class IndexingNotReadyError(IndexingError):
"""Activation was refused because not every expected assignment is READY."""
class IndexingProviderError(IndexingError):
"""An unexpected provider failure was converted to a safe worker error."""
@dataclass(frozen=True, slots=True)
class IndexingItem:
chunk_id: uuid.UUID
ordinal: int
embedding_text: str
embedding_text_sha256: str
assignment_status: AssignmentStatus
@dataclass(frozen=True, slots=True)
class ApprovedIndexingPlan:
knowledge_base_id: uuid.UUID
document_version_id: uuid.UUID
review_state: str
outbound_manifest_sha256: str
expected_count: int
profile: ActiveEmbeddingProfile
items: tuple[IndexingItem, ...]
@dataclass(frozen=True, slots=True)
class CachedEmbedding:
cache_key: str
profile_hash: str
embedding_text_sha256: str
resolved_model: str
dimension: int
@dataclass(frozen=True, slots=True)
class EmbeddingCacheLookup:
"""Derived key plus the indexed database key needed for an efficient lookup."""
cache_key: str
profile_hash: str
embedding_text_sha256: str
@dataclass(frozen=True, slots=True)
class EmbeddingWrite:
chunk_id: uuid.UUID
batch_index: int
cache_key: str
profile_hash: str
embedding_text_sha256: str
source: WriteSource
embedding: tuple[float, ...] | None
resolved_model: str
provider_request_id: str | None
usage: ProviderUsage
elapsed_ms: float
@dataclass(frozen=True, slots=True)
class AssignmentProgress:
expected_count: int
ready_count: int
@dataclass(frozen=True, slots=True)
class IndexingResult:
document_version_id: uuid.UUID
profile_hash: str
expected_count: int
ready_count: int
cache_hit_count: int
newly_embedded_count: int
provider_call_count: int
activated: bool
class IndexingRepository(Protocol):
"""Short-operation persistence port; implementations must not retain transactions."""
def load_approved_plan(
self,
*,
lease: JobLease,
document_version_id: uuid.UUID,
) -> ApprovedIndexingPlan: ...
def lookup_cache(
self,
*,
lease: JobLease,
lookups: Sequence[EmbeddingCacheLookup],
) -> Mapping[str, CachedEmbedding]: ...
def begin_model_invocation(
self,
*,
lease: JobLease,
trace_id: uuid.UUID,
profile_hash: str,
model: str,
item_count: int,
) -> uuid.UUID: ...
def finish_model_invocation(
self,
*,
lease: JobLease,
invocation_id: uuid.UUID,
status: InvocationStatus,
provider_request_id: str | None,
usage: ProviderUsage,
elapsed_ms: float,
error_code: str | None,
) -> None: ...
def fenced_persist_batch(
self,
*,
lease: JobLease,
document_version_id: uuid.UUID,
profile_hash: str,
writes: Sequence[EmbeddingWrite],
) -> AssignmentProgress: ...
def fenced_activate(
self,
*,
lease: JobLease,
document_version_id: uuid.UUID,
profile_hash: str,
expected_count: int,
) -> bool: ...
class DocumentIndexingService:
"""Resolve cache hits, embed misses, and activate only a complete projection."""
def __init__(
self,
*,
repository: IndexingRepository,
embedding_provider: EmbeddingProvider,
synthetic_embedding_provider: EmbeddingProvider | None = None,
) -> None:
self._repository = repository
self._embedding_provider = embedding_provider
self._synthetic_embedding_provider = synthetic_embedding_provider
async def index_document_version(
self,
*,
lease: JobLease,
document_version_id: uuid.UUID,
trace_id: uuid.UUID,
) -> IndexingResult:
plan = await asyncio.to_thread(
self._repository.load_approved_plan,
lease=lease,
document_version_id=document_version_id,
)
_validate_plan(plan, document_version_id)
provider = self._provider_for(plan.profile)
initial_ready = sum(item.assignment_status == "READY" for item in plan.items)
progress = AssignmentProgress(plan.expected_count, initial_ready)
pending = tuple(item for item in plan.items if item.assignment_status != "READY")
items_by_cache_key = _group_by_cache_key(pending, plan.profile.profile_hash)
cached_by_key: dict[str, CachedEmbedding] = {}
cache_keys = tuple(items_by_cache_key)
for key_batch in _batches(cache_keys, EMBEDDING_BATCH_SIZE):
lookups = tuple(
EmbeddingCacheLookup(
cache_key=key,
profile_hash=plan.profile.profile_hash,
embedding_text_sha256=items_by_cache_key[key][0].embedding_text_sha256,
)
for key in key_batch
)
found = await asyncio.to_thread(
self._repository.lookup_cache,
lease=lease,
lookups=lookups,
)
cached_by_key.update(_validated_cache(found, lookups, plan.profile))
cache_writes = tuple(
_cache_write(item, cached_by_key[key], plan.profile)
for key, items in items_by_cache_key.items()
if key in cached_by_key
for item in items
)
for write_batch in _batches(cache_writes, EMBEDDING_BATCH_SIZE):
progress = await self._persist(
lease=lease,
plan=plan,
writes=write_batch,
previous=progress,
)
missing_keys = tuple(key for key in items_by_cache_key if key not in cached_by_key)
provider_call_count = 0
newly_embedded_count = 0
for key_batch in _batches(missing_keys, EMBEDDING_BATCH_SIZE):
batch_items = tuple(items_by_cache_key[key][0] for key in key_batch)
invocation_id = await asyncio.to_thread(
self._repository.begin_model_invocation,
lease=lease,
trace_id=trace_id,
profile_hash=plan.profile.profile_hash,
model=plan.profile.model,
item_count=len(batch_items),
)
provider_call_count += 1
result, validated = await self._call_provider(
provider=provider,
items=batch_items,
profile=plan.profile,
lease=lease,
invocation_id=invocation_id,
)
newly_embedded_count += len(validated)
provider_writes = tuple(
_provider_write(
item,
cache_key=cache_key,
batch_index=validated_index,
vector=validated[validated_index],
result=result,
profile=plan.profile,
)
for validated_index, cache_key in enumerate(key_batch)
for item in items_by_cache_key[cache_key]
)
for write_batch in _batches(provider_writes, EMBEDDING_BATCH_SIZE):
progress = await self._persist(
lease=lease,
plan=plan,
writes=write_batch,
previous=progress,
)
if (
progress.expected_count != plan.expected_count
or progress.ready_count != plan.expected_count
):
raise IndexingNotReadyError("not every expected embedding assignment is READY")
activated = await asyncio.to_thread(
self._repository.fenced_activate,
lease=lease,
document_version_id=plan.document_version_id,
profile_hash=plan.profile.profile_hash,
expected_count=plan.expected_count,
)
if not activated:
raise IndexingNotReadyError("fenced activation rejected an incomplete projection")
return IndexingResult(
document_version_id=plan.document_version_id,
profile_hash=plan.profile.profile_hash,
expected_count=plan.expected_count,
ready_count=progress.ready_count,
cache_hit_count=len(cache_writes),
newly_embedded_count=newly_embedded_count,
provider_call_count=provider_call_count,
activated=True,
)
def _provider_for(self, profile: ActiveEmbeddingProfile) -> EmbeddingProvider:
if not profile.synthetic:
return self._embedding_provider
if self._synthetic_embedding_provider is None:
raise InvalidIndexingPlanError("synthetic profile has no local embedding provider")
return self._synthetic_embedding_provider
async def _call_provider(
self,
*,
provider: EmbeddingProvider,
items: tuple[IndexingItem, ...],
profile: ActiveEmbeddingProfile,
lease: JobLease,
invocation_id: uuid.UUID,
) -> tuple[EmbeddingResult, tuple[tuple[float, ...], ...]]:
try:
result = await provider.embed_documents(tuple(item.embedding_text for item in items))
except ModelProviderError as exc:
await asyncio.to_thread(
self._repository.finish_model_invocation,
lease=lease,
invocation_id=invocation_id,
status=_provider_failure_status(exc.kind),
provider_request_id=_safe_request_id(exc.request_id),
usage=ProviderUsage(),
elapsed_ms=0.0,
error_code=f"EMBEDDING_{exc.kind.value.upper()}",
)
raise
except Exception:
await asyncio.to_thread(
self._repository.finish_model_invocation,
lease=lease,
invocation_id=invocation_id,
status="UNKNOWN",
provider_request_id=None,
usage=ProviderUsage(),
elapsed_ms=0.0,
error_code="EMBEDDING_PROVIDER_UNEXPECTED",
)
raise IndexingProviderError("embedding provider failed unexpectedly") from None
try:
validated = _validated_result(result, items, profile)
except InvalidEmbeddingResponseError:
await asyncio.to_thread(
self._repository.finish_model_invocation,
lease=lease,
invocation_id=invocation_id,
status="FAILED",
provider_request_id=_safe_request_id(result.request_id),
usage=_safe_usage(result.usage),
elapsed_ms=_safe_elapsed(result.elapsed_ms),
error_code="INVALID_EMBEDDING_RESPONSE",
)
raise
await asyncio.to_thread(
self._repository.finish_model_invocation,
lease=lease,
invocation_id=invocation_id,
status="SUCCEEDED",
provider_request_id=result.request_id,
usage=result.usage,
elapsed_ms=result.elapsed_ms,
error_code=None,
)
return result, validated
async def _persist(
self,
*,
lease: JobLease,
plan: ApprovedIndexingPlan,
writes: tuple[EmbeddingWrite, ...],
previous: AssignmentProgress,
) -> AssignmentProgress:
if not writes or len(writes) > EMBEDDING_BATCH_SIZE:
raise InvalidIndexingPlanError(
"persistence batches must contain between 1 and 10 items"
)
progress = await asyncio.to_thread(
self._repository.fenced_persist_batch,
lease=lease,
document_version_id=plan.document_version_id,
profile_hash=plan.profile.profile_hash,
writes=writes,
)
if (
progress.expected_count != plan.expected_count
or not previous.ready_count <= progress.ready_count <= plan.expected_count
):
raise IndexingNotReadyError("assignment progress violated the expected READY count")
return progress
def embedding_cache_key(embedding_text_sha256: str, profile_hash: str) -> str:
"""Return the deterministic application cache key required by the indexing contract."""
if not _valid_sha256(embedding_text_sha256) or not _valid_sha256(profile_hash):
raise InvalidIndexingPlanError("cache key inputs must be lowercase SHA-256 values")
return hashlib.sha256(f"{embedding_text_sha256}{profile_hash}".encode()).hexdigest()
def _validate_plan(plan: ApprovedIndexingPlan, requested_version_id: uuid.UUID) -> None:
profile = plan.profile
if plan.document_version_id != requested_version_id:
raise InvalidIndexingPlanError("repository returned a different document version")
if plan.review_state != "CLOUD_APPROVED":
raise InvalidIndexingPlanError("document version is not cloud approved")
if not _valid_sha256(plan.outbound_manifest_sha256):
raise InvalidIndexingPlanError("approved manifest hash is invalid")
if (
not _valid_sha256(profile.profile_hash)
or not profile.model.strip()
or profile.dimension != EMBEDDING_DIMENSION
):
raise InvalidIndexingPlanError("embedding profile is invalid or unsupported")
if (
isinstance(plan.expected_count, bool)
or plan.expected_count < 0
or plan.expected_count != len(plan.items)
):
raise InvalidIndexingPlanError("expected chunk count does not match the approved plan")
chunk_ids: set[uuid.UUID] = set()
ordinals: set[int] = set()
valid_statuses = {"PENDING", "EMBEDDING", "READY", "FAILED", "STALE"}
for item in plan.items:
if item.chunk_id in chunk_ids or item.ordinal in ordinals:
raise InvalidIndexingPlanError("approved indexing items must be unique")
if isinstance(item.ordinal, bool) or item.ordinal < 0:
raise InvalidIndexingPlanError("chunk ordinal is invalid")
if not item.embedding_text:
raise InvalidIndexingPlanError("approved embedding text must not be empty")
if item.assignment_status not in valid_statuses:
raise InvalidIndexingPlanError("embedding assignment status is invalid")
actual_text_hash = hashlib.sha256(item.embedding_text.encode()).hexdigest()
if item.embedding_text_sha256 != actual_text_hash:
raise InvalidIndexingPlanError("approved embedding text hash does not match its text")
chunk_ids.add(item.chunk_id)
ordinals.add(item.ordinal)
if ordinals != set(range(plan.expected_count)):
raise InvalidIndexingPlanError("approved chunk ordinals must be contiguous")
def _group_by_cache_key(
items: tuple[IndexingItem, ...],
profile_hash: str,
) -> dict[str, list[IndexingItem]]:
grouped: dict[str, list[IndexingItem]] = {}
for item in items:
key = embedding_cache_key(item.embedding_text_sha256, profile_hash)
grouped.setdefault(key, []).append(item)
return grouped
def _validated_cache(
found: Mapping[str, CachedEmbedding],
lookups: tuple[EmbeddingCacheLookup, ...],
profile: ActiveEmbeddingProfile,
) -> dict[str, CachedEmbedding]:
requested = {lookup.cache_key: lookup for lookup in lookups}
if any(key not in requested for key in found):
raise InvalidIndexingPlanError("cache lookup returned an unrequested key")
validated: dict[str, CachedEmbedding] = {}
for key, record in found.items():
expected_key = embedding_cache_key(record.embedding_text_sha256, record.profile_hash)
if (
key != record.cache_key
or key != expected_key
or record.profile_hash != profile.profile_hash
or record.profile_hash != requested[key].profile_hash
or record.embedding_text_sha256 != requested[key].embedding_text_sha256
or record.resolved_model != profile.model
or record.dimension != profile.dimension
):
raise InvalidIndexingPlanError("cache record does not match the active profile")
validated[key] = record
return validated
def _validated_result(
result: EmbeddingResult,
items: tuple[IndexingItem, ...],
profile: ActiveEmbeddingProfile,
) -> tuple[tuple[float, ...], ...]:
if result.model != profile.model:
raise InvalidEmbeddingResponseError("embedding model did not match the active profile")
if len(result.vectors) != len(items):
raise InvalidEmbeddingResponseError("embedding result count did not match the batch")
if (
isinstance(result.elapsed_ms, bool)
or not isinstance(result.elapsed_ms, (int, float))
or not math.isfinite(float(result.elapsed_ms))
or result.elapsed_ms < 0
):
raise InvalidEmbeddingResponseError("embedding elapsed time is invalid")
if result.request_id is not None and _safe_request_id(result.request_id) is None:
raise InvalidEmbeddingResponseError("embedding request identifier is invalid")
if _safe_usage(result.usage) != result.usage:
raise InvalidEmbeddingResponseError("embedding usage metadata is invalid")
vectors: list[tuple[float, ...]] = []
for index, vector in enumerate(result.vectors):
if index >= len(items):
raise InvalidEmbeddingResponseError("embedding index exceeded the requested batch")
if len(vector) != profile.dimension:
raise InvalidEmbeddingResponseError("embedding dimension did not match the profile")
normalized: list[float] = []
for component in vector:
if (
isinstance(component, bool)
or not isinstance(component, (int, float))
or not math.isfinite(float(component))
):
raise InvalidEmbeddingResponseError("embedding contains a non-finite component")
normalized.append(float(component))
if math.hypot(*normalized) <= 0:
raise InvalidEmbeddingResponseError("embedding vector must have a nonzero norm")
vectors.append(tuple(normalized))
return tuple(vectors)
def _cache_write(
item: IndexingItem,
cached: CachedEmbedding,
profile: ActiveEmbeddingProfile,
) -> EmbeddingWrite:
return EmbeddingWrite(
chunk_id=item.chunk_id,
batch_index=0,
cache_key=cached.cache_key,
profile_hash=profile.profile_hash,
embedding_text_sha256=item.embedding_text_sha256,
source="cache",
embedding=None,
resolved_model=cached.resolved_model,
provider_request_id=None,
usage=ProviderUsage(),
elapsed_ms=0.0,
)
def _provider_write(
item: IndexingItem,
*,
cache_key: str,
batch_index: int,
vector: tuple[float, ...],
result: EmbeddingResult,
profile: ActiveEmbeddingProfile,
) -> EmbeddingWrite:
return EmbeddingWrite(
chunk_id=item.chunk_id,
batch_index=batch_index,
cache_key=cache_key,
profile_hash=profile.profile_hash,
embedding_text_sha256=item.embedding_text_sha256,
source="provider",
embedding=vector,
resolved_model=result.model,
provider_request_id=result.request_id,
usage=result.usage,
elapsed_ms=result.elapsed_ms,
)
def _provider_failure_status(kind: ProviderErrorKind) -> InvocationStatus:
if kind in {ProviderErrorKind.TIMEOUT, ProviderErrorKind.TRANSPORT}:
return "UNKNOWN"
return "FAILED"
def _safe_request_id(value: str | None) -> str | None:
if value is None:
return None
if not value.strip() or len(value) > 512 or any(character.isspace() for character in value):
return None
return value
def _safe_usage(value: ProviderUsage) -> ProviderUsage:
def valid(component: int | None) -> int | None:
if component is not None and (
isinstance(component, bool) or not isinstance(component, int) or component < 0
):
return None
return component
return ProviderUsage(
input_tokens=valid(value.input_tokens),
output_tokens=valid(value.output_tokens),
total_tokens=valid(value.total_tokens),
)
def _safe_elapsed(value: float) -> float:
if (
isinstance(value, bool)
or not isinstance(value, (int, float))
or not math.isfinite(float(value))
or value < 0
):
return 0.0
return float(value)
def _valid_sha256(value: str) -> bool:
return bool(_SHA256_PATTERN.fullmatch(value))
def _batches[T](values: Sequence[T], size: int) -> tuple[tuple[T, ...], ...]:
return tuple(tuple(values[index : index + size]) for index in range(0, len(values), size))