From ee5e119aeda33ee0cf2fbffd9357b01113742086 Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Thu, 18 Jun 2026 19:11:11 +0530 Subject: [PATCH 1/6] =?UTF-8?q?week=5F3:=20Module=20C=20C.1=20=E2=80=94=20?= =?UTF-8?q?candidate=20retriever=20(in-memory=20+=20pgvector)=20+=20pipeli?= =?UTF-8?q?ne=20switch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The semantic search step. For sections with no explicit CRE id, embed the text and cosine-rank the CRE vector hub, returning the top-K (default 20) candidates as an RFC RetrievalAudit (reranked[] empty until W4). Two interchangeable backends behind one retrieve() seam, selected by CRE_LIBRARIAN_RETRIEVER_BACKEND: - CandidateRetriever: in-memory sklearn cosine (SQLite/CI/harness) - PgVectorRetriever: Postgres-side <=> cosine over embedding_vec build_retriever() is the factory. Reuses OpenCRE's embedding stack (PromptHandler.get_text_embeddings, db.get_embeddings_by_doc_type). Dim gate fails loudly on empty/ragged hubs and query/hub width mismatch. CLI switch (--run_librarian / --librarian_dry_run) runs the pipeline dry-run; harness --use_live_embeddings measures retrieval recall@k. The pgvector schema migration is a W8 deliverable (validated against real Postgres there); W3 ships only the code path, unit-tested via a fake connection. 82 librarian tests pass. --- .env.example | 14 +- application/cmd/cre_main.py | 114 +++++++ .../librarian/candidate_retriever_test.py | 213 +++++++++++++ .../utils/librarian/candidate_retriever.py | 284 ++++++++++++++++++ application/utils/librarian/config_loader.py | 13 + cre.py | 18 ++ scripts/benchmark_retriever.py | 121 ++++++++ scripts/evaluate_librarian.py | 77 +++++ 8 files changed, 849 insertions(+), 5 deletions(-) create mode 100644 application/tests/librarian/candidate_retriever_test.py create mode 100644 application/utils/librarian/candidate_retriever.py create mode 100644 scripts/benchmark_retriever.py 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..8a3588117 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,114 @@ 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. + """ + 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) + # 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/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/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/evaluate_librarian.py b/scripts/evaluate_librarian.py index aac3effed..46c71d4b1 100644 --- a/scripts/evaluate_librarian.py +++ b/scripts/evaluate_librarian.py @@ -94,6 +94,63 @@ 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, + ) + + database = db_connect(path=cache_file) + ph = prompt_client.PromptHandler(database=database) + pool = CandidatePool.from_mapping( + database.get_embeddings_by_doc_type(cre_defs.Credoctypes.CRE.value) + ) + 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 +168,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 +235,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 From 56691a475e6ad7fc08f02105ee20dbdf63013952 Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Thu, 25 Jun 2026 13:41:52 +0530 Subject: [PATCH 2/6] week_3: address CodeRabbit review on #937 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolve the CodeRabbit review comments on the Module C PR: - schemas.py: enforce RFC format parity — AnyUrl for url fields, datetime for committed_at/filtered_at/classified_at/created_at - section_validator.py: replace assert with a typed MalformedKnowledgeItemError guard (survives python -O) - knowledge_source.py: skip+log malformed JSONL rows instead of aborting the whole iteration on ValidationError - config_loader_test.py: isolate os.environ and assert the specific FrozenInstanceError - dataset_test.py: add a timeout to the determinism subprocess call - build_golden_dataset.py: rename unused loop var node_id -> _node_id - section_validator_test.py: assert committed_at as the typed datetime --- .../tests/librarian/config_loader_test.py | 6 ++++-- application/tests/librarian/dataset_test.py | 1 + .../tests/librarian/section_validator_test.py | 6 +++++- application/utils/librarian/knowledge_source.py | 17 +++++++++++++++-- application/utils/librarian/schemas.py | 15 ++++++++------- .../utils/librarian/section_validator.py | 5 +++-- scripts/build_golden_dataset.py | 2 +- 7 files changed, 37 insertions(+), 15 deletions(-) diff --git a/application/tests/librarian/config_loader_test.py b/application/tests/librarian/config_loader_test.py index d74deb1b4..cadcf9d20 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 @@ -18,8 +19,9 @@ 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] 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/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/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: From e6585716d6fdbaf425b51571550339a5356f0b45 Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Thu, 2 Jul 2026 11:37:01 +0530 Subject: [PATCH 3/6] =?UTF-8?q?week=5F3:=20address=20review=20=E2=80=94=20?= =?UTF-8?q?warn=20loudly=20when=20pgvector=20backend=20runs=20on=20non-Pos?= =?UTF-8?q?tgres?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Maintainer nit on #937: CRE_LIBRARIAN_RETRIEVER_BACKEND=pgvector is selectable via env, but the embedding_vec column doesn't land until W8. Emit a loud startup warning at pipeline-switch time when the backend is pgvector on a non-Postgres (e.g. SQLite) database, since retrieve() would otherwise fail with an opaque SQL error. --- application/cmd/cre_main.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/application/cmd/cre_main.py b/application/cmd/cre_main.py index 8a3588117..486bd0c76 100644 --- a/application/cmd/cre_main.py +++ b/application/cmd/cre_main.py @@ -1105,6 +1105,16 @@ def run_librarian( 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). From f71fbcc56f55e6b2622283c583d956c78f7061d6 Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Fri, 3 Jul 2026 07:48:40 +0530 Subject: [PATCH 4/6] week_3: address CodeRabbit findings on #937 - config_loader_test: assert retriever_backend in both the defaults and the override tests (was the only LibrarianConfig field left unchecked). - __init__.py: refresh the stale 'No linking logic yet' scope note to reflect W1 contracts + W2 C.0 boundary + W3 C.1 retriever. Other #937 CodeRabbit findings were already addressed on this branch (knowledge_source model_validate_json error handling, schemas AnyUrl/datetime fields, section_validator assert->if, test_config_is_frozen isolation + FrozenInstanceError, dataset_test subprocess timeout, build_golden_dataset unused loop var). --- application/tests/librarian/config_loader_test.py | 3 +++ application/utils/librarian/__init__.py | 4 +++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/application/tests/librarian/config_loader_test.py b/application/tests/librarian/config_loader_test.py index cadcf9d20..1b7abb294 100644 --- a/application/tests/librarian/config_loader_test.py +++ b/application/tests/librarian/config_loader_test.py @@ -10,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) @@ -27,6 +28,7 @@ def test_config_is_frozen(self): 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", @@ -39,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/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). From df36899bb7a20898d25c564f66c4e19b224ec416 Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Fri, 3 Jul 2026 08:09:23 +0530 Subject: [PATCH 5/6] week_3: document ops posture for the librarian CLI (production note) Address northdpole's production note on #937: record at the run_librarian entrypoint that it is opt-in CLI only (not on Procfile, not wired into web or worker until W8) and that the paid embedding API cost is incurred only on manual runs, never by the running deployment. --- application/cmd/cre_main.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/application/cmd/cre_main.py b/application/cmd/cre_main.py index 486bd0c76..4b814640a 100644 --- a/application/cmd/cre_main.py +++ b/application/cmd/cre_main.py @@ -1077,6 +1077,11 @@ def run_librarian( 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, From 87103b54d2b3572a31629c4e6434880baae8e3db Mon Sep 17 00:00:00 2001 From: PRAteek-singHWY Date: Tue, 7 Jul 2026 11:43:20 +0530 Subject: [PATCH 6/6] week_3: fix recall harness id-space (UUID->external_id) so live recall@20 is measurable The embeddings pool is keyed by the CRE internal UUID, but the golden dataset expects external_ids (e.g. 616-305). Without translating, every comparison missed and recall@20 read 0%. Map pool keys via cre.id-> external_id (pass-through for DBs already keyed by external_id). --- scripts/evaluate_librarian.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/scripts/evaluate_librarian.py b/scripts/evaluate_librarian.py index 46c71d4b1..7c2ff3b03 100644 --- a/scripts/evaluate_librarian.py +++ b/scripts/evaluate_librarian.py @@ -118,11 +118,30 @@ def report_retrieval_recall( 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( - database.get_embeddings_by_doc_type(cre_defs.Credoctypes.CRE.value) + { + 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,