From 65c6d26008315c36fa66fe860beaf419c6863545 Mon Sep 17 00:00:00 2001 From: Stefan Hudici Date: Fri, 3 Jul 2026 12:54:51 +0300 Subject: [PATCH] feat: keep embeddings fresh after builds and purge orphaned vectors Embeddings were write-once: nothing refreshed them after a build, so semantic search silently decayed as code changed, and vectors for deleted or renamed nodes kept surfacing as ghosts. Two additions: - EmbeddingStore.purge_orphans() deletes vectors whose qualified_name no longer exists in the nodes table (both tables share one SQLite file). embed_all_nodes() now calls it, so a manual embed also cleans up ghosts. - refresh_embeddings(graph_store) runs after every build-time post-process: the shared run_post_processing() pipeline (watch mode, eval runner) gains it as a fifth step, and the tool pipeline used by the CLI build command and the MCP build_or_update_graph / run_postprocess tools invokes it at the "full" postprocess level. run_postprocess grows an embeddings flag (default True) alongside its existing flows/communities/fts knobs. The refresh is opt-in by construction: it only acts when the graph already contains embeddings (the user ran embed at least once), and only with the exact provider identity that produced them (stored in the per-row provider tag). If the tagged provider cannot be resolved (missing extras or env vars) or resolves to a different model/endpoint, the refresh is skipped, so a build can never trigger a surprise full re-embed or cloud spend. The hash-incremental embed_nodes() keeps the refresh cheap: only new or changed nodes are re-encoded. Every wiring point downgrades provider or transport errors to a build warning, matching the existing step-isolation contract. Co-Authored-By: Claude Fable 5 --- code_review_graph/embeddings.py | 87 +++++++++++++++- code_review_graph/main.py | 4 + code_review_graph/postprocessing.py | 28 ++++- code_review_graph/tools/build.py | 28 +++++ tests/test_embeddings.py | 152 ++++++++++++++++++++++++++++ tests/test_postprocessing.py | 46 +++++++++ tests/test_tools.py | 92 +++++++++++++++++ 7 files changed, 435 insertions(+), 2 deletions(-) diff --git a/code_review_graph/embeddings.py b/code_review_graph/embeddings.py index ccb13508..3af7bdbd 100644 --- a/code_review_graph/embeddings.py +++ b/code_review_graph/embeddings.py @@ -966,15 +966,41 @@ def remove_node(self, qualified_name: str) -> None: ) self._conn.commit() + def purge_orphans(self) -> int: + """Delete vectors whose node no longer exists in the graph. + + The embeddings table lives in the same SQLite file as the graph, + so deleted or renamed nodes leave stale vectors behind that keep + surfacing in semantic search. Returns the number of rows removed; + no-op when there is no ``nodes`` table (standalone embedding DBs). + """ + has_nodes = self._conn.execute( + "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'nodes'" + ).fetchone() + if not has_nodes: + return 0 + cur = self._conn.execute( + "DELETE FROM embeddings WHERE qualified_name NOT IN " + "(SELECT qualified_name FROM nodes)" + ) + self._conn.commit() + return cur.rowcount + def count(self) -> int: return self._conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0] def embed_all_nodes(graph_store: GraphStore, embedding_store: EmbeddingStore) -> int: - """Embed all non-file nodes in the graph.""" + """Embed all non-file nodes in the graph. + + Also drops vectors whose node no longer exists, so deleted or renamed + symbols stop surfacing in semantic search. + """ if not embedding_store.available: return 0 + embedding_store.purge_orphans() + all_files = graph_store.get_all_files() all_nodes: list[GraphNode] = [] for f in all_files: @@ -983,6 +1009,65 @@ def embed_all_nodes(graph_store: GraphStore, embedding_store: EmbeddingStore) -> return embedding_store.embed_nodes(all_nodes) +def refresh_embeddings(graph_store: GraphStore) -> dict[str, int] | None: + """Incrementally refresh an existing embedding set after a build. + + Opt-in by construction: this runs only when the graph already contains + embeddings (the user ran ``embed`` at least once), and only with the + provider that produced them. If that provider cannot be resolved + (missing extras or env vars) or resolves to a different identity + (other model or endpoint), the refresh is skipped — a build must never + trigger a surprise full re-embed, or cloud spend, under a new identity. + + Returns: + ``{"embedded": n, "purged": m}`` when a refresh ran, else ``None``. + """ + try: + row = graph_store._conn.execute( + "SELECT provider, COUNT(*) AS n FROM embeddings " + "GROUP BY provider ORDER BY n DESC LIMIT 1" + ).fetchone() + except sqlite3.OperationalError: + return None # embeddings table absent -> the user never embedded + if row is None: + return None + + tag = row["provider"] + key, _, model = tag.partition(":") + if key == "openai" and "@" in model: + # openai tags are "openai:@"; the host part comes from + # CRG_OPENAI_BASE_URL and is re-derived by the provider itself. + model = model.split("@", 1)[0] + + try: + embedding_store = EmbeddingStore( + graph_store.db_path, provider=key, model=model or None, + ) + except ValueError as exc: + logger.info("Embedding refresh skipped: %s", exc) + return None + try: + if not embedding_store.available or embedding_store.provider.name != tag: + resolved = ( + embedding_store.provider.name + if embedding_store.available else "unavailable" + ) + logger.info( + "Embedding refresh skipped: existing vectors were made by %r " + "but the resolved provider is %r", + tag, + resolved, + ) + return None + # embed_all_nodes purges again on entry; that second pass is a + # no-op DELETE — this call exists to capture the purge count. + purged = embedding_store.purge_orphans() + embedded = embed_all_nodes(graph_store, embedding_store) + return {"embedded": embedded, "purged": purged} + finally: + embedding_store.close() + + def semantic_search( query: str, graph_store: GraphStore, diff --git a/code_review_graph/main.py b/code_review_graph/main.py index 629b1674..57ec5f64 100644 --- a/code_review_graph/main.py +++ b/code_review_graph/main.py @@ -137,6 +137,7 @@ async def run_postprocess_tool( flows: bool = True, communities: bool = True, fts: bool = True, + embeddings: bool = True, repo_root: Optional[str] = None, ) -> dict: """Run post-processing on existing graph (flows, communities, FTS index). @@ -152,11 +153,14 @@ async def run_postprocess_tool( flows: Run flow detection. Default: True. communities: Run community detection. Default: True. fts: Rebuild FTS index. Default: True. + embeddings: Refresh semantic embeddings when the graph already + has them (no-op otherwise). Default: True. repo_root: Repository root path. Auto-detected if omitted. """ return await asyncio.to_thread( run_postprocess, flows=flows, communities=communities, fts=fts, + embeddings=embeddings, repo_root=_resolve_repo_root(repo_root), ) diff --git a/code_review_graph/postprocessing.py b/code_review_graph/postprocessing.py index c7dec597..dee4ed8f 100644 --- a/code_review_graph/postprocessing.py +++ b/code_review_graph/postprocessing.py @@ -1,12 +1,13 @@ """Shared post-build processing pipeline. -After the core Tree-sitter parse (full_build or incremental_update), four +After the core Tree-sitter parse (full_build or incremental_update), five post-processing steps must run to populate derived tables: 1. Compute node signatures 2. Rebuild FTS5 search index 3. Trace execution flows 4. Detect code communities +5. Refresh semantic embeddings (only when the graph was embedded before) This module extracts that pipeline so every entry point — MCP tool, CLI commands, and watch mode — produces identical results. @@ -43,6 +44,7 @@ def run_post_processing(store: GraphStore) -> dict[str, Any]: _rebuild_fts_index(store, result, warnings) _trace_flows(store, result, warnings) _detect_communities(store, result, warnings) + _refresh_embeddings(store, result, warnings) if warnings: result["warnings"] = warnings @@ -132,3 +134,27 @@ def _detect_communities( except (sqlite3.OperationalError, ImportError) as e: logger.warning("Community detection failed: %s", e) warnings.append(f"Community detection failed: {type(e).__name__}: {e}") + + +def _refresh_embeddings( + store: GraphStore, + result: dict[str, Any], + warnings: list[str], +) -> None: + """Incrementally refresh semantic embeddings after a build. + + Runs only when the graph already has embeddings (see + :func:`code_review_graph.embeddings.refresh_embeddings`). This step can + hit an embedding model or a network endpoint, so any failure is + downgraded to a warning — a build must never fail on embedding refresh. + """ + try: + from .embeddings import refresh_embeddings + + refreshed = refresh_embeddings(store) + if refreshed is not None: + result["embeddings_refreshed"] = refreshed["embedded"] + result["embeddings_purged"] = refreshed["purged"] + except Exception as e: + logger.warning("Embedding refresh failed: %s", e) + warnings.append(f"Embedding refresh failed: {type(e).__name__}: {e}") diff --git a/code_review_graph/tools/build.py b/code_review_graph/tools/build.py index 11e05e93..39a4fefd 100644 --- a/code_review_graph/tools/build.py +++ b/code_review_graph/tools/build.py @@ -113,6 +113,18 @@ def _run_postprocess( logger.warning("Community detection failed: %s", e) warnings.append(f"Community detection failed: {type(e).__name__}: {e}") + # -- Refresh semantic embeddings (no-op unless the graph was embedded) -- + try: + from code_review_graph.embeddings import refresh_embeddings + + refreshed = refresh_embeddings(store) + if refreshed is not None: + build_result["embeddings_refreshed"] = refreshed.get("embedded", 0) + build_result["embeddings_purged"] = refreshed.get("purged", 0) + except Exception as e: # provider/transport errors vary by backend + logger.warning("Embedding refresh failed: %s", e) + warnings.append(f"Embedding refresh failed: {type(e).__name__}: {e}") + # -- Compute pre-computed summary tables -- try: _compute_summaries(store) @@ -441,6 +453,7 @@ def run_postprocess( flows: bool = True, communities: bool = True, fts: bool = True, + embeddings: bool = True, repo_root: str | None = None, ) -> dict[str, Any]: """Run post-processing steps on an existing graph. @@ -453,6 +466,8 @@ def run_postprocess( flows: Run flow detection. Default: True. communities: Run community detection. Default: True. fts: Rebuild FTS index. Default: True. + embeddings: Refresh semantic embeddings when the graph already + has them (no-op otherwise). Default: True. repo_root: Repository root path. Auto-detected if omitted. Returns: @@ -529,6 +544,19 @@ def run_postprocess( logger.warning("Community detection failed: %s", e) warnings.append(f"Community detection failed: {type(e).__name__}: {e}") + if embeddings: + try: + from code_review_graph.embeddings import refresh_embeddings + + refreshed = refresh_embeddings(store) + if refreshed is not None: + result["embeddings_refreshed"] = refreshed.get("embedded", 0) + result["embeddings_purged"] = refreshed.get("purged", 0) + except Exception as e: # provider/transport errors vary by backend + store.rollback() + logger.warning("Embedding refresh failed: %s", e) + warnings.append(f"Embedding refresh failed: {type(e).__name__}: {e}") + store.set_metadata( "last_postprocessed_at", time.strftime("%Y-%m-%dT%H:%M:%S"), diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py index 740c3c9e..ca9eb444 100644 --- a/tests/test_embeddings.py +++ b/tests/test_embeddings.py @@ -1131,3 +1131,155 @@ def test_subdomain_spoof_triggers_warning(self, capsys): get_provider("openai") captured = capsys.readouterr() assert "cloud" in captured.err.lower() + + +class _StubProvider: + """Deterministic in-memory provider for refresh tests.""" + + name = "stub:mini" + dimension = 2 + + def embed(self, texts): + return [[float(len(t)), 1.0] for t in texts] + + def embed_query(self, text): + return [1.0, 0.0] + + +class TestPurgeOrphans: + def test_removes_vectors_for_deleted_nodes(self, tmp_path): + from code_review_graph.graph import GraphStore + from code_review_graph.parser import NodeInfo + + db = tmp_path / "graph.db" + gstore = GraphStore(db) + gstore.upsert_node(NodeInfo( + kind="Function", name="keep", file_path="/r/a.py", + line_start=1, line_end=2, language="python", + )) + gstore.commit() + + with patch("code_review_graph.embeddings.get_provider", + return_value=_StubProvider()): + estore = EmbeddingStore(db) + for qn in ("/r/a.py::keep", "/r/a.py::ghost"): + estore._conn.execute( + "INSERT INTO embeddings (qualified_name, vector, text_hash," + " provider) VALUES (?, ?, ?, ?)", + (qn, b"\x00\x00\x00\x00", "h", "stub:mini"), + ) + estore._conn.commit() + + assert estore.purge_orphans() == 1 + rows = estore._conn.execute( + "SELECT qualified_name FROM embeddings" + ).fetchall() + assert [r["qualified_name"] for r in rows] == ["/r/a.py::keep"] + estore.close() + gstore.close() + + def test_noop_without_nodes_table(self, tmp_path): + with patch("code_review_graph.embeddings.get_provider", + return_value=_StubProvider()): + estore = EmbeddingStore(tmp_path / "standalone.db") + assert estore.purge_orphans() == 0 + estore.close() + + +class TestRefreshEmbeddings: + def _graph_with_node(self, tmp_path): + from code_review_graph.graph import GraphStore + from code_review_graph.parser import NodeInfo + + gstore = GraphStore(tmp_path / "graph.db") + # embed_all_nodes discovers nodes via File nodes (get_all_files), + # so seed one like the parser always does. + gstore.upsert_node(NodeInfo( + kind="File", name="/r/a.py", file_path="/r/a.py", + line_start=1, line_end=20, language="python", + )) + gstore.upsert_node(NodeInfo( + kind="Function", name="existing", file_path="/r/a.py", + line_start=1, line_end=2, language="python", + )) + gstore.commit() + return gstore + + def test_returns_none_when_never_embedded(self, tmp_path): + from code_review_graph.embeddings import refresh_embeddings + + gstore = self._graph_with_node(tmp_path) + try: + assert refresh_embeddings(gstore) is None + finally: + gstore.close() + + def test_skips_when_stored_provider_cannot_be_resolved(self, tmp_path): + """Legacy rows are tagged 'unknown'; the provider must not be guessed.""" + from code_review_graph.embeddings import refresh_embeddings + + gstore = self._graph_with_node(tmp_path) + with patch("code_review_graph.embeddings.get_provider", + return_value=_StubProvider()): + estore = EmbeddingStore(gstore.db_path) + estore._conn.execute( + "INSERT INTO embeddings (qualified_name, vector, text_hash," + " provider) VALUES (?, ?, ?, ?)", + ("/r/a.py::existing", b"\x00\x00\x00\x00", "h", "unknown"), + ) + estore._conn.commit() + estore.close() + + try: + # Real get_provider raises ValueError for the 'unknown' name. + assert refresh_embeddings(gstore) is None + finally: + gstore.close() + + def test_skips_when_provider_identity_differs(self, tmp_path): + from code_review_graph.embeddings import refresh_embeddings + + gstore = self._graph_with_node(tmp_path) + with patch("code_review_graph.embeddings.get_provider", + return_value=_StubProvider()): + estore = EmbeddingStore(gstore.db_path) + estore._conn.execute( + "INSERT INTO embeddings (qualified_name, vector, text_hash," + " provider) VALUES (?, ?, ?, ?)", + ("/r/a.py::existing", b"\x00\x00\x00\x00", "h", "stub:other"), + ) + estore._conn.commit() + estore.close() + + # get_provider resolves to "stub:mini" != stored "stub:other". + assert refresh_embeddings(gstore) is None + gstore.close() + + def test_embeds_new_nodes_and_purges_orphans(self, tmp_path): + from code_review_graph.embeddings import refresh_embeddings + from code_review_graph.parser import NodeInfo + + gstore = self._graph_with_node(tmp_path) + stub = _StubProvider() + with patch("code_review_graph.embeddings.get_provider", + return_value=stub): + estore = EmbeddingStore(gstore.db_path) + estore.embed_nodes(gstore.get_nodes_by_file("/r/a.py")) + estore._conn.execute( + "INSERT INTO embeddings (qualified_name, vector, text_hash," + " provider) VALUES (?, ?, ?, ?)", + ("/r/a.py::ghost", b"\x00\x00\x00\x00", "stale", stub.name), + ) + estore._conn.commit() + estore.close() + + gstore.upsert_node(NodeInfo( + kind="Function", name="added_later", file_path="/r/a.py", + line_start=5, line_end=8, language="python", + )) + gstore.commit() + + out = refresh_embeddings(gstore) + + assert out == {"embedded": 1, "purged": 1} + gstore.close() diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index f9b0f946..c1a6cac1 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -325,3 +325,49 @@ def run_watch(): callback.assert_not_called() finally: store.close() + + +class TestEmbeddingRefreshStep: + def _store_with_node(self, tmp_path): + store = GraphStore(tmp_path / "graph.db") + store.upsert_node(NodeInfo( + kind="Function", name="fn", file_path="/repo/a.py", + line_start=1, line_end=5, language="python", + )) + store.commit() + return store + + def test_skipped_when_never_embedded(self, tmp_path): + store = self._store_with_node(tmp_path) + try: + result = run_post_processing(store) + finally: + store.close() + assert "embeddings_refreshed" not in result + assert "warnings" not in result + + def test_reports_counts_when_refresh_runs(self, tmp_path): + store = self._store_with_node(tmp_path) + try: + with patch( + "code_review_graph.embeddings.refresh_embeddings", + return_value={"embedded": 3, "purged": 2}, + ): + result = run_post_processing(store) + finally: + store.close() + assert result["embeddings_refreshed"] == 3 + assert result["embeddings_purged"] == 2 + + def test_refresh_failure_is_warning_not_fatal(self, tmp_path): + store = self._store_with_node(tmp_path) + try: + with patch( + "code_review_graph.embeddings.refresh_embeddings", + side_effect=RuntimeError("provider exploded"), + ): + result = run_post_processing(store) + finally: + store.close() + assert result["signatures_computed"] == 1 + assert any("Embedding refresh" in w for w in result["warnings"]) diff --git a/tests/test_tools.py b/tests/test_tools.py index 578536d4..001dee01 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -1194,6 +1194,98 @@ def test_postprocess_full_matches_default(self): assert "communities_detected" in result +class TestBuildRefreshesEmbeddings: + """The build and postprocess tools must run the embedding refresh. + + The refresh itself is opt-in by construction (it acts only when the + graph already has embeddings), so these tests pin the *wiring*, not + the policy: the tool paths must invoke it and surface its counts. + """ + + def setup_method(self): + self.tmp = tempfile.mkdtemp() + self.root = Path(self.tmp) + (self.root / ".git").mkdir() + (self.root / "sample.py").write_text( + "def hello():\n pass\n\nclass Foo:\n pass\n" + ) + + def teardown_method(self): + import shutil + shutil.rmtree(self.tmp, ignore_errors=True) + + def _build(self, postprocess: str, refresh_mock): + from unittest.mock import patch + + from code_review_graph.tools.build import build_or_update_graph + + with patch( + "code_review_graph.incremental.get_all_tracked_files", + return_value=["sample.py"], + ), patch( + "code_review_graph.embeddings.refresh_embeddings", + refresh_mock, + ): + return build_or_update_graph( + full_rebuild=True, repo_root=str(self.root), + postprocess=postprocess, + ) + + def test_build_full_invokes_embedding_refresh(self): + refresh = MagicMock(return_value={"embedded": 2, "purged": 1}) + result = self._build("full", refresh) + + refresh.assert_called_once() + assert result["embeddings_refreshed"] == 2 + assert result["embeddings_purged"] == 1 + + def test_build_minimal_skips_embedding_refresh(self): + refresh = MagicMock(return_value={"embedded": 2, "purged": 1}) + result = self._build("minimal", refresh) + + refresh.assert_not_called() + assert "embeddings_refreshed" not in result + + def test_build_without_prior_embeddings_reports_nothing(self): + refresh = MagicMock(return_value=None) # never-embedded graph + result = self._build("full", refresh) + + refresh.assert_called_once() + assert "embeddings_refreshed" not in result + assert "warnings" not in result + + def test_refresh_failure_is_a_warning_not_a_build_failure(self): + refresh = MagicMock(side_effect=RuntimeError("provider exploded")) + result = self._build("full", refresh) + + assert result["status"] == "ok" + assert any("Embedding refresh failed" in w for w in result["warnings"]) + + def test_run_postprocess_embeddings_flag(self): + from unittest.mock import patch + + from code_review_graph.tools.build import run_postprocess + + refresh = MagicMock(return_value={"embedded": 3, "purged": 0}) + self._build("none", MagicMock(return_value=None)) + + with patch( + "code_review_graph.embeddings.refresh_embeddings", refresh + ): + skipped = run_postprocess( + flows=False, communities=False, fts=False, + embeddings=False, repo_root=str(self.root), + ) + refresh.assert_not_called() + ran = run_postprocess( + flows=False, communities=False, fts=False, + repo_root=str(self.root), + ) + refresh.assert_called_once() + assert "embeddings_refreshed" not in skipped + assert ran["embeddings_refreshed"] == 3 + + class TestComputeSummaries: """Tests for _compute_summaries: pins the contents of the three summary tables so that the batch-aggregate refactor can't silently