shira-hermes/api/services/kb/ingest_jobs.py

"""Async ingest-job tracking for the KB upload UI (Phase 2 / Task #13).

When a user uploads a PDF/DOCX/TXT through the KB management panel, we don't
want to make the browser wait while we parse → chunk → embed (that's tens
of seconds for a large PDF). Instead:

  1. The HTTP handler writes the bytes to MinIO under
     `inbox/<topic_slug>/<kind>/<uuid>-<filename>`.
  2. Inserts a row in `kb_ingest_job` with `status='queued'`.
  3. Schedules `asyncio.create_task(process_job(job_id))` and returns
     immediately with `{job_id, status:'queued'}`.
  4. The background task transitions queued → processing → done|failed,
     calling `kb_ingest.ingest_source` with the user-supplied metadata.
  5. The UI polls `/admin/kb/jobs/<id>` every 2s until it sees a terminal
     status.

The legacy `/admin/kb/scan-inbox` cron path is unchanged and unaware of
`kb_ingest_job` — it still reads from the flat `inbox/<kind>/` layout. The
new upload path uses a topic-scoped layout (`inbox/<topic_slug>/<kind>/`)
so the cron can't accidentally pick up an upload-form file. Single-source
of truth: when an item is in `kb_ingest_job`, only `process_job` ever
touches it.
"""

from __future__ import annotations

import asyncio
import datetime as _dt
import json
import logging

from api.services.kb import ingest as kb_ingest
from api.services.kb import s3 as kb_s3
from api.services.kb.db import get_pool

logger = logging.getLogger("shira.kb.jobs")


