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 classifier as kb_classifier
from api.services.kb import ingest as kb_ingest
from api.services.kb import labels as kb_labels
from api.services.kb import s3 as kb_s3
from api.services.kb import topics as kb_topics
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,
    batch_id: str | None = 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))` (legacy single-stage) or
    `asyncio.create_task(process_classify_stage(job_id, topic_name))`
    (Phase 6 two-phase) after this returns.

    `batch_id` groups multiple files uploaded in one drag-drop drop.
    The review screen uses it to load the whole batch in one round trip.
    """
    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, batch_id)
            VALUES ($1, $2, $3, $4, $5::jsonb, 'queued', $6, $7)
            RETURNING id
            """,
            original_filename, s3_key, kind, topic_id,
            json.dumps(metadata, ensure_ascii=False),
            requested_by_user, batch_id,
        )
    job_id = row["id"]
    logger.info(
        "[kb.jobs] queued job_id=%s kind=%s topic_id=%s file=%r s3=%s batch=%s",
        job_id, kind, topic_id, original_filename[:80], s3_key, batch_id,
    )
    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_by_batch(batch_id: str) -> list[dict]:
    """All jobs in one drag-drop batch, oldest first. Used by the
    review screen for a single round-trip load.
    """
    pool = await get_pool()
    async with pool.acquire() as conn:
        rows = await conn.fetch(
            """
            SELECT * FROM kb_ingest_job
            WHERE batch_id = $1
            ORDER BY created_at, id
            """,
            batch_id,
        )
    return [_row_to_dict(r) for r in rows]


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],
            )


# ── Phase 6 (Task #20): two-phase classify → await_review → embed ─────────

async def process_classify_stage(job_id: int) -> None:
    """First half of the new two-phase pipeline.

    queued → processing(stage=classifying)
           → awaiting_review (with ai_suggestions populated, possibly {})

    Called via asyncio.create_task from the new /admin/kb/upload-batch
    route. Fail-soft: the AI call timing out, returning malformed
    output, or the file being unparseable all transition to
    awaiting_review with empty suggestions — the user fills metadata
    manually and proceeds normally. Hard fail only on S3-unreachable.
    """
    pool = await get_pool()

    # queued → processing(classifying), guarded against double-entry
    async with pool.acquire() as conn:
        updated = await conn.fetchrow(
            """
            UPDATE kb_ingest_job
               SET status = 'processing',
                   processing_stage = 'classifying',
                   started_at = now()
             WHERE id = $1 AND status = 'queued'
            RETURNING *
            """,
            job_id,
        )
    if not updated:
        logger.warning("[classify] job %s not in queued state — skipping", job_id)
        return

    job = _row_to_dict(updated)
    logger.info("[classify] job_id=%s starting", job_id)

    suggestions: dict = {}
    try:
        # Pull just enough of the file to classify. Hard-fail here if S3
        # is unreachable — that's an infra issue, not a classifier issue.
        data = kb_s3.fetch(job["s3_key"])
        if data:
            text = kb_ingest.parse_first_pages(data, job["original_filename"], n=3)
            topic = await kb_topics.get_topic(job["topic_id"])
            topic_name = (topic or {}).get("name")
            suggestions = await kb_classifier.classify(
                text=text,
                filename=job["original_filename"],
                topic_name=topic_name,
            )
    except Exception as e:
        # Note we don't `raise` — empty suggestions are an acceptable
        # outcome. Any unexpected exception means the user fills the
        # form manually instead of getting a head start.
        logger.warning("[classify] job %s suggestions failed: %s", job_id, e)
        suggestions = {}

    # Whatever the outcome, transition to awaiting_review so the UI
    # never gets stuck on `processing`.
    async with pool.acquire() as conn:
        await conn.execute(
            """
            UPDATE kb_ingest_job
               SET status = 'awaiting_review',
                   processing_stage = NULL,
                   ai_suggestions = $2::jsonb
             WHERE id = $1
            """,
            job_id,
            json.dumps(suggestions, ensure_ascii=False),
        )
    logger.info(
        "[classify] job_id=%s awaiting_review (suggestions=%d fields)",
        job_id, len(suggestions),
    )


async def commit_job(
    *,
    job_id: int,
    fields: dict,
    label_specs: list[dict],
    requested_by: str | None,
) -> dict:
    """User confirmed metadata on the review screen. Resolve labels,
    then transition awaiting_review → processing(embedding) and
    schedule the embed task.

    `fields` is whitelisted in the route handler — kind, title,
    identifier, source_url, published_at, effective_at, summary.
    `label_specs` is a list of {slug, name, topic_id?} — existing labels
    by slug, new ones by display name. lookup_or_create_batch resolves
    them all to ids race-safely.
    """
    pool = await get_pool()

    async with pool.acquire() as conn:
        job_row = await conn.fetchrow(
            "SELECT * FROM kb_ingest_job WHERE id = $1", job_id
        )
    if not job_row:
        raise LookupError(f"job {job_id} not found")
    job = _row_to_dict(job_row)
    if job["status"] != "awaiting_review":
        raise ValueError(
            f"job {job_id} is in status {job['status']!r}, expected 'awaiting_review'"
        )

    # Resolve labels — both existing and new — to ids before we kick off
    # the embed task. If label resolution fails the user can re-commit.
    label_id_map = await kb_labels.lookup_or_create_batch(
        label_specs, created_by=requested_by,
    )
    label_ids = list(label_id_map.values())

    # Transition awaiting_review → processing(embedding)
    async with pool.acquire() as conn:
        await conn.execute(
            """
            UPDATE kb_ingest_job
               SET status = 'processing',
                   processing_stage = 'embedding',
                   started_at = COALESCE(started_at, now()),
                   metadata_json = $2::jsonb
             WHERE id = $1 AND status = 'awaiting_review'
            """,
            job_id,
            json.dumps(fields, ensure_ascii=False),
        )

    asyncio.create_task(process_embed_stage(job_id, fields, label_ids))
    return {"job_id": job_id, "status": "processing", "label_ids": label_ids}


async def process_embed_stage(
    job_id: int, fields: dict, label_ids: list[int],
) -> None:
    """Second half: ingest_source + apply_labels.

    Failures here are real failures — the file's been confirmed by the
    user but parse/embed broke. Surface as `failed` with error_message.
    """
    pool = await get_pool()
    async with pool.acquire() as conn:
        job_row = await conn.fetchrow(
            "SELECT * FROM kb_ingest_job WHERE id = $1", job_id
        )
    if not job_row:
        logger.warning("[embed] job %s vanished", job_id)
        return
    job = _row_to_dict(job_row)

    try:
        data = kb_s3.fetch(job["s3_key"])
        if not data:
            raise kb_ingest.IngestError("S3 object is empty")

        title = (fields.get("title") or "").strip() or _filename_stem(job["original_filename"])
        result = await kb_ingest.ingest_source(
            kind=fields.get("kind") or job["kind"],
            title=title,
            identifier=fields.get("identifier"),
            content=data,
            filename=job["original_filename"],
            published_at=fields.get("published_at"),
            effective_at=fields.get("effective_at"),
            source_url=fields.get("source_url"),
            description=fields.get("summary"),
            original_path=f"s3://{kb_s3._bucket()}/{job['s3_key']}",
            topic_id=job["topic_id"],
            ai_classified=bool(job.get("ai_suggestions")),
        )
        source_id = result["source_id"]

        if label_ids:
            await kb_labels.apply_to_source(source_id, label_ids)

        async with pool.acquire() as conn:
            await conn.execute(
                """
                UPDATE kb_ingest_job
                   SET status = 'done',
                       processing_stage = NULL,
                       source_id = $2,
                       chunks_created = $3,
                       completed_at = now()
                 WHERE id = $1
                """,
                job_id, source_id, result["chunks_created"],
            )
        logger.info(
            "[embed] done job_id=%s source_id=%s chunks=%d labels=%d",
            job_id, source_id, result["chunks_created"], len(label_ids),
        )
    except Exception as exc:
        logger.exception("[embed] failed job_id=%s", job_id)
        async with pool.acquire() as conn:
            await conn.execute(
                """
                UPDATE kb_ingest_job
                   SET status = 'failed',
                       processing_stage = NULL,
                       error_message = $2,
                       completed_at = now()
                 WHERE id = $1
                """,
                job_id, str(exc)[:4000],
            )


async def discard_job(job_id: int, requested_by: str | None) -> dict:
    """User removed a job from the review screen. Status → failed,
    S3 file deleted (best-effort), label assignments aren't an issue
    yet (commit hasn't happened).
    """
    pool = await get_pool()
    async with pool.acquire() as conn:
        job_row = await conn.fetchrow(
            "SELECT * FROM kb_ingest_job WHERE id = $1", job_id
        )
    if not job_row:
        raise LookupError(f"job {job_id} not found")
    job = _row_to_dict(job_row)
    if job["status"] in ("done",):
        raise ValueError(f"job {job_id} already done — cannot discard")

    # Best-effort S3 cleanup; the row stays as audit trail.
    try:
        kb_s3._client().delete_object(Bucket=kb_s3._bucket(), Key=job["s3_key"])
    except Exception as e:
        logger.warning("[discard] S3 delete failed for %s: %s", job["s3_key"], e)

    reason = f"discarded by {requested_by or 'user'}"
    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, reason,
        )
    logger.info("[discard] job_id=%s by=%s", job_id, requested_by)
    return {"job_id": job_id, "status": "failed", "reason": reason}


# ── 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"