| """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() |
| |
| self.protected_doc_ids: Set[str] = set() |
| self.rebuild_lexical() |
|
|
| |
| @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: |
| |
| |
| self.vector_store.drop() |
| self.manifest.close() |
|
|
| |
| @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), |
| } |
|
|
| |
| 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() |
|
|
| |
| 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) |
|
|
| |
| 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 |
| 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: |
| |
| 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) |
| |
|
|
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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: |
| 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: |
| logger.error("Rollback: could not delete vectors: %s", exc) |
| try: |
| self.storage.delete(storage_name) |
| except Exception as exc: |
| 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) |
|
|
| |
| 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) |
| |
| 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) |
|
|
| |
| 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) |
|
|
