ai-resume/corpus/pipeline.py

"""The pipeline: content/ -> validate -> chunk -> enrich -> embed -> upsert.

One entry point, `run`, with two modes:

* **rebuild** — drop the collection, empty the store, embed everything. This is
  the "one command rebuilds the corpus from an empty volume" path.
* **sync** — the default. Only documents whose content hash, embedding model,
  embedding dimension or chunker version changed are re-chunked and
  re-embedded. Everything else is left alone.

The order matters. Payloads are built and validated for *every* changed document
before a single vector is written, so an invalid payload stops the run with an
empty diff rather than leaving half an index behind.

Deterministic ids (`sha256(source_path + chunk_index + content_hash)`) mean
re-running updates points instead of duplicating them, and the id is a valid
UUID so Qdrant accepts it as a point id without a mapping table.
"""

from __future__ import annotations

import hashlib
import json
import time
import uuid
from dataclasses import asdict, dataclass, field

from . import chunker, embeddings, loader, payloads, store as store_module
from .config import CHUNKER_VERSION
from .qdrant import Qdrant, QdrantError


class PipelineError(RuntimeError):
    """The run cannot continue without corrupting the index."""


def chunk_id(source_path: str, chunk_index: int, content_hash: str) -> str:
    """Deterministic point id: sha256(source_path + chunk_index + content_hash).

    Rendered as a UUID because Qdrant only accepts unsigned integers or UUIDs as
    point ids; the first 128 bits of the digest keep collisions improbable.
    """
    digest = hashlib.sha256(
        f"{source_path}\x00{chunk_index}\x00{content_hash}".encode("utf-8")
    ).hexdigest()
    return str(uuid.UUID(hex=digest[:32]))


@dataclass
class RunStats:
    mode: str = "sync"
    documents: int = 0
    added: int = 0
    updated: int = 0
    unchanged: int = 0
    removed: int = 0
    chunks_written: int = 0
    chunks_deleted: int = 0
    embed_calls: int = 0
    texts_embedded: int = 0
    points_sent: int = 0
    collection_points: int = 0
    store_chunks: int = 0
    duration_s: float = 0.0
    collection: str = ""
    provider: str = ""
    model: str = ""
    dim: int = 0
    chunker_version: str = CHUNKER_VERSION
    notes: list = field(default_factory=list)

    def summary(self) -> str:
        rows = [
            f"mode={self.mode} documents={self.documents} chunks={self.store_chunks}",
            f"documents added={self.added} updated={self.updated} unchanged={self.unchanged} removed={self.removed}",
            f"chunks written={self.chunks_written} deleted={self.chunks_deleted}",
            f"embedding calls={self.embed_calls} texts={self.texts_embedded}",
            f"points sent={self.points_sent} collection points={self.collection_points}",
            f"provider={self.provider} model={self.model} dim={self.dim} chunker={self.chunker_version}",
            f"collection={self.collection} duration={self.duration_s:.2f}s",
        ]
        rows.extend(self.notes)
        return "\n".join(rows)

    def to_dict(self) -> dict:
        return asdict(self)


