src/knowledge_ingest/service.py

"""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)