"""The real `PipelinePort` - one document, blob to persisted rows (E1-E4). Stages, ordered by what they cost: download blob -> a temp file free parse Mistral OCR, per page PAID (cheap) filter GLiNER + regex free cluster grouping mentions into term clusters free estimate builds the exact prompts, sends none free ---------------------------------------- the spend gate ---------------- extract ~90 LLM calls for nine pages PAID (the real cost) persist artifact + run + entries free The gate sits where it does because `estimate_cost()` needs a parsed, clustered document to be exact - it builds the real prompts and sends nothing. So a job can refuse *before* the expensive stage while still quoting a real number, which is the difference between a spend ceiling and a guess. This module owns no extraction logic. Parsing is `src/knowledge_parsing/` (Sofhia) and extraction is `src/knowledge_extraction/` (Rifqi); both are called through their public entry points and neither is reimplemented here. **`doc_id` is `documents.id`, deliberately.** Every `entity_id` downstream is derived from it (`ids.py`), so a filename-derived key means renaming a file mints all-new ids and orphans every expert approval recorded against the old ones. `parse_one` accepts the key but cannot enforce it - passing the stable one is this caller's job, and is the whole point of the parsing-side change. """ from __future__ import annotations import asyncio import tempfile from concurrent.futures import ThreadPoolExecutor from functools import partial from pathlib import Path from src.db.postgres.models import Document from src.knowledge_extraction.genre import classify as classify_genre from src.knowledge_extraction.service import ( build_clusters, estimate_cost, extract_all, run_filters, ) from src.knowledge_ingest.knowledge_store import KnowledgeStore from src.knowledge_ingest.models import DocumentResult from src.middlewares.logging import get_logger logger = get_logger("knowledge_pipeline") # Refuse a document whose estimate exceeds this. Deliberately generous - it is a # runaway guard, not a budget: the reference nine-page standard estimates around # 150k prompt tokens, so this stops a 2,000-page handbook being ingested by # accident without getting in the way of anything real. MAX_PROMPT_TOKENS = 2_000_000 # The SECOND ceiling, and the one that tracks what actually scales. # # Spend is driven by CLUSTER COUNT, not pages: the pipeline makes one call per # term cluster, so a dense chapter costs more than a long simple one. Two things # the token ceiling alone cannot see: # # 1. **Per-call cost is not linear in prompt tokens.** Each call carries its # own completion tokens and its own latency, and the fixed cached prefix # becomes a smaller share of each prompt as evidence grows - so the CACHE # HIT RATE degrades as a run gets bigger, and the real bill rises faster # than the token estimate does. # 2. **The estimate is partly an approximation** (see `estimate_cost`), while # cluster count is exact and known for free before anything is spent. # # This became urgent rather than theoretical on 2026-09-10: retiring the fuzzy # clustering pass was correct, and it produces MORE, SMALLER clusters - the Open # Pit textbook's single 46-page chapter went from 56% to 71% of the token # ceiling on that change alone. A ceiling that only counts tokens would have let # the next such improvement through unnoticed. # # Measured free on the three v2 artifacts, 2026-09-11: # # clusters calls est. prompt tokens % of the token ceiling # BUMA 83 106 201,751 10.1% # Komatsu 93 95 227,329 11.4% # Open Pit 427 494 1,560,548 78.0% # # Read those honestly: at this corpus's density (~3,200 prompt tokens per # cluster) the TOKEN ceiling binds first, at roughly 630 clusters. So this # ceiling is not the tighter of the two for documents that look like these - it # is the guard for the other shape, a document with many cheap clusters, where # the cost is call COUNT rather than prompt size: 494 calls already runs about a # quarter of an hour, and a batch is twenty documents. # # 750 is ~1.75x the largest real document measured. Generous enough not to # obstruct anything observed, tight enough that one document cannot occupy the # single-flight worker for an hour. Raise it with a measurement, not a hunch. MAX_CLUSTERS = 750 # Extraction is synchronous and blocks for ~170 s on nine pages, so it cannot run # on the event loop. Its own pool, following the #49 `dbexec` precedent: sharing # the default `to_thread` pool means a long parse competes with every other # `to_thread` caller in the process, notably the tabular Parquet loader. # # One worker, not several: the point is to get the work OFF the loop, not to run # batches concurrently. Two concurrent extractions would hold two ~1.2 GB GLiNER # models resident, which is how this falls over on a Space. _EXECUTOR = ThreadPoolExecutor(max_workers=1, thread_name_prefix="knowledge") # Which branches a run covers, in `extract_all`'s own vocabulary. Recorded on # the run row so a later comparison knows what was actually ATTEMPTED rather # than inferring it from what happened to appear. # # Note "summary" is the branch that produces the DocumentBrief - the branch kept # its original name through the v3 rename, while the entity and the entry `kind` # are both "document". Passing "document" here silently extracts nothing. BRANCHES = ("glossary", "rule", "formula", "summary") class SpendCeilingExceededError(RuntimeError): """The pre-flight estimate exceeded a ceiling; nothing was spent. Names WHICH ceiling bound and by how much. An operator's next move differs completely between the two - a cluster refusal means the document is dense and wants splitting, a token refusal means it is simply long - and "over budget" tells them neither. """ def __init__(self, estimate: dict, ceiling: str) -> None: self.estimate = estimate self.ceiling = ceiling if ceiling == "clusters": measured, limit, unit = estimate.get("clusters", 0), MAX_CLUSTERS, "term clusters" else: measured, limit, unit = ( estimate.get("estimated_prompt_tokens", 0), MAX_PROMPT_TOKENS, "prompt tokens", ) self.measured, self.limit = measured, limit super().__init__( f"estimated {measured:,} {unit} exceeds the ceiling of {limit:,} " f"({measured / limit:.0%} of it) — nothing was spent; " f"{estimate.get('total_calls', 0):,} LLM calls were declined" ) def check_ceilings(estimate: dict) -> None: """Raise if either ceiling binds. Clusters are checked FIRST because they are exact where the token figure is partly approximate.""" if (estimate.get("clusters") or 0) > MAX_CLUSTERS: raise SpendCeilingExceededError(estimate, "clusters") if (estimate.get("estimated_prompt_tokens") or 0) > MAX_PROMPT_TOKENS: raise SpendCeilingExceededError(estimate, "prompt_tokens") def _extractor(): """The Azure client the extraction branches call. Built here rather than injected so the pipeline has one obvious wiring point; `AzureExtractor` raises if its deployment quad is unset, which is a better failure than a run that starts and then fails on its first call. """ from src.knowledge_extraction.extract import AzureExtractor return AzureExtractor() def _storage(): """The configured blob backend, chosen the same way `TabularExecutor` does.""" from src.config.settings import settings provider = (settings.storage_provider or "").strip().lower() if provider == "supabase_s3": from src.storage.object_storage.supabase_s3 import SupabaseS3Storage return SupabaseS3Storage() from src.storage.az_blob.az_blob import AzureBlobStorage return AzureBlobStorage() class KnowledgePipeline: """Run one document end to end and report its outcome. Returns a failed `DocumentResult` rather than raising for expected failures (an unreadable blob, a spend refusal); the runner records either shape identically, so neither can take a batch down. """ def __init__(self, store: KnowledgeStore | None = None) -> None: self._store = store or KnowledgeStore() # The most recent pre-flight estimate, so `run_document` can report it # on the SUCCESS path too (a refusal carries its own on the exception). # Safe as instance state only because documents are run one at a time — # `IngestService` holds a single-flight lock and `_EXECUTOR` has one # worker. If either ever changes, pass it back instead of stashing it. self.last_estimate: dict | None = None async def run_document( self, document: Document, scope_id: str, user_id: str ) -> DocumentResult: result = DocumentResult( document_id=document.id, filename=document.filename, status="pending" ) try: artifact = await self._parse(document) except Exception as e: logger.error( "knowledge_parse_failed", extra={"document_id": document.id, "error": repr(e)}, ) result.status = "failed" result.error = f"parse failed: {e!r}" return result doc_row = await self._store.save_document( artifact, scope_id=scope_id, user_id=user_id, document_id=document.id ) result.knowledge_document_id = doc_row.id try: n_entries, run_id = await self.extract_and_persist( artifact, doc_row.id, scope_id=scope_id, user_id=user_id ) except SpendCeilingExceededError as e: logger.warning( "knowledge_spend_ceiling", extra={ "document_id": document.id, "ceiling": e.ceiling, "measured": e.measured, "limit": e.limit, "estimate": e.estimate, }, ) result.status = "failed" result.error = str(e) result.estimate = e.estimate result.blocked_by = e.ceiling return result except Exception as e: logger.error( "knowledge_extract_failed", extra={"document_id": document.id, "error": repr(e)}, ) result.status = "failed" result.error = f"extraction failed: {e!r}" return result result.run_id = run_id result.n_entries = n_entries result.estimate = self.last_estimate result.status = "succeeded" return result async def _parse(self, document: Document): """Blob -> `ParsedDocument`, with `documents.id` as the stable key.""" from src.knowledge_parsing.config import PipelineConfig from src.knowledge_parsing.run import parse_one storage = _storage() content = await storage.download_file(document.blob_name) suffix = Path(document.filename or "").suffix or ".pdf" with tempfile.TemporaryDirectory(prefix="knowledge-ingest-") as tmp: # Named for the original file so anything reading `source_path` or # `source_title` sees something a human recognises; identity comes # from `doc_id`, which is passed explicitly. source = Path(tmp) / f"{Path(document.filename or 'document').stem}{suffix}" source.write_bytes(content) cfg = PipelineConfig(input_dir=Path(tmp)) loop = asyncio.get_running_loop() return await loop.run_in_executor( _EXECUTOR, parse_one, source, cfg, document.id ) async def extract_and_persist( self, artifact, knowledge_document_id: str, *, scope_id: str, user_id: str, active_glossary: list[dict] | None = None, ) -> tuple[int, str]: """Filter, cluster, gate on cost, extract, persist. Returns (rows, run_id). Also the entry point for re-extracting a stored artifact, which needs no parsing at all - prompt retuning is the common case and must not re-pay the parser for a byte-identical result. `active_glossary` defaults to **the tenant's approved entries, fetched here**. Passing it was previously the caller's job and the ingest path forgot, so every entry came back `diff_status="new"`: the diff ran against an empty list, and a second document never showed as duplicating or contradicting what an expert had already approved - which is exactly the cross-document dedup that makes the queue finishable. Fetching it at the single point that needs it means no caller can omit it again. Pass an explicit `[]` to diff against nothing. """ from src.knowledge_extraction.adapter import parsed_doc_from_artifact if active_glossary is None: active_glossary = await self._store.approved_glossary(scope_id) doc = parsed_doc_from_artifact(artifact) loop = asyncio.get_running_loop() # Free stages, but CPU-bound (GLiNER), so still off the loop. filtered = await loop.run_in_executor(_EXECUTOR, run_filters, doc) clustered = await loop.run_in_executor(_EXECUTOR, build_clusters, doc, filtered) estimate = estimate_cost(doc, clustered, filtered) # What KIND of document this is, decided before anything is spent. # Advisory today — see `genre.ROUTE_LABELS` / `ROUTE_BRANCHES`, both # off — but recorded, because "104 calls, 93 abstentions" is a fact an # operator should be able to see the reason for. # # A never-throw seam, like every other observability path here: this # describes the run, it does not decide it, and a classifier defect must # not take down an extraction that was about to work. `repr(e)` because # an empty `str(e)` is how the Fernet mismatch hid for a day. try: verdict = classify_genre(doc, filtered) estimate["genre"] = verdict.as_dict() for warning in verdict.warnings: logger.warning( "knowledge_genre_warning", extra={ "doc_id": doc.doc_id, "genre": verdict.genre, "warning": warning, }, ) except Exception as e: # noqa: BLE001 - described above logger.warning( "knowledge_genre_failed", extra={"doc_id": doc.doc_id, "error": repr(e)}, ) self.last_estimate = estimate # Logged on EVERY run, not only on a refusal: an operator who can only # see the number when it is already too late cannot plan a batch, and # "why was this refused" is answerable only against what a run that # succeeded looked like. logger.info( "knowledge_pre_flight_estimate", extra={ "doc_id": doc.doc_id, "clusters": estimate.get("clusters"), "total_calls": estimate.get("total_calls"), "estimated_prompt_tokens": estimate.get("estimated_prompt_tokens"), "cluster_ceiling_pct": round( 100 * (estimate.get("clusters") or 0) / MAX_CLUSTERS, 1 ), "token_ceiling_pct": round( 100 * (estimate.get("estimated_prompt_tokens") or 0) / MAX_PROMPT_TOKENS, 1, ), }, ) check_ceilings(estimate) # Built before the run row so a missing Azure deployment fails without # leaving an orphan `running` row behind. extractor = _extractor() run = await self._store.start_run( knowledge_document_id=knowledge_document_id, doc_id=doc.doc_id, scope_id=scope_id, user_id=user_id, branches=list(BRANCHES), model_deployment=getattr(extractor, "deployment", None), ) # EVERYTHING after `start_run` is inside this block. It used to cover only # `extract_all`, so a failure in persistence left the run row `running` # with n_calls=0 - a row claiming nothing happened for a run that had # already spent ~90 LLM calls. An expensive failure must leave an honest # trace, which is the whole reason the row is opened before the spend. try: out = await loop.run_in_executor( _EXECUTOR, partial( extract_all, doc, clustered, filtered, extractor, active_glossary=active_glossary or None, branches=BRANCHES, ), ) entries = { "glossary": out.glossary, "rule": out.rules, "formula": out.formulas, "document": [out.brief] if out.brief else [], } n = await self._store.save_entries( run_id=run.id, knowledge_document_id=knowledge_document_id, doc_id=doc.doc_id, scope_id=scope_id, user_id=user_id, entries_by_kind=entries, ) except Exception as e: # The run row is closed even on failure: an expensive extraction # that crashed must leave a trace, not just an absence. Persistence # is inside the block for the same reason - a UNIQUE violation on # entry ids used to leave the row `running` with n_calls=0, claiming # nothing had happened for a run that had already been paid for. await self._store.finish_run( run.id, status="failed", error_message=repr(e) ) raise await self._store.finish_run( run.id, status="succeeded", usages=[u.model_dump() if hasattr(u, "model_dump") else dict(u) for u in out.usages], rejected=[ r.model_dump() if hasattr(r, "model_dump") else dict(r) for r in out.rejected ], counts={ "entries": n, "glossary": len(out.glossary), "rules": len(out.rules), "formulas": len(out.formulas), "queue": len(out.review_queue), **{f"link_{k}": v for k, v in (out.links or {}).items()}, }, ) return n, run.id