"""KnowledgeStore - persisting artifacts, runs and candidate entries (N4). Writes the three tables the pipeline produces: knowledge_documents one parsed artifact version (content-addressed) knowledge_extraction_runs one row per extraction of one document knowledge_entries one row per candidate, all four kinds Schema of record: `docs/knowledge/KNOWLEDGE_PERSISTENCE_CONTRACT.md` section 2. Go owns the dedorch migration (CLAUDE.md 2.2); this module only INSERTs and SELECTs and executes no DDL, ever. **Deliberately NOT a never-throw seam.** The seams listed in CLAUDE.md 5.4 exist so a *chat turn* degrades to soft output rather than 500ing at the user. An ingestion job is the opposite case: it is asynchronous, the caller polls for its outcome, and a job that silently lost its rows while reporting success is strictly worse than one that reports `failed`. Failures here propagate and the job settles `failed`. Do not "fix" this into the house pattern. Two invariants worth stating because they are easy to break later: - **`payload` is the source of truth.** The promoted columns (`label`, `mention_count`, `extraction_status`, `diff_status`, `definition_conflict`) are a read optimisation so the review queue can be ordered in SQL without opening jsonb. A writer sets both from the same object in the same insert. - **Entries are written per run, never upserted across runs.** Two extractions of one document produce two full sets of rows sharing `entity_id`s. That is the point: it is what makes a quality change attributable. """ from __future__ import annotations import uuid from datetime import UTC, datetime from typing import Any from sqlalchemy import select from src.db.postgres.connection import AsyncSessionLocal from src.db.postgres.models import ( KnowledgeDocumentRow, KnowledgeEntryRow, KnowledgeExtractionRunRow, KnowledgeReviewRow, ) from src.middlewares.logging import get_logger logger = get_logger("knowledge_store") # How each branch's output maps onto a `knowledge_entries` row. The id field # differs per kind because each is derived from different content (`ids.py`), # and the label is whatever a reviewer would recognise the row by in a list. _ENTRY_KINDS: dict[str, tuple[str, str]] = { "glossary": ("term_id", "term"), "rule": ("rule_id", "statement"), "formula": ("formula_id", "name"), # The per-document brief. `domain` is reserved for the scope-level # semantic model composed in `src/knowledge_domain/` (v4). "document": ("brief_id", "title"), # Scope-level, not produced by a run - see `save_declaration`. "domain": ("domain_id", "domain_name"), } def _label(payload: dict[str, Any], key: str) -> str | None: """The listing label, trimmed. A rule's `statement` can be a paragraph, and a column meant for scanning a queue should not carry one.""" value = payload.get(key) if not isinstance(value, str): return None value = value.strip() return (value[:197] + "...") if len(value) > 200 else value or None class KnowledgeStore: """Insert and read the pipeline's persisted output.""" async def find_document( self, doc_id: str, content_hash: str, version: int = 1 ) -> KnowledgeDocumentRow | None: """The content-addressing lookup: has this exact artifact been stored? This is what makes a second click on unchanged documents free. The key is `(doc_id, content_hash, version)` and it is UNIQUE in the schema, so a race between two jobs surfaces as an integrity error rather than a duplicate row. """ async with AsyncSessionLocal() as session: result = await session.execute( select(KnowledgeDocumentRow).where( KnowledgeDocumentRow.doc_id == doc_id, KnowledgeDocumentRow.content_hash == content_hash, KnowledgeDocumentRow.version == version, ) ) return result.scalars().first() async def save_document( self, artifact: Any, *, scope_id: str, user_id: str, document_id: str | None = None, ) -> KnowledgeDocumentRow: """Persist one `ParsedDocument`, or return the row that already holds it. `artifact` is the parsing half's model; it is stored whole in `artifact` (jsonb) with its provenance lifted into columns. The four fingerprints are kept separate on purpose - `parser_config` covers the parser's settings, `normaliser_version` our post-processing, `schema_version` the artifact's shape, `parser_version` the parser build - because none of them moves when the others do, and a stale artifact once went unnoticed for a day precisely because they were conflated. """ payload = ( artifact.model_dump(mode="json") if hasattr(artifact, "model_dump") else dict(artifact) ) doc_id = payload["doc_id"] content_hash = payload["content_hash"] version = int(payload.get("version", 1)) existing = await self.find_document(doc_id, content_hash, version) if existing is not None: logger.info( "knowledge_document_reused", extra={"doc_id": doc_id, "knowledge_document_id": existing.id}, ) return existing row = KnowledgeDocumentRow( id=str(uuid.uuid4()), doc_id=doc_id, document_id=document_id, scope_id=scope_id, user_id=user_id, source_path=payload.get("source_path") or "", content_hash=content_hash, n_pages=int(payload.get("n_pages") or 0), version=version, schema_version=payload.get("schema_version") or "", parser_name=payload.get("parser_name") or "mistral", parser_version=payload.get("parser_version"), parser_backend=payload.get("parser_backend"), parser_config=payload.get("parser_config"), normaliser_version=payload.get("normaliser_version"), raw_output_dir=payload.get("raw_output_dir"), artifact=payload, ) async with AsyncSessionLocal() as session: session.add(row) await session.commit() await session.refresh(row) logger.info( "knowledge_document_saved", extra={ "doc_id": doc_id, "knowledge_document_id": row.id, "chunks": len(payload.get("chunks") or []), }, ) return row async def start_run( self, *, knowledge_document_id: str, doc_id: str, scope_id: str, user_id: str, branches: list[str], model_deployment: str | None = None, ) -> KnowledgeExtractionRunRow: """Open a run row before extraction begins. Opened first, not written at the end: a run that crashes mid-way must still be visible, otherwise the only trace of an expensive failed extraction is its absence. """ row = KnowledgeExtractionRunRow( id=str(uuid.uuid4()), knowledge_document_id=knowledge_document_id, doc_id=doc_id, scope_id=scope_id, user_id=user_id, status="running", branches=branches, model_deployment=model_deployment, ) async with AsyncSessionLocal() as session: session.add(row) await session.commit() await session.refresh(row) return row async def finish_run( self, run_id: str, *, status: str, usages: list[dict] | None = None, rejected: list[dict] | None = None, counts: dict | None = None, error_message: str | None = None, ) -> None: """Close a run with its usage, rejections and counts. `model_version` and `system_fingerprint` come from the API response, not from config: the deployment name stays stable while the model behind it moves, and a run without them recorded is a measurement of an unknown. """ usages = usages or [] async with AsyncSessionLocal() as session: row = await session.get(KnowledgeExtractionRunRow, run_id) if row is None: raise LookupError(f"no extraction run {run_id}") row.status = status row.usage = usages row.rejected = rejected or [] row.counts = counts or {} row.error_message = error_message row.n_calls = len(usages) row.prompt_tokens = sum(u.get("prompt_tokens", 0) or 0 for u in usages) row.cached_tokens = sum(u.get("cached_tokens", 0) or 0 for u in usages) row.completion_tokens = sum( u.get("completion_tokens", 0) or 0 for u in usages ) # Whatever actually served the calls. Distinct values would mean one # run spanned a model change, which is worth seeing rather than # averaging away, so the first non-null is recorded and the set is # visible in `usage`. for key in ("model_version", "system_fingerprint"): value = next((u.get(key) for u in usages if u.get(key)), None) if value: setattr(row, key, value) row.completed_at = datetime.now(UTC) await session.commit() async def save_entries( self, *, run_id: str, knowledge_document_id: str, doc_id: str, scope_id: str, user_id: str, entries_by_kind: dict[str, list[dict]], ) -> int: """Write every candidate entry for one run. Returns the row count. One flush for the whole run: a half-written entry set is not a useful thing to leave behind, and the caller settles the job on the exception. """ rows: list[KnowledgeEntryRow] = [] # `(run_id, kind, entity_id)` is UNIQUE, and ids are content-derived, so # two entries CAN collide inside one run: `formula_id` falls back # name -> latex -> chunk_id, and two formulas sharing a name share an id. # Writing both is impossible; dropping one silently is worse than saying # so, hence the warning and the count on the run row. seen: set[tuple[str, str]] = set() collisions: list[str] = [] for kind, payloads in entries_by_kind.items(): if kind not in _ENTRY_KINDS: raise ValueError(f"unknown entry kind {kind!r}") id_field, label_field = _ENTRY_KINDS[kind] for payload in payloads: entity_id = payload.get(id_field) if not entity_id: # Not recoverable by guessing: the id IS the thing an # expert's approval hangs on, and inventing one would make # the row unreviewable in a way nothing downstream detects. raise ValueError(f"{kind} entry has no {id_field}") if (kind, entity_id) in seen: collisions.append(f"{kind}:{entity_id}") continue seen.add((kind, entity_id)) rows.append( KnowledgeEntryRow( id=str(uuid.uuid4()), run_id=run_id, knowledge_document_id=knowledge_document_id, doc_id=doc_id, scope_id=scope_id, user_id=user_id, kind=kind, entity_id=entity_id, label=_label(payload, label_field), mention_count=int(payload.get("mention_count") or 0), extraction_status=payload.get("extraction_status"), diff_status=payload.get("diff_status"), definition_conflict=bool( payload.get("definition_conflict", False) ), payload=payload, ) ) if not rows: return 0 async with AsyncSessionLocal() as session: session.add_all(rows) await session.commit() if collisions: logger.warning( "knowledge_entry_id_collision", extra={"run_id": run_id, "doc_id": doc_id, "dropped": len(collisions), "ids": collisions[:10]}, ) logger.info( "knowledge_entries_saved", extra={"run_id": run_id, "doc_id": doc_id, "count": len(rows), "collisions": len(collisions)}, ) self.last_collisions = collisions return len(rows) async def list_entries( self, *, scope_id: str, doc_id: str | None = None, kind: str | None = None, run_id: str | None = None, limit: int = 500, ) -> list[KnowledgeEntryRow]: """Read candidates back. Always filtered by tenant. `scope_id` is not optional: every read filters on it, so a missing filter is a cross-tenant read rather than a broader result set. """ stmt = select(KnowledgeEntryRow).where(KnowledgeEntryRow.scope_id == scope_id) if doc_id: stmt = stmt.where(KnowledgeEntryRow.doc_id == doc_id) if kind: stmt = stmt.where(KnowledgeEntryRow.kind == kind) if run_id: stmt = stmt.where(KnowledgeEntryRow.run_id == run_id) stmt = stmt.order_by( KnowledgeEntryRow.definition_conflict.desc(), KnowledgeEntryRow.mention_count.desc(), ).limit(limit) async with AsyncSessionLocal() as session: result = await session.execute(stmt) return list(result.scalars().all()) async def latest_run_for( self, scope_id: str, doc_id: str ) -> KnowledgeExtractionRunRow | None: """The most recent extraction of one document. Reads default to this rather than unioning every historical run's entries, which is never what a caller wants. """ async with AsyncSessionLocal() as session: result = await session.execute( select(KnowledgeExtractionRunRow) .where( KnowledgeExtractionRunRow.scope_id == scope_id, KnowledgeExtractionRunRow.doc_id == doc_id, ) .order_by(KnowledgeExtractionRunRow.created_at.desc()) .limit(1) ) return result.scalars().first() async def list_documents( self, scope_id: str, limit: int = 50 ) -> list[KnowledgeDocumentRow]: """Stored artifacts for one tenant, newest first. What the re-extract route is picked from: prompt retuning runs against a stored artifact and must never re-pay the parser for a byte-identical result. """ async with AsyncSessionLocal() as session: result = await session.execute( select(KnowledgeDocumentRow) .where(KnowledgeDocumentRow.scope_id == scope_id) .order_by(KnowledgeDocumentRow.created_at.desc()) .limit(limit) ) return list(result.scalars().all()) async def get_document( self, knowledge_document_id: str, scope_id: str ) -> KnowledgeDocumentRow | None: """One artifact, tenant-checked. `scope_id` is a predicate rather than a post-hoc assertion: filtering in SQL means a wrong tenant gets "not found" instead of a row it then has to be trusted not to use. """ async with AsyncSessionLocal() as session: result = await session.execute( select(KnowledgeDocumentRow).where( KnowledgeDocumentRow.id == knowledge_document_id, KnowledgeDocumentRow.scope_id == scope_id, ) ) return result.scalars().first() async def find_entry( self, entity_id: str, scope_id: str ) -> KnowledgeEntryRow | None: """The most recent entry row for one durable entity id. A decision is recorded against `entity_id`, which survives re-extraction; `entry_id` does not. So a review posted after a re-run must attach to the NEWEST row for that entity, which is what this returns. """ async with AsyncSessionLocal() as session: result = await session.execute( select(KnowledgeEntryRow) .where( KnowledgeEntryRow.entity_id == entity_id, KnowledgeEntryRow.scope_id == scope_id, ) .order_by(KnowledgeEntryRow.created_at.desc()) .limit(1) ) return result.scalars().first() async def save_review( self, *, entity_id: str, scope_id: str, decision: str, reviewer_id: str, edited_payload: dict | None = None, note: str | None = None, ) -> KnowledgeReviewRow: """Record one expert decision. APPEND-ONLY. A changed mind is a NEW row, never an update: the point of the table is that it records what was ruled and when, and an update destroys the history the whole span-verification apparatus exists to produce. `content_hash` is derived server-side from the artifact the entry came from, never accepted from the caller — a stale client would otherwise attribute a decision to the wrong artifact version, which is exactly the confusion the field exists to prevent. """ entry = await self.find_entry(entity_id, scope_id) if entry is None: raise LookupError(f"no entry {entity_id} in this scope") async with AsyncSessionLocal() as session: document = await session.get( KnowledgeDocumentRow, entry.knowledge_document_id ) content_hash = document.content_hash if document else "" row = KnowledgeReviewRow( id=str(uuid.uuid4()), entry_id=entry.id, entity_id=entity_id, doc_id=entry.doc_id, scope_id=scope_id, content_hash=content_hash, decision=decision, edited_payload=edited_payload, note=note, reviewer_id=reviewer_id, ) session.add(row) await session.commit() await session.refresh(row) logger.info( "knowledge_review_saved", extra={ "entity_id": entity_id, "decision": decision, "content_hash": content_hash[:12], }, ) return row async def latest_reviews( self, scope_id: str, entity_ids: list[str] | None = None ) -> dict[str, KnowledgeReviewRow]: """Current decision per entity — the latest row wins. Returned as a map so a queue read can annotate rows in one pass rather than a query per row. """ stmt = select(KnowledgeReviewRow).where( KnowledgeReviewRow.scope_id == scope_id ) if entity_ids: stmt = stmt.where(KnowledgeReviewRow.entity_id.in_(entity_ids)) stmt = stmt.order_by(KnowledgeReviewRow.created_at.asc()) async with AsyncSessionLocal() as session: result = await session.execute(stmt) # Ascending, so a later row overwrites an earlier one and the map # ends up holding the current decision. return {r.entity_id: r for r in result.scalars().all()} async def approved_glossary(self, scope_id: str) -> list[dict]: """Approved glossary entries — the baseline the next run diffs against. THIS IS WHAT CLOSES THE LOOP: an expert's approvals become the active glossary a later extraction classifies new candidates against, which is what turns a batch extractor into a knowledge base that improves as it is worked. An `edited` decision contributes the expert's corrected payload, not the model's original claim. """ reviews = await self.latest_reviews(scope_id) approved = { entity_id: review for entity_id, review in reviews.items() if review.decision in ("approved", "edited") } if not approved: return [] async with AsyncSessionLocal() as session: result = await session.execute( select(KnowledgeEntryRow).where( KnowledgeEntryRow.scope_id == scope_id, KnowledgeEntryRow.kind == "glossary", KnowledgeEntryRow.entity_id.in_(list(approved)), ) ) rows = list(result.scalars().all()) out: list[dict] = [] seen: set[str] = set() for row in sorted(rows, key=lambda r: r.created_at or 0, reverse=True): if row.entity_id in seen: continue seen.add(row.entity_id) review = approved[row.entity_id] out.append( review.edited_payload if review.decision == "edited" and review.edited_payload else row.payload ) return out async def save_declaration( self, *, scope_id: str, user_id: str, payload: dict ) -> KnowledgeEntryRow: """Persist a scope-level domain declaration as a `kind='domain'` entry. Deliberately an ENTRY rather than its own table: it has to be reviewable, and `knowledge_reviews` keys on `entity_id`. Making it an entry means the expert approves, edits or rejects a domain declaration through exactly the same route as a glossary term - no second review mechanism to build, and no second one to keep consistent. Parentless by nature: it is composed across every document in the scope, so it belongs to no run and no document. A DB CHECK allows that only for this kind. Re-proposing REPLACES the standing draft rather than accumulating versions, because a draft is not a finding. The expert's decisions live in `knowledge_reviews` and are keyed on `entity_id`, so they survive the replacement - which is the whole reason the id is deterministic. """ entity_id = payload["domain_id"] async with AsyncSessionLocal() as session: existing = await session.execute( select(KnowledgeEntryRow).where( KnowledgeEntryRow.scope_id == scope_id, KnowledgeEntryRow.kind == "domain", KnowledgeEntryRow.entity_id == entity_id, ) ) row = existing.scalars().first() if row is None: row = KnowledgeEntryRow( id=str(uuid.uuid4()), run_id=None, knowledge_document_id=None, doc_id=scope_id, # scope-level; there is no document scope_id=scope_id, user_id=user_id, kind="domain", entity_id=entity_id, ) session.add(row) row.label = _label(payload, "domain_name") row.payload = payload row.extraction_status = "ok" await session.commit() await session.refresh(row) logger.info( "knowledge_declaration_saved", extra={"scope_id": scope_id, "entity_id": entity_id}, ) return row async def get_declaration(self, scope_id: str) -> KnowledgeEntryRow | None: """The standing domain declaration for a scope, approved or not.""" async with AsyncSessionLocal() as session: result = await session.execute( select(KnowledgeEntryRow).where( KnowledgeEntryRow.scope_id == scope_id, KnowledgeEntryRow.kind == "domain", ) ) return result.scalars().first()