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

"""Source management for the KB admin panel (Phase 3 / Task #14).

Today the only way to clean up or correct a source is `DELETE FROM
kb_source WHERE id=...` on the shared PG plus an mc rm in MinIO. This
module backs the new sources table in the EspoCRM "ניהול" tab:

  list_sources       — table view with chunk count and last-ingest status
  update_source      — patch the user-editable metadata
                       (title, identifier, dates, source_url)
  delete_source      — hard delete (cascades chunks; drops S3 object)
  start_reingest     — re-parse the existing file in place; preserves
                       source_id and metadata, replaces chunks
  process_reingest_job — async background worker that does the actual
                         parse→chunk→embed and swaps chunks in a single
                         transaction (UI polls the same /admin/kb/jobs
                         endpoint as upload jobs)

The reingest path differs from ingest_source: ingest_source creates a
NEW kb_source row and supersedes the previous version. Reingest keeps
the same row id (so cached _lastSearch chunk_index references in the
browser still point at a valid source), wipes its chunks, and inserts
fresh ones from the same file. Hand-edited metadata is preserved.
"""

from __future__ import annotations

import datetime as _dt
import json
import logging
import re
from typing import Optional

from api.services.kb import ingest as kb_ingest
from api.services.kb import s3 as kb_s3
from api.services.kb import voyage
from api.services.kb.chunker import chunk as chunk_text
from api.services.kb.db import get_pool

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

# Whitelist of fields the user can edit through PUT /admin/kb/sources/{id}.
# Anything else (kind, topic_id, checksum, source_id, ...) is structural
# and must not be flipped from a metadata edit.
_EDITABLE_FIELDS = ("title", "identifier", "source_url", "published_at",
                    "effective_at", "description")


# ── Listing ─────────────────────────────────────────────────────────────────

async def list_sources(
    *,
    topic_id: Optional[int] = None,
    kind: Optional[str] = None,
    label_id: Optional[int] = None,
    limit: int = 200,
) -> list[dict]:
    """Source rows enriched with chunk_count, last_ingest_status, and labels.

    The last-ingest column is a LATERAL join against kb_ingest_job — picks
    the most recent job row for the source (regardless of status), so the
    UI can flag sources whose last re-ingest failed.

    Labels are aggregated as a JSON array via a correlated subquery —
    keeps the row count clean (no N×labels Cartesian product).
    """
    limit = max(1, min(500, limit))
    where: list[str] = ["s.superseded_by IS NULL"]
    args: list = []
    if topic_id is not None:
        args.append(topic_id)
        where.append(f"s.topic_id = ${len(args)}")
    if kind:
        args.append(kind)
        where.append(f"s.kind = ${len(args)}")
    if label_id is not None:
        args.append(label_id)
        where.append(f"EXISTS (SELECT 1 FROM kb_source_label sl "
                     f"WHERE sl.source_id = s.id AND sl.label_id = ${len(args)})")
    args.append(limit)
    sql = f"""
        SELECT
            s.id,
            s.kind,
            s.title,
            s.identifier,
            s.published_at,
            s.effective_at,
            s.source_url,
            s.original_path,
            s.checksum,
            s.topic_id,
            s.description,
            s.ai_classified_at,
            s.created_at,
            s.updated_at,
            (SELECT COUNT(*) FROM kb_chunk c WHERE c.source_id = s.id) AS chunk_count,
            (
                SELECT json_build_object(
                    'job_id', j.id,
                    'status', j.status,
                    'created_at', j.created_at,
                    'completed_at', j.completed_at,
                    'error_message', j.error_message
                )
                FROM kb_ingest_job j
                WHERE j.source_id = s.id
                ORDER BY j.created_at DESC
                LIMIT 1
            ) AS last_ingest,
            COALESCE(
                (
                    SELECT json_agg(json_build_object(
                        'id', l.id, 'slug', l.slug, 'name', l.name,
                        'topic_id', l.topic_id
                    ) ORDER BY l.name)
                    FROM kb_label l
                    JOIN kb_source_label sl ON sl.label_id = l.id
                    WHERE sl.source_id = s.id
                ),
                '[]'::json
            ) AS labels
        FROM kb_source s
        WHERE {' AND '.join(where)}
        ORDER BY s.kind, s.title
        LIMIT ${len(args)}
    """
    pool = await get_pool()
    async with pool.acquire() as conn:
        rows = await conn.fetch(sql, *args)
    return [_row_to_dict(r) for r in rows]


async def get_source(source_id: int) -> Optional[dict]:
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT * FROM kb_source WHERE id = $1",
            source_id,
        )
    return _row_to_dict(row) if row else None


# ── Editing ─────────────────────────────────────────────────────────────────

