diff --git a/.env.example b/.env.example index 06d177493..c4eb07f93 100644 --- a/.env.example +++ b/.env.example @@ -70,14 +70,18 @@ CRE_NOISE_FILTER_CONFIDENCE_THRESHOLD=0.8 OpenCRE_gspread_Auth=path/to/credentials.json -# Module C — Librarian -# Loaded by application/utils/librarian/config_loader.py. Nothing consumes these -# yet; defaults match the OIE design doc. Constraints enforced at load time: -# top_k_rerank <= top_k_retrieval; thresholds in [0.0, 1.0]; counts > 0. +# Module C — The Librarian (defaults match the OIE design doc; see +# application/utils/librarian/config_loader.py for validation) -CRE_LIBRARIAN_CROSSENCODER_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2 +# Retrieval backend: in_memory (sklearn cosine; SQLite dev/CI/harness) or +# pgvector (Postgres vector column — pending prod extension + mentor OK). +CRE_LIBRARIAN_RETRIEVER_BACKEND=in_memory +# Candidate shortlist size produced by C.1 and reranked by C.2 (W4). CRE_LIBRARIAN_TOP_K_RETRIEVAL=20 CRE_LIBRARIAN_TOP_K_RERANK=5 +# Cross-encoder reranker model (W4), pinned. +CRE_LIBRARIAN_CROSSENCODER_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2 +# Auto-link confidence threshold; below this -> human review (W5). CRE_LIBRARIAN_LINK_THRESHOLD=0.8 CRE_LIBRARIAN_BATCH_SIZE=32 CRE_LIBRARIAN_ECE_TARGET=0.10 diff --git a/application/cmd/cre_main.py b/application/cmd/cre_main.py index 216f7c7b6..4b814640a 100644 --- a/application/cmd/cre_main.py +++ b/application/cmd/cre_main.py @@ -1005,6 +1005,12 @@ def run(args: argparse.Namespace) -> None: # pragma: no cover ) if args.upstream_sync: download_graph_from_upstream(args.cache_file) + if args.run_librarian or args.librarian_dry_run: + run_librarian( + cache_file=args.cache_file, + dry_run=args.librarian_dry_run or not args.run_librarian, + source_jsonl=args.librarian_source, + ) def ai_client_init(database: db.Node_collection): @@ -1051,6 +1057,129 @@ def generate_embeddings(db_url: str) -> None: prompt_client.PromptHandler(database, load_all_embeddings=True) +_DEFAULT_LIBRARIAN_SOURCE = os.path.join( + os.path.dirname(os.path.dirname(__file__)), + "tests", + "librarian", + "fixtures", + "sample_knowledge_queue.jsonl", +) + + +def run_librarian( + cache_file: str, dry_run: bool = True, source_jsonl: Optional[str] = None +) -> None: + """Module C entrypoint — the pipeline switch (W3). + + For each knowledge-queue section: try the deterministic explicit-CRE fast + path (C.0.5); on no/ambiguous reference, run the semantic retriever (C.1) + and log the top-K candidate shortlist. The cross-encoder rerank (C.2, W4), + decision/threshold routing (C.3-C.4, W5) and graph writes (W8) are not + built yet, so this is dry-run only: it never writes a link. ``--run_librarian`` + without writes behaves identically and warns. + + Ops note: this is opt-in CLI only — it is not on the ``Procfile`` and is + wired into neither the web app nor the background worker (that lands W8). + It calls the paid embedding API, so cost is incurred only when someone runs + the command manually; the running deployment never triggers it on its own. + """ + from application.utils.librarian.candidate_retriever import ( + CandidatePool, + RetrieverBackend, + build_retriever, + ) + from application.utils.librarian.config_loader import load_config + from application.utils.librarian.explicit_link_resolver import ( + ResolutionOutcome, + resolve, + ) + from application.utils.librarian.knowledge_source import FixtureKnowledgeSource + from application.utils.librarian.section_validator import ( + SectionValidationError, + section_from_queue_row, + ) + + if not dry_run: + logger.warning( + "the Librarian cannot write links yet (DecisionEngine + graph write " + "land W8); running in dry-run mode" + ) + + cfg = load_config() + database = db_connect(path=cache_file) + ph = prompt_client.PromptHandler(database=database) + + backend = RetrieverBackend(cfg.retriever_backend) + if backend is RetrieverBackend.pgvector: + dialect = database.session.connection().dialect.name + if dialect != "postgresql": + logger.warning( + "CRE_LIBRARIAN_RETRIEVER_BACKEND=pgvector selected on a %r " + "database, but the pgvector backend needs Postgres with the " + "embedding_vec column (lands W8) and will fail at retrieve() " + "time here. Set the backend to in_memory until then.", + dialect, + ) + # The CRE ids present in the hub are exactly the known ids the explicit + # resolver may auto-link to (W2 seeded this from the golden set; here it is + # the real DB-backed registry). + cre_embeddings = database.get_embeddings_by_doc_type(defs.Credoctypes.CRE.value) + known_ids = set(cre_embeddings.keys()) + # in_memory loads the hub matrix; pgvector ranks in the DB over the + # embedding_vec column (no in-RAM pool). Both honor the same retrieve(). + pool = ( + CandidatePool.from_mapping(cre_embeddings) + if backend is RetrieverBackend.in_memory + else None + ) + retriever = build_retriever( + backend, + embed_fn=ph.get_text_embeddings, + top_k=cfg.top_k_retrieval, + threshold=cfg.link_threshold, + pool=pool, + connection=database.session.connection(), + ) + + source = FixtureKnowledgeSource(source_jsonl or _DEFAULT_LIBRARIAN_SOURCE) + sections = explicit = semantic = rejected = 0 + for item in source.items(): + try: + section = section_from_queue_row(item) + except SectionValidationError as exc: + rejected += 1 + logger.warning("section rejected at C.0 boundary: %s", exc) + continue + sections += 1 + + resolution = resolve(section.text, known_ids) + if resolution.outcome == ResolutionOutcome.resolved: + explicit += 1 + logger.info("[explicit] %s -> %s", section.chunk_id, resolution.cre_ids[0]) + continue + + semantic += 1 + audit = retriever.retrieve(section.text) + top = ", ".join( + f"{c.cre_id}:{c.score_vector:.3f}" for c in audit.candidates[:5] + ) + logger.info( + "[semantic] %s -> %d candidates (top5: %s)", + section.chunk_id, + len(audit.candidates), + top or "", + ) + + logger.info( + "librarian dry-run complete: %d sections (%d explicit, %d semantic), " + "%d rejected at boundary", + sections, + explicit, + semantic, + rejected, + ) + + def regenerate_embeddings(db_url: str) -> None: """Wipe all embedding rows, then rebuild (CRE + every node type) like ``--generate_embeddings``.""" database = db_connect(path=db_url) diff --git a/application/tests/librarian/candidate_retriever_test.py b/application/tests/librarian/candidate_retriever_test.py new file mode 100644 index 000000000..2d44dc4a8 --- /dev/null +++ b/application/tests/librarian/candidate_retriever_test.py @@ -0,0 +1,213 @@ +"""Tests for the C.1 semantic candidate retriever (Week 3). + +Hermetic: the embedding function and the CRE vector hub are injected as +controlled vectors, so cosine ordering, top-K truncation, the dim gate, and +the RetrievalAudit shape are all assertable without an LLM or a DB. +""" + +import unittest + +from application.utils.librarian.candidate_retriever import ( + PGVECTOR_RETRIEVER_NAME, + RETRIEVER_NAME, + CandidatePool, + CandidateRetriever, + DimensionMismatchError, + EmptyPoolError, + PgVectorRetriever, + RetrieverBackend, + RetrieverError, + build_retriever, + to_pgvector_literal, +) + +# A controlled hub. Query [1,0,0] -> cosine a=1.0, c=0.707, b=0.0 -> order a,c,b. +HUB = { + "111-111": [1.0, 0.0, 0.0], # "a" + "222-222": [0.0, 1.0, 0.0], # "b" + "333-333": [1.0, 1.0, 0.0], # "c" +} + +# Deterministic embedder: text -> a fixed query vector. +_VECTORS = { + "about-a": [1.0, 0.0, 0.0], + "about-b": [0.0, 1.0, 0.0], + "wrong-width": [1.0, 0.0], +} + + +def fake_embed(text): + return _VECTORS[text] + + +def make_retriever(top_k=20, threshold=0.8, cre_names=None): + return CandidateRetriever( + embed_fn=fake_embed, + pool=CandidatePool.from_mapping(HUB), + top_k=top_k, + threshold=threshold, + cre_names=cre_names, + ) + + +class CandidatePoolTest(unittest.TestCase): + def test_from_mapping_builds_aligned_matrix(self) -> None: + pool = CandidatePool.from_mapping(HUB) + self.assertEqual(pool.dim, 3) + self.assertEqual(set(pool.cre_ids), set(HUB)) + self.assertEqual(pool.matrix.shape, (3, 3)) + + def test_empty_pool_rejected(self) -> None: + with self.assertRaises(EmptyPoolError): + CandidatePool.from_mapping({}) + + def test_ragged_vectors_rejected(self) -> None: + with self.assertRaises(DimensionMismatchError): + CandidatePool.from_mapping({"a": [1.0, 0.0], "b": [1.0]}) + + def test_zero_width_vectors_rejected(self) -> None: + with self.assertRaises(DimensionMismatchError): + CandidatePool.from_mapping({"a": [], "b": []}) + + +class RetrieveTest(unittest.TestCase): + def test_candidates_rank_ordered_by_cosine(self) -> None: + audit = make_retriever().retrieve("about-a") + self.assertEqual( + [c.cre_id for c in audit.candidates], + ["111-111", "333-333", "222-222"], + ) + # score_vector is the cosine; populated and descending. + scores = [c.score_vector for c in audit.candidates] + self.assertAlmostEqual(scores[0], 1.0) + self.assertEqual(scores, sorted(scores, reverse=True)) + + def test_top_k_truncates_and_caps_at_pool_size(self) -> None: + # top_k larger than the hub returns the whole hub, no error. + self.assertEqual( + len(make_retriever(top_k=20).retrieve("about-a").candidates), 3 + ) + # top_k smaller than the hub returns exactly the best k. + top1 = make_retriever(top_k=1).retrieve("about-a") + self.assertEqual([c.cre_id for c in top1.candidates], ["111-111"]) + + def test_audit_shape_for_w3(self) -> None: + audit = make_retriever(threshold=0.8).retrieve("about-b") + self.assertEqual(audit.retriever, RETRIEVER_NAME) + self.assertEqual(audit.reranked, []) # cross-encoder lands W4 + self.assertEqual(audit.threshold, 0.8) + # The closest to [0,1,0] is b, then c (0.707), then a (0.0). + self.assertEqual(audit.candidates[0].cre_id, "222-222") + + def test_cre_names_populated_when_provided(self) -> None: + audit = make_retriever(cre_names={"111-111": "Authentication"}).retrieve( + "about-a" + ) + self.assertEqual(audit.candidates[0].cre_name, "Authentication") + # Unnamed candidates stay None, never KeyError. + self.assertIsNone(audit.candidates[-1].cre_name) + + def test_query_dim_mismatch_is_caught(self) -> None: + with self.assertRaises(DimensionMismatchError): + make_retriever().retrieve("wrong-width") + + +class ConstructionTest(unittest.TestCase): + def test_non_positive_top_k_rejected(self) -> None: + with self.assertRaises(RetrieverError): + make_retriever(top_k=0) + + +# --- pgvector backend (hermetic: the DB is a fake recording connection) --- + + +class _FakeRow: + def __init__(self, cre_id, score): + self.cre_id = cre_id + self.score = score + + +class _FakeResult: + def __init__(self, rows): + self._rows = rows + + def fetchall(self): + return self._rows + + +class _FakeConnection: + """Records the last execute() call and returns canned, pre-ordered rows.""" + + def __init__(self, rows): + self._rows = rows + self.last_sql = None + self.last_params = None + + def execute(self, sql, params): + self.last_sql = str(sql) + self.last_params = params + return _FakeResult(self._rows) + + +class PgVectorRetrieverTest(unittest.TestCase): + def test_pgvector_literal_format(self) -> None: + self.assertEqual(to_pgvector_literal([1, 2.5, 0]), "[1.0,2.5,0.0]") + + def test_parses_db_rows_in_order(self) -> None: + # The DB does the ranking; the retriever preserves ORDER BY order. + conn = _FakeConnection([_FakeRow("111-111", 0.97), _FakeRow("333-333", 0.71)]) + retriever = PgVectorRetriever( + embed_fn=fake_embed, connection=conn, top_k=20, threshold=0.8 + ) + audit = retriever.retrieve("about-a") + self.assertEqual(audit.retriever, PGVECTOR_RETRIEVER_NAME) + self.assertEqual([c.cre_id for c in audit.candidates], ["111-111", "333-333"]) + self.assertAlmostEqual(audit.candidates[0].score_vector, 0.97) + self.assertEqual(audit.reranked, []) + + def test_binds_query_vector_doctype_and_limit(self) -> None: + conn = _FakeConnection([]) + PgVectorRetriever( + embed_fn=fake_embed, connection=conn, top_k=7, threshold=0.8 + ).retrieve("about-a") + self.assertEqual(conn.last_params["q"], "[1.0,0.0,0.0]") + self.assertEqual(conn.last_params["doc_type"], "CRE") + self.assertEqual(conn.last_params["k"], 7) + # Cosine via the <=> operator, scored as similarity (1 - distance). + self.assertIn("<=>", conn.last_sql) + + +class BuildRetrieverTest(unittest.TestCase): + def test_in_memory_requires_pool(self) -> None: + with self.assertRaises(RetrieverError): + build_retriever( + RetrieverBackend.in_memory, fake_embed, top_k=20, threshold=0.8 + ) + + def test_pgvector_requires_connection(self) -> None: + with self.assertRaises(RetrieverError): + build_retriever( + RetrieverBackend.pgvector, fake_embed, top_k=20, threshold=0.8 + ) + + def test_factory_routes_to_each_backend(self) -> None: + in_mem = build_retriever( + RetrieverBackend.in_memory, + fake_embed, + top_k=20, + threshold=0.8, + pool=CandidatePool.from_mapping(HUB), + ) + self.assertIsInstance(in_mem, CandidateRetriever) + pg = build_retriever( + RetrieverBackend.pgvector, + fake_embed, + top_k=20, + threshold=0.8, + connection=_FakeConnection([]), + ) + self.assertIsInstance(pg, PgVectorRetriever) + + +if __name__ == "__main__": + unittest.main() diff --git a/application/tests/librarian/config_loader_test.py b/application/tests/librarian/config_loader_test.py index d74deb1b4..1b7abb294 100644 --- a/application/tests/librarian/config_loader_test.py +++ b/application/tests/librarian/config_loader_test.py @@ -1,5 +1,6 @@ import os import unittest +from dataclasses import FrozenInstanceError from unittest import mock from application.utils.librarian.config_loader import LibrarianConfig, load_config @@ -9,6 +10,7 @@ class TestConfigLoaderDefaults(unittest.TestCase): def test_defaults_when_env_unset(self): with mock.patch.dict(os.environ, {}, clear=True): cfg = load_config() + self.assertEqual(cfg.retriever_backend, "in_memory") self.assertEqual(cfg.crossencoder_model, "cross-encoder/ms-marco-MiniLM-L-6-v2") self.assertEqual(cfg.top_k_retrieval, 20) self.assertEqual(cfg.top_k_rerank, 5) @@ -18,13 +20,15 @@ def test_defaults_when_env_unset(self): self.assertEqual(cfg.conformal_alpha, 0.10) def test_config_is_frozen(self): - cfg = load_config() - with self.assertRaises(Exception): + with mock.patch.dict(os.environ, {}, clear=True): + cfg = load_config() + with self.assertRaises(FrozenInstanceError): cfg.link_threshold = 0.5 # type: ignore[misc] class TestConfigLoaderOverrides(unittest.TestCase): OVERRIDES = { + "CRE_LIBRARIAN_RETRIEVER_BACKEND": "pgvector", "CRE_LIBRARIAN_CROSSENCODER_MODEL": "cross-encoder/other", "CRE_LIBRARIAN_TOP_K_RETRIEVAL": "50", "CRE_LIBRARIAN_TOP_K_RERANK": "10", @@ -37,6 +41,7 @@ class TestConfigLoaderOverrides(unittest.TestCase): def test_env_overrides_apply(self): with mock.patch.dict(os.environ, self.OVERRIDES, clear=True): cfg = load_config() + self.assertEqual(cfg.retriever_backend, "pgvector") self.assertEqual(cfg.crossencoder_model, "cross-encoder/other") self.assertEqual(cfg.top_k_retrieval, 50) self.assertEqual(cfg.top_k_rerank, 10) diff --git a/application/tests/librarian/dataset_test.py b/application/tests/librarian/dataset_test.py index bf37c858d..0dab83f17 100644 --- a/application/tests/librarian/dataset_test.py +++ b/application/tests/librarian/dataset_test.py @@ -111,6 +111,7 @@ def test_build_check_matches_committed_dataset(self): [sys.executable, _BUILD_SCRIPT, "--check"], capture_output=True, text=True, + timeout=120, ) self.assertEqual( result.returncode, diff --git a/application/tests/librarian/section_validator_test.py b/application/tests/librarian/section_validator_test.py index a8daeadbd..f9dbff76e 100644 --- a/application/tests/librarian/section_validator_test.py +++ b/application/tests/librarian/section_validator_test.py @@ -6,6 +6,7 @@ """ import unittest +from datetime import datetime, timezone from pydantic import ValidationError @@ -86,7 +87,10 @@ def test_valid_row_builds_section_with_synthesized_identity(self) -> None: section.artifact_id, "art:OWASP/ASVS:4.0/en/0x11-V2-Authentication.md" ) self.assertEqual(section.source.repo, "OWASP/ASVS") - self.assertEqual(section.source.committed_at, "2026-05-25T02:25:00Z") + self.assertEqual( + section.source.committed_at, + datetime(2026, 5, 25, 2, 25, tzinfo=timezone.utc), + ) self.assertEqual(section.locator.path, "4.0/en/0x11-V2-Authentication.md") self.assertEqual(section.language, "en") diff --git a/application/utils/librarian/__init__.py b/application/utils/librarian/__init__.py index e0581f28a..f6b8440a5 100644 --- a/application/utils/librarian/__init__.py +++ b/application/utils/librarian/__init__.py @@ -13,7 +13,9 @@ W1 (C.-1): contracts + config + eval harness + golden dataset. W2 (C.0): input boundary — SectionValidator (validate/adapt without re-normalizing text) and ExplicitLinkResolver (fail-safe - explicit-link resolution). No retrieval/ranking logic yet. + explicit-link resolution). + W3 (C.1): candidate retriever (in-memory + pgvector) + pipeline switch. +Cross-encoder rerank (C.2, W4) onward is not built yet. Vendored RFC JSON schemas live under ``_rfc_schemas/``. They are pinned to upstream/owasp-graph @ 2b1437987768d5ed20fe9ee721ab9a898c4b84af (PR #734). diff --git a/application/utils/librarian/candidate_retriever.py b/application/utils/librarian/candidate_retriever.py new file mode 100644 index 000000000..d4e36c12a --- /dev/null +++ b/application/utils/librarian/candidate_retriever.py @@ -0,0 +1,284 @@ +"""Module C.1 — semantic candidate retriever (Week 3). The search step. + +For sections with no explicit CRE id (the majority — most chunks just +describe a security concept in plain English), C must *find* which of +OpenCRE's nodes the text is about. That is a search problem: embed the +section text, cosine-match it against the CRE-node vector hub, and return +the top-K rank-ordered candidates — the shortlist W4's cross-encoder +reranks and later weeks turn into a confident yes/no. + +The retriever is a thin, dependency-injected seam over two collaborators: + + - ``embed_fn(text) -> Sequence[float]`` — turns text into a vector. Prod + wires ``PromptHandler.get_text_embeddings``; the harness and tests inject + a deterministic stub. C never embeds directly, so it stays import-light + and hermetically testable (mirrors W2's resolver taking an injected + known-id set). + - a ``CandidatePool`` — the ``{cre_id -> vector}`` hub for every CRE node. + Prod loads ``db.get_embeddings_by_doc_type(CRE)``; the pool is prevalidated + to one common width so the dim gate below is meaningful. + +Two interchangeable backends behind one ``retrieve()`` seam, selected by +``CRE_LIBRARIAN_RETRIEVER_BACKEND``: + + - ``in_memory`` — sklearn cosine over an in-RAM matrix. The only backend + that works on SQLite, so it is what CI, the golden-dataset harness, and + SQLite dev use. Loads the whole hub into RAM. + - ``pgvector`` — pushes the cosine into Postgres via the ``<=>`` operator + over an ``embedding_vec vector(dim)`` column; never loads the hub into + RAM. Requires the ``vector`` extension + that column (the Alembic + migration lands W8, during live integration, where it is validated + against real Postgres). Unavailable on SQLite. + +The RFC is silent on retrieval tech — it mandates only the +``candidates[]``/``reranked[]`` audit trail — so the backend choice is ours; +both emit the same ``RetrievalAudit``. + +Gate (PR 3): dim assertion — the query-vector width must equal the stored +CRE-vector width, or every cosine score is silently meaningless. (Enforced +in-process for in_memory; Postgres enforces it structurally for pgvector via +the fixed-width ``vector(dim)`` column.) +""" + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Callable, List, Mapping, Optional, Sequence, Tuple + +import numpy as np +from sklearn.metrics.pairwise import cosine_similarity + +from application.utils.librarian.schemas import CreCandidate, RetrievalAudit + +# A function that turns one piece of text into a single dense vector. +EmbedFn = Callable[[str], Sequence[float]] + +# Identify the retriever in the RFC audit trail (RetrievalAudit.retriever). +# Bumped when the matching algorithm changes so a stored proposal is traceable +# to the code that produced it. +RETRIEVER_NAME = "in-memory-cosine/0.1.0" +PGVECTOR_RETRIEVER_NAME = "pgvector-cosine/0.1.0" + + +class RetrieverBackend(str, Enum): + # sklearn cosine over an in-RAM matrix — SQLite dev, CI, and the harness. + in_memory = "in_memory" + # Postgres-side cosine via pgvector's ``<=>`` operator (needs the vector + # extension + an embedding_vec column; migration lands W8). + pgvector = "pgvector" + + +class RetrieverError(ValueError): + """Base class for retriever construction/usage failures.""" + + +class EmptyPoolError(RetrieverError): + """No CRE vectors to search against — retrieval cannot run.""" + + +class DimensionMismatchError(RetrieverError): + """Query- and CRE-vector widths differ — cosine scores would be meaningless.""" + + +@dataclass(frozen=True) +class CandidatePool: + """An immutable ``{cre_id -> vector}`` hub, prevalidated to one width. + + ``matrix`` is ``(n_cre, dim)`` row-aligned with ``cre_ids`` so a single + cosine call scores the whole hub at once. + """ + + cre_ids: Tuple[str, ...] + matrix: np.ndarray + dim: int + + @classmethod + def from_mapping(cls, embeddings: Mapping[str, Sequence[float]]) -> "CandidatePool": + """Build a pool from ``db.get_embeddings_by_doc_type``-shaped data. + + Rejects an empty hub and any ragged vector — both are silent-failure + traps if they reach the cosine step. + """ + if not embeddings: + raise EmptyPoolError("candidate pool is empty; no CRE vectors to search") + cre_ids = tuple(embeddings.keys()) + rows = [list(embeddings[cre_id]) for cre_id in cre_ids] + widths = {len(r) for r in rows} + if len(widths) != 1: + raise DimensionMismatchError( + f"CRE vectors have inconsistent widths {sorted(widths)}; " + "the hub must be a single embedding model/dimension" + ) + dim = widths.pop() + if dim == 0: + raise DimensionMismatchError("CRE vectors are zero-width") + return cls(cre_ids=cre_ids, matrix=np.asarray(rows, dtype=float), dim=dim) + + +class CandidateRetriever: + """Embed a section, cosine-rank the CRE hub, return the top-K shortlist. + + ``top_k`` is ``CRE_LIBRARIAN_TOP_K_RETRIEVAL`` (default 20). ``threshold`` + is the configured link threshold, carried verbatim into the audit so a + stored proposal is self-describing; the retriever itself does **not** + threshold — it always returns the top-K so W4 can rerank a full shortlist. + """ + + def __init__( + self, + embed_fn: EmbedFn, + pool: CandidatePool, + top_k: int, + *, + threshold: float, + cre_names: Optional[Mapping[str, str]] = None, + ) -> None: + if top_k <= 0: + raise RetrieverError(f"top_k must be > 0, got {top_k}") + self._embed_fn = embed_fn + self._pool = pool + self._top_k = top_k + self._threshold = threshold + self._cre_names = dict(cre_names or {}) + + def retrieve(self, text: str) -> RetrievalAudit: + """Return the top-K CRE candidates for ``text`` as a RetrievalAudit. + + ``reranked`` is empty — the cross-encoder lands W4. Candidates are + rank-ordered (highest cosine first) with ``score_vector`` populated. + """ + query = np.asarray(list(self._embed_fn(text)), dtype=float) + if query.shape[0] != self._pool.dim: + raise DimensionMismatchError( + f"query vector width {query.shape[0]} != CRE hub width " + f"{self._pool.dim}; check the embedding model matches the hub" + ) + + scores = cosine_similarity(query.reshape(1, -1), self._pool.matrix)[0] + # Top-K by descending score. argsort is ascending, so take the tail + # and reverse; cap at pool size when the hub is smaller than K. + k = min(self._top_k, len(self._pool.cre_ids)) + top_idx = np.argsort(scores)[-k:][::-1] + + candidates: List[CreCandidate] = [ + CreCandidate( + cre_id=self._pool.cre_ids[i], + cre_name=self._cre_names.get(self._pool.cre_ids[i]), + score_vector=float(scores[i]), + ) + for i in top_idx + ] + return RetrievalAudit( + retriever=RETRIEVER_NAME, + candidates=candidates, + reranked=[], + threshold=self._threshold, + ) + + +def to_pgvector_literal(vector: Sequence[float]) -> str: + """Render a vector in pgvector's text input format: ``[1.0,2.0,3.0]``.""" + return "[" + ",".join(repr(float(x)) for x in vector) + "]" + + +class PgVectorRetriever: + """Postgres-side top-K cosine via pgvector's ``<=>`` operator. + + The similarity is computed and ranked in the database against the + ``embedding_vec`` column, so the hub is never loaded into RAM — the win + over ``CandidateRetriever`` on a large CRE corpus. ``<=>`` is cosine + *distance* (0 = identical), so the score is ``1 - distance`` to match the + in-memory backend's cosine *similarity*. + + Needs the ``vector`` extension and an ``embedding_vec vector(dim)`` column + (migration lands W8). Unavailable on SQLite — the factory routes SQLite to + ``in_memory``. + """ + + # Parameterized; :q is bound as a pgvector text literal and cast in-SQL. + _SQL = ( + "SELECT cre_id, 1 - (embedding_vec <=> CAST(:q AS vector)) AS score " + "FROM embeddings " + "WHERE doc_type = :doc_type AND cre_id IS NOT NULL " + "AND embedding_vec IS NOT NULL " + "ORDER BY embedding_vec <=> CAST(:q AS vector) " + "LIMIT :k" + ) + + def __init__( + self, + embed_fn: EmbedFn, + connection: Any, + top_k: int, + *, + threshold: float, + doc_type: str = "CRE", + cre_names: Optional[Mapping[str, str]] = None, + ) -> None: + if top_k <= 0: + raise RetrieverError(f"top_k must be > 0, got {top_k}") + self._embed_fn = embed_fn + self._conn = connection + self._top_k = top_k + self._threshold = threshold + self._doc_type = doc_type + self._cre_names = dict(cre_names or {}) + + def retrieve(self, text: str) -> RetrievalAudit: + """Return the top-K CRE candidates for ``text`` as a RetrievalAudit. + + Rows arrive already rank-ordered by the SQL ``ORDER BY``; we preserve + that order. ``sqlalchemy.text`` is imported lazily so the in-memory + backend (and CI) never needs a DB driver loaded. + """ + from sqlalchemy import text as sql_text + + query = to_pgvector_literal(list(self._embed_fn(text))) + rows = self._conn.execute( + sql_text(self._SQL), + {"q": query, "doc_type": self._doc_type, "k": self._top_k}, + ).fetchall() + + candidates = [ + CreCandidate( + cre_id=row.cre_id, + cre_name=self._cre_names.get(row.cre_id), + score_vector=float(row.score), + ) + for row in rows + ] + return RetrievalAudit( + retriever=PGVECTOR_RETRIEVER_NAME, + candidates=candidates, + reranked=[], + threshold=self._threshold, + ) + + +def build_retriever( + backend: RetrieverBackend, + embed_fn: EmbedFn, + *, + top_k: int, + threshold: float, + pool: Optional[CandidatePool] = None, + connection: Any = None, + cre_names: Optional[Mapping[str, str]] = None, +) -> Any: + """Construct the retriever for ``backend`` behind the shared ``retrieve()``. + + ``in_memory`` needs ``pool``; ``pgvector`` needs ``connection``. Mismatches + fail loudly rather than silently doing nothing. + """ + if backend is RetrieverBackend.in_memory: + if pool is None: + raise RetrieverError("in_memory backend requires a CandidatePool") + return CandidateRetriever( + embed_fn, pool, top_k, threshold=threshold, cre_names=cre_names + ) + if backend is RetrieverBackend.pgvector: + if connection is None: + raise RetrieverError("pgvector backend requires a DB connection") + return PgVectorRetriever( + embed_fn, connection, top_k, threshold=threshold, cre_names=cre_names + ) + raise RetrieverError(f"unknown retriever backend {backend!r}") diff --git a/application/utils/librarian/config_loader.py b/application/utils/librarian/config_loader.py index a78452e2e..945a210c1 100644 --- a/application/utils/librarian/config_loader.py +++ b/application/utils/librarian/config_loader.py @@ -7,10 +7,16 @@ import os from dataclasses import dataclass +# Retrieval backends (see candidate_retriever.RetrieverBackend). Kept as a +# plain set here so the loader stays dependency-free; the retriever owns the +# enum it maps to. +_RETRIEVER_BACKENDS = frozenset({"in_memory", "pgvector"}) + @dataclass(frozen=True) class LibrarianConfig: crossencoder_model: str + retriever_backend: str top_k_retrieval: int top_k_rerank: int link_threshold: float @@ -23,6 +29,7 @@ def load_config() -> LibrarianConfig: crossencoder_model = os.getenv( "CRE_LIBRARIAN_CROSSENCODER_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2" ) + retriever_backend = os.getenv("CRE_LIBRARIAN_RETRIEVER_BACKEND", "in_memory") top_k_retrieval = int(os.getenv("CRE_LIBRARIAN_TOP_K_RETRIEVAL", "20")) top_k_rerank = int(os.getenv("CRE_LIBRARIAN_TOP_K_RERANK", "5")) link_threshold = float(os.getenv("CRE_LIBRARIAN_LINK_THRESHOLD", "0.8")) @@ -30,6 +37,11 @@ def load_config() -> LibrarianConfig: ece_target = float(os.getenv("CRE_LIBRARIAN_ECE_TARGET", "0.10")) conformal_alpha = float(os.getenv("CRE_LIBRARIAN_CONFORMAL_ALPHA", "0.10")) + if retriever_backend not in _RETRIEVER_BACKENDS: + raise ValueError( + f"CRE_LIBRARIAN_RETRIEVER_BACKEND must be one of " + f"{sorted(_RETRIEVER_BACKENDS)}, got {retriever_backend!r}" + ) if top_k_retrieval <= 0: raise ValueError( f"CRE_LIBRARIAN_TOP_K_RETRIEVAL must be > 0, got {top_k_retrieval}" @@ -58,6 +70,7 @@ def load_config() -> LibrarianConfig: return LibrarianConfig( crossencoder_model=crossencoder_model, + retriever_backend=retriever_backend, top_k_retrieval=top_k_retrieval, top_k_rerank=top_k_rerank, link_threshold=link_threshold, diff --git a/application/utils/librarian/knowledge_source.py b/application/utils/librarian/knowledge_source.py index 76ad493a1..9f04aece7 100644 --- a/application/utils/librarian/knowledge_source.py +++ b/application/utils/librarian/knowledge_source.py @@ -6,11 +6,16 @@ from each row at processing time (master guide §1.2). """ +import logging from abc import ABC, abstractmethod from typing import Iterator +from pydantic import ValidationError + from application.utils.librarian.schemas import KnowledgeQueueItem +logger = logging.getLogger(__name__) + class KnowledgeSource(ABC): @abstractmethod @@ -27,7 +32,15 @@ def __init__(self, jsonl_path: str) -> None: def items(self) -> Iterator[KnowledgeQueueItem]: with open(self._path, encoding="utf-8") as fh: - for line in fh: + for line_no, line in enumerate(fh, start=1): line = line.strip() if line: - yield KnowledgeQueueItem.model_validate_json(line) + try: + yield KnowledgeQueueItem.model_validate_json(line) + except ValidationError as exc: + logger.warning( + "Skipping malformed knowledge_queue row at line %d: %s", + line_no, + exc, + ) + continue diff --git a/application/utils/librarian/schemas.py b/application/utils/librarian/schemas.py index d632e2859..d0a41983d 100644 --- a/application/utils/librarian/schemas.py +++ b/application/utils/librarian/schemas.py @@ -12,10 +12,11 @@ from __future__ import annotations import re +from datetime import datetime from enum import Enum from typing import List, Literal, Optional -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import AnyUrl, BaseModel, ConfigDict, Field, model_validator SCHEMA_VERSION = "0.2.0" _SCHEMA_VERSION_RE = re.compile(r"^0\.\d+\.\d+$") @@ -64,10 +65,10 @@ class SourceRef(BaseModel): model_config = ConfigDict(extra="forbid") type: SourceType repo: Optional[str] = None - url: Optional[str] = None + url: Optional[AnyUrl] = None commit_sha: Optional[str] = Field(default=None, min_length=7) commit_message: Optional[str] = None - committed_at: str + committed_at: datetime author_login: Optional[str] = None @model_validator(mode="after") @@ -84,7 +85,7 @@ class Locator(BaseModel): kind: LocatorKind id: str = Field(min_length=1) path: Optional[str] = None - url: Optional[str] = None + url: Optional[AnyUrl] = None title: Optional[str] = None @model_validator(mode="after") @@ -201,7 +202,7 @@ class KnowledgeItem(BaseModel): artifact_id: str event_id: str pipeline_run_id: str - filtered_at: str + filtered_at: datetime status: KnowledgeStatus source: SourceRef locator: Locator @@ -228,7 +229,7 @@ class LinkProposal(BaseModel): chunk_id: str artifact_id: str pipeline_run_id: str - classified_at: str + classified_at: datetime status: Literal["linked"] = "linked" knowledge: KnowledgeSnapshot retrieval: RetrievalAudit @@ -251,7 +252,7 @@ class ReviewItem(BaseModel): chunk_id: str artifact_id: str pipeline_run_id: str - created_at: str + created_at: datetime status: Literal["review_required"] = "review_required" reason_code: ReasonCode knowledge: KnowledgeSnapshot diff --git a/application/utils/librarian/section_validator.py b/application/utils/librarian/section_validator.py index 610a91788..8d92eb082 100644 --- a/application/utils/librarian/section_validator.py +++ b/application/utils/librarian/section_validator.py @@ -164,8 +164,9 @@ def section_from_knowledge_item( raise NotKnowledgeError( f"status={item.status.value!r}; only 'accepted' items may be linked" ) - # status=accepted guarantees content (enforced by the KnowledgeItem model). - assert item.content is not None + # Keep boundary behavior typed even for pre-built/mutated model instances. + if item.content is None: + raise MalformedKnowledgeItemError("status='accepted' requires content") _require_text(item.content.text) language = _require_language(item.content.language) diff --git a/cre.py b/cre.py index 559e595e4..98b0198f2 100644 --- a/cre.py +++ b/cre.py @@ -209,6 +209,24 @@ def main() -> None: action="store_true", help="delete all embedding rows then rebuild embeddings for every CRE and node (use after smart-extract or model changes)", ) + parser.add_argument( + "--run_librarian", + action="store_true", + help="run Module C (the Librarian): for each knowledge-queue section, " + "resolve explicit CRE ids or retrieve the top-K semantic CRE candidates", + ) + parser.add_argument( + "--librarian_dry_run", + action="store_true", + help="run the Librarian without writing any links (the only supported " + "mode pre-W8; logs the candidate shortlist per section)", + ) + parser.add_argument( + "--librarian_source", + default=None, + help="path to a knowledge_queue JSONL for --run_librarian " + "(defaults to the bundled sample fixture)", + ) parser.add_argument( "--populate_neo4j_db", action="store_true", diff --git a/scripts/benchmark_retriever.py b/scripts/benchmark_retriever.py new file mode 100644 index 000000000..e6bcdc474 --- /dev/null +++ b/scripts/benchmark_retriever.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python +"""Benchmark the two C.1 retriever backends (Week 3, PR 3). + +Times in-memory cosine vs pgvector over the real CRE hub for a set of probe +queries, and checks they agree on the top-1. The in-memory backend always +runs (it only needs the embeddings already loaded); the pgvector backend runs +only when the DB is Postgres with the ``embedding_vec`` column present — +otherwise it is reported as skipped, never silently passed. + +Usage: + python scripts/benchmark_retriever.py --cache_file standards_cache.sqlite \\ + [--queries "password storage" "access control" ...] [--runs 5] +""" + +import argparse +import os +import sys +import time +from typing import List + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from application.cmd.cre_main import db_connect +from application.defs import cre_defs +from application.prompt_client import prompt_client +from application.utils.librarian.candidate_retriever import ( + CandidatePool, + PgVectorRetriever, + build_retriever, + RetrieverBackend, +) +from application.utils.librarian.config_loader import load_config + +_DEFAULT_QUERIES = [ + "Verify that user-set passwords are at least 12 characters in length.", + "Enforce least privilege for all access control decisions.", + "Do not use broken cryptographic algorithms such as MD5 or SHA-1.", + "Protect against cross-site scripting in all rendered output.", +] + + +def _time_backend(retriever, queries: List[str], runs: int) -> float: + # Warm up (model/index load), then take the best wall-clock of `runs`. + retriever.retrieve(queries[0]) + best = float("inf") + for _ in range(runs): + start = time.perf_counter() + for q in queries: + retriever.retrieve(q) + best = min(best, time.perf_counter() - start) + return best + + +def _pgvector_available(database) -> bool: + """True only on Postgres with the embedding_vec column present.""" + conn = database.session.connection() + if conn.dialect.name != "postgresql": + return False + from sqlalchemy import text + + row = conn.execute( + text( + "SELECT 1 FROM information_schema.columns " + "WHERE table_name = 'embeddings' AND column_name = 'embedding_vec'" + ) + ).fetchone() + return row is not None + + +def main(argv: List[str]) -> int: + cfg = load_config() + parser = argparse.ArgumentParser(description="Benchmark C.1 retriever backends") + parser.add_argument("--cache_file", default="standards_cache.sqlite") + parser.add_argument("--queries", nargs="*", default=_DEFAULT_QUERIES) + parser.add_argument("--runs", type=int, default=5) + parser.add_argument("--top_k", type=int, default=cfg.top_k_retrieval) + args = parser.parse_args(argv) + + database = db_connect(path=args.cache_file) + ph = prompt_client.PromptHandler(database=database) + cre_embeddings = database.get_embeddings_by_doc_type(cre_defs.Credoctypes.CRE.value) + print(f"CRE hub: {len(cre_embeddings)} vectors; {len(args.queries)} probe queries") + + in_mem = build_retriever( + RetrieverBackend.in_memory, + embed_fn=ph.get_text_embeddings, + top_k=args.top_k, + threshold=cfg.link_threshold, + pool=CandidatePool.from_mapping(cre_embeddings), + ) + in_mem_time = _time_backend(in_mem, args.queries, args.runs) + print(f"in_memory : {in_mem_time * 1000:8.1f} ms / {len(args.queries)} queries") + + if not _pgvector_available(database): + print( + "pgvector : SKIPPED (needs Postgres + the embedding_vec column; " + "lands with the W8 pgvector migration + backfill)" + ) + return 0 + + pg = PgVectorRetriever( + embed_fn=ph.get_text_embeddings, + connection=database.session.connection(), + top_k=args.top_k, + threshold=cfg.link_threshold, + ) + pg_time = _time_backend(pg, args.queries, args.runs) + print(f"pgvector : {pg_time * 1000:8.1f} ms / {len(args.queries)} queries") + + # Agreement check: do the backends pick the same top-1 per query? + agree = sum( + in_mem.retrieve(q).candidates[0].cre_id == pg.retrieve(q).candidates[0].cre_id + for q in args.queries + if in_mem.retrieve(q).candidates and pg.retrieve(q).candidates + ) + print(f"top-1 agreement: {agree}/{len(args.queries)}") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/scripts/build_golden_dataset.py b/scripts/build_golden_dataset.py index 0905bd265..7a371c836 100644 --- a/scripts/build_golden_dataset.py +++ b/scripts/build_golden_dataset.py @@ -258,7 +258,7 @@ def build_positive_multilink(conn: sqlite3.Connection) -> List[Dict]: """ ).fetchall() out = [] - for node_id, name, section_id, text, cre_concat in rows: + for _node_id, name, section_id, text, cre_concat in rows: cre_ids = sorted(set(cre_concat.split("|"))) std = "OTHER" if "Top 10" in name: diff --git a/scripts/evaluate_librarian.py b/scripts/evaluate_librarian.py index aac3effed..7c2ff3b03 100644 --- a/scripts/evaluate_librarian.py +++ b/scripts/evaluate_librarian.py @@ -94,6 +94,82 @@ def predict(section: Section, registry: Set[str], hub: List[HubRep]) -> List[str return [] +def report_retrieval_recall( + rows: List[GoldenDatasetRow], cache_file: str, top_k: int, threshold: float +) -> None: + """Measure C.1 retrieval recall@k over the positive slice, live. + + Recall@k is the W3 metric — not the Jaccard link rule (that grades the + final auto-link, which needs the W4 reranker). It asks the only question + the search step controls: does the expected CRE id make it into the top-K + shortlist the reranker will later see? A miss here is unrecoverable + downstream, so it is the right thing to gate retrieval on. + + Live-only: there is no honest way to compute this offline. The candidate + pool must be the real CRE-node vectors, and seeding it from the golden + text itself is exactly the leakage the hub firewall strips. + """ + # Live deps are imported lazily so the offline harness needs neither a DB + # nor an embedding model. + from application.cmd.cre_main import db_connect + from application.defs import cre_defs + from application.prompt_client import prompt_client + from application.utils.librarian.candidate_retriever import ( + CandidatePool, + CandidateRetriever, + ) + from sqlalchemy import text as _sql_text + + database = db_connect(path=cache_file) + ph = prompt_client.PromptHandler(database=database) + + # The embeddings pool is keyed by the CRE's internal UUID (add_embedding + # stores cre_id=db_object.id), but the golden dataset speaks external_ids + # ("616-305"). Translate the pool keys to external_id so recall compares in + # the same id-space. Keys not in the map (a DB that already stores + # external_ids) pass through unchanged. + id_to_ext = { + row[0]: row[1] + for row in database.session.execute( + _sql_text("SELECT id, external_id FROM cre") + ) + if row[1] + } + pool = CandidatePool.from_mapping( + { + id_to_ext.get(k, k): v + for k, v in database.get_embeddings_by_doc_type( + cre_defs.Credoctypes.CRE.value + ).items() + } + ) + retriever = CandidateRetriever( + embed_fn=ph.get_text_embeddings, + pool=pool, + top_k=top_k, + threshold=threshold, + ) + + positives = [r for r in rows if r.slice.value == "positive" and r.expected.cre_ids] + if not positives: + print("retrieval recall: no positive rows with expected ids in this selection") + return + any_hit = all_hit = 0 + for row in positives: + retrieved = {c.cre_id for c in retriever.retrieve(row.input.text).candidates} + expected = set(row.expected.cre_ids or []) + if expected & retrieved: + any_hit += 1 + if expected <= retrieved: + all_hit += 1 + n = len(positives) + print( + f"retrieval recall@{top_k} (live, {n} positive rows): " + f"any-hit {any_hit}/{n} ({any_hit / n:.0%}), " + f"all-hit {all_hit}/{n} ({all_hit / n:.0%})" + ) + + def main(argv: List[str]) -> int: cfg = load_config() parser = argparse.ArgumentParser(description="Module C eval harness (W2: C.0)") @@ -111,6 +187,17 @@ def main(argv: List[str]) -> int: action="store_true", help="disable the leakage firewall (firewall is ON by default)", ) + parser.add_argument( + "--use_live_embeddings", + action="store_true", + help="connect to the OpenCRE DB + embedding model and measure semantic " + "retrieval recall@k over the positive slice (needs an LLM + populated DB)", + ) + parser.add_argument( + "--cache_file", + default="standards_cache.sqlite", + help="OpenCRE cache DB path for --use_live_embeddings", + ) args = parser.parse_args(argv) rows = load_dataset(args.dataset) @@ -167,6 +254,15 @@ def main(argv: List[str]) -> int: ) if not gate_ok: return 1 + if args.use_live_embeddings: + report_retrieval_recall( + rows, args.cache_file, args.top_k_retrieval, args.threshold + ) + else: + print( + "semantic retriever (C.1): wired; recall@k needs --use_live_embeddings " + "(no CRE vectors offline — seeding from golden text would be leakage)" + ) print(f"correct overall (semantic path still stubbed): {correct}/{len(rows)}") return 0