"""IngestService — the job lifecycle: select, admit, run, settle. Owns *orchestration only*. Parsing lives in `src/knowledge_parsing/` (Sofhia) and extraction in `src/knowledge_extraction/` (Rifqi); this module knows neither, and reaches both through `PipelinePort`. That indirection is not decoration — the single-document parsing entry point does not exist yet (N9 in KNOWLEDGE_SURFACE_PLAN.md: `run.py` is folder-oriented and its per-document glue is exactly what N9 extracts), so the port is what lets the job lifecycle be built and verified before it lands. Two rules this module exists to enforce: - **Single-flight per scope.** Two concurrent runs mean two 1.2 GB GLiNER models resident at once. A second click while a job is in flight returns the job already running, rather than queueing a second one. - **One bad document does not fail the batch.** Each document is attempted independently and its outcome recorded; the job settles `partial` when some succeeded and some did not. """ from __future__ import annotations import asyncio from datetime import UTC, datetime from typing import Protocol, runtime_checkable from src.db.postgres.models import Document from src.knowledge_ingest.knowledge_store import KnowledgeStore from src.knowledge_ingest.models import DocumentResult, IngestJob, SkippedDocument from src.knowledge_ingest.pipeline import KnowledgePipeline from src.knowledge_ingest.selector import ( resolve_scope_id, select_documents, ) from src.knowledge_ingest.store import JobStore, PostgresJobStore from src.middlewares.logging import get_logger logger = get_logger("knowledge_ingest") def _roll_up_estimates(results: list[DocumentResult]) -> dict | None: """The batch's pre-flight shape, per document and in total. `IngestJob.estimated_cost` has existed since the job model was written and was never populated, so the field a caller would poll to answer "what is this going to cost, and why was that document refused" always came back null. Spend is driven by CLUSTER COUNT — one LLM call per term cluster — so the per-document breakdown is what an operator needs to see; a total alone cannot tell them which document to split. """ estimates = [(r.document_id, r.estimate) for r in results if r.estimate] if not estimates: return None return { "clusters": sum(e.get("clusters") or 0 for _, e in estimates), "total_calls": sum(e.get("total_calls") or 0 for _, e in estimates), "estimated_prompt_tokens": sum( e.get("estimated_prompt_tokens") or 0 for _, e in estimates ), "per_document": { doc_id: { "clusters": e.get("clusters"), "total_calls": e.get("total_calls"), "estimated_prompt_tokens": e.get("estimated_prompt_tokens"), } for doc_id, e in estimates }, "blocked": { r.document_id: r.blocked_by for r in results if r.blocked_by }, } @runtime_checkable class PipelinePort(Protocol): """Run the whole pipeline over ONE document and report its outcome. Implementations own everything between the blob and the persisted rows: download, parse, filter, cluster, estimate, extract, persist. They report failure by returning a `DocumentResult` with `status="failed"` — a raised exception is caught by the runner and recorded the same way, so neither style can take the batch down. """ async def run_document( self, document: Document, scope_id: str, user_id: str ) -> DocumentResult: ... class UnavailablePipeline: """Placeholder until the parsing entry point (N9) and the tables land. Fails every document with a legible reason instead of pretending to work. Deliberately not a silent no-op: a job that reports success having done nothing is the single most expensive way to be wrong here. """ REASON = ( "pipeline unavailable: single-document parsing entry point not built (N9) " "and knowledge_* tables not migrated (contract pending Harry)" ) async def run_document( self, document: Document, scope_id: str, user_id: str ) -> DocumentResult: return DocumentResult( document_id=document.id, filename=document.filename, status="failed", error=self.REASON, ) class IngestService: """Create and drive ingestion jobs.""" def __init__( self, store: JobStore | None = None, pipeline: PipelinePort | None = None ) -> None: self._store = store or PostgresJobStore() self._pipeline = pipeline or KnowledgePipeline() self._admit = asyncio.Lock() # Admission is per scope; EXECUTION is serialised process-wide. Two # scopes may each have a job admitted, but they must not extract at the # same time — the GLiNER span model is ~1.2 GB resident and two copies # is how this falls over on a Space. self._knowledge = KnowledgeStore() self._run_lock = asyncio.Lock() self._tasks: set[asyncio.Task] = set() async def get(self, job_id: str) -> IngestJob | None: return await self._store.get(job_id) async def list_for_user(self, user_id: str, limit: int = 20) -> list[IngestJob]: return await self._store.list_for_user(user_id, limit) async def start( self, user_id: str, document_ids: list[str] | None = None, review_mode: str = "expert", ) -> tuple[IngestJob, bool]: """Admit a batch. Returns (job, reused). `reused=True` means a job for this scope was already in flight and is being returned instead of a second one being started. The button asks for a state, not a command, so a double click is not an error. """ scope_id = await resolve_scope_id(user_id) # Admission is serialised so two simultaneous clicks cannot both find # the scope idle and each start a run. async with self._admit: active = await self._store.find_active(scope_id) if active is not None: logger.info( "ingest_job_reused", extra={"job_id": active.id, "scope_id": scope_id}, ) return active, True selected, skipped = await select_documents(user_id, document_ids) job = IngestJob( scope_id=scope_id, user_id=user_id, review_mode=review_mode, document_ids=[d.id for d in selected], n_requested=len(selected), results=[ DocumentResult(document_id=d.id, filename=d.filename) for d in selected ], skipped=[ SkippedDocument( document_id=d.id, filename=d.filename, reason=reason ) for d, reason in skipped ], ) # Nothing to do is a terminal state, not a run. Settling here keeps # the caller from polling a job that will never move. if not selected: job.status = "succeeded" if not skipped else "failed" job.error_message = ( None if not skipped else "No parseable documents in this batch." ) job.completed_at = datetime.now(UTC) await self._store.create(job) return job, False await self._store.create(job) self._spawn(job, selected) return job, False def _spawn(self, job: IngestJob, documents: list[Document]) -> None: """Run the batch in the background; the request does not wait on it.""" task = asyncio.create_task(self._run(job, documents)) # asyncio holds only a weak reference to a running task, so without this # a batch can be garbage-collected mid-flight. Discarded on completion. self._tasks.add(task) task.add_done_callback(self._tasks.discard) async def _run(self, job: IngestJob, documents: list[Document]) -> None: async with self._run_lock: await self._run_serialised(job, documents) async def _run_serialised(self, job: IngestJob, documents: list[Document]) -> None: job.status = "running" job.started_at = datetime.now(UTC) await self._store.save(job) by_id = {r.document_id: r for r in job.results} for document in documents: try: result = await self._pipeline.run_document( document, job.scope_id, job.user_id ) except Exception as e: # one document must not take the batch down logger.error( "ingest_document_failed", extra={ "job_id": job.id, "document_id": document.id, "error": repr(e), }, ) result = DocumentResult( document_id=document.id, filename=document.filename, status="failed", error=repr(e), ) # r1: in `llm` mode there is no reviewer step - the run's # entries are approved as soon as they land, recorded against # the person who chose the mode (r2/r3). Never-throw: an # auto-approval that fails leaves the entries REVIEWABLE, which # is the safe direction, and must not fail the document. if job.review_mode == "llm" and result.status == "succeeded" and result.run_id: try: await self._knowledge.auto_approve_run( run_id=result.run_id, scope_id=job.scope_id, authorised_by=job.user_id, ) except Exception as e: # noqa: BLE001 logger.error( "ingest_auto_approve_failed", extra={"job_id": job.id, "run_id": result.run_id, "error": repr(e)}, ) by_id[result.document_id] = result job.results = list(by_id.values()) job.estimated_cost = _roll_up_estimates(job.results) await self._store.save(job) job.settle() await self._store.save(job) logger.info( "ingest_job_settled", extra={ "job_id": job.id, "status": job.status, "succeeded": job.n_succeeded, "failed": job.n_failed, }, ) async def start_reextract( self, user_id: str, knowledge_document_id: str ) -> tuple[IngestJob, bool]: """Admit a re-extraction of one stored artifact. Returns (job, reused). Shares the job machinery with `/ingest` — same single-flight rule, same polling route — but skips parsing entirely, which is the point: prompt retuning is the common case and re-parsing a byte-identical artifact costs money for an identical result. """ scope_id = await resolve_scope_id(user_id) async with self._admit: active = await self._store.find_active(scope_id) if active is not None: logger.info( "ingest_job_reused", extra={"job_id": active.id, "scope_id": scope_id}, ) return active, True document = await self._knowledge.get_document( knowledge_document_id, scope_id ) job = IngestJob(scope_id=scope_id, user_id=user_id) # Tenant-checked in SQL, so a wrong scope is indistinguishable from a # wrong id — deliberately: neither should confirm the row exists. if document is None: job.status = "failed" job.error_message = ( f"No stored artifact {knowledge_document_id} in this scope." ) job.completed_at = datetime.now(UTC) await self._store.create(job) return job, False # `document_id` is the source `documents` row; it is null for an # artifact parsed by the CLI rather than ingested, and the result # falls back to the artifact's own key so the row is still legible. job.document_ids = [knowledge_document_id] job.n_requested = 1 job.results = [ DocumentResult( document_id=document.document_id or document.doc_id, filename=None, knowledge_document_id=document.id, ) ] await self._store.create(job) self._spawn_reextract(job, document) return job, False def _spawn_reextract(self, job: IngestJob, document) -> None: task = asyncio.create_task(self._run_reextract(job, document)) self._tasks.add(task) task.add_done_callback(self._tasks.discard) async def _run_reextract(self, job: IngestJob, document) -> None: async with self._run_lock: job.status = "running" job.stage = "extracting" job.started_at = datetime.now(UTC) await self._store.save(job) result = job.results[0] try: # The approved set is fetched inside `extract_and_persist` for # The approved set is fetched inside `extract_and_persist` for # BOTH paths - it used to be passed here and forgotten on the # ingest path, which silently disabled cross-document diffing. # This is the edge that closes the loop: an expert's approvals # become the baseline the next run classifies against. n, run_id = await self._pipeline.extract_and_persist( document.artifact, document.id, scope_id=job.scope_id, user_id=job.user_id, ) result.run_id = run_id result.n_entries = n result.status = "succeeded" except Exception as e: logger.error( "reextract_failed", extra={"job_id": job.id, "error": repr(e)}, ) result.status = "failed" result.error = repr(e) job.results = [result] job.settle() await self._store.save(job)