async def create_job(
    *,
    topic_id: int,
    topic_slug: str,
    kind: str,
    original_filename: str,
    file_bytes: bytes,
    metadata: dict,
    requested_by_user: str | None,
) -> int:
    """Write the file to S3 and insert a queued kb_ingest_job row.

    Returns the new job id. The caller is expected to schedule
    `asyncio.create_task(process_job(job_id))` after this returns.
    """
    import uuid as _uuid

    # UUID prefix prevents collisions when two users upload the same filename.
    safe_filename = original_filename.replace("/", "_")
    s3_key = f"inbox/{topic_slug}/{kind}/{_uuid.uuid4().hex[:8]}-{safe_filename}"

    # Upload to MinIO BEFORE the DB insert, so we never have a queued row
    # pointing at a non-existent S3 object. If the upload fails the user
    # gets a 5xx and no job is created.
    s3 = kb_s3._client()
    s3.put_object(
        Bucket=kb_s3._bucket(),
        Key=s3_key,
        Body=file_bytes,
        ContentType=_guess_content_type(original_filename),
    )

    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            INSERT INTO kb_ingest_job
                (original_filename, s3_key, kind, topic_id, metadata_json,
                 status, requested_by_user)
            VALUES ($1, $2, $3, $4, $5::jsonb, 'queued', $6)
            RETURNING id
            """,
            original_filename, s3_key, kind, topic_id,
            json.dumps(metadata, ensure_ascii=False),
            requested_by_user,
        )
    job_id = row["id"]
    logger.info(
        "[kb.jobs] queued job_id=%s kind=%s topic_id=%s file=%r s3=%s",
        job_id, kind, topic_id, original_filename[:80], s3_key,
    )
    return job_id


async def get_job(job_id: int) -> dict | None:
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT * FROM kb_ingest_job WHERE id = $1",
            job_id,
        )
    return _row_to_dict(row) if row else None


async def list_jobs(
    *,
    topic_id: int | None = None,
    status: str | None = None,
    limit: int = 50,
) -> list[dict]:
    limit = max(1, min(200, limit))
    pool = await get_pool()
    where: list[str] = []
    args: list = []
    if topic_id is not None:
        args.append(topic_id)
        where.append(f"topic_id = ${len(args)}")
    if status:
        args.append(status)
        where.append(f"status = ${len(args)}")
    where_sql = ("WHERE " + " AND ".join(where)) if where else ""
    args.append(limit)
    sql = f"""
        SELECT * FROM kb_ingest_job
        {where_sql}
        ORDER BY created_at DESC
        LIMIT ${len(args)}
    """
    async with pool.acquire() as conn:
        rows = await conn.fetch(sql, *args)
    return [_row_to_dict(r) for r in rows]


async def process_job(job_id: int) -> None:
    """Run a single job through ingest_source, recording status transitions.

    Called via `asyncio.create_task(...)` from the upload route. Any
    exception here is caught and recorded as `status='failed'` so the
    polling UI never gets stuck on `processing`.
    """
    pool = await get_pool()

    # 1) queued → processing
    async with pool.acquire() as conn:
        # Guard against double-processing if process_job is somehow scheduled
        # twice for the same id (e.g. retry on container restart later).
        updated = await conn.fetchrow(
            """
            UPDATE kb_ingest_job
               SET status = 'processing', started_at = now()
             WHERE id = $1 AND status = 'queued'
            RETURNING *
            """,
            job_id,
        )
    if not updated:
        logger.warning("[kb.jobs] process_job %s: not in queued state, skipping", job_id)
        return

    job = _row_to_dict(updated)
    logger.info("[kb.jobs] processing job_id=%s s3=%s", job_id, job["s3_key"])

    try:
        data = kb_s3.fetch(job["s3_key"])
        if not data:
            raise kb_ingest.IngestError("S3 object is empty")
        meta = job.get("metadata_json") or {}
        # Title is required by ingest_source; fall back to filename stem.
        title = meta.get("title") or _filename_stem(job["original_filename"])

        result = await kb_ingest.ingest_source(
            kind=job["kind"],
            title=title,
            identifier=meta.get("identifier"),
            content=data,
            filename=job["original_filename"],
            published_at=meta.get("published_at"),
            effective_at=meta.get("effective_at"),
            source_url=meta.get("source_url"),
            original_path=f"s3://{kb_s3._bucket()}/{job['s3_key']}",
            topic_id=job["topic_id"],
        )

        async with pool.acquire() as conn:
            await conn.execute(
                """
                UPDATE kb_ingest_job
                   SET status = 'done',
                       source_id = $2,
                       chunks_created = $3,
                       completed_at = now()
                 WHERE id = $1
                """,
                job_id,
                result["source_id"],
                result["chunks_created"],
            )
        logger.info(
            "[kb.jobs] done job_id=%s source_id=%s chunks=%d",
            job_id, result["source_id"], result["chunks_created"],
        )

    except Exception as exc:
        logger.exception("[kb.jobs] failed job_id=%s", job_id)
        async with pool.acquire() as conn:
            await conn.execute(
                """
                UPDATE kb_ingest_job
                   SET status = 'failed',
                       error_message = $2,
                       completed_at = now()
                 WHERE id = $1
                """,
                job_id,
                # asyncpg can't store >1GB; clamp the message defensively.
                str(exc)[:4000],
            )


# ── helpers ─────────────────────────────────────────────────────────────────

def _row_to_dict(row) -> dict:
    """Convert an asyncpg.Record to a plain dict the API can return as JSON."""
    if row is None:
        return None  # type: ignore[return-value]
    out = {}
    for k, v in dict(row).items():
        if isinstance(v, _dt.datetime):
            out[k] = v.isoformat()
        elif k == "metadata_json" and isinstance(v, str):
            try:
                out[k] = json.loads(v)
            except Exception:
                out[k] = {}
        else:
            out[k] = v
    return out


def _filename_stem(filename: str) -> str:
    stem = filename.rsplit("/", 1)[-1]
    for ext in (".pdf", ".docx", ".txt", ".md"):
        if stem.lower().endswith(ext):
            return stem[: -len(ext)]
    return stem


def _guess_content_type(filename: str) -> str:
    lower = filename.lower()
    if lower.endswith(".pdf"):
        return "application/pdf"
    if lower.endswith(".docx"):
        return "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    if lower.endswith(".txt"):
        return "text/plain; charset=utf-8"
    return "application/octet-stream"