def run(cfg, mode: str = "sync", log=None) -> RunStats:
    """Build or refresh the corpus. Returns what happened; raises on any error."""
    if mode not in ("sync", "rebuild"):
        raise PipelineError(f"unknown mode {mode!r}; expected 'sync' or 'rebuild'")
    started = time.time()
    log = log or (lambda _msg: None)
    stats = RunStats(
        mode=mode,
        collection=cfg.collection,
        provider=cfg.provider,
        model=cfg.model,
        dim=cfg.dim,
        chunker_version=cfg.chunker_version,
    )

    documents = loader.load_documents(cfg)
    stats.documents = len(documents)
    log(f"validated {len(documents)} document(s)")

    provider = embeddings.build_provider(cfg)
    qdrant = Qdrant(cfg.qdrant_url, timeout=cfg.timeout)
    if not qdrant.health():
        raise PipelineError(
            f"Qdrant is not answering at {cfg.qdrant_url}. Start it "
            "(deploy/qdrant.yml) or point RESUME_QDRANT_URL at a running instance."
        )

    with store_module.Store(cfg.db_path) as store:
        contract = _contract(cfg)
        mode = _check_contract(store, qdrant, cfg, contract, mode)

        _ensure_collection(qdrant, cfg)

        manifest = store.manifest()
        changed, removed = _classify(documents, manifest, cfg, mode)
        stats.unchanged = len(documents) - len(changed)
        stats.removed = len(removed)
        log(
            f"incremental plan: {len(changed)} to embed, "
            f"{stats.unchanged} unchanged, {len(removed)} removed"
        )

        for path in removed:
            stats.chunks_deleted += store.delete_document(path)
            qdrant.delete_by_source(cfg.collection, path)
            log(f"removed {path}")

        for doc in changed:
            payloads_for_doc = _prepare(doc, cfg, manifest)
            stats.chunks_written += len(payloads_for_doc)

            vectors = provider.embed([p["text"] for p in payloads_for_doc])
            if len(vectors) != len(payloads_for_doc):
                raise PipelineError(
                    f"{doc.path}: provider returned {len(vectors)} vectors for "
                    f"{len(payloads_for_doc)} chunks"
                )

            points = []
            for payload, vector in zip(payloads_for_doc, vectors):
                if len(vector) != cfg.dim:
                    raise PipelineError(
                        f"{doc.path}: vector dimension {len(vector)} != configured {cfg.dim}"
                    )
                points.append({"id": payload["id"], "vector": vector, "payload": payload})

            # Order: remove what this document used to have (chunk counts change),
            # then write the new rows and points. A crash between the two leaves an
            # index that reports as incomplete, never one that claims to be whole.
            stats.chunks_deleted += store.delete_document(doc.path)
            qdrant.delete_by_source(cfg.collection, doc.path)
            stats.points_sent += qdrant.upsert(cfg.collection, points)
            store.write_document(doc, payloads_for_doc)

            if doc.path in manifest:
                stats.updated += 1
                log(f"updated {doc.path} ({len(points)} chunk(s))")
            else:
                stats.added += 1
                log(f"added {doc.path} ({len(points)} chunk(s))")

        provider_stats = provider.stats()
        stats.embed_calls = provider_stats["calls"]
        stats.texts_embedded = provider_stats["texts"]
        stats.store_chunks = store.chunk_count()
        stats.collection_points = qdrant.point_count(cfg.collection)

        if stats.collection_points != stats.store_chunks:
            raise PipelineError(
                "index is inconsistent after the run: "
                f"{stats.collection_points} point(s) in Qdrant but "
                f"{stats.store_chunks} chunk(s) in SQLite. Run `python3 -m corpus.cli verify`."
            )

        store.set_meta("collection", contract)
        store.set_meta("last_run", {**stats.to_dict(), "finished_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())})
        _write_collection_file(cfg, store, stats)

    stats.duration_s = round(time.time() - started, 3)
    return stats


def _contract(cfg) -> dict:
    """What this collection *is*. Recorded so a change is a migration, not drift."""
    return {
        "provider": cfg.provider,
        "model": cfg.model,
        "dim": cfg.dim,
        "chunker_version": cfg.chunker_version,
        "distance": "Cosine",
    }


def _check_contract(store, qdrant, cfg, contract: dict, mode: str) -> str:
    stored = store.get_meta("collection")
    if mode == "rebuild":
        qdrant.drop_collection(cfg.collection)
        store.reset()
        return mode
    if stored is None:
        # No contract yet: either a fresh store, or one whose database was lost.
        if qdrant.collection_exists(cfg.collection) and qdrant.point_count(cfg.collection) > 0:
            raise PipelineError(
                f"collection {cfg.collection!r} holds points but this store has no "
                "collection contract (was data/corpus/corpus.db deleted?). "
                "Confirm you want to overwrite it: python3 -m corpus.cli rebuild"
            )
        return mode
    if stored != contract:
        raise PipelineError(
            "the collection contract does not match this configuration — refusing to "
            "mix vectors in one collection (docs/phase-4-corpus.md §4.3).\n"
            f"  stored: {json.dumps(stored, sort_keys=True)}\n"
            f"  config: {json.dumps(contract, sort_keys=True)}\n"
            "A provider or chunker change is a governed migration: run "
            "`python3 -m corpus.cli rebuild` once Sam has approved the change."
        )
    return mode


def _ensure_collection(qdrant: Qdrant, cfg) -> None:
    if not qdrant.collection_exists(cfg.collection):
        qdrant.create_collection(
            cfg.collection, cfg.dim, vectors_on_disk=cfg.vectors_on_disk
        )
        return
    info = qdrant.collection_info(cfg.collection)
    vectors = ((info.get("config") or {}).get("params") or {}).get("vectors")
    size = None
    if isinstance(vectors, dict):
        size = vectors.get("size")
    if size is not None and int(size) != int(cfg.dim):
        raise PipelineError(
            f"collection {cfg.collection!r} has vectors of size {size} but this run "
            f"is configured for {cfg.dim}. That is a governed migration: rebuild."
        )


def _classify(documents: list, manifest: dict, cfg, mode: str):
    """Which documents need work, and which are gone."""
    changed = []
    on_disk = set()
    for doc in documents:
        on_disk.add(doc.path)
        if mode == "rebuild":
            changed.append(doc)
            continue
        entry = manifest.get(doc.path)
        if entry is None:
            changed.append(doc)
            continue
        if (
            entry["content_hash"] != doc.digest
            or entry["embedding_model"] != cfg.model
            or int(entry["embedding_dim"]) != int(cfg.dim)
            or entry["chunker_version"] != cfg.chunker_version
        ):
            changed.append(doc)
    removed = sorted(path for path in manifest if path not in on_disk)
    return changed, removed


def _prepare(doc, cfg, manifest: dict) -> list:
    """Chunk one document and build validated payloads. No network, no writes."""
    try:
        chunks = chunker.chunk_document(doc.body)
    except chunker.ChunkError as exc:
        raise PipelineError(f"{doc.path}: {exc}") from exc
    if not chunks:
        raise PipelineError(f"{doc.path}: produced no chunks (empty body?)")
    doc.chunk_count = len(chunks)

    out = []
    for chunk in chunks:
        payload = payloads.build_payload(doc, chunk, cfg)
        payload["id"] = chunk_id(doc.path, chunk.index, doc.digest)
        try:
            payloads.validate_payload(payload, where=f"{doc.path}#chunk-{chunk.index}")
        except payloads.PayloadError as exc:
            raise PipelineError(str(exc)) from exc
        out.append(payload)
    return out


def _write_collection_file(cfg, store, stats: RunStats) -> None:
    """Mirror the contract to data/corpus/collection.json for the operator."""
    body = {
        "collection": cfg.collection,
        "contract": _contract(cfg),
        "qdrant_url": cfg.qdrant_url,
        "chunks": store.chunk_count(),
        "sources": store.source_count(),
        "last_run": store.get_meta("last_run"),
    }
    with open(cfg.meta_path, "w", encoding="utf-8") as handle:
        json.dump(body, handle, indent=2, sort_keys=True)
        handle.write("\n")


# -- verification and demonstration ---------------------------------------


def verify(cfg, log=None) -> list:
    """Return a list of problems. Empty means the index agrees with content/."""
    log = log or (lambda _msg: None)
    problems = []
    documents = loader.load_documents(cfg)
    qdrant = Qdrant(cfg.qdrant_url, timeout=cfg.timeout)
    if not qdrant.health():
        return [f"Qdrant is not answering at {cfg.qdrant_url}"]

    contract = _contract(cfg)
    with store_module.Store(cfg.db_path) as store:
        stored = store.get_meta("collection")
        if stored != contract:
            problems.append(f"collection contract mismatch: stored={stored} config={contract}")

        if not qdrant.collection_exists(cfg.collection):
            problems.append(f"collection {cfg.collection!r} does not exist")
            return problems

        info = qdrant.collection_info(cfg.collection)
        size = (((info.get("config") or {}).get("params") or {}).get("vectors") or {}).get("size")
        if size is not None and int(size) != cfg.dim:
            problems.append(f"collection vector size {size} != configured {cfg.dim}")

        points_by_source = qdrant.scroll_sources(cfg.collection)
        rows = store.chunk_payloads()
        by_source = {}
        for row in rows:
            by_source.setdefault(row["source_path"], []).append(row)
            try:
                payloads.validate_payload(row, where=f"{row['source_path']}#chunk-{row['chunk_index']}")
            except payloads.PayloadError as exc:
                problems.append(str(exc))

        for doc in documents:
            payload_rows = by_source.get(doc.path, [])
            if not payload_rows:
                problems.append(f"{doc.path}: no chunks in the index (run sync)")
                continue
            if payload_rows[0]["content_hash"] != doc.digest:
                problems.append(f"{doc.path}: source changed since it was embedded")
            expected = len(payload_rows)
            actual = points_by_source.get(doc.path, 0)
            if actual != expected:
                problems.append(
                    f"{doc.path}: {expected} chunk row(s) but {actual} vector(s) in Qdrant"
                )

        for source in sorted(set(by_source) - {doc.path for doc in documents}):
            problems.append(f"{source}: indexed but no longer in content/ (run sync)")

        total_points = sum(points_by_source.values())
        if total_points != store.chunk_count():
            problems.append(
                f"{total_points} vector(s) in Qdrant but {store.chunk_count()} chunk(s) in SQLite"
            )
        log(
            f"verify: {store.chunk_count()} chunk(s) across {store.source_count()} source(s), "
            f"{total_points} vector(s) in Qdrant"
        )
    return problems


def demo_area_filter(cfg, query: str, area: str = None, limit: int = 5) -> dict:
    """Same query with and without an `area` filter — the DoD's filter evidence."""
    provider = embeddings.build_provider(cfg)
    qdrant = Qdrant(cfg.qdrant_url, timeout=cfg.timeout)
    vector = provider.embed([query])[0]
    unfiltered = qdrant.search(cfg.collection, vector, limit=limit)
    filtered = qdrant.search(cfg.collection, vector, limit=limit, area=area) if area else []
    return {
        "query": query,
        "area": area,
        "unfiltered": [_hit(row) for row in unfiltered],
        "filtered": [_hit(row) for row in filtered],
    }


def _hit(row: dict) -> dict:
    payload = row.get("payload") or {}
    return {
        "score": round(row["score"], 4),
        "area": payload.get("area"),
        "confidence": payload.get("confidence"),
        "title": payload.get("title"),
        "url": payload.get("url"),
        "source_path": payload.get("source_path"),
        "chunk_index": payload.get("chunk_index"),
        "preview": (payload.get("text") or "").replace("\n", " ")[:120],
    }


def hybrid_search(cfg, query: str, area: str = None, limit: int = 5) -> list:
    """Vectors + FTS5 fused with reciprocal rank fusion (decision D10).

    Phase 5 owns the query graph; this is here because Phase 4 is the phase that
    has to prove the two halves line up — same ids, same metadata, one ranking.
    """
    provider = embeddings.build_provider(cfg)
    qdrant = Qdrant(cfg.qdrant_url, timeout=cfg.timeout)
    vector = provider.embed([query])[0]
    ranked = {}
    for rank, hit in enumerate(qdrant.search(cfg.collection, vector, limit=limit * 4, area=area)):
        entry = ranked.setdefault(hit["id"], {"id": hit["id"], "payload": hit["payload"], "score": 0.0})
        entry["score"] += 1.0 / (60 + rank)
    with store_module.Store(cfg.db_path) as store:
        for rank, hit in enumerate(store.keyword_search(query, area=area, limit=limit * 4)):
            entry = ranked.setdefault(hit["id"], {"id": hit["id"], "payload": {}, "score": 0.0})
            entry.setdefault("keyword", hit["snippet"])
            entry["score"] += 1.0 / (60 + rank)
    ordered = sorted(ranked.values(), key=lambda row: row["score"], reverse=True)[:limit]
    for row in ordered:
        payload = row.get("payload") or {}
        row["title"] = payload.get("title")
        row["url"] = payload.get("url")
    return ordered