Rifqi Hafizuddin
Claude Opus 5
[NOTICKET] fix: gate spend in cluster units and surface the estimate (Step 6)
cdb5fe8 | """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 | |