"""Document lifecycle service: ingest, list, inspect, delete. Supported process model: one API process owning one corpus. Writes are serialized; the manifest is authoritative; the lexical index is rebuilt from the manifest after every committed change; vectors live in Chroma when an embedding provider is configured. Ingestion is staged: file -> manifest (INDEXING) -> chunks -> vectors -> commit (READY). Any failure rolls the staged state back and, for new documents, leaves a FAILED record that is never searchable. Replacement keeps the previous version searchable until the new one is committed. """ from __future__ import annotations import asyncio import json import logging import time import uuid from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Optional, Set from ..core.chunking import Chunk, chunk_text from ..core.config import Settings from ..core.embeddings import EmbeddingProvider, build_embedding_provider from ..core.lexical import LexicalEntry, LexicalIndex from ..core.manifest import Manifest, make_chunk_id, utcnow_iso from ..core.types import ChunkRecord, DocumentRecord, DocumentStatus, IngestionError, NotFoundError from ..core.vector_store import VectorStore from ..utils.document_loader import ( SUPPORTED_EXTENSIONS, UploadStorage, content_hash, detect_extension, extract_document, normalize_display_filename, ) from .cache import AnswerCache logger = logging.getLogger(__name__) RESERVED_METADATA_KEYS = { "doc_id", "version", "chunk_id", "ordinal", "filename", "status", "content_hash", "source", "storage_name", "location", "owner", } def validate_user_metadata( metadata: Optional[Dict[str, Any]], *, max_keys: int, max_value_length: int ) -> Dict[str, Any]: """Accept only flat, scalar user metadata that does not shadow reserved fields.""" if metadata is None: return {} if not isinstance(metadata, dict): raise IngestionError("metadata must be a JSON object", code="invalid_metadata") if len(metadata) > max_keys: raise IngestionError(f"metadata may have at most {max_keys} keys", code="invalid_metadata") clean: Dict[str, Any] = {} for key, value in metadata.items(): if not isinstance(key, str) or not key or len(key) > 64: raise IngestionError("metadata keys must be non-empty strings", code="invalid_metadata") if key in RESERVED_METADATA_KEYS: raise IngestionError(f"metadata key '{key}' is reserved", code="invalid_metadata") if isinstance(value, bool) or value is None: clean[key] = value elif isinstance(value, (int, float)): clean[key] = value elif isinstance(value, str): if len(value) > max_value_length: raise IngestionError( f"metadata value for '{key}' exceeds {max_value_length} characters", code="invalid_metadata", ) clean[key] = value else: raise IngestionError( f"metadata value for '{key}' must be a string, number, boolean or null", code="invalid_metadata", ) return clean @dataclass class IngestResult: record: DocumentRecord created: bool duplicate_of: Optional[str] = None warnings: List[str] = field(default_factory=list) timings_ms: Dict[str, float] = field(default_factory=dict) evicted: List[str] = field(default_factory=list) @dataclass class SeedResult: doc_ids: List[str] errors: List[str] @dataclass class DeleteResult: doc_id: str removed_chunks: int removed_vectors: int file_removed: bool class DocumentService: def __init__( self, settings: Settings, *, manifest: Manifest, storage: UploadStorage, lexical: LexicalIndex, answer_cache: AnswerCache, embedder: Optional[EmbeddingProvider], vector_store: Optional[VectorStore], ): self.settings = settings self.manifest = manifest self.storage = storage self.lexical = lexical self.answer_cache = answer_cache self.embedder = embedder self.vector_store = vector_store self._write_lock = asyncio.Lock() # Documents ingested at startup from ``demo_seed_dir``; never evicted. self.protected_doc_ids: Set[str] = set() self.rebuild_lexical() # -- construction ------------------------------------------------------ @classmethod def from_settings(cls, settings: Settings) -> "DocumentService": embedder = build_embedding_provider(settings) if settings.storage_mode == "ephemeral": manifest = Manifest.in_memory() persist_dir = None else: manifest = Manifest.open(settings.manifest_path) persist_dir = settings.index_dir storage = UploadStorage(settings.uploads_dir) vector_store = None if embedder is not None: vector_store = VectorStore( embedding_identity=embedder.identity, collection_name=settings.chroma_collection_name, persist_dir=persist_dir, ) return cls( settings, manifest=manifest, storage=storage, lexical=LexicalIndex(), answer_cache=AnswerCache( settings.answer_cache_ttl_seconds, settings.answer_cache_max_entries ), embedder=embedder, vector_store=vector_store, ) def close(self) -> None: if self.manifest.path is None and self.vector_store is not None: # Ephemeral Chroma clients share process-wide state; an ephemeral # corpus must not outlive the service that owns it. self.vector_store.drop() self.manifest.close() # -- capabilities -------------------------------------------------------- @property def dense_available(self) -> bool: return self.embedder is not None and self.vector_store is not None def capabilities(self) -> Dict[str, Any]: return { "storage_mode": self.settings.storage_mode, "persistent": self.manifest.path is not None, "embedding_provider": self.settings.resolved_embedding_provider, "embedding_identity": self.embedder.identity if self.embedder else None, "embedding_semantic": bool(self.embedder and self.embedder.semantic), "dense_retrieval": self.dense_available, "lexical_retrieval": True, "reranker_mode": self.settings.reranker_mode, "generation_provider": self.settings.resolved_generation_provider, "chunking_config": self.settings.chunking_config_id, "supported_extensions": sorted(SUPPORTED_EXTENSIONS), "max_upload_size": self.settings.max_upload_size, "max_documents": self.settings.max_documents, } def stats(self) -> Dict[str, Any]: docs = self.manifest.list_documents() return { "documents": len(docs), "documents_ready": sum(1 for d in docs if d.status == DocumentStatus.READY), "active_chunks": self.manifest.active_chunk_count(), "lexical_entries": self.lexical.size, "vectors": self.vector_store.count() if self.vector_store else None, "corpus_generation": self.manifest.corpus_generation(), "answer_cache_entries": len(self.answer_cache), } # -- lexical ------------------------------------------------------------- def rebuild_lexical(self) -> None: entries = [ LexicalEntry(chunk_id=c.chunk_id, doc_id=c.doc_id, text=c.text) for c in self.manifest.iter_active_chunks() ] self.lexical.rebuild(entries, self.manifest.corpus_generation()) def _commit_corpus_change(self) -> None: self.manifest.bump_generation() self.rebuild_lexical() self.answer_cache.clear() # -- reads ----------------------------------------------------------------- def get_document(self, doc_id: str) -> DocumentRecord: record = self.manifest.get_document(doc_id) if record is None or record.status == DocumentStatus.DELETED: raise NotFoundError(doc_id) return record def list_documents(self) -> List[DocumentRecord]: return [d for d in self.manifest.list_documents() if d.status != DocumentStatus.DELETED] def get_chunks(self, doc_id: str, *, offset: int = 0, limit: int = 50) -> List[ChunkRecord]: record = self.get_document(doc_id) return self.manifest.get_chunks(doc_id, record.version, offset=offset, limit=limit) # -- ingestion --------------------------------------------------------------- async def ingest( self, filename: Optional[str], content: bytes, user_metadata: Optional[Dict[str, Any]] = None, *, replace_doc_id: Optional[str] = None, ) -> IngestResult: settings = self.settings display_name = normalize_display_filename(filename, settings.max_filename_length) extension = detect_extension(display_name) if len(content) > settings.max_upload_size: raise IngestionError( f"File exceeds the {settings.max_upload_size} byte limit", code="too_large", status=413, ) metadata = validate_user_metadata( user_metadata, max_keys=settings.max_metadata_keys, max_value_length=settings.max_metadata_value_length, ) digest = content_hash(content) async with self._write_lock: existing: Optional[DocumentRecord] = None if replace_doc_id is not None: existing = self.manifest.get_document(replace_doc_id) if existing is None or existing.status == DocumentStatus.DELETED: raise NotFoundError(replace_doc_id) if existing.status == DocumentStatus.READY and existing.content_hash == digest: return IngestResult(existing, created=False, duplicate_of=existing.doc_id) else: for candidate in self.manifest.find_by_hash(digest): if candidate.status == DocumentStatus.READY: return IngestResult(candidate, created=False, duplicate_of=candidate.doc_id) if candidate.status in (DocumentStatus.FAILED, DocumentStatus.INDEXING): existing = candidate # retry under the same identity break evicted: List[str] = [] if existing is None and settings.max_documents > 0: evicted = await self._make_room_locked(settings.max_documents) timings: Dict[str, float] = {} t0 = time.perf_counter() extracted = await asyncio.to_thread( extract_document, display_name, content, max_extracted_chars=settings.max_extracted_chars, max_pdf_pages=settings.max_pdf_pages, ) timings["extract"] = (time.perf_counter() - t0) * 1000 t0 = time.perf_counter() chunks: List[Chunk] = await asyncio.to_thread( chunk_text, extracted.text, chunk_size=settings.chunk_size, chunk_overlap=settings.chunk_overlap, max_chunks=settings.max_chunks_per_document, pages=extracted.pages, extension=extension, ) timings["chunk"] = (time.perf_counter() - t0) * 1000 if not chunks: raise IngestionError("Document produced no chunks", code="no_text") if existing is None: doc_id = uuid.uuid4().hex version = 1 previous_version = None previous_storage = None else: doc_id = existing.doc_id version = existing.version + 1 previous_version = ( existing.version if existing.status == DocumentStatus.READY else None ) previous_storage = existing.storage_name if previous_version is None: # Failed or interrupted earlier attempt: clear its remains. self._purge_versions(doc_id, keep_version=None) storage_name = await asyncio.to_thread( self.storage.save, doc_id, version, extension, content ) now = utcnow_iso() record = DocumentRecord( doc_id=doc_id, version=version, display_filename=display_name, extension=extension, content_hash=digest, size_bytes=len(content), status=DocumentStatus.INDEXING, chunk_count=0, created_at=existing.created_at if existing else now, updated_at=now, storage_name=storage_name, chunking_config=settings.chunking_config_id, embedding_identity=self.embedder.identity if self.embedder else None, user_metadata=metadata, page_count=extracted.page_count, ) if existing is None or previous_version is None: self.manifest.upsert_document(record) # For replacement, the row keeps the old READY version until commit. chunk_records = [ ChunkRecord( chunk_id=make_chunk_id(doc_id, version, c.ordinal), doc_id=doc_id, version=version, ordinal=c.ordinal, text=c.text, location=c.location, text_hash=c.text_hash, ) for c in chunks ] try: self.manifest.replace_chunks(doc_id, version, chunk_records) if self.embedder is not None and self.vector_store is not None: t0 = time.perf_counter() vectors = await asyncio.to_thread( self.embedder.embed_documents, [c.text for c in chunk_records] ) timings["embed"] = (time.perf_counter() - t0) * 1000 if len(vectors) != len(chunk_records): raise RuntimeError("embedding count mismatch") t0 = time.perf_counter() await asyncio.to_thread( self.vector_store.upsert, [c.chunk_id for c in chunk_records], vectors, [ {"doc_id": doc_id, "version": version, "ordinal": c.ordinal} for c in chunk_records ], ) timings["index_vectors"] = (time.perf_counter() - t0) * 1000 present = self.vector_store.has_ids([c.chunk_id for c in chunk_records]) missing = [cid for cid, ok in present.items() if not ok] if missing: raise RuntimeError(f"{len(missing)} vectors missing after write") except Exception as exc: logger.error("Ingestion of %s failed: %s", display_name, exc) self._rollback_version(doc_id, version, storage_name) if previous_version is None: self.manifest.set_status( doc_id, DocumentStatus.FAILED, error=str(exc)[:500], chunk_count=0 ) raise IngestionError( f"Indexing failed: {exc}", code="indexing_failed", status=500 ) from exc # Commit record.status = DocumentStatus.READY record.chunk_count = len(chunk_records) record.updated_at = utcnow_iso() self.manifest.upsert_document(record) if previous_version is not None: self._purge_versions(doc_id, keep_version=version) if previous_storage and previous_storage != storage_name: self.storage.delete(previous_storage) t0 = time.perf_counter() self._commit_corpus_change() timings["index_lexical"] = (time.perf_counter() - t0) * 1000 logger.info( "Indexed %s as %s v%d (%d chunks)", display_name, doc_id, version, len(chunk_records), ) return IngestResult( record, created=True, warnings=extracted.warnings, timings_ms=timings, evicted=evicted, ) async def _make_room_locked(self, max_documents: int) -> List[str]: """Evict the oldest unprotected documents until one more fits. Caller holds the lock.""" evicted: List[str] = [] while True: docs = self.list_documents() if len(docs) < max_documents: return evicted candidates = sorted( (d for d in docs if d.doc_id not in self.protected_doc_ids), key=lambda d: (d.created_at, d.doc_id), ) if not candidates: raise IngestionError( f"The corpus holds its maximum of {max_documents} documents and none " "can be evicted", code="corpus_full", status=409, ) victim = candidates[0] await self._delete_document_locked(victim.doc_id) logger.info( "Evicted %s (%s) to stay within max_documents", victim.doc_id, victim.display_filename, ) evicted.append(victim.doc_id) # -- seeding ---------------------------------------------------------------- async def seed_directory(self, directory: Path) -> SeedResult: """Ingest every supported file in ``directory`` (sorted) and protect the results. Files that fail to ingest are reported, not raised, so one bad sample does not take the whole service down. Re-running on a persistent corpus is idempotent: identical bytes resolve to the existing record. """ directory = Path(directory) doc_ids: List[str] = [] errors: List[str] = [] if not directory.is_dir(): return SeedResult([], [f"seed directory not found: {directory}"]) for path in sorted(directory.iterdir()): if not path.is_file() or path.suffix.lower() not in SUPPORTED_EXTENSIONS: continue try: result = await self.ingest(path.name, path.read_bytes()) except Exception as exc: errors.append(f"{path.name}: {exc}") continue doc_ids.append(result.record.doc_id) self.protected_doc_ids.add(result.record.doc_id) return SeedResult(doc_ids, errors) def _rollback_version(self, doc_id: str, version: int, storage_name: Optional[str]) -> None: try: self.manifest.delete_chunks(doc_id, version) except Exception as exc: # pragma: no cover logger.error("Rollback: could not delete chunks: %s", exc) if self.vector_store is not None: try: self.vector_store.delete_document(doc_id, version) except Exception as exc: # pragma: no cover logger.error("Rollback: could not delete vectors: %s", exc) try: self.storage.delete(storage_name) except Exception as exc: # pragma: no cover logger.error("Rollback: could not delete file: %s", exc) def _purge_versions(self, doc_id: str, keep_version: Optional[int]) -> None: """Remove chunks/vectors of every version except ``keep_version``.""" if self.vector_store is not None: stale = [ cid for cid in self.vector_store.ids_for_document(doc_id) if keep_version is None or f":v{keep_version}:" not in cid ] if stale: self.vector_store.delete_ids(stale) if keep_version is None: self.manifest.delete_chunks(doc_id) return versions = { int(cid.split(":v")[1].split(":")[0]) for cid in self.manifest.chunk_ids(doc_id) if f":v{keep_version}:" not in cid } for version in versions: self.manifest.delete_chunks(doc_id, version) # -- deletion ------------------------------------------------------------- async def delete_document(self, doc_id: str) -> DeleteResult: async with self._write_lock: return await self._delete_document_locked(doc_id) async def _delete_document_locked(self, doc_id: str) -> DeleteResult: record = self.manifest.get_document(doc_id) if record is None: raise NotFoundError(doc_id) # Mark first so retrieval excludes the document even if a later step fails. self.manifest.set_status(doc_id, DocumentStatus.DELETED) self.answer_cache.clear() removed_vectors = 0 try: if self.vector_store is not None: removed_vectors = await asyncio.to_thread(self.vector_store.delete_document, doc_id) remaining = self.vector_store.ids_for_document(doc_id) if remaining: raise RuntimeError(f"{len(remaining)} vectors still present") except Exception as exc: self.manifest.set_status( doc_id, DocumentStatus.DELETED, error=f"vector deletion incomplete: {exc}" ) self._commit_corpus_change() raise IngestionError( "Deletion incomplete: vector index removal failed; retry the request", code="deletion_incomplete", status=500, ) from exc removed_chunks = self.manifest.delete_chunks(doc_id) file_removed = self.storage.delete(record.storage_name) self.manifest.delete_document_rows(doc_id) self.protected_doc_ids.discard(doc_id) self._commit_corpus_change() logger.info("Deleted document %s (%d chunks)", doc_id, removed_chunks) return DeleteResult(doc_id, removed_chunks, removed_vectors, file_removed) async def clear_all(self) -> int: async with self._write_lock: records = self.manifest.list_documents() for record in records: self.manifest.set_status(record.doc_id, DocumentStatus.DELETED) self.answer_cache.clear() if self.vector_store is not None: await asyncio.to_thread(self.vector_store.reset) for record in records: self.storage.delete(record.storage_name) self.manifest.clear() self.protected_doc_ids.clear() self._commit_corpus_change() return len(records) # -- consistency --------------------------------------------------------------- def consistency_report(self) -> Dict[str, Any]: """Compare manifest, vector index and lexical index counts.""" active = self.manifest.active_chunk_count() report: Dict[str, Any] = { "manifest_active_chunks": active, "lexical_entries": self.lexical.size, "lexical_in_sync": self.lexical.size == active and self.lexical.generation == self.manifest.corpus_generation(), } if self.vector_store is not None: report["vector_count"] = self.vector_store.count() report["vectors_in_sync"] = report["vector_count"] == active return report def describe_config(self) -> str: return json.dumps(self.capabilities(), sort_keys=True)