async def update_source(
    source_id: int,
    fields: dict,
    *,
    label_ids: Optional[list[int]] = None,
) -> dict:
    """UPDATE only whitelisted columns. Returns the post-update row.

    If `label_ids` is provided (not None), the source's label set is
    REPLACED to exactly those ids — usage_count adjusted both ways.
    Pass an empty list to clear all labels; pass None to leave labels
    untouched.
    """
    from api.services.kb import labels as kb_labels  # local import: cycle

    clean: dict = {}
    for k in _EDITABLE_FIELDS:
        if k not in fields:
            continue
        v = fields[k]
        if k in ("published_at", "effective_at"):
            clean[k] = kb_ingest._coerce_date(v)
        elif v is None or v == "":
            clean[k] = None
        else:
            clean[k] = str(v).strip()
    if "title" in clean and not clean["title"]:
        raise ValueError("title cannot be empty")

    pool = await get_pool()

    if clean:
        set_pieces = []
        args: list = [source_id]
        for k, v in clean.items():
            args.append(v)
            set_pieces.append(f"{k} = ${len(args)}")
        sql = f"""
            UPDATE kb_source
               SET {', '.join(set_pieces)}
             WHERE id = $1
            RETURNING id
        """
        async with pool.acquire() as conn:
            updated = await conn.fetchrow(sql, *args)
        if not updated:
            raise LookupError(f"source {source_id} not found")
    else:
        # Caller passed nothing editable; verify the source still exists
        # so we can return a sensible 404 vs blindly returning empty.
        async with pool.acquire() as conn:
            existing = await conn.fetchrow(
                "SELECT id FROM kb_source WHERE id = $1", source_id
            )
        if not existing:
            raise LookupError(f"source {source_id} not found")

    if label_ids is not None:
        await kb_labels.replace_for_source(source_id, label_ids)

    # Re-read with labels via list_sources([id])-style query — reuse
    # get_source() for simplicity, then fetch labels separately.
    row = await get_source(source_id)
    if row is not None:
        row["labels"] = await kb_labels.get_labels_for_source(source_id)
    return row


# ── Deletion ────────────────────────────────────────────────────────────────

async def delete_source(source_id: int) -> dict:
    """Hard-delete a source. kb_chunk rows cascade via FK; the underlying
    S3 object is deleted on a best-effort basis (a stray file is harmless,
    but cluttered).
    """
    src = await get_source(source_id)
    if not src:
        raise LookupError(f"source {source_id} not found")

    pool = await get_pool()
    async with pool.acquire() as conn:
        result = await conn.execute(
            "DELETE FROM kb_source WHERE id = $1", source_id
        )
    deleted = result.endswith(" 1")

    # Best-effort S3 cleanup — never fail the whole delete on a stray
    # object that might already be gone.
    s3_deleted = False
    bucket, key = _parse_s3_uri(src.get("original_path") or "")
    if key:
        try:
            kb_s3._client().delete_object(Bucket=bucket or kb_s3._bucket(), Key=key)
            s3_deleted = True
        except Exception as exc:
            logger.warning("[admin.delete] S3 delete failed for %s: %s", key, exc)

    return {"deleted": deleted, "source_id": source_id, "s3_deleted": s3_deleted}


# ── Re-ingest ───────────────────────────────────────────────────────────────

