Agentic-Service-Data-Eyond-Catalog

Running

Agentic-Service-Data-Eyond-Catalog / src /knowledge_ingest /service.py

/fix parsing and term extract (#21)

f07443e about 1 month ago

14.8 kB

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

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