From 13b86f08445697db2dfba082c69a1680c8cde9c1 Mon Sep 17 00:00:00 2001 From: Spyros Date: Mon, 13 Jul 2026 11:48:27 +0100 Subject: [PATCH 1/4] feat: pgvector embedding_vec migration + chat similarity cutover Merge dual Alembic heads into one revision that adds embedding_vec, backfills from CSV, dual-writes on add_embedding, and prefers Postgres <=> for chat/import matching. Closes #977. Co-authored-by: Cursor --- .env.example | 2 +- application/cmd/cre_main.py | 14 +- application/database/db.py | 73 ++++++++ application/database/pgvector_utils.py | 157 ++++++++++++++++++ application/prompt_client/prompt_client.py | 34 ++++ application/tests/db_test.py | 26 +++ application/tests/pgvector_utils_test.py | 112 +++++++++++++ .../prompt_client_pgvector_similarity_test.py | 69 ++++++++ .../utils/librarian/candidate_retriever.py | 17 +- ...c7d8e9f0a1b2_add_embedding_vec_pgvector.py | 63 +++++++ 10 files changed, 548 insertions(+), 19 deletions(-) create mode 100644 application/database/pgvector_utils.py create mode 100644 application/tests/pgvector_utils_test.py create mode 100644 application/tests/prompt_client_pgvector_similarity_test.py create mode 100644 migrations/versions/c7d8e9f0a1b2_add_embedding_vec_pgvector.py diff --git a/.env.example b/.env.example index c4eb07f93..17ad95d27 100644 --- a/.env.example +++ b/.env.example @@ -74,7 +74,7 @@ OpenCRE_gspread_Auth=path/to/credentials.json # application/utils/librarian/config_loader.py for validation) # Retrieval backend: in_memory (sklearn cosine; SQLite dev/CI/harness) or -# pgvector (Postgres vector column — pending prod extension + mentor OK). +# pgvector (Postgres ``embedding_vec`` after Alembic c7d8e9f0a1b2 / #977). 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 diff --git a/application/cmd/cre_main.py b/application/cmd/cre_main.py index f42a4a5c9..14015cafe 100644 --- a/application/cmd/cre_main.py +++ b/application/cmd/cre_main.py @@ -1115,14 +1115,14 @@ def run_librarian( backend = RetrieverBackend(cfg.retriever_backend) if backend is RetrieverBackend.pgvector: - dialect = database.session.connection().dialect.name - if dialect != "postgresql": + # Readiness must include the embedding_vec column (not dialect alone): + # Postgres pre-migration would otherwise fail on every retrieve(). + if not database.can_use_pgvector_similarity(): 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, + "CRE_LIBRARIAN_RETRIEVER_BACKEND=pgvector selected, but " + "Postgres is not ready for pgvector similarity (need Alembic " + "c7d8e9f0a1b2 / #977 embedding_vec). Set the backend to " + "in_memory until then." ) # 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 diff --git a/application/database/db.py b/application/database/db.py index e3290e5bc..bb62bad60 100644 --- a/application/database/db.py +++ b/application/database/db.py @@ -191,6 +191,9 @@ class Embeddings(BaseModel): # type: ignore embeddings_content = sqla.Column(sqla.String, nullable=True, default=None) embedding_model_id = sqla.Column(sqla.String, nullable=True, default=None) embedding_dim = sqla.Column(sqla.Integer, nullable=True, default=None) + # Postgres: real ``vector(N)`` via Alembic (c7d8e9f0a1b2). SQLite tests map + # this as Text storing a pgvector literal for dual-write coverage. + embedding_vec = sqla.Column(sqla.Text, nullable=True, default=None) class GapAnalysisResults(BaseModel): @@ -2486,6 +2489,12 @@ def add_embedding( ) existing = self.get_embedding(db_object.id) embeddings_str = ",".join([str(e) for e in embeddings]) + # Dual-write CSV + pgvector literal. On Postgres after migration the + # column is ``vector(N)``; on SQLite tests it is Text. Never LLM + # re-embed here — see application.database.pgvector_utils. + from application.database.pgvector_utils import to_pgvector_literal + + embedding_vec_literal = to_pgvector_literal(embeddings) resolved_node_url: Optional[str] = None if doctype != cre_defs.Credoctypes.CRE: resolved_node_url = ( @@ -2502,6 +2511,7 @@ def add_embedding( embeddings_content=embedding_text, embedding_model_id=embedding_model_id, embedding_dim=embedding_dim, + embedding_vec=embedding_vec_literal, ) else: emb = Embeddings( @@ -2512,6 +2522,7 @@ def add_embedding( embeddings_url=resolved_node_url, embedding_model_id=embedding_model_id, embedding_dim=embedding_dim, + embedding_vec=embedding_vec_literal, ) self.session.add(emb) self.session.commit() @@ -2522,6 +2533,7 @@ def add_embedding( existing[0].embeddings_content = embedding_text existing[0].embedding_model_id = embedding_model_id existing[0].embedding_dim = embedding_dim + existing[0].embedding_vec = embedding_vec_literal if doctype != cre_defs.Credoctypes.CRE: if embeddings_url is not None: existing[0].embeddings_url = embeddings_url @@ -2531,6 +2543,67 @@ def add_embedding( return existing + def can_use_pgvector_similarity(self) -> bool: + """True when Postgres has ``embeddings.embedding_vec`` for SQL cosine.""" + cached = getattr(self, "_pgvector_similarity_ready", None) + if cached is not None: + return bool(cached) + try: + bind = self.session.get_bind() + from application.database.pgvector_utils import embedding_vec_column_exists + + ready = getattr( + bind.dialect, "name", "" + ) == "postgresql" and embedding_vec_column_exists(bind) + except Exception: + ready = False + self._pgvector_similarity_ready = ready + return ready + + def find_most_similar_embedding_id( + self, + query_embedding: List[float], + *, + doc_type: str, + id_column: str, + similarity_threshold: float, + ) -> Tuple[Optional[str], Optional[float]]: + """Top-1 cosine match via pgvector ``<=>`` (Postgres only). + + Returns ``(object_id, score)`` or ``(None, None)`` below threshold / + when no rows match. + """ + from application.database.pgvector_utils import ( + most_similar_id_sql, + to_pgvector_literal, + ) + from sqlalchemy import text as sql_text + + sql = most_similar_id_sql(id_column) + bind = self.session.get_bind() + try: + row = bind.execute( + sql_text(sql), + { + "q": to_pgvector_literal(query_embedding), + "doc_type": doc_type, + }, + ).fetchone() + if not row: + return None, None + score = float(row.score) + if score < similarity_threshold: + return None, None + return str(row.object_id), score + except Exception as exc: + # Match prior sklearn-path resilience: chat/import should degrade + # to "no match" rather than 500 on a transient driver/query error. + logger.warning( + "pgvector similarity query failed (%s); returning no match", + exc, + ) + return None, None + def assert_embedding_contract( self, *, diff --git a/application/database/pgvector_utils.py b/application/database/pgvector_utils.py new file mode 100644 index 000000000..7f1074a73 --- /dev/null +++ b/application/database/pgvector_utils.py @@ -0,0 +1,157 @@ +"""Helpers for the ``embeddings.embedding_vec`` pgvector column. + +Do **not** re-embed the full corpus on Heroku/production dynos — likely +OOM-killed. Re-embedding (fix missing / wrong-dim rows via the LLM) must run +only where there is enough RAM (local or a high-resource worker), then +sync/backfill vectors. Alembic on prod only: enable extension, add column, +copy CSV→vector for matching dims. + +HNSW/IVFFlat indexes are deferred (Essential-1 capacity); see migration +``c7d8e9f0a1b2`` docstring for the follow-up plan. +""" + +from __future__ import annotations + +import os +from typing import Any, Iterable, Optional, Sequence, Set, Tuple + +from sqlalchemy import text + +HEROKU_REEMBED_WARNING = ( + "Do not re-embed the full corpus on Heroku/production dynos — likely " + "OOM-killed. Re-embedding must run only on a machine with enough RAM, " + "then sync/backfill vectors. Alembic on prod only copies CSV→vector." +) + + +class MultiDimEmbeddingError(RuntimeError): + """Raised when more than one embedding dimension is present in the DB.""" + + +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) + "]" + + +def csv_embeddings_to_literal(csv: str) -> str: + """Wrap a comma-separated CSV embedding string as a pgvector literal.""" + return f"[{(csv or '').strip()}]" + + +def parse_csv_embedding_dim(csv: str) -> int: + """Return the number of floats implied by a CSV embeddings string.""" + parts = [p for p in (csv or "").split(",") if p.strip() != ""] + return len(parts) + + +def collect_distinct_dims( + rows: Iterable[Tuple[Optional[int], Optional[str]]], +) -> Set[int]: + """Collect distinct dims from ``(embedding_dim, embeddings_csv)`` rows.""" + dims: Set[int] = set() + for embedding_dim, csv in rows: + if embedding_dim is not None: + dims.add(int(embedding_dim)) + continue + if csv is None or str(csv).strip() == "": + continue + dims.add(parse_csv_embedding_dim(str(csv))) + return dims + + +def require_single_embedding_dim(connection: Any) -> int: + """Return the single project-wide embedding dim or raise. + + Prefers ``embedding_dim`` metadata; falls back to CSV length. If the table + is empty, uses ``CRE_EMBED_EXPECTED_DIM`` when set. + """ + rows = connection.execute( + text("SELECT embedding_dim, embeddings FROM embeddings") + ).fetchall() + dims = collect_distinct_dims((row[0], row[1]) for row in rows) + if len(dims) > 1: + raise MultiDimEmbeddingError( + f"multiple embedding dimensions detected in DB: {sorted(dims)}. " + "Refuse to add embedding_vec — reconcile dims offline (do not " + "re-embed on Heroku). " + HEROKU_REEMBED_WARNING + ) + if len(dims) == 1: + return next(iter(dims)) + + env_dim = (os.environ.get("CRE_EMBED_EXPECTED_DIM", "") or "").strip() + if env_dim: + return int(env_dim) + raise RuntimeError( + "embeddings table has no vectors to infer dimension; set " + "CRE_EMBED_EXPECTED_DIM before running this migration. " + + HEROKU_REEMBED_WARNING + ) + + +def embedding_vec_column_exists(connection: Any) -> bool: + """True when ``embeddings.embedding_vec`` exists (Postgres information_schema).""" + if getattr(getattr(connection, "dialect", None), "name", None) != "postgresql": + # SQLite / others: probe via PRAGMA or SQLAlchemy inspector if needed. + dialect_name = getattr(getattr(connection, "dialect", None), "name", "") + if dialect_name == "sqlite": + rows = connection.execute(text("PRAGMA table_info(embeddings)")).fetchall() + return any(r[1] == "embedding_vec" for r in rows) + return False + row = connection.execute( + text( + "SELECT 1 FROM information_schema.columns " + "WHERE table_name = 'embeddings' AND column_name = 'embedding_vec' " + "LIMIT 1" + ) + ).fetchone() + return row is not None + + +def backfill_embedding_vec(connection: Any, dim: int) -> int: + """Idempotent CSV → ``embedding_vec`` copy for rows matching ``dim``. + + Returns the number of rows reported updated (driver-dependent). + """ + result = connection.execute( + text( + """ + UPDATE embeddings + SET embedding_vec = CAST(('[' || embeddings || ']') AS vector) + WHERE embedding_vec IS NULL + AND embeddings IS NOT NULL + AND btrim(embeddings) <> '' + AND ( + embedding_dim = :dim + OR ( + embedding_dim IS NULL + AND ( + length(embeddings) + - length(replace(embeddings, ',', '')) + + 1 + ) = :dim + ) + ) + """ + ), + {"dim": dim}, + ) + return int(getattr(result, "rowcount", 0) or 0) + + +def most_similar_id_sql(id_column: str) -> str: + """SQL for top-1 cosine similarity over ``embedding_vec``. + + ``id_column`` must be ``cre_id`` or ``node_id`` (caller-validated). + Score is ``1 - (embedding_vec <=> query)`` (cosine similarity). + """ + if id_column not in ("cre_id", "node_id"): + raise ValueError(f"id_column must be cre_id or node_id, got {id_column!r}") + return ( + f"SELECT {id_column} AS object_id, " + f"1 - (embedding_vec <=> CAST(:q AS vector)) AS score " + f"FROM embeddings " + f"WHERE doc_type = :doc_type AND {id_column} IS NOT NULL " + f"AND embedding_vec IS NOT NULL " + f"ORDER BY embedding_vec <=> CAST(:q AS vector) " + f"LIMIT 1" + ) diff --git a/application/prompt_client/prompt_client.py b/application/prompt_client/prompt_client.py index 112114034..d2ec1fdb3 100644 --- a/application/prompt_client/prompt_client.py +++ b/application/prompt_client/prompt_client.py @@ -1060,6 +1060,15 @@ def get_id_of_most_similar_cre(self, item_embedding: List[float]) -> Optional[st Returns: str: _description_ """ + if self.database.can_use_pgvector_similarity(): + match_id, _score = self.database.find_most_similar_embedding_id( + item_embedding, + doc_type=cre_defs.Credoctypes.CRE.value, + id_column="cre_id", + similarity_threshold=SIMILARITY_THRESHOLD, + ) + return match_id + if not hasattr(self, "existing_cre_embeddings"): ( self.existing_cre_embeddings, @@ -1102,6 +1111,15 @@ def get_id_of_most_similar_node( Returns: str: the database id of the closest database standard """ + if self.database.can_use_pgvector_similarity(): + match_id, _score = self.database.find_most_similar_embedding_id( + standard_text_embedding, + doc_type=cre_defs.Credoctypes.Standard.value, + id_column="node_id", + similarity_threshold=SIMILARITY_THRESHOLD, + ) + return match_id + if not hasattr(self, "existing_node_embeddings"): ( self.existing_node_embeddings, @@ -1145,6 +1163,14 @@ def get_id_of_most_similar_cre_paginated( Returns: str: the ID of the CRE with the closest cosine_similarity """ + if self.database.can_use_pgvector_similarity(): + return self.database.find_most_similar_embedding_id( + item_embedding, + doc_type=cre_defs.Credoctypes.CRE.value, + id_column="cre_id", + similarity_threshold=similarity_threshold, + ) + embedding_array = sparse.csr_matrix( np.array(item_embedding).reshape(1, -1) ) # convert embedding into a 1-dimentional numpy array @@ -1197,6 +1223,14 @@ def get_id_of_most_similar_node_paginated( Returns: str: the db id of the most similar object """ + if self.database.can_use_pgvector_similarity(): + return self.database.find_most_similar_embedding_id( + question_embedding, + doc_type=cre_defs.Credoctypes.Standard.value, + id_column="node_id", + similarity_threshold=similarity_threshold, + ) + embedding_array = sparse.csr_matrix( np.array(question_embedding).reshape(1, -1) ) # convert embedding into a 1-dimentional numpy array diff --git a/application/tests/db_test.py b/application/tests/db_test.py index b75f46730..4f8c85efd 100644 --- a/application/tests/db_test.py +++ b/application/tests/db_test.py @@ -2390,6 +2390,32 @@ def test_add_embedding_persists_embedding_contract_metadata(self): finally: os.environ.pop("CRE_EMBED_MODEL", None) + def test_add_embedding_dual_writes_embedding_vec_literal(self): + """CSV + embedding_vec literal are persisted together (#977).""" + from application.database.pgvector_utils import to_pgvector_literal + + dbsa = db.Node( + subsection="", + section="Sec", + name="DualWriteStd", + link="https://example.com/dw", + ntype=defs.Credoctypes.Standard.value, + ) + self.collection.session.add(dbsa) + self.collection.session.commit() + vec = [0.1, 0.2, 0.3] + self.collection.add_embedding( + db_object=dbsa, + doctype=defs.Credoctypes.Standard.value, + embeddings=vec, + embedding_text="dual-write", + ) + row = self.collection.get_embedding(dbsa.id)[0] + self.assertEqual(row.embeddings, "0.1,0.2,0.3") + self.assertEqual(row.embeddings_content, "dual-write") + self.assertEqual(row.embedding_vec, to_pgvector_literal(vec)) + self.assertEqual(row.embedding_dim, 3) + def test_assert_embedding_contract_fails_on_mixed_dimensions(self): n1 = db.Node( subsection="", diff --git a/application/tests/pgvector_utils_test.py b/application/tests/pgvector_utils_test.py new file mode 100644 index 000000000..c66463e95 --- /dev/null +++ b/application/tests/pgvector_utils_test.py @@ -0,0 +1,112 @@ +"""Tests for pgvector embedding helpers (dim gate, literal, backfill SQL).""" + +import os +import unittest +from unittest.mock import MagicMock + +from application.database.pgvector_utils import ( + MultiDimEmbeddingError, + backfill_embedding_vec, + collect_distinct_dims, + csv_embeddings_to_literal, + most_similar_id_sql, + parse_csv_embedding_dim, + require_single_embedding_dim, + to_pgvector_literal, +) + + +class PgvectorLiteralTest(unittest.TestCase): + def test_to_pgvector_literal(self) -> None: + self.assertEqual(to_pgvector_literal([1, 2.5, 0]), "[1.0,2.5,0.0]") + + def test_csv_to_literal(self) -> None: + self.assertEqual(csv_embeddings_to_literal("1.0,2.0"), "[1.0,2.0]") + + def test_parse_csv_dim(self) -> None: + self.assertEqual(parse_csv_embedding_dim("0.1,0.2,0.3"), 3) + self.assertEqual(parse_csv_embedding_dim(""), 0) + + +class DistinctDimsTest(unittest.TestCase): + def test_single_dim_from_metadata(self) -> None: + self.assertEqual(collect_distinct_dims([(3, "1,2,3"), (3, "4,5,6")]), {3}) + + def test_csv_fallback_when_metadata_missing(self) -> None: + self.assertEqual(collect_distinct_dims([(None, "1,2"), (None, "3,4")]), {2}) + + def test_mixed_dims_detected(self) -> None: + self.assertEqual( + collect_distinct_dims([(2, "1,2"), (3, "1,2,3")]), + {2, 3}, + ) + + +class RequireSingleDimTest(unittest.TestCase): + def test_raises_on_multiple_dims(self) -> None: + conn = MagicMock() + conn.execute.return_value.fetchall.return_value = [ + (2, "1,2"), + (3, "1,2,3"), + ] + with self.assertRaises(MultiDimEmbeddingError) as cm: + require_single_embedding_dim(conn) + self.assertIn("multiple embedding dimensions", str(cm.exception)) + self.assertIn("Heroku", str(cm.exception)) + + def test_returns_single_dim(self) -> None: + conn = MagicMock() + conn.execute.return_value.fetchall.return_value = [ + (768, "1," * 767 + "1"), + (768, None), + ] + self.assertEqual(require_single_embedding_dim(conn), 768) + + def test_empty_table_uses_env_dim(self) -> None: + conn = MagicMock() + conn.execute.return_value.fetchall.return_value = [] + os.environ["CRE_EMBED_EXPECTED_DIM"] = "4" + try: + self.assertEqual(require_single_embedding_dim(conn), 4) + finally: + os.environ.pop("CRE_EMBED_EXPECTED_DIM", None) + + def test_empty_table_without_env_raises(self) -> None: + conn = MagicMock() + conn.execute.return_value.fetchall.return_value = [] + os.environ.pop("CRE_EMBED_EXPECTED_DIM", None) + with self.assertRaises(RuntimeError): + require_single_embedding_dim(conn) + + +class BackfillTest(unittest.TestCase): + def test_backfill_runs_idempotent_update(self) -> None: + conn = MagicMock() + result = MagicMock() + result.rowcount = 2 + conn.execute.return_value = result + self.assertEqual(backfill_embedding_vec(conn, 3), 2) + args, kwargs = conn.execute.call_args + sql = str(args[0]) + self.assertIn("embedding_vec IS NULL", sql) + self.assertIn("CAST", sql) + # bound dim + bound = args[1] if len(args) > 1 else kwargs.get("parameters") or kwargs + if isinstance(bound, dict): + self.assertEqual(bound.get("dim"), 3) + + +class MostSimilarSqlTest(unittest.TestCase): + def test_sql_uses_id_column_and_cosine(self) -> None: + sql = most_similar_id_sql("node_id") + self.assertIn("node_id", sql) + self.assertIn("<=>", sql) + self.assertIn("LIMIT 1", sql) + + def test_rejects_unknown_column(self) -> None: + with self.assertRaises(ValueError): + most_similar_id_sql("id") + + +if __name__ == "__main__": + unittest.main() diff --git a/application/tests/prompt_client_pgvector_similarity_test.py b/application/tests/prompt_client_pgvector_similarity_test.py new file mode 100644 index 000000000..6e11ce0b3 --- /dev/null +++ b/application/tests/prompt_client_pgvector_similarity_test.py @@ -0,0 +1,69 @@ +"""Chat/import similarity prefers pgvector when the DB reports it ready.""" + +import unittest +from unittest.mock import MagicMock, patch + +from application.prompt_client import prompt_client as prompt_client_mod + + +class PgvectorSimilarityCutoverTest(unittest.TestCase): + def test_node_paginated_uses_pgvector_when_available(self) -> None: + database = MagicMock() + database.can_use_pgvector_similarity.return_value = True + database.find_most_similar_embedding_id.return_value = ("node-1", 0.91) + + handler = prompt_client_mod.PromptHandler.__new__( + prompt_client_mod.PromptHandler + ) + handler.database = database + + result = handler.get_id_of_most_similar_node_paginated( + [0.1, 0.2, 0.3], similarity_threshold=0.7 + ) + self.assertEqual(result, ("node-1", 0.91)) + database.find_most_similar_embedding_id.assert_called_once() + kwargs = database.find_most_similar_embedding_id.call_args.kwargs + self.assertEqual(kwargs["id_column"], "node_id") + database.get_embeddings_by_doc_type_paginated.assert_not_called() + + def test_cre_paginated_uses_pgvector_when_available(self) -> None: + database = MagicMock() + database.can_use_pgvector_similarity.return_value = True + database.find_most_similar_embedding_id.return_value = ("cre-1", 0.88) + + handler = prompt_client_mod.PromptHandler.__new__( + prompt_client_mod.PromptHandler + ) + handler.database = database + + result = handler.get_id_of_most_similar_cre_paginated( + [0.1, 0.2, 0.3], similarity_threshold=0.7 + ) + self.assertEqual(result, ("cre-1", 0.88)) + kwargs = database.find_most_similar_embedding_id.call_args.kwargs + self.assertEqual(kwargs["id_column"], "cre_id") + + +class FindMostSimilarEmbeddingIdResilienceTest(unittest.TestCase): + def test_query_error_returns_no_match(self) -> None: + from application.database.db import Node_collection + + database = Node_collection.__new__(Node_collection) + session = MagicMock() + session.execute.side_effect = RuntimeError("driver boom") + session.get_bind.return_value = MagicMock() + database.session = session + + with patch("application.database.db.logger"): + match_id, score = database.find_most_similar_embedding_id( + [0.1, 0.2, 0.3], + doc_type="CRE", + id_column="cre_id", + similarity_threshold=0.7, + ) + self.assertIsNone(match_id) + self.assertIsNone(score) + + +if __name__ == "__main__": + unittest.main() diff --git a/application/utils/librarian/candidate_retriever.py b/application/utils/librarian/candidate_retriever.py index 7e1c6f93c..63c6d8b6d 100644 --- a/application/utils/librarian/candidate_retriever.py +++ b/application/utils/librarian/candidate_retriever.py @@ -26,9 +26,8 @@ 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. + RAM. Requires the ``vector`` extension + that column (Alembic + ``c7d8e9f0a1b2`` / #977). Unavailable on SQLite. The RFC is silent on retrieval tech — it mandates only the ``candidates[]``/``reranked[]`` audit trail — so the backend choice is ours; @@ -47,6 +46,7 @@ import numpy as np from sklearn.metrics.pairwise import cosine_similarity +from application.database.pgvector_utils import to_pgvector_literal from application.utils.librarian.schemas import CreCandidate, RetrievalAudit # A function that turns one piece of text into a single dense vector. @@ -63,7 +63,7 @@ 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). + # extension + an embedding_vec column; Alembic c7d8e9f0a1b2 / #977). pgvector = "pgvector" @@ -177,11 +177,6 @@ def retrieve(self, text: str) -> RetrievalAudit: ) -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. @@ -192,8 +187,8 @@ class PgVectorRetriever: 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``. + (Alembic ``c7d8e9f0a1b2`` / #977). Unavailable on SQLite — the factory + routes SQLite to ``in_memory``. """ # Parameterized; :q is bound as a pgvector text literal and cast in-SQL. diff --git a/migrations/versions/c7d8e9f0a1b2_add_embedding_vec_pgvector.py b/migrations/versions/c7d8e9f0a1b2_add_embedding_vec_pgvector.py new file mode 100644 index 000000000..605bd408f --- /dev/null +++ b/migrations/versions/c7d8e9f0a1b2_add_embedding_vec_pgvector.py @@ -0,0 +1,63 @@ +""" +Add embeddings.embedding_vec (pgvector) and merge dual Alembic heads. + +Revision ID: c7d8e9f0a1b2 +Revises: 9b1c2d3e4f50, ab12cd34ef56 +Create Date: 2026-07-13 + +Do **not** re-embed the full corpus on Heroku/production dynos — likely +OOM-killed. This migration only: CREATE EXTENSION vector, ADD COLUMN +embedding_vec, and copy matching-dim CSV → vector. Fix missing / wrong-dim +rows offline on a high-RAM machine, then re-run / backfill. + +ANN index (HNSW/IVFFlat) is intentionally deferred: opencreorg runs on +Heroku Postgres Essential-1; add a follow-up migration with +``CREATE INDEX ... USING hnsw (embedding_vec vector_cosine_ops)`` once +dims/backfill are settled and plan capacity allows. +""" + +from alembic import op +import sqlalchemy as sa + +from application.database.pgvector_utils import ( + HEROKU_REEMBED_WARNING, + backfill_embedding_vec, + require_single_embedding_dim, +) + +# revision identifiers, used by Alembic. +revision = "c7d8e9f0a1b2" +down_revision = ("9b1c2d3e4f50", "ab12cd34ef56") +branch_labels = None +depends_on = None + + +def upgrade(): + """Postgres-only: enable pgvector, add embedding_vec, backfill from CSV. + + SQLite/CI: no-op (chat/Librarian keep in-memory / CSV paths). + """ + conn = op.get_bind() + if conn.dialect.name != "postgresql": + return + + # See module docstring / HEROKU_REEMBED_WARNING — no LLM re-embed here. + _ = HEROKU_REEMBED_WARNING + + op.execute(sa.text("CREATE EXTENSION IF NOT EXISTS vector")) + dim = require_single_embedding_dim(conn) + op.execute( + sa.text( + f"ALTER TABLE embeddings " + f"ADD COLUMN IF NOT EXISTS embedding_vec vector({int(dim)})" + ) + ) + backfill_embedding_vec(conn, dim) + + +def downgrade(): + conn = op.get_bind() + if conn.dialect.name != "postgresql": + return + op.execute(sa.text("ALTER TABLE embeddings DROP COLUMN IF EXISTS embedding_vec")) + # Intentionally leave the `vector` extension installed (may be shared). From 495037a1faf829a37a848a65403baaa38f560cce Mon Sep 17 00:00:00 2001 From: Spyros Date: Mon, 13 Jul 2026 12:02:21 +0100 Subject: [PATCH 2/4] fix: use SQLAlchemy 2 Engine-safe execute for pgvector helpers session.get_bind() returns an Engine which no longer has .execute; open a Connection so can_use_pgvector_similarity and similarity queries work on Postgres. --- application/database/db.py | 18 +++++++----- application/database/pgvector_utils.py | 39 ++++++++++++++++++-------- 2 files changed, 39 insertions(+), 18 deletions(-) diff --git a/application/database/db.py b/application/database/db.py index bb62bad60..0d799cae4 100644 --- a/application/database/db.py +++ b/application/database/db.py @@ -2581,14 +2581,18 @@ def find_most_similar_embedding_id( sql = most_similar_id_sql(id_column) bind = self.session.get_bind() + params = { + "q": to_pgvector_literal(query_embedding), + "doc_type": doc_type, + } try: - row = bind.execute( - sql_text(sql), - { - "q": to_pgvector_literal(query_embedding), - "doc_type": doc_type, - }, - ).fetchone() + # Prefer the Session so we stay inside the app transaction; fall + # back to a Connection from the Engine for SQLAlchemy 2. + if hasattr(self.session, "execute"): + row = self.session.execute(sql_text(sql), params).fetchone() + else: + with bind.connect() as conn: + row = conn.execute(sql_text(sql), params).fetchone() if not row: return None, None score = float(row.score) diff --git a/application/database/pgvector_utils.py b/application/database/pgvector_utils.py index 7f1074a73..817cccb87 100644 --- a/application/database/pgvector_utils.py +++ b/application/database/pgvector_utils.py @@ -16,6 +16,7 @@ from typing import Any, Iterable, Optional, Sequence, Set, Tuple from sqlalchemy import text +from sqlalchemy.engine import Engine HEROKU_REEMBED_WARNING = ( "Do not re-embed the full corpus on Heroku/production dynos — likely " @@ -28,6 +29,20 @@ class MultiDimEmbeddingError(RuntimeError): """Raised when more than one embedding dimension is present in the DB.""" +def _execute(connection: Any, statement: Any, params: Optional[dict] = None): + """Run ``statement`` on a Connection/Session, or open one from an Engine. + + SQLAlchemy 2 removed ``Engine.execute``; Alembic upgrade passes a + Connection, while runtime code often has ``session.get_bind()`` → Engine. + """ + if isinstance(connection, Engine): + with connection.connect() as conn: + result = conn.execute(statement, params or {}) + conn.commit() + return result + return connection.execute(statement, params or {}) + + 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) + "]" @@ -65,8 +80,8 @@ def require_single_embedding_dim(connection: Any) -> int: Prefers ``embedding_dim`` metadata; falls back to CSV length. If the table is empty, uses ``CRE_EMBED_EXPECTED_DIM`` when set. """ - rows = connection.execute( - text("SELECT embedding_dim, embeddings FROM embeddings") + rows = _execute( + connection, text("SELECT embedding_dim, embeddings FROM embeddings") ).fetchall() dims = collect_distinct_dims((row[0], row[1]) for row in rows) if len(dims) > 1: @@ -90,19 +105,20 @@ def require_single_embedding_dim(connection: Any) -> int: def embedding_vec_column_exists(connection: Any) -> bool: """True when ``embeddings.embedding_vec`` exists (Postgres information_schema).""" - if getattr(getattr(connection, "dialect", None), "name", None) != "postgresql": - # SQLite / others: probe via PRAGMA or SQLAlchemy inspector if needed. - dialect_name = getattr(getattr(connection, "dialect", None), "name", "") - if dialect_name == "sqlite": - rows = connection.execute(text("PRAGMA table_info(embeddings)")).fetchall() - return any(r[1] == "embedding_vec" for r in rows) + dialect_name = getattr(getattr(connection, "dialect", None), "name", "") or "" + + if dialect_name == "sqlite": + rows = _execute(connection, text("PRAGMA table_info(embeddings)")).fetchall() + return any(r[1] == "embedding_vec" for r in rows) + if dialect_name != "postgresql": return False - row = connection.execute( + row = _execute( + connection, text( "SELECT 1 FROM information_schema.columns " "WHERE table_name = 'embeddings' AND column_name = 'embedding_vec' " "LIMIT 1" - ) + ), ).fetchone() return row is not None @@ -112,7 +128,8 @@ def backfill_embedding_vec(connection: Any, dim: int) -> int: Returns the number of rows reported updated (driver-dependent). """ - result = connection.execute( + result = _execute( + connection, text( """ UPDATE embeddings From 2dc4062313c81ce049f03ffca941a85e322ed7b0 Mon Sep 17 00:00:00 2001 From: Spyros Date: Mon, 13 Jul 2026 12:19:25 +0100 Subject: [PATCH 3/4] chore: use pgvector/pgvector for all local Postgres containers make docker-postgres (and benchmark/staging seed helpers) default to pgvector/pgvector:pg16 and ensure CREATE EXTENSION vector on start. --- AGENTS.md | 5 +++-- Makefile | 38 +++++++++++++++++++++++++++++++-- scripts/run_benchmark.sh | 3 ++- scripts/setup-heroku-staging.sh | 4 +++- 4 files changed, 44 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 75b31bd40..656d4b7c3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,12 +25,13 @@ Use Makefile targets — do not hand-roll `docker run`: ```bash make docker-neo4j # Neo4j (7474/7687) make docker-redis # Redis Stack (6379/8001) -make docker-postgres # local Postgres (5432, cre/password) +make docker-postgres # local Postgres+pgvector (5432, cre/password) make start-containers # neo4j + redis only make migrate-upgrade # after postgres is up ``` -Reset volumes: `make docker-neo4j-rm` / `make docker-redis-rm`. +Reset volumes: `make docker-neo4j-rm` / `make docker-redis-rm` / `make docker-postgres-rm`. +Local Postgres is **`pgvector/pgvector:pg16`** by default (`POSTGRES_IMAGE` override). Plain `postgres` images are not supported for app DB use — `embedding_vec` / chat / Librarian need the `vector` extension. ### Imports and Neo4j populate diff --git a/Makefile b/Makefile index ad76a3f7d..00d349b0f 100644 --- a/Makefile +++ b/Makefile @@ -25,9 +25,43 @@ docker-redis: docker start cre-redis-stack 2>/dev/null ||\ docker run -d --name cre-redis-stack -p 6379:6379 -p 8001:8001 redis/redis-stack:latest +# Local app DB — always pgvector (embedding_vec / Librarian / chat similarity). +POSTGRES_IMAGE ?= pgvector/pgvector:pg16 + +docker-postgres-rm: + -docker stop cre-postgres + -docker rm -f cre-postgres + docker-postgres: - docker start cre-postgres 2>/dev/null ||\ - docker run -d --name cre-postgres -e POSTGRES_PASSWORD=password -e POSTGRES_USER=cre -e POSTGRES_DB=cre -p 5432:5432 postgres + @wanted="$(POSTGRES_IMAGE)"; \ + if docker inspect cre-postgres >/dev/null 2>&1; then \ + have=$$(docker inspect -f '{{.Config.Image}}' cre-postgres); \ + if [ "$$have" = "$$wanted" ]; then \ + docker start cre-postgres >/dev/null; \ + else \ + echo "Recreating cre-postgres (was $$have → $$wanted)"; \ + docker stop cre-postgres >/dev/null 2>&1 || true; \ + docker rm -f cre-postgres >/dev/null 2>&1 || true; \ + docker run -d --name cre-postgres \ + -e POSTGRES_PASSWORD=password \ + -e POSTGRES_USER=cre \ + -e POSTGRES_DB=cre \ + -p 5432:5432 \ + "$$wanted"; \ + fi; \ + else \ + docker run -d --name cre-postgres \ + -e POSTGRES_PASSWORD=password \ + -e POSTGRES_USER=cre \ + -e POSTGRES_DB=cre \ + -p 5432:5432 \ + "$$wanted"; \ + fi; \ + for i in 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20; do \ + docker exec cre-postgres pg_isready -U cre -d cre >/dev/null 2>&1 && break; \ + sleep 1; \ + done; \ + docker exec cre-postgres psql -U cre -d cre -c "CREATE EXTENSION IF NOT EXISTS vector;" >/dev/null start-containers: docker-neo4j docker-redis diff --git a/scripts/run_benchmark.sh b/scripts/run_benchmark.sh index b530308b9..0e3b660e3 100755 --- a/scripts/run_benchmark.sh +++ b/scripts/run_benchmark.sh @@ -39,12 +39,13 @@ setup_db() { echo "Setting up DB in container $container on port $port" docker stop "$container" 2>/dev/null || true docker rm -v "$container" 2>/dev/null || true - docker run -d --name "$container" -e POSTGRES_PASSWORD=password -e POSTGRES_USER=cre -e POSTGRES_DB=cre -p "$port:5432" postgres + docker run -d --name "$container" -e POSTGRES_PASSWORD=password -e POSTGRES_USER=cre -e POSTGRES_DB=cre -p "$port:5432" "${POSTGRES_IMAGE:-pgvector/pgvector:pg16}" echo "Waiting for $container to be ready..." until docker exec "$container" pg_isready -U cre -d cre; do sleep 1 done + docker exec "$container" psql -U cre -d cre -c "CREATE EXTENSION IF NOT EXISTS vector;" >/dev/null echo "Running migrations for $container..." [ -d "./venv" ] && . ./venv/bin/activate diff --git a/scripts/setup-heroku-staging.sh b/scripts/setup-heroku-staging.sh index 19186a53e..a42768314 100755 --- a/scripts/setup-heroku-staging.sh +++ b/scripts/setup-heroku-staging.sh @@ -237,7 +237,7 @@ ensure_local_postgres_if_needed() { -e POSTGRES_USER="${LOCAL_PG_USER}" \ -e POSTGRES_DB="${LOCAL_PG_DB}" \ -p "${LOCAL_PG_PORT}:5432" \ - postgres:16 >/dev/null + "${LOCAL_PG_IMAGE:-pgvector/pgvector:pg16}" >/dev/null fi local ready=0 @@ -249,6 +249,8 @@ ensure_local_postgres_if_needed() { sleep 1 done [[ "${ready}" == "1" ]] || die "Local postgres docker did not become ready in time" + docker exec "${LOCAL_PG_CONTAINER}" psql -U "${LOCAL_PG_USER}" -d "${LOCAL_PG_DB}" \ + -c "CREATE EXTENSION IF NOT EXISTS vector;" >/dev/null } copy_env_from_prod() { From 01d3a4b948d25b32bd145097b7caf87a90517124 Mon Sep 17 00:00:00 2001 From: Spyros Date: Wed, 15 Jul 2026 10:39:46 +0100 Subject: [PATCH 4/4] feat: make embedding_vec the sole store and hard-fail pgvector-on-SQLite Drop the legacy CSV embeddings column in c7d8e9f0a1b2 after backfill, port readers/writers (including Module C) to embedding_vec, and SystemExit when pgvector similarity is requested without Postgres. Co-authored-by: Cursor --- application/cmd/cre_main.py | 17 +-- application/database/db.py | 70 +++++++--- application/database/pgvector_utils.py | 132 +++++++++++++++++- application/tests/db_test.py | 14 +- .../librarian/candidate_retriever_test.py | 14 ++ application/tests/pci_dss_parser_test.py | 11 +- application/tests/pgvector_utils_test.py | 67 +++++++++ .../prompt_client_pgvector_similarity_test.py | 19 +++ .../parsers/cloud_native_security_controls.py | 2 +- .../parsers/juiceshop.py | 2 +- .../parsers/pci_dss.py | 9 +- .../utils/librarian/candidate_retriever.py | 22 ++- ...c7d8e9f0a1b2_add_embedding_vec_pgvector.py | 47 ++++++- scripts/benchmark_retriever.py | 18 +-- scripts/link_pci_dss_cre.py | 12 +- scripts/rewrite_sqlite_embeddings_to_vec.py | 116 +++++++++++++++ scripts/sync_embeddings_table.py | 100 +++++++++++-- 17 files changed, 592 insertions(+), 80 deletions(-) create mode 100644 scripts/rewrite_sqlite_embeddings_to_vec.py diff --git a/application/cmd/cre_main.py b/application/cmd/cre_main.py index 14015cafe..c85b5d2e5 100644 --- a/application/cmd/cre_main.py +++ b/application/cmd/cre_main.py @@ -1115,14 +1115,12 @@ def run_librarian( backend = RetrieverBackend(cfg.retriever_backend) if backend is RetrieverBackend.pgvector: - # Readiness must include the embedding_vec column (not dialect alone): - # Postgres pre-migration would otherwise fail on every retrieve(). + # Hard-fail: do not silently fall back to in_memory / sklearn CSV paths. + from application.database.pgvector_utils import fail_pgvector_unavailable + if not database.can_use_pgvector_similarity(): - logger.warning( - "CRE_LIBRARIAN_RETRIEVER_BACKEND=pgvector selected, but " - "Postgres is not ready for pgvector similarity (need Alembic " - "c7d8e9f0a1b2 / #977 embedding_vec). Set the backend to " - "in_memory until then." + fail_pgvector_unavailable( + context="CRE_LIBRARIAN_RETRIEVER_BACKEND=pgvector" ) # 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 @@ -1136,13 +1134,16 @@ def run_librarian( if backend is RetrieverBackend.in_memory else None ) + pg_connection = None + if backend is RetrieverBackend.pgvector: + pg_connection = database.session.connection() 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(), + connection=pg_connection, ) # C.2 reranker: reads each (section, candidate-CRE) pair together and diff --git a/application/database/db.py b/application/database/db.py index 0d799cae4..6748bb0a3 100644 --- a/application/database/db.py +++ b/application/database/db.py @@ -174,7 +174,6 @@ class Embeddings(BaseModel): # type: ignore __tablename__ = "embeddings" id = sqla.Column(sqla.String, primary_key=True, default=generate_uuid) - embeddings = sqla.Column(sqla.String, nullable=False) doc_type = sqla.Column(sqla.String, nullable=False) cre_id = sqla.Column( sqla.String, @@ -191,9 +190,10 @@ class Embeddings(BaseModel): # type: ignore embeddings_content = sqla.Column(sqla.String, nullable=True, default=None) embedding_model_id = sqla.Column(sqla.String, nullable=True, default=None) embedding_dim = sqla.Column(sqla.Integer, nullable=True, default=None) - # Postgres: real ``vector(N)`` via Alembic (c7d8e9f0a1b2). SQLite tests map - # this as Text storing a pgvector literal for dual-write coverage. - embedding_vec = sqla.Column(sqla.Text, nullable=True, default=None) + # Postgres: real ``vector(N)`` via Alembic (c7d8e9f0a1b2) — sole store for + # vectors (legacy CSV ``embeddings`` column is dropped in that migration). + # SQLite/tests: Text holding a pgvector literal ``[1.0,2.0,...]``. + embedding_vec = sqla.Column(sqla.Text, nullable=False) class GapAnalysisResults(BaseModel): @@ -2367,16 +2367,27 @@ def health_check(self) -> Dict[str, Any]: } def get_embeddings_by_doc_type(self, doc_type: str) -> Dict[str, List[float]]: + from application.database.pgvector_utils import ( + parse_stored_embedding_vec, + require_embedding_vec_store, + ) + + require_embedding_vec_store( + self.session.get_bind(), context="get_embeddings_by_doc_type" + ) res = {} embeddings = ( self.session.query(Embeddings).filter(Embeddings.doc_type == doc_type).all() ) if embeddings: for entry in embeddings: + vec = parse_stored_embedding_vec(entry.embedding_vec) + if not vec: + continue if doc_type == cre_defs.Credoctypes.CRE.value: - res[entry.cre_id] = [float(e) for e in entry.embeddings.split(",")] + res[entry.cre_id] = vec else: - res[entry.node_id] = [float(e) for e in entry.embeddings.split(",")] + res[entry.node_id] = vec return res def get_embedding_contents_by_doc_type(self, doc_type: str) -> Dict[str, str]: @@ -2403,6 +2414,14 @@ def get_embedding_contents_by_doc_type(self, doc_type: str) -> Dict[str, str]: def get_embeddings_by_doc_type_paginated( self, doc_type: str, page: int = 1, per_page: int = 100 ) -> Tuple[Dict[str, List[float]], int, int]: + from application.database.pgvector_utils import ( + parse_stored_embedding_vec, + require_embedding_vec_store, + ) + + require_embedding_vec_store( + self.session.get_bind(), context="get_embeddings_by_doc_type_paginated" + ) res = {} embeddings = ( self.session.query(Embeddings) @@ -2412,13 +2431,21 @@ def get_embeddings_by_doc_type_paginated( total_pages = embeddings.pages if embeddings.items: for entry in embeddings.items: + vec = parse_stored_embedding_vec(entry.embedding_vec) + if not vec: + continue if doc_type == cre_defs.Credoctypes.CRE.value: - res[entry.cre_id] = [float(e) for e in entry.embeddings.split(",")] + res[entry.cre_id] = vec else: - res[entry.node_id] = [float(e) for e in entry.embeddings.split(",")] + res[entry.node_id] = vec return res, total_pages, page def get_embeddings_for_doc(self, doc: cre_defs.Node | cre_defs.CRE) -> Embeddings: + from application.database.pgvector_utils import require_embedding_vec_store + + require_embedding_vec_store( + self.session.get_bind(), context="get_embeddings_for_doc" + ) if doc.doctype == cre_defs.Credoctypes.CRE: obj = self.session.query(CRE).filter(CRE.external_id == doc.id).first() return ( @@ -2487,13 +2514,15 @@ def add_embedding( f"embedding dimension mismatch for {db_object.id}: " f"expected {expected_dim}, got {len(embeddings)}" ) - existing = self.get_embedding(db_object.id) - embeddings_str = ",".join([str(e) for e in embeddings]) - # Dual-write CSV + pgvector literal. On Postgres after migration the - # column is ``vector(N)``; on SQLite tests it is Text. Never LLM - # re-embed here — see application.database.pgvector_utils. - from application.database.pgvector_utils import to_pgvector_literal + from application.database.pgvector_utils import ( + require_embedding_vec_store, + to_pgvector_literal, + ) + require_embedding_vec_store(self.session.get_bind(), context="add_embedding") + existing = self.get_embedding(db_object.id) + # Sole store is ``embedding_vec`` (pgvector on Postgres; Text literal + # on SQLite). Never LLM re-embed here — see pgvector_utils. embedding_vec_literal = to_pgvector_literal(embeddings) resolved_node_url: Optional[str] = None if doctype != cre_defs.Credoctypes.CRE: @@ -2505,7 +2534,6 @@ def add_embedding( emb = None if doctype == cre_defs.Credoctypes.CRE: emb = Embeddings( - embeddings=embeddings_str, cre_id=db_object.id, doc_type=cre_defs.Credoctypes.CRE.value, embeddings_content=embedding_text, @@ -2515,7 +2543,6 @@ def add_embedding( ) else: emb = Embeddings( - embeddings=embeddings_str, node_id=db_object.id, doc_type=db_object.ntype, embeddings_content=embedding_text, @@ -2529,7 +2556,6 @@ def add_embedding( return emb else: logger.debug(f"knew of embedding for object {db_object.id} ,updating") - existing[0].embeddings = embeddings_str existing[0].embeddings_content = embedding_text existing[0].embedding_model_id = embedding_model_id existing[0].embedding_dim = embedding_dim @@ -2571,14 +2597,20 @@ def find_most_similar_embedding_id( """Top-1 cosine match via pgvector ``<=>`` (Postgres only). Returns ``(object_id, score)`` or ``(None, None)`` below threshold / - when no rows match. + when no rows match. Refuses SQLite / missing ``embedding_vec`` with + ``SystemExit`` — callers that need a soft fallback must gate on + ``can_use_pgvector_similarity()`` first and use sklearn instead. """ from application.database.pgvector_utils import ( + fail_pgvector_unavailable, most_similar_id_sql, to_pgvector_literal, ) from sqlalchemy import text as sql_text + if not self.can_use_pgvector_similarity(): + fail_pgvector_unavailable(context="find_most_similar_embedding_id") + sql = most_similar_id_sql(id_column) bind = self.session.get_bind() params = { @@ -2599,6 +2631,8 @@ def find_most_similar_embedding_id( if score < similarity_threshold: return None, None return str(row.object_id), score + except SystemExit: + raise except Exception as exc: # Match prior sklearn-path resilience: chat/import should degrade # to "no match" rather than 500 on a transient driver/query error. diff --git a/application/database/pgvector_utils.py b/application/database/pgvector_utils.py index 817cccb87..506a7f0e5 100644 --- a/application/database/pgvector_utils.py +++ b/application/database/pgvector_utils.py @@ -3,8 +3,8 @@ Do **not** re-embed the full corpus on Heroku/production dynos — likely OOM-killed. Re-embedding (fix missing / wrong-dim rows via the LLM) must run only where there is enough RAM (local or a high-resource worker), then -sync/backfill vectors. Alembic on prod only: enable extension, add column, -copy CSV→vector for matching dims. +sync/backfill vectors. Alembic on prod: enable extension, add column, copy +CSV→vector for matching dims, then DROP the legacy CSV ``embeddings`` column. HNSW/IVFFlat indexes are deferred (Essential-1 capacity); see migration ``c7d8e9f0a1b2`` docstring for the follow-up plan. @@ -21,7 +21,26 @@ HEROKU_REEMBED_WARNING = ( "Do not re-embed the full corpus on Heroku/production dynos — likely " "OOM-killed. Re-embedding must run only on a machine with enough RAM, " - "then sync/backfill vectors. Alembic on prod only copies CSV→vector." + "then sync/backfill vectors. Alembic on prod only copies CSV→vector " + "then drops the CSV column." +) + +PGVECTOR_UNAVAILABLE_EXIT_MSG = ( + "pgvector embeddings are required but unavailable on this database. " + "Need Postgres with embeddings.embedding_vec (Alembic c7d8e9f0a1b2). " + "SQLite cannot serve pgvector similarity (``<=>``). For Module C / " + "Librarian on SQLite or CI, set CRE_LIBRARIAN_RETRIEVER_BACKEND=in_memory " + "(reads ``embedding_vec`` Text literals into an in-RAM hub). For " + "Postgres-side similarity, use local Postgres (make docker-postgres) or " + "a remote Postgres URL with the vector extension." +) + +EMBEDDING_VEC_REQUIRED_EXIT_MSG = ( + "embeddings.embedding_vec is required; the legacy CSV ``embeddings`` " + "column is no longer a supported store. Postgres: run Alembic upgrade " + "to c7d8e9f0a1b2. SQLite caches: rewrite with " + "`python scripts/rewrite_sqlite_embeddings_to_vec.py --db PATH` or " + "re-export from Postgres after the migration." ) @@ -29,6 +48,96 @@ class MultiDimEmbeddingError(RuntimeError): """Raised when more than one embedding dimension is present in the DB.""" +class PgVectorUnavailableError(RuntimeError): + """Raised when a pgvector operation is requested without Postgres+vector.""" + + +def _exit_with_message(msg: str) -> None: + import logging + + logging.getLogger(__name__).error(msg) + raise SystemExit(msg) + + +def fail_pgvector_unavailable(*, context: str = "") -> None: + """Log and exit — callers must not silently fall back to CSV/sklearn. + + Intentionally raises ``SystemExit`` so CLI/import paths stop with a clear + message when ``CRE_LIBRARIAN_RETRIEVER_BACKEND=pgvector`` (or other + pgvector-only code) runs against SQLite / pre-migration Postgres. + """ + detail = f" ({context})" if context else "" + _exit_with_message(PGVECTOR_UNAVAILABLE_EXIT_MSG + detail) + + +def fail_embedding_vec_required(*, context: str = "") -> None: + """Exit when the vector store column is missing (legacy CSV-only schema).""" + detail = f" ({context})" if context else "" + _exit_with_message(EMBEDDING_VEC_REQUIRED_EXIT_MSG + detail) + + +def connection_dialect_name(connection: Any) -> str: + """Best-effort dialect name from a SQLAlchemy Connection/Engine/Session bind.""" + dialect = getattr(connection, "dialect", None) + if dialect is not None: + return getattr(dialect, "name", "") or "" + bind = getattr(connection, "get_bind", None) + if callable(bind): + return connection_dialect_name(bind()) + engine = getattr(connection, "engine", None) + if engine is not None: + return connection_dialect_name(engine) + return "" + + +def require_pgvector_connection(connection: Any, *, context: str = "") -> None: + """Refuse SQLite (and unknown non-Postgres) when loading pgvector embeddings.""" + dialect = connection_dialect_name(connection) + if dialect == "sqlite": + fail_pgvector_unavailable(context=context or "sqlite connection") + # Hermetic test fakes often omit dialect — allow those through. + if dialect and dialect != "postgresql": + fail_pgvector_unavailable(context=context or f"unsupported dialect {dialect!r}") + + +def _embeddings_table_columns(connection: Any) -> Optional[Set[str]]: + """Return column names for ``embeddings``, or None when the table is absent.""" + dialect = connection_dialect_name(connection) + if dialect == "sqlite": + rows = _execute(connection, text("PRAGMA table_info(embeddings)")).fetchall() + if not rows: + return None + return {str(r[1]) for r in rows} + if dialect != "postgresql": + return None + rows = _execute( + connection, + text( + "SELECT column_name FROM information_schema.columns " + "WHERE table_name = 'embeddings'" + ), + ).fetchall() + if not rows: + return None + return {str(r[0]) for r in rows} + + +def require_embedding_vec_store(connection: Any, *, context: str = "") -> None: + """Refuse legacy CSV-only schemas; require ``embedding_vec``. + + No-ops for hermetic fakes (empty dialect) and when the ``embeddings`` + table does not exist yet (caller will hit a normal SQLAlchemy error). + """ + dialect = connection_dialect_name(connection) + if not dialect: + return + cols = _embeddings_table_columns(connection) + if cols is None: + return + if "embedding_vec" not in cols: + fail_embedding_vec_required(context=context or f"{dialect} embeddings schema") + + def _execute(connection: Any, statement: Any, params: Optional[dict] = None): """Run ``statement`` on a Connection/Session, or open one from an Engine. @@ -53,6 +162,23 @@ def csv_embeddings_to_literal(csv: str) -> str: return f"[{(csv or '').strip()}]" +def parse_stored_embedding_vec(value: Any) -> list[float]: + """Parse a stored ``embedding_vec`` (pgvector / Text literal) to floats. + + Accepts ``[1.0, 2.0]``, ``1.0,2.0``, or a sequence of numbers. + """ + if value is None: + return [] + if isinstance(value, (list, tuple)): + return [float(x) for x in value] + s = str(value).strip() + if not s: + return [] + if s.startswith("[") and s.endswith("]"): + s = s[1:-1] + return [float(part.strip()) for part in s.split(",") if part.strip() != ""] + + def parse_csv_embedding_dim(csv: str) -> int: """Return the number of floats implied by a CSV embeddings string.""" parts = [p for p in (csv or "").split(",") if p.strip() != ""] diff --git a/application/tests/db_test.py b/application/tests/db_test.py index 4f8c85efd..33eb18479 100644 --- a/application/tests/db_test.py +++ b/application/tests/db_test.py @@ -2390,8 +2390,8 @@ def test_add_embedding_persists_embedding_contract_metadata(self): finally: os.environ.pop("CRE_EMBED_MODEL", None) - def test_add_embedding_dual_writes_embedding_vec_literal(self): - """CSV + embedding_vec literal are persisted together (#977).""" + def test_add_embedding_writes_embedding_vec_literal(self): + """Vectors are stored only in embedding_vec (#977 option C).""" from application.database.pgvector_utils import to_pgvector_literal dbsa = db.Node( @@ -2408,13 +2408,17 @@ def test_add_embedding_dual_writes_embedding_vec_literal(self): db_object=dbsa, doctype=defs.Credoctypes.Standard.value, embeddings=vec, - embedding_text="dual-write", + embedding_text="vec-only", ) row = self.collection.get_embedding(dbsa.id)[0] - self.assertEqual(row.embeddings, "0.1,0.2,0.3") - self.assertEqual(row.embeddings_content, "dual-write") + self.assertNotIn("embeddings", type(row).__table__.c) + self.assertEqual(row.embeddings_content, "vec-only") self.assertEqual(row.embedding_vec, to_pgvector_literal(vec)) self.assertEqual(row.embedding_dim, 3) + by_type = self.collection.get_embeddings_by_doc_type( + defs.Credoctypes.Standard.value + ) + self.assertEqual(by_type[dbsa.id], vec) def test_assert_embedding_contract_fails_on_mixed_dimensions(self): n1 = db.Node( diff --git a/application/tests/librarian/candidate_retriever_test.py b/application/tests/librarian/candidate_retriever_test.py index 2d44dc4a8..5e9ea359b 100644 --- a/application/tests/librarian/candidate_retriever_test.py +++ b/application/tests/librarian/candidate_retriever_test.py @@ -6,6 +6,7 @@ """ import unittest +from unittest.mock import MagicMock from application.utils.librarian.candidate_retriever import ( PGVECTOR_RETRIEVER_NAME, @@ -208,6 +209,19 @@ def test_factory_routes_to_each_backend(self) -> None: ) self.assertIsInstance(pg, PgVectorRetriever) + def test_pgvector_backend_exits_on_sqlite_connection(self) -> None: + conn = MagicMock() + conn.dialect.name = "sqlite" + with self.assertRaises(SystemExit) as cm: + build_retriever( + RetrieverBackend.pgvector, + fake_embed, + top_k=20, + threshold=0.8, + connection=conn, + ) + self.assertIn("pgvector embeddings are required", str(cm.exception)) + if __name__ == "__main__": unittest.main() diff --git a/application/tests/pci_dss_parser_test.py b/application/tests/pci_dss_parser_test.py index 01815e5b9..8469b4714 100644 --- a/application/tests/pci_dss_parser_test.py +++ b/application/tests/pci_dss_parser_test.py @@ -62,7 +62,16 @@ def test_best_cre_via_bridge_standard_picks_highest_similarity_linked_node( low_cre = defs.CRE(id="111-111", name="Low", description="") high_cre = defs.CRE(id="222-222", name="High", description="") cache.get_nodes.return_value = [low_node, high_node] - cache.get_embeddings_for_doc.side_effect = [[0.0, 1.0], [1.0, 0.0]] + + def _emb_row(vec): + row = Mock() + row.embedding_vec = "[" + ",".join(str(x) for x in vec) + "]" + return row + + cache.get_embeddings_for_doc.side_effect = [ + _emb_row([0.0, 1.0]), + _emb_row([1.0, 0.0]), + ] cache.find_cres_of_node.side_effect = [ [Mock(id="low-db")], [Mock(id="high-db")], diff --git a/application/tests/pgvector_utils_test.py b/application/tests/pgvector_utils_test.py index c66463e95..0dfb1c53c 100644 --- a/application/tests/pgvector_utils_test.py +++ b/application/tests/pgvector_utils_test.py @@ -27,6 +27,73 @@ def test_parse_csv_dim(self) -> None: self.assertEqual(parse_csv_embedding_dim("0.1,0.2,0.3"), 3) self.assertEqual(parse_csv_embedding_dim(""), 0) + def test_parse_stored_embedding_vec(self) -> None: + from application.database.pgvector_utils import parse_stored_embedding_vec + + self.assertEqual(parse_stored_embedding_vec("[1.0, 2.0]"), [1.0, 2.0]) + self.assertEqual(parse_stored_embedding_vec("1,2,3"), [1.0, 2.0, 3.0]) + self.assertEqual(parse_stored_embedding_vec(None), []) + self.assertEqual(parse_stored_embedding_vec([1, 2]), [1.0, 2.0]) + + def test_sqlite_connection_refuses_pgvector(self) -> None: + from application.database.pgvector_utils import ( + PGVECTOR_UNAVAILABLE_EXIT_MSG, + require_pgvector_connection, + ) + + conn = MagicMock() + conn.dialect.name = "sqlite" + with self.assertRaises(SystemExit) as cm: + require_pgvector_connection(conn, context="unit-test") + self.assertIn("pgvector embeddings are required", str(cm.exception)) + self.assertIn(PGVECTOR_UNAVAILABLE_EXIT_MSG[:32], str(cm.exception)) + + def test_postgres_connection_accepted(self) -> None: + from application.database.pgvector_utils import require_pgvector_connection + + conn = MagicMock() + conn.dialect.name = "postgresql" + require_pgvector_connection(conn) # does not raise + + def test_fake_connection_without_dialect_allowed(self) -> None: + from application.database.pgvector_utils import require_pgvector_connection + + class NoDialect: + pass + + require_pgvector_connection(NoDialect()) # hermetic fakes OK + + +class EmbeddingVecStoreGateTest(unittest.TestCase): + def test_legacy_sqlite_schema_exits(self) -> None: + from application.database.pgvector_utils import ( + EMBEDDING_VEC_REQUIRED_EXIT_MSG, + require_embedding_vec_store, + ) + + conn = MagicMock() + conn.dialect.name = "sqlite" + # PRAGMA table_info rows: (cid, name, type, notnull, dflt, pk) + conn.execute.return_value.fetchall.return_value = [ + (0, "embeddings", "TEXT", 1, None, 0), + (1, "doc_type", "TEXT", 1, None, 0), + ] + with self.assertRaises(SystemExit) as cm: + require_embedding_vec_store(conn, context="unit-test") + self.assertIn("embedding_vec is required", str(cm.exception)) + self.assertIn(EMBEDDING_VEC_REQUIRED_EXIT_MSG[:40], str(cm.exception)) + + def test_vec_column_present_ok(self) -> None: + from application.database.pgvector_utils import require_embedding_vec_store + + conn = MagicMock() + conn.dialect.name = "sqlite" + conn.execute.return_value.fetchall.return_value = [ + (0, "embedding_vec", "TEXT", 1, None, 0), + (1, "doc_type", "TEXT", 1, None, 0), + ] + require_embedding_vec_store(conn) + class DistinctDimsTest(unittest.TestCase): def test_single_dim_from_metadata(self) -> None: diff --git a/application/tests/prompt_client_pgvector_similarity_test.py b/application/tests/prompt_client_pgvector_similarity_test.py index 6e11ce0b3..845f8000d 100644 --- a/application/tests/prompt_client_pgvector_similarity_test.py +++ b/application/tests/prompt_client_pgvector_similarity_test.py @@ -53,6 +53,7 @@ def test_query_error_returns_no_match(self) -> None: session.execute.side_effect = RuntimeError("driver boom") session.get_bind.return_value = MagicMock() database.session = session + database._pgvector_similarity_ready = True with patch("application.database.db.logger"): match_id, score = database.find_most_similar_embedding_id( @@ -64,6 +65,24 @@ def test_query_error_returns_no_match(self) -> None: self.assertIsNone(match_id) self.assertIsNone(score) + def test_sqlite_refuses_pgvector_similarity_with_systemexit(self) -> None: + from application.database.db import Node_collection + from application.database.pgvector_utils import PGVECTOR_UNAVAILABLE_EXIT_MSG + + database = Node_collection.__new__(Node_collection) + database.session = MagicMock() + database._pgvector_similarity_ready = False + + with self.assertRaises(SystemExit) as cm: + database.find_most_similar_embedding_id( + [0.1, 0.2, 0.3], + doc_type="CRE", + id_column="cre_id", + similarity_threshold=0.7, + ) + self.assertIn("pgvector embeddings are required", str(cm.exception)) + self.assertIn(PGVECTOR_UNAVAILABLE_EXIT_MSG[:40], str(cm.exception)) + if __name__ == "__main__": unittest.main() diff --git a/application/utils/external_project_parsers/parsers/cloud_native_security_controls.py b/application/utils/external_project_parsers/parsers/cloud_native_security_controls.py index 64089cf93..2d050ceee 100644 --- a/application/utils/external_project_parsers/parsers/cloud_native_security_controls.py +++ b/application/utils/external_project_parsers/parsers/cloud_native_security_controls.py @@ -58,7 +58,7 @@ def parse(self, cache: db.Node_collection, ph: prompt_client.PromptHandler): ) if existing: embeddings = cache.get_embeddings_for_doc(existing[0]) - if embeddings: + if embeddings and getattr(embeddings, "embedding_vec", None): logger.info( f"Node {cnsc.todict()} already exists and has embeddings, skipping" ) diff --git a/application/utils/external_project_parsers/parsers/juiceshop.py b/application/utils/external_project_parsers/parsers/juiceshop.py index ca6c89e9d..c001f5b16 100644 --- a/application/utils/external_project_parsers/parsers/juiceshop.py +++ b/application/utils/external_project_parsers/parsers/juiceshop.py @@ -68,7 +68,7 @@ def parse( ) if existing: embeddings = cache.get_embeddings_for_doc(existing[0]) - if embeddings: + if embeddings and getattr(embeddings, "embedding_vec", None): logger.info( f"Node {chal.todict()} already exists and has embeddings, skipping" ) diff --git a/application/utils/external_project_parsers/parsers/pci_dss.py b/application/utils/external_project_parsers/parsers/pci_dss.py index 555404b3a..487c8790a 100644 --- a/application/utils/external_project_parsers/parsers/pci_dss.py +++ b/application/utils/external_project_parsers/parsers/pci_dss.py @@ -101,8 +101,13 @@ def best_cre_via_bridge_standard( best_similarity = -1.0 best_cre: Optional[defs.CRE] = None + from application.database.pgvector_utils import parse_stored_embedding_vec + for node in cache.get_nodes(name=standard_name) or []: - node_embedding = cache.get_embeddings_for_doc(node) + emb_row = cache.get_embeddings_for_doc(node) + if not emb_row or not emb_row.embedding_vec: + continue + node_embedding = parse_stored_embedding_vec(emb_row.embedding_vec) if not node_embedding: continue node_array = sparse.csr_matrix( @@ -264,7 +269,7 @@ def __parse( ) if existing: embeddings = cache.get_embeddings_for_doc(existing[0]) - if embeddings: + if embeddings and getattr(embeddings, "embedding_vec", None): logger.info( f"Node {pci_control.todict()} already exists and has embeddings, skipping" ) diff --git a/application/utils/librarian/candidate_retriever.py b/application/utils/librarian/candidate_retriever.py index 63c6d8b6d..9db70e9c8 100644 --- a/application/utils/librarian/candidate_retriever.py +++ b/application/utils/librarian/candidate_retriever.py @@ -21,13 +21,15 @@ 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. + - ``in_memory`` — sklearn cosine over an in-RAM matrix. Works on SQLite + (CI / harness) by reading stored ``embedding_vec`` Text literals into a + hub matrix. Does not use the pgvector ``<=>`` operator. - ``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 (Alembic - ``c7d8e9f0a1b2`` / #977). Unavailable on SQLite. + ``c7d8e9f0a1b2`` / #977). On SQLite (or Postgres without the column) + construction/use raises ``SystemExit`` with a clear message — never + silently falls back. The RFC is silent on retrieval tech — it mandates only the ``candidates[]``/``reranked[]`` audit trail — so the backend choice is ours; @@ -187,8 +189,8 @@ class PgVectorRetriever: in-memory backend's cosine *similarity*. Needs the ``vector`` extension and an ``embedding_vec vector(dim)`` column - (Alembic ``c7d8e9f0a1b2`` / #977). Unavailable on SQLite — the factory - routes SQLite to ``in_memory``. + (Alembic ``c7d8e9f0a1b2`` / #977). SQLite is refused at construct/execute time + with ``SystemExit`` — never silently routed to ``in_memory``. """ # Parameterized; :q is bound as a pgvector text literal and cast in-SQL. @@ -213,6 +215,11 @@ def __init__( ) -> None: if top_k <= 0: raise RetrieverError(f"top_k must be > 0, got {top_k}") + from application.database.pgvector_utils import require_pgvector_connection + + require_pgvector_connection( + connection, context="PgVectorRetriever / Module C.1" + ) self._embed_fn = embed_fn self._conn = connection self._top_k = top_k @@ -229,6 +236,9 @@ def retrieve(self, text: str) -> RetrievalAudit: """ from sqlalchemy import text as sql_text + from application.database.pgvector_utils import require_pgvector_connection + + require_pgvector_connection(self._conn, context="PgVectorRetriever.retrieve") query = to_pgvector_literal(list(self._embed_fn(text))) rows = self._conn.execute( sql_text(self._SQL), diff --git a/migrations/versions/c7d8e9f0a1b2_add_embedding_vec_pgvector.py b/migrations/versions/c7d8e9f0a1b2_add_embedding_vec_pgvector.py index 605bd408f..de12f2bb7 100644 --- a/migrations/versions/c7d8e9f0a1b2_add_embedding_vec_pgvector.py +++ b/migrations/versions/c7d8e9f0a1b2_add_embedding_vec_pgvector.py @@ -1,5 +1,5 @@ """ -Add embeddings.embedding_vec (pgvector) and merge dual Alembic heads. +Add embeddings.embedding_vec (pgvector), backfill from CSV, drop CSV column. Revision ID: c7d8e9f0a1b2 Revises: 9b1c2d3e4f50, ab12cd34ef56 @@ -7,13 +7,14 @@ Do **not** re-embed the full corpus on Heroku/production dynos — likely OOM-killed. This migration only: CREATE EXTENSION vector, ADD COLUMN -embedding_vec, and copy matching-dim CSV → vector. Fix missing / wrong-dim -rows offline on a high-RAM machine, then re-run / backfill. +embedding_vec, copy matching-dim CSV → vector, then DROP the legacy CSV +``embeddings`` column. Fix missing / wrong-dim rows offline on a high-RAM +machine before deploying. ANN index (HNSW/IVFFlat) is intentionally deferred: opencreorg runs on Heroku Postgres Essential-1; add a follow-up migration with ``CREATE INDEX ... USING hnsw (embedding_vec vector_cosine_ops)`` once -dims/backfill are settled and plan capacity allows. +plan capacity allows. """ from alembic import op @@ -33,9 +34,9 @@ def upgrade(): - """Postgres-only: enable pgvector, add embedding_vec, backfill from CSV. + """Postgres-only: pgvector column, CSV backfill, drop legacy CSV column. - SQLite/CI: no-op (chat/Librarian keep in-memory / CSV paths). + SQLite/CI: no-op (ORM maps ``embedding_vec`` as Text for create_all). """ conn = op.get_bind() if conn.dialect.name != "postgresql": @@ -54,10 +55,44 @@ def upgrade(): ) backfill_embedding_vec(conn, dim) + # Refuse to drop CSV if any row still lacks a vector (should be empty after + # a clean backfill; wrong-dim / empty CSV rows must be fixed offline). + leftover = conn.execute( + sa.text("SELECT COUNT(*) FROM embeddings WHERE embedding_vec IS NULL") + ).scalar() + if leftover: + raise RuntimeError( + f"{leftover} embeddings row(s) still have NULL embedding_vec after " + "backfill — refuse to DROP CSV column. Reconcile offline " + "(do not re-embed on Heroku). " + HEROKU_REEMBED_WARNING + ) + + op.execute(sa.text("ALTER TABLE embeddings DROP COLUMN IF EXISTS embeddings")) + op.execute( + sa.text("ALTER TABLE embeddings ALTER COLUMN embedding_vec SET NOT NULL") + ) + def downgrade(): conn = op.get_bind() if conn.dialect.name != "postgresql": return + # Restore CSV from vector text form ``[1,2,3]`` → ``1,2,3``. + op.execute( + sa.text("ALTER TABLE embeddings ADD COLUMN IF NOT EXISTS embeddings TEXT") + ) + op.execute( + sa.text( + """ + UPDATE embeddings + SET embeddings = trim(both '[]' FROM replace(embedding_vec::text, ' ', '')) + WHERE embedding_vec IS NOT NULL + AND (embeddings IS NULL OR btrim(embeddings) = '') + """ + ) + ) + op.execute( + sa.text("ALTER TABLE embeddings ALTER COLUMN embedding_vec DROP NOT NULL") + ) op.execute(sa.text("ALTER TABLE embeddings DROP COLUMN IF EXISTS embedding_vec")) # Intentionally leave the `vector` extension installed (may be shared). diff --git a/scripts/benchmark_retriever.py b/scripts/benchmark_retriever.py index e6bcdc474..8c46b04e2 100644 --- a/scripts/benchmark_retriever.py +++ b/scripts/benchmark_retriever.py @@ -53,18 +53,7 @@ def _time_backend(retriever, queries: List[str], runs: int) -> float: 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 + return bool(database.can_use_pgvector_similarity()) def main(argv: List[str]) -> int: @@ -93,8 +82,9 @@ def main(argv: List[str]) -> int: if not _pgvector_available(database): print( - "pgvector : SKIPPED (needs Postgres + the embedding_vec column; " - "lands with the W8 pgvector migration + backfill)" + "pgvector : SKIPPED (needs Postgres + embedding_vec; SQLite cannot " + "run <=> — use CRE_LIBRARIAN_RETRIEVER_BACKEND=in_memory or " + "postgresql://… after Alembic c7d8e9f0a1b2)" ) return 0 diff --git a/scripts/link_pci_dss_cre.py b/scripts/link_pci_dss_cre.py index 0b5a94b19..df908991c 100644 --- a/scripts/link_pci_dss_cre.py +++ b/scripts/link_pci_dss_cre.py @@ -85,12 +85,20 @@ def main() -> int: .filter(db.Embeddings.node_id == db_node.id) .first() ) - if not emb_row or not emb_row.embeddings: + if not emb_row or not emb_row.embedding_vec: logger.warning("No embedding for PCI node %s; skipping", db_node.section_id) skipped_no_match += 1 continue - control_embeddings = [float(e) for e in emb_row.embeddings.split(",")] + from application.database.pgvector_utils import parse_stored_embedding_vec + + control_embeddings = parse_stored_embedding_vec(emb_row.embedding_vec) + if not control_embeddings: + logger.warning( + "Empty embedding for PCI node %s; skipping", db_node.section_id + ) + skipped_no_match += 1 + continue cre_id = ph.get_id_of_most_similar_cre(control_embeddings) if not cre_id: standard_id = ph.get_id_of_most_similar_node(control_embeddings) diff --git a/scripts/rewrite_sqlite_embeddings_to_vec.py b/scripts/rewrite_sqlite_embeddings_to_vec.py new file mode 100644 index 000000000..f9d236e1c --- /dev/null +++ b/scripts/rewrite_sqlite_embeddings_to_vec.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Rewrite a SQLite embeddings cache from legacy CSV to ``embedding_vec`` Text. + +After Alembic ``c7d8e9f0a1b2``, OpenCRE stores vectors only in ``embedding_vec`` +(Postgres: real ``vector(N)``; SQLite / CI: Text literals ``[1.0,2.0,...]``). +Old ``standards_cache.sqlite`` files still have a CSV ``embeddings`` column — +the ORM will refuse them. This script converts in place: + + * ADD ``embedding_vec`` TEXT + * copy ``'[' || embeddings || ']'`` + * DROP ``embeddings`` + +Usage:: + + python scripts/rewrite_sqlite_embeddings_to_vec.py --db standards_cache.sqlite +""" + +from __future__ import annotations + +import argparse +import sqlite3 +import sys + + +def rewrite(path: str) -> int: + conn = sqlite3.connect(path) + try: + cols = {r[1] for r in conn.execute("PRAGMA table_info(embeddings)").fetchall()} + if not cols: + print(f"error: no embeddings table in {path!r}", file=sys.stderr) + return 2 + if "embedding_vec" in cols and "embeddings" not in cols: + print(f"{path}: already embedding_vec-only; nothing to do") + return 0 + if "embedding_vec" not in cols: + conn.execute("ALTER TABLE embeddings ADD COLUMN embedding_vec TEXT") + conn.execute( + """ + UPDATE embeddings + SET embedding_vec = '[' || embeddings || ']' + WHERE embeddings IS NOT NULL + AND trim(embeddings) <> '' + AND (embedding_vec IS NULL OR trim(embedding_vec) = '') + """ + ) + nulls = conn.execute( + "SELECT COUNT(*) FROM embeddings WHERE embedding_vec IS NULL OR trim(embedding_vec) = ''" + ).fetchone()[0] + if nulls: + print( + f"error: {nulls} row(s) still lack embedding_vec after copy; " + "refusing to drop CSV column", + file=sys.stderr, + ) + conn.rollback() + return 3 + conn.commit() + # SQLite lacks reliable DROP COLUMN on older builds; recreate table. + src_cols = { + r[1] for r in conn.execute("PRAGMA table_info(embeddings)").fetchall() + } + select_id = "id" if "id" in src_cols else "NULL" + select_model = ( + "embedding_model_id" if "embedding_model_id" in src_cols else "NULL" + ) + select_dim = "embedding_dim" if "embedding_dim" in src_cols else "NULL" + select_url = "embeddings_url" if "embeddings_url" in src_cols else "NULL" + select_content = ( + "embeddings_content" if "embeddings_content" in src_cols else "NULL" + ) + conn.execute( + """ + CREATE TABLE embeddings_new ( + id TEXT PRIMARY KEY, + doc_type TEXT NOT NULL, + cre_id TEXT, + node_id TEXT, + embeddings_url TEXT, + embeddings_content TEXT, + embedding_model_id TEXT, + embedding_dim INTEGER, + embedding_vec TEXT NOT NULL + ) + """ + ) + conn.execute( + f""" + INSERT INTO embeddings_new ( + id, doc_type, cre_id, node_id, embeddings_url, + embeddings_content, embedding_model_id, embedding_dim, embedding_vec + ) + SELECT + {select_id}, doc_type, cre_id, node_id, {select_url}, + {select_content}, {select_model}, {select_dim}, embedding_vec + FROM embeddings + """ + ) + conn.execute("DROP TABLE embeddings") + conn.execute("ALTER TABLE embeddings_new RENAME TO embeddings") + conn.commit() + n = conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0] + print(f"{path}: rewrote {n} row(s) to embedding_vec-only") + return 0 + finally: + conn.close() + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--db", required=True, help="Path to SQLite database file") + args = p.parse_args() + return rewrite(args.db) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/sync_embeddings_table.py b/scripts/sync_embeddings_table.py index f61790a9b..8457a30a3 100644 --- a/scripts/sync_embeddings_table.py +++ b/scripts/sync_embeddings_table.py @@ -2,7 +2,14 @@ """ Copy only the ``embeddings`` table between databases. -Typical flow (embeddings live in local SQLite): +After ``c7d8e9f0a1b2``, Postgres stores vectors in ``embedding_vec`` (pgvector), +not the legacy CSV ``embeddings`` column. This script: + + * From SQLite: prefers ``embedding_vec`` if present; else converts legacy CSV + ``embeddings`` into a pgvector literal and writes ``embedding_vec``. + * From/to Postgres: copies ``embedding_vec`` (+ content/url metadata). + +Typical flow: 1) Push into local Postgres (same ``cre`` / ``node`` ids as in the sqlite file): @@ -34,6 +41,11 @@ from psycopg2 import extras +# Keep sync usable without importing the full app package on slim hosts. +def _csv_to_literal(csv: str) -> str: + return f"[{(csv or '').strip()}]" + + def _blank_to_none(value: Any) -> Any: if isinstance(value, str) and value == "": return None @@ -44,9 +56,9 @@ def _normalize_embedding_row(row: Tuple[Any, ...]) -> Tuple[Any, ...]: # SQLite snapshots often store "missing" FK fields as empty strings. # Postgres FK checks treat empty string as a real value (not NULL), # so convert blanks to NULL for cre_id/node_id before insert. - embeddings, doc_type, cre_id, node_id, embeddings_content, embeddings_url = row + embedding_vec, doc_type, cre_id, node_id, embeddings_content, embeddings_url = row return ( - embeddings, + embedding_vec, doc_type, _blank_to_none(cre_id), _blank_to_none(node_id), @@ -67,17 +79,66 @@ def _pg_host_is_loopback(url: str) -> bool: return h in ("127.0.0.1", "localhost", "::1", "0.0.0.0") or h == "" +def _sqlite_table_columns(conn: sqlite3.Connection, table: str) -> set[str]: + cur = conn.execute(f"PRAGMA table_info({table})") + return {r[1] for r in cur.fetchall()} + + +def _as_vec_literal(value: Any) -> str: + s = str(value or "").strip() + if not s: + return "[]" + if s.startswith("[") and s.endswith("]"): + return s + return _csv_to_literal(s) + + def _fetch_sqlite_embeddings(path: str) -> Tuple[List[str], List[Tuple[Any, ...]]]: conn = sqlite3.connect(path) conn.row_factory = sqlite3.Row - cur = conn.execute( - "SELECT embeddings, doc_type, cre_id, node_id, embeddings_content, embeddings_url " - "FROM embeddings" - ) - rows = [_normalize_embedding_row(tuple(r)) for r in cur.fetchall()] + cols_present = _sqlite_table_columns(conn, "embeddings") + if "embedding_vec" in cols_present: + cur = conn.execute( + "SELECT embedding_vec, doc_type, cre_id, node_id, " + "embeddings_content, embeddings_url FROM embeddings" + ) + raw_rows = [ + ( + _as_vec_literal(r[0]), + r[1], + r[2], + r[3], + r[4], + r[5], + ) + for r in cur.fetchall() + ] + elif "embeddings" in cols_present: + # Legacy CSV column from pre-c7d8e9f0a1b2 SQLite caches. + cur = conn.execute( + "SELECT embeddings, doc_type, cre_id, node_id, " + "embeddings_content, embeddings_url FROM embeddings" + ) + raw_rows = [ + ( + _csv_to_literal(str(r[0] or "")), + r[1], + r[2], + r[3], + r[4], + r[5], + ) + for r in cur.fetchall() + ] + else: + conn.close() + raise RuntimeError( + "sqlite embeddings table has neither embedding_vec nor embeddings" + ) + rows = [_normalize_embedding_row(tuple(r)) for r in raw_rows] conn.close() cols = [ - "embeddings", + "embedding_vec", "doc_type", "cre_id", "node_id", @@ -89,7 +150,7 @@ def _fetch_sqlite_embeddings(path: str) -> Tuple[List[str], List[Tuple[Any, ...] def _fetch_postgres_embeddings(pg_url: str) -> Tuple[List[str], List[Tuple[Any, ...]]]: cols = [ - "embeddings", + "embedding_vec", "doc_type", "cre_id", "node_id", @@ -101,7 +162,19 @@ def _fetch_postgres_embeddings(pg_url: str) -> Tuple[List[str], List[Tuple[Any, try: cur = conn.cursor() cur.execute(f"SELECT {quoted} FROM public.embeddings") - rows = [_normalize_embedding_row(tuple(r)) for r in cur.fetchall()] + rows = [ + _normalize_embedding_row( + ( + _as_vec_literal(r[0]), + r[1], + r[2], + r[3], + r[4], + r[5], + ) + ) + for r in cur.fetchall() + ] cur.close() return cols, rows finally: @@ -120,8 +193,9 @@ def _replace_postgres_embeddings(pg_url: str, rows: Sequence[Tuple[Any, ...]]) - cur, """ INSERT INTO public.embeddings ( - embeddings, doc_type, cre_id, node_id, embeddings_content, embeddings_url - ) VALUES (%s, %s, %s, %s, %s, %s) + embedding_vec, doc_type, cre_id, node_id, + embeddings_content, embeddings_url + ) VALUES (%s::vector, %s, %s, %s, %s, %s) """, list(rows), page_size=500,