feat(llm): incremental pgvector indexer + stack llm index CLI (refs #567)

src/cli/__init__.py

@@ -20,6 +20,7 @@ from cli.docs import app as docs_app
 from cli.generate import app as generate_app
 from cli.health import health
 from cli.lake import app as lake_app
+from cli.llm import app as llm_app
 from cli.load import app as load_app
 from cli.mail import app as mail_app
 from cli.perf import app as perf_app
@@ -47,6 +48,7 @@ app.add_typer(
     help="CMS rulemaking comment text extraction & analysis.",
 )
 app.add_typer(lake_app, name="lake", help="Lakehouse schema and data.")
+app.add_typer(llm_app, name="llm", help="Local RAG — index, search, tag.")
 app.add_typer(db_app, name="db", help="DuckDB utilities.")
 app.add_typer(docs_app, name="docs", help="Documentation generation.")
 app.add_typer(api_app, name="api", help="API server.")

src/cli/llm.py

@@ -0,0 +1,46 @@
+"""stack llm — local RAG over comments + corpus (P33: index only)."""
+
+import typer
+
+app = typer.Typer(no_args_is_help=True)
+
+
+@app.command()
+def index(
+    collection: str = typer.Option(
+        "comments", help="Which collection: comments | corpus."
+    ),
+    docket: str = typer.Option("", help="Limit comments to one docket id."),
+    force: bool = typer.Option(False, help="Re-embed even when unchanged."),
+    limit: int = typer.Option(0, help="Stop after N docs (0 = all)."),
+) -> None:
+    """Embed comments/corpus into pgvector (incremental, resumable)."""
+    import itertools
+
+    from conf.connect import bib
+    from llm import config as llm_config
+    from llm.index import index_docs
+    from llm.pool import HostPool
+    from llm.source import iter_comment_docs, iter_corpus_docs
+
+    cfg = llm_config.load()
+    store = bib()
+    if collection == "comments":
+        docs = iter_comment_docs(store, docket=docket)
+    elif collection == "corpus":
+        docs = iter_corpus_docs(store)
+    else:
+        raise typer.BadParameter("collection must be 'comments' or 'corpus'")
+    if limit:
+        docs = itertools.islice(docs, limit)
+    stats = index_docs(
+        docs,
+        collection=collection,
+        cfg=cfg,
+        pool=HostPool.from_config(cfg),
+        force=force,
+    )
+    typer.echo(
+        f"indexed={stats['indexed']} skipped={stats['skipped']} "
+        f"chunks={stats['chunks']}"
+    )

src/llm/index.py

@@ -0,0 +1,126 @@
+"""Incremental, resumable embedding indexer.
+
+Per doc: compare ``content_hash`` against ``index_state``; skip when
+unchanged, else delete the item's old chunks (by ``cmetadata->>
+'item_key'``), embed via the host pool, upsert with deterministic ids,
+and record the new hash — one transaction per doc, so a killed run
+resumes exactly where it stopped.
+
+The metadata delete runs after ``vectorstore()`` has already constructed
+the ``PGVector`` store: the installed ``langchain-postgres==0.0.17``
+creates ``langchain_pg_embedding``/``langchain_pg_collection`` and the
+collection row synchronously in ``PGVector.__post_init__`` (sync mode),
+so the table always exists by the time any ``DELETE`` runs — but the
+delete is still wrapped defensively in case a future version defers
+table creation to first ``add_embeddings()``.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Iterable
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import ProgrammingError
+
+from llm.chunk import Doc, chunk_doc, content_hash
+from llm.config import LlmConfig, pg_url
+from llm.migrate import ensure_hnsw, migrate
+from llm.pool import HostPool, PoolEmbeddings, embed_texts
+
+log = logging.getLogger(__name__)
+
+
+def _engine(cfg: LlmConfig) -> Engine:
+    return create_engine(pg_url(cfg))
+
+
+def vectorstore(collection: str, cfg: LlmConfig, pool: HostPool):
+    from langchain_postgres import PGVector
+
+    return PGVector(
+        embeddings=PoolEmbeddings(pool, cfg.embed_model),
+        collection_name=collection,
+        connection=pg_url(cfg),
+        embedding_length=cfg.embed_dim,
+        use_jsonb=True,
+    )
+
+
+def _state(engine: Engine, collection: str) -> dict[str, str]:
+    with engine.begin() as conn:
+        rows = conn.execute(
+            text(
+                "SELECT item_key, content_hash FROM index_state WHERE collection = :c"
+            ),
+            {"c": collection},
+        ).fetchall()
+    return dict(rows)
+
+
+def _delete_old_chunks(engine: Engine, item_key: str) -> None:
+    try:
+        with engine.begin() as conn:
+            conn.execute(
+                text(
+                    "DELETE FROM langchain_pg_embedding "
+                    "WHERE cmetadata->>'item_key' = :k"
+                ),
+                {"k": item_key},
+            )
+    except ProgrammingError:
+        # Fresh database, store tables not created yet — nothing to delete.
+        log.debug("langchain_pg_embedding not present yet; skipping delete")
+
+
+def index_docs(
+    docs: Iterable[Doc],
+    *,
+    collection: str,
+    cfg: LlmConfig,
+    pool: HostPool,
+    force: bool = False,
+) -> dict:
+    engine = _engine(cfg)
+    migrate(engine)
+    pool.check(cfg.embed_model)
+    store = vectorstore(collection, cfg, pool)
+    seen = _state(engine, collection)
+    stats = {"indexed": 0, "skipped": 0, "chunks": 0}
+
+    for doc in docs:
+        chunks = chunk_doc(doc)
+        if not chunks:
+            stats["skipped"] += 1
+            continue
+        h = content_hash(doc.text)
+        if not force and seen.get(doc.key) == h:
+            stats["skipped"] += 1
+            continue
+        vectors = embed_texts(pool, cfg.embed_model, [c.text for c in chunks])
+        _delete_old_chunks(engine, doc.key)
+        store.add_embeddings(
+            texts=[c.text for c in chunks],
+            embeddings=vectors,
+            metadatas=[c.metadata for c in chunks],
+            ids=[c.id for c in chunks],
+        )
+        with engine.begin() as conn:
+            conn.execute(
+                text(
+                    "INSERT INTO index_state "
+                    "(item_key, collection, content_hash, chunk_count) "
+                    "VALUES (:k, :c, :h, :n) "
+                    "ON CONFLICT (item_key, collection) DO UPDATE SET "
+                    "content_hash = :h, chunk_count = :n, indexed_at = now()"
+                ),
+                {"k": doc.key, "c": collection, "h": h, "n": len(chunks)},
+            )
+        stats["indexed"] += 1
+        stats["chunks"] += len(chunks)
+        if stats["indexed"] % 100 == 0:
+            log.info("indexed %(indexed)s (+%(chunks)s chunks)", stats)
+
+    ensure_hnsw(engine, cfg.embed_dim)
+    return stats

tests/llm/test_index.py

@@ -0,0 +1,74 @@
+"""llm.index — incremental, resumable embedding indexer."""
+
+from unittest.mock import MagicMock, patch
+
+from llm.chunk import Doc, content_hash
+from llm.config import LlmConfig
+from llm.index import index_docs
+
+CFG = LlmConfig(
+    ollama_hosts=("http://h1:11434",),
+    embed_model="m",
+    instruct_model="g",
+    embed_dim=768,
+    pg_host="x",
+    pg_port=5432,
+    pg_db="llm",
+    pg_user="llm",
+)
+
+DOC = Doc(key="K1", text="Some body text.", metadata={"docket": "D"})
+
+
+def _run(docs, state_rows, force=False):
+    """Run index_docs with everything external mocked; return mocks."""
+    store = MagicMock()
+    engine = MagicMock()
+    conn = engine.begin.return_value.__enter__.return_value
+    conn.execute.return_value.fetchall.return_value = state_rows
+    with (
+        patch("llm.index._engine", return_value=engine),
+        patch("llm.index.vectorstore", return_value=store),
+        patch("llm.index.embed_texts", return_value=[[0.0] * 3]),
+        patch("llm.index.HostPool") as MockPool,
+    ):
+        MockPool.return_value.check.return_value = ["http://h1:11434"]
+        stats = index_docs(
+            docs,
+            collection="comments",
+            cfg=CFG,
+            pool=MockPool.return_value,
+            force=force,
+        )
+    return stats, store, conn
+
+
+class TestIndexDocs:
+    def test_new_doc_embedded_and_recorded(self):
+        stats, store, conn = _run([DOC], state_rows=[])
+        assert stats == {"indexed": 1, "skipped": 0, "chunks": 1}
+        store.add_embeddings.assert_called_once()
+        kwargs = store.add_embeddings.call_args.kwargs
+        assert kwargs["ids"][0].startswith("K1:")
+
+    def test_unchanged_doc_skipped(self):
+        h = content_hash(DOC.text)
+        stats, store, _ = _run([DOC], state_rows=[("K1", h)])
+        assert stats["skipped"] == 1
+        store.add_embeddings.assert_not_called()
+
+    def test_changed_doc_deletes_old_chunks_first(self):
+        stats, store, conn = _run([DOC], state_rows=[("K1", "stalehash")])
+        assert stats["indexed"] == 1
+        deletes = " ".join(str(c.args[0]) for c in conn.execute.call_args_list)
+        assert "item_key" in deletes  # DELETE ... cmetadata->>'item_key'
+
+    def test_force_reembeds_unchanged(self):
+        h = content_hash(DOC.text)
+        stats, store, _ = _run([DOC], state_rows=[("K1", h)], force=True)
+        assert stats["indexed"] == 1
+
+    def test_empty_doc_counts_skipped(self):
+        empty = Doc(key="K2", text="   ", metadata={})
+        stats, store, _ = _run([empty], state_rows=[])
+        assert stats == {"indexed": 0, "skipped": 1, "chunks": 0}