Skip to content

feat: pgvector embedding_vec migration + chat similarity cutover #979

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
northdpole wants to merge 4 commits into
base: main
Choose a base branch
from
Open

feat: pgvector embedding_vec migration + chat similarity cutover #979

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
13b86f0
feat: pgvector embedding_vec migration + chat similarity cutover
northdpole Jul 13, 2026
495037a
fix: use SQLAlchemy 2 Engine-safe execute for pgvector helpers
northdpole Jul 13, 2026
2dc4062
chore: use pgvector/pgvector for all local Postgres containers
northdpole Jul 13, 2026
01d3a4b
feat: make embedding_vec the sole store and hard-fail pgvector-on-SQLite
northdpole Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .env.example
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 3 additions & 2 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
38 changes: 36 additions & 2 deletions Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +60 to +64

@coderabbitai coderabbitai Bot Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Readiness loops for Docker Postgres lack explicit failure on timeout.

The Docker Postgres readiness loops in the Makefile and benchmark script don't fail explicitly when the container doesn't become ready — the Makefile falls through to psql with a confusing error, and the benchmark script's unbounded until loop can hang indefinitely. scripts/setup-heroku-staging.sh (L243–L251) demonstrates the correct pattern with a bounded loop and explicit die on timeout.

  • Makefile#L60-L64: After the 20-iteration loop, verify readiness with an explicit failure before running psql, so the target fails with a clear message instead of a confusing connection error.
  • scripts/run_benchmark.sh#L44-L47: Replace the unbounded until loop with a bounded iteration loop and explicit failure, matching the pattern in scripts/setup-heroku-staging.sh.
🔧 Proposed fix for Makefile readiness check 
 	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
+	docker exec cre-postgres pg_isready -U cre -d cre >/dev/null 2>&1 || { \
+		echo "ERROR: cre-postgres did not become ready in 20s"; exit 1; \
+	}; \
+	docker exec cre-postgres psql -U cre -d cre -c "CREATE EXTENSION IF NOT EXISTS vector;" >/dev/null
🔧 Proposed fix for run_benchmark.sh readiness check 
     echo "Waiting for $container to be ready..."
-    until docker exec "$container" pg_isready -U cre -d cre; do
-      sleep 1
-    done
+    local ready=0
+    for _ in $(seq 1 30); do
+      if docker exec "$container" pg_isready -U cre -d cre >/dev/null 2>&1; then
+        ready=1
+        break
+      fi
+      sleep 1
+    done
+    [[ "$ready" == "1" ]] || { echo "ERROR: $container did not become ready in 30s"; exit 1; }
     docker exec "$container" psql -U cre -d cre -c "CREATE EXTENSION IF NOT EXISTS vector;" >/dev/null
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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
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 pg_isready -U cre -d cre >/dev/null 2>&1 || { \
echo "ERROR: cre-postgres did not become ready in 20s"; exit 1; \
}; \
docker exec cre-postgres psql -U cre -d cre -c "CREATE EXTENSION IF NOT EXISTS vector;" >/dev/null
📍 Affects 2 files
  • Makefile#L60-L64 (this comment)
  • scripts/run_benchmark.sh#L44-L47
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@Makefile` around lines 60 - 64, Makefile lines 60-64: update the Docker
Postgres readiness loop to verify readiness after all 20 attempts and fail with
a clear timeout message before running psql. scripts/run_benchmark.sh lines
44-47: replace the unbounded readiness until loop with a bounded iteration loop
that exits successfully when pg_isready succeeds and explicitly fails with a
timeout message when all attempts are exhausted, matching the pattern used by
setup-heroku-staging.sh.


start-containers: docker-neo4j docker-redis

Expand Down
19 changes: 10 additions & 9 deletions application/cmd/cre_main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1115,14 +1115,12 @@ def run_librarian(

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,
# 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():
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
Expand All @@ -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
Expand Down
129 changes: 120 additions & 9 deletions application/database/db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -191,6 +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) — 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)
Comment on lines +193 to +196

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win

Restore the legacy embeddings CSV column.

The PR objectives and linked Issue #977 explicitly require preserving the CSV embeddings column for backwards compatibility and dual-writing. Removing the legacy column and making embedding_vec the sole store violates this requirement.

💡 Proposed fix 
     embedding_dim = sqla.Column(sqla.Integer, nullable=True, default=None)
+    embeddings = sqla.Column(sqla.String, 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)
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
# 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)
embedding_dim = sqla.Column(sqla.Integer, nullable=True, default=None)
embeddings = sqla.Column(sqla.String, 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)
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@application/database/db.py` around lines 193 - 196, Restore the legacy
nullable `embeddings` column alongside `embedding_vec` in the model, preserving
`embedding_vec` for the pgvector/SQLite representation. Update the related
persistence paths to dual-write both columns so existing CSV consumers remain
compatible.



class GapAnalysisResults(BaseModel):
Expand Down Expand Up @@ -2364,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)

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Fallback to legacy embeddings if embedding_vec is unavailable.

To safely support the transition period before the backfill completes, and to honor the PR requirement to preserve the CSV column, the embedding retrieval methods should fall back to reading the legacy embeddings column when embedding_vec is empty.

  • application/database/db.py#L2384-L2384: In get_embeddings_by_doc_type, change to vec = parse_stored_embedding_vec(getattr(entry, "embedding_vec", None) or getattr(entry, "embeddings", None)).
  • application/database/db.py#L2434-L2434: In get_embeddings_by_doc_type_paginated, change to vec = parse_stored_embedding_vec(getattr(entry, "embedding_vec", None) or getattr(entry, "embeddings", None)).
📍 Affects 1 file
  • application/database/db.py#L2384-L2384 (this comment)
  • application/database/db.py#L2434-L2434
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@application/database/db.py` at line 2384, Update get_embeddings_by_doc_type
and get_embeddings_by_doc_type_paginated to pass embedding_vec to
parse_stored_embedding_vec when available, otherwise fall back to the legacy
embeddings value via safe attribute access. Apply this change at
application/database/db.py lines 2384 and 2434.

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]:
Expand All @@ -2400,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)
Expand All @@ -2409,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 (
Expand Down Expand Up @@ -2484,8 +2514,16 @@ def add_embedding(
f"embedding dimension mismatch for {db_object.id}: "
f"expected {expected_dim}, got {len(embeddings)}"
)
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)
embeddings_str = ",".join([str(e) for e in embeddings])
# 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:
resolved_node_url = (
Expand All @@ -2496,32 +2534,32 @@ 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,
embedding_model_id=embedding_model_id,
embedding_dim=embedding_dim,
embedding_vec=embedding_vec_literal,
)
else:
emb = Embeddings(
embeddings=embeddings_str,
node_id=db_object.id,
doc_type=db_object.ntype,
embeddings_content=embedding_text,
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()
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
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
Expand All @@ -2531,6 +2569,79 @@ 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. 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 = {
"q": to_pgvector_literal(query_embedding),
"doc_type": doc_type,
}
try:
# 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)
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.
logger.warning(
"pgvector similarity query failed (%s); returning no match",
exc,
)
return None, None

Comment thread
coderabbitai[bot] marked this conversation as resolved.
def assert_embedding_contract(
self,
*,
Expand Down
Loading
Loading