async def start_reingest(source_id: int, requested_by_user: Optional[str]) -> int:
    """Create a kb_ingest_job tied to the existing source_id.

    The actual work is in `process_reingest_job`, scheduled via
    `asyncio.create_task` from the route. Returns the new job_id so the
    browser can poll using the same /admin/kb/jobs/{id} endpoint as
    upload jobs.
    """
    src = await get_source(source_id)
    if not src:
        raise LookupError(f"source {source_id} not found")
    if not src.get("original_path"):
        raise ValueError(
            f"source {source_id} has no original_path — re-ingest needs the original file"
        )

    # Pre-flight: confirm the file actually exists in S3 before queuing
    # the job. Saves the user from waiting 30s on a doomed job.
    bucket, key = _parse_s3_uri(src["original_path"])
    if not key:
        raise ValueError(f"unparseable original_path: {src['original_path']}")
    try:
        kb_s3._client().head_object(Bucket=bucket or kb_s3._bucket(), Key=key)
    except Exception as exc:
        raise ValueError(f"original file unavailable: {exc}")

    # Reuse the kb_ingest_job table; mark this row as a re-ingest by
    # pre-filling source_id (uploads start with source_id NULL).
    metadata = {
        "title": src.get("title"),
        "identifier": src.get("identifier"),
        "published_at": _date_to_iso(src.get("published_at")),
        "effective_at": _date_to_iso(src.get("effective_at")),
        "source_url": src.get("source_url"),
        "reingest": True,
    }
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            INSERT INTO kb_ingest_job
                (source_id, original_filename, s3_key, kind, topic_id,
                 metadata_json, status, requested_by_user)
            VALUES ($1, $2, $3, $4, $5, $6::jsonb, 'queued', $7)
            RETURNING id
            """,
            source_id,
            src.get("title") or _filename_from_key(key),
            key,
            src["kind"],
            src["topic_id"],
            json.dumps(metadata, ensure_ascii=False),
            requested_by_user,
        )
    job_id = row["id"]
    logger.info("[admin.reingest] queued job_id=%s source_id=%s", job_id, source_id)
    return job_id


async def process_reingest_job(job_id: int) -> None:
    """Background worker for re-ingest. Replaces chunks in-place; preserves
    the kb_source row and its metadata.

    Errors are written to the job row (status='failed', error_message=...)
    so the polling UI never gets stuck on 'processing'.
    """
    pool = await get_pool()
    async with pool.acquire() as conn:
        job_row = await conn.fetchrow(
            """
            UPDATE kb_ingest_job
               SET status = 'processing', started_at = now()
             WHERE id = $1 AND status = 'queued'
            RETURNING *
            """,
            job_id,
        )
    if not job_row:
        logger.warning("[reingest] job %s not in queued state, skipping", job_id)
        return

    job = _row_to_dict(job_row)
    source_id = job["source_id"]
    if source_id is None:
        await _fail_job(pool, job_id, "re-ingest job has no source_id")
        return

    try:
        bucket, key = _parse_s3_uri(job["s3_key"])  # s3_key here is just the key
        # If s3_key is a bare key (no s3:// prefix), _parse_s3_uri returns
        # ("", "") — fall back to the bare key.
        if not key:
            key = job["s3_key"]
        data = kb_s3._client().get_object(
            Bucket=bucket or kb_s3._bucket(),
            Key=key,
        )["Body"].read()
        if not data:
            raise kb_ingest.IngestError(f"empty file at {key}")

        text = kb_ingest.parse(data, job.get("original_filename") or key)
        if not text:
            raise kb_ingest.IngestError("no text extracted")

        chunks = chunk_text(text, job["kind"])
        if not chunks:
            raise kb_ingest.IngestError("chunker produced zero chunks")

        # Embed first (slow, network-bound) so the actual DB swap is fast.
        texts = [c["content"] for c in chunks]
        vectors = await voyage.embed(texts, input_type="document")

        import hashlib
        new_checksum = hashlib.sha256(text.encode("utf-8")).hexdigest()

        async with pool.acquire() as conn:
            async with conn.transaction():
                # Drop existing chunks and reinsert in one tx so we never
                # leave the source partially-chunked.
                await conn.execute(
                    "DELETE FROM kb_chunk WHERE source_id = $1", source_id
                )
                rows = [
                    (
                        source_id,
                        c["chunk_index"],
                        c["heading_path"],
                        c["section_ref"],
                        c["content"],
                        vectors[i],
                        c["token_count"],
                        c.get("page_number"),
                    )
                    for i, c in enumerate(chunks)
                ]
                await conn.executemany(
                    """
                    INSERT INTO kb_chunk
                        (source_id, chunk_index, heading_path, section_ref,
                         content, embedding, token_count, page_number)
                    VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
                    """,
                    rows,
                )
                # Update only the structural fields that depend on the file
                # (checksum). Hand-edited metadata stays as-is.
                await conn.execute(
                    "UPDATE kb_source SET checksum = $2 WHERE id = $1",
                    source_id, new_checksum,
                )
                await conn.execute(
                    """
                    UPDATE kb_ingest_job
                       SET status = 'done',
                           chunks_created = $2,
                           completed_at = now()
                     WHERE id = $1
                    """,
                    job_id, len(chunks),
                )
        logger.info(
            "[reingest] done job_id=%s source_id=%s chunks=%d",
            job_id, source_id, len(chunks),
        )
    except Exception as exc:
        logger.exception("[reingest] failed job_id=%s", job_id)
        await _fail_job(pool, job_id, str(exc))


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

async def _fail_job(pool, job_id: int, msg: str) -> None:
    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, msg[:4000],
        )


def _row_to_dict(row) -> Optional[dict]:
    if row is None:
        return None
    out: dict = {}
    for k, v in dict(row).items():
        if isinstance(v, _dt.datetime):
            out[k] = v.isoformat()
        elif isinstance(v, _dt.date):
            out[k] = v.isoformat()
        elif k == "metadata_json" and isinstance(v, str):
            try:
                out[k] = json.loads(v)
            except Exception:
                out[k] = {}
        elif k == "last_ingest" and isinstance(v, str):
            try:
                out[k] = json.loads(v)
            except Exception:
                out[k] = None
        elif k == "labels" and isinstance(v, str):
            try:
                out[k] = json.loads(v)
            except Exception:
                out[k] = []
        else:
            out[k] = v
    return out


_S3_URI = re.compile(r"^s3://([^/]+)/(.+)$")

def _parse_s3_uri(uri: str) -> tuple[str, str]:
    """('bucket', 'key/path') or ('', '') if not an s3:// URI."""
    if not uri:
        return ("", "")
    m = _S3_URI.match(uri)
    if not m:
        return ("", "")
    return (m.group(1), m.group(2))


def _filename_from_key(key: str) -> str:
    return key.rsplit("/", 1)[-1] if key else "source"


def _date_to_iso(value) -> Optional[str]:
    if value is None:
        return None
    if isinstance(value, str):
        return value[:10] if value else None
    if isinstance(value, _dt.date):
        return value.isoformat()
    return None