diff --git a/application/graphrag/store.py b/application/graphrag/store.py new file mode 100644 index 00000000..b1c6c456 --- /dev/null +++ b/application/graphrag/store.py @@ -0,0 +1,637 @@ +"""Per-source knowledge-graph store co-located with the pgvector ``documents`` table. + +GraphRAG is pgvector-only (D29): the graph tables live in the same DB as the +pgvector store and are created on-demand (``CREATE TABLE IF NOT EXISTS`` + +``CREATE EXTENSION IF NOT EXISTS vector``), mirroring +``PGVectorStore._ensure_table_exists`` rather than going through app-DB Alembic. +That DB may be a separate cluster (e.g. Neon) from the app DB where ``sources`` +lives, so ``source_id`` is a plain indexed UUID column with no cross-DB FK and +all ids are generated in Python. +""" + +from __future__ import annotations + +import logging +import uuid +from typing import Any, Dict, List, Optional + +from psycopg.types.json import Jsonb + +from application.core.settings import settings + +DEFAULT_NAME_EMBEDDING_DIM = 768 + +MAX_SUBGRAPH_NODES = 500 +MAX_SUBGRAPH_EDGES = 2000 + + +class GraphStore: + """Stores and queries a per-source knowledge graph in the pgvector DB.""" + + def __init__(self, connection_string: Optional[str] = None): + self._connection_string = connection_string or getattr( + settings, "PGVECTOR_CONNECTION_STRING", None + ) + + if not self._connection_string and getattr(settings, "POSTGRES_URI", None): + from application.core.db_uri import normalize_pgvector_connection_string + + self._connection_string = normalize_pgvector_connection_string( + settings.POSTGRES_URI + ) + + if not self._connection_string: + raise ValueError( + "PostgreSQL connection string is required. " + "Set PGVECTOR_CONNECTION_STRING or POSTGRES_URI in settings, " + "or pass connection_string parameter." + ) + + try: + import psycopg + from pgvector.psycopg import register_vector + except ImportError: + raise ImportError( + "Could not import required packages. " + "Please install with `pip install 'psycopg[binary,pool]' pgvector`." + ) + + self._psycopg = psycopg + self._register_vector = register_vector + self._connection = None + self._ensure_tables() + + def _get_connection(self): + if self._connection is None or self._connection.closed: + self._connection = self._psycopg.connect(self._connection_string) + self._register_vector(self._connection) + return self._connection + + def _embedding_dim(self) -> int: + """Dimension of the configured embeddings model, matching ``PGVectorStore``. + + Falls back to ``DEFAULT_NAME_EMBEDDING_DIM`` so the graph table and the + pgvector ``documents`` table always agree on the configured model. + """ + from application.vectorstore.base import EmbeddingsSingleton + + embedding = EmbeddingsSingleton.get_instance( + settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY + ) + return getattr(embedding, "dimension", DEFAULT_NAME_EMBEDDING_DIM) + + def _ensure_tables(self): + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;") + + embedding_dim = self._embedding_dim() + + cursor.execute( + f""" + CREATE TABLE IF NOT EXISTS graph_nodes ( + id UUID PRIMARY KEY, + source_id UUID NOT NULL, + name TEXT, + normalized_name TEXT, + type TEXT, + description TEXT, + degree INT DEFAULT 0, + doc_freq INT DEFAULT 0, + name_embedding vector({embedding_dim}), + UNIQUE (source_id, normalized_name) + ); + """ + ) + + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS graph_edges ( + id UUID PRIMARY KEY, + source_id UUID NOT NULL, + src_node_id UUID, + dst_node_id UUID, + type TEXT, + description TEXT, + weight REAL DEFAULT 1.0, + source_chunk_ids JSONB + ); + """ + ) + + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS graph_node_chunks ( + source_id UUID NOT NULL, + node_id UUID NOT NULL, + chunk_id TEXT NOT NULL, + PRIMARY KEY (source_id, node_id, chunk_id) + ); + """ + ) + + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS graph_ingest_progress ( + source_id UUID NOT NULL, + chunk_id TEXT NOT NULL, + status TEXT, + PRIMARY KEY (source_id, chunk_id) + ); + """ + ) + + cursor.execute( + "CREATE INDEX IF NOT EXISTS graph_nodes_source_id_idx " + "ON graph_nodes (source_id);" + ) + cursor.execute( + "CREATE INDEX IF NOT EXISTS graph_edges_source_id_idx " + "ON graph_edges (source_id);" + ) + cursor.execute( + "CREATE INDEX IF NOT EXISTS graph_edges_src_node_id_idx " + "ON graph_edges (src_node_id);" + ) + cursor.execute( + "CREATE INDEX IF NOT EXISTS graph_edges_dst_node_id_idx " + "ON graph_edges (dst_node_id);" + ) + cursor.execute( + "CREATE INDEX IF NOT EXISTS graph_node_chunks_node_id_idx " + "ON graph_node_chunks (node_id);" + ) + cursor.execute( + "CREATE INDEX IF NOT EXISTS graph_nodes_name_embedding_idx " + "ON graph_nodes USING ivfflat (name_embedding vector_cosine_ops) " + "WITH (lists = 100);" + ) + + conn.commit() + except Exception as e: + conn.rollback() + logging.error(f"Error creating graph tables: {e}") + raise + finally: + cursor.close() + + def upsert_node( + self, + source_id: str, + name: str, + normalized_name: str, + type: Optional[str] = None, + description: Optional[str] = None, + name_embedding: Optional[List[float]] = None, + ) -> str: + """Insert a node or merge into the existing one for ``(source_id, normalized_name)``. + + On conflict the description is concatenated (de-duped), ``doc_freq`` is + incremented, the type is refreshed if previously empty, and the + embedding is refreshed when provided. Returns the node id either way. + """ + node_id = str(uuid.uuid4()) + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + """ + INSERT INTO graph_nodes + (id, source_id, name, normalized_name, type, description, + doc_freq, name_embedding) + VALUES (%s, %s, %s, %s, %s, %s, 1, %s) + ON CONFLICT (source_id, normalized_name) DO UPDATE SET + description = CASE + WHEN EXCLUDED.description IS NULL + OR EXCLUDED.description = '' THEN graph_nodes.description + WHEN graph_nodes.description IS NULL + OR graph_nodes.description = '' THEN EXCLUDED.description + WHEN position(EXCLUDED.description IN graph_nodes.description) > 0 + THEN graph_nodes.description + ELSE graph_nodes.description || ' ' || EXCLUDED.description + END, + type = CASE + WHEN graph_nodes.type IS NULL + OR graph_nodes.type = '' THEN EXCLUDED.type + ELSE graph_nodes.type + END, + name = COALESCE(graph_nodes.name, EXCLUDED.name), + doc_freq = graph_nodes.doc_freq + 1, + name_embedding = COALESCE( + EXCLUDED.name_embedding, graph_nodes.name_embedding + ) + RETURNING id; + """, + ( + node_id, + source_id, + name, + normalized_name, + type, + description, + name_embedding, + ), + ) + returned_id = cursor.fetchone()[0] + conn.commit() + return str(returned_id) + except Exception as e: + conn.rollback() + logging.error(f"Error upserting node: {e}") + raise + finally: + cursor.close() + + def add_edge( + self, + source_id: str, + src_node_id: str, + dst_node_id: str, + type: Optional[str] = None, + description: Optional[str] = None, + weight: float = 1.0, + source_chunk_ids: Optional[List[str]] = None, + ) -> str: + """Insert an edge and bump the degree of both endpoints. Returns its id.""" + edge_id = str(uuid.uuid4()) + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + """ + INSERT INTO graph_edges + (id, source_id, src_node_id, dst_node_id, type, description, + weight, source_chunk_ids) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s); + """, + ( + edge_id, + source_id, + src_node_id, + dst_node_id, + type, + description, + weight, + Jsonb(source_chunk_ids or []), + ), + ) + cursor.execute( + "UPDATE graph_nodes SET degree = degree + 1 " + "WHERE source_id = %s AND id IN (%s, %s);", + (source_id, src_node_id, dst_node_id), + ) + conn.commit() + return edge_id + except Exception as e: + conn.rollback() + logging.error(f"Error adding edge: {e}") + raise + finally: + cursor.close() + + def link_node_chunk(self, source_id: str, node_id: str, chunk_id: str): + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + """ + INSERT INTO graph_node_chunks (source_id, node_id, chunk_id) + VALUES (%s, %s, %s) + ON CONFLICT (source_id, node_id, chunk_id) DO NOTHING; + """, + (source_id, node_id, str(chunk_id)), + ) + conn.commit() + except Exception as e: + conn.rollback() + logging.error(f"Error linking node chunk: {e}") + raise + finally: + cursor.close() + + def get_node_by_normalized( + self, source_id: str, normalized_name: str + ) -> Optional[Dict[str, Any]]: + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + """ + SELECT id, name, normalized_name, type, description, degree, doc_freq + FROM graph_nodes + WHERE source_id = %s AND normalized_name = %s; + """, + (source_id, normalized_name), + ) + row = cursor.fetchone() + if row is None: + return None + return { + "id": str(row[0]), + "name": row[1], + "normalized_name": row[2], + "type": row[3], + "description": row[4], + "degree": row[5], + "doc_freq": row[6], + } + except Exception as e: + logging.error(f"Error getting node by normalized name: {e}") + return None + finally: + cursor.close() + + def count_nodes(self, source_id: str) -> int: + """Number of nodes for a source. Zero drives the ClassicRAG fallback (G5).""" + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + "SELECT count(*) FROM graph_nodes WHERE source_id = %s;", + (source_id,), + ) + return int(cursor.fetchone()[0]) + except Exception as e: + logging.error(f"Error counting nodes: {e}") + return 0 + finally: + cursor.close() + + def search_nodes_by_embedding( + self, source_id: str, query_embedding: List[float], k: int = 10 + ) -> List[Dict[str, Any]]: + """Cosine NN over ``graph_nodes.name_embedding`` scoped to a source.""" + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + """ + SELECT id, name, description, + (name_embedding <=> %s::vector) AS distance + FROM graph_nodes + WHERE source_id = %s AND name_embedding IS NOT NULL + ORDER BY name_embedding <=> %s::vector + LIMIT %s; + """, + (query_embedding, source_id, query_embedding, k), + ) + rows = cursor.fetchall() + return [ + { + "id": str(row[0]), + "name": row[1], + "description": row[2], + "distance": row[3], + } + for row in rows + ] + except Exception as e: + logging.error(f"Error searching nodes by embedding: {e}") + return [] + finally: + cursor.close() + + def get_subgraph( + self, source_id: str, node_ids: List[str], hops: int = 1 + ) -> Dict[str, List[Dict[str, Any]]]: + """Bounded 1-2-hop neighborhood of ``node_ids`` via indexed joins. + + Expands the seed set one hop at a time over edges (no recursive PageRank + in SQL), capping node and edge counts so a hub never explodes the fetch. + """ + if not node_ids: + return {"nodes": [], "edges": []} + + conn = self._get_connection() + cursor = conn.cursor() + try: + frontier = set(str(n) for n in node_ids) + visited = set(frontier) + for _ in range(max(1, hops)): + if not frontier or len(visited) >= MAX_SUBGRAPH_NODES: + break + cursor.execute( + """ + SELECT src_node_id, dst_node_id + FROM graph_edges + WHERE source_id = %s + AND (src_node_id = ANY(%s) OR dst_node_id = ANY(%s)) + LIMIT %s; + """, + ( + source_id, + list(frontier), + list(frontier), + MAX_SUBGRAPH_EDGES, + ), + ) + next_frontier = set() + for src, dst in cursor.fetchall(): + for neighbor in (str(src), str(dst)): + if neighbor not in visited: + next_frontier.add(neighbor) + if len(visited) + len(next_frontier) > MAX_SUBGRAPH_NODES: + allowed = MAX_SUBGRAPH_NODES - len(visited) + next_frontier = set(sorted(next_frontier)[:allowed]) + visited |= next_frontier + frontier = next_frontier + + node_id_list = list(visited) + cursor.execute( + """ + SELECT id, name, type, description, degree, doc_freq + FROM graph_nodes + WHERE source_id = %s AND id = ANY(%s); + """, + (source_id, node_id_list), + ) + nodes = [ + { + "id": str(row[0]), + "name": row[1], + "type": row[2], + "description": row[3], + "degree": row[4], + "doc_freq": row[5], + } + for row in cursor.fetchall() + ] + + cursor.execute( + """ + SELECT id, src_node_id, dst_node_id, type, weight + FROM graph_edges + WHERE source_id = %s + AND src_node_id = ANY(%s) AND dst_node_id = ANY(%s) + LIMIT %s; + """, + (source_id, node_id_list, node_id_list, MAX_SUBGRAPH_EDGES), + ) + edges = [ + { + "id": str(row[0]), + "src_node_id": str(row[1]), + "dst_node_id": str(row[2]), + "type": row[3], + "weight": row[4], + } + for row in cursor.fetchall() + ] + return {"nodes": nodes, "edges": edges} + except Exception as e: + logging.error(f"Error getting subgraph: {e}") + return {"nodes": [], "edges": []} + finally: + cursor.close() + + def get_chunk_ids_for_nodes( + self, source_id: str, node_ids: List[str] + ) -> Dict[str, List[str]]: + if not node_ids: + return {} + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + """ + SELECT node_id, chunk_id + FROM graph_node_chunks + WHERE source_id = %s AND node_id = ANY(%s); + """, + (source_id, [str(n) for n in node_ids]), + ) + result: Dict[str, List[str]] = {} + for node_id, chunk_id in cursor.fetchall(): + result.setdefault(str(node_id), []).append(chunk_id) + return result + except Exception as e: + logging.error(f"Error getting chunk ids for nodes: {e}") + return {} + finally: + cursor.close() + + def set_node_degrees(self, source_id: str): + """Recompute every node's degree from its incident edges for a source. + + A self-loop counts once, matching ``add_edge``'s incremental update + (``WHERE id IN (src, dst)`` bumps the endpoint a single time when + ``src == dst``). ``UNION`` deduplicates the two endpoints of each edge. + """ + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + """ + UPDATE graph_nodes n + SET degree = COALESCE(d.deg, 0) + FROM ( + SELECT node_id, count(*) AS deg + FROM ( + SELECT id, src_node_id AS node_id FROM graph_edges + WHERE source_id = %s + UNION + SELECT id, dst_node_id AS node_id FROM graph_edges + WHERE source_id = %s + ) incident + GROUP BY node_id + ) d + WHERE n.source_id = %s AND n.id = d.node_id; + """, + (source_id, source_id, source_id), + ) + conn.commit() + except Exception as e: + conn.rollback() + logging.error(f"Error setting node degrees: {e}") + raise + finally: + cursor.close() + + def mark_chunk(self, source_id: str, chunk_id: str, status: str): + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + """ + INSERT INTO graph_ingest_progress (source_id, chunk_id, status) + VALUES (%s, %s, %s) + ON CONFLICT (source_id, chunk_id) DO UPDATE SET status = EXCLUDED.status; + """, + (source_id, str(chunk_id), status), + ) + conn.commit() + except Exception as e: + conn.rollback() + logging.error(f"Error marking chunk: {e}") + raise + finally: + cursor.close() + + def pending_chunks(self, source_id: str, all_chunk_ids: List[str]) -> List[str]: + """Chunk ids from ``all_chunk_ids`` not yet marked ``done`` for the source.""" + if not all_chunk_ids: + return [] + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + """ + SELECT chunk_id FROM graph_ingest_progress + WHERE source_id = %s AND status = 'done'; + """, + (source_id,), + ) + done = {row[0] for row in cursor.fetchall()} + return [str(c) for c in all_chunk_ids if str(c) not in done] + except Exception as e: + logging.error(f"Error getting pending chunks: {e}") + return [str(c) for c in all_chunk_ids] + finally: + cursor.close() + + def get_progress(self, source_id: str) -> Dict[str, str]: + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + "SELECT chunk_id, status FROM graph_ingest_progress " + "WHERE source_id = %s;", + (source_id,), + ) + return {row[0]: row[1] for row in cursor.fetchall()} + except Exception as e: + logging.error(f"Error getting progress: {e}") + return {} + finally: + cursor.close() + + def delete_by_source(self, source_id: str): + """Remove every graph row for a source (no FK cascade across clusters).""" + conn = self._get_connection() + cursor = conn.cursor() + try: + for table in ( + "graph_node_chunks", + "graph_edges", + "graph_nodes", + "graph_ingest_progress", + ): + cursor.execute( + f"DELETE FROM {table} WHERE source_id = %s;", (source_id,) + ) + conn.commit() + except Exception as e: + conn.rollback() + logging.error(f"Error deleting graph by source: {e}") + raise + finally: + cursor.close() + + def __del__(self): + if ( + hasattr(self, "_connection") + and self._connection + and not self._connection.closed + ): + self._connection.close() diff --git a/tests/graphrag/test_store.py b/tests/graphrag/test_store.py new file mode 100644 index 00000000..64ae1858 --- /dev/null +++ b/tests/graphrag/test_store.py @@ -0,0 +1,365 @@ +"""Tests for the GraphRAG GraphStore (on-demand tables in the pgvector DB). + +Two layers: + +* A live-pg integration test that exercises the real DDL + SQL against the + pgvector store DB (same connection-string source as ``PGVectorStore``). It + uses a unique temp ``source_id`` and tears down every row it creates. +* A mock-cursor test that asserts the parameterized SQL shapes — ``source_id`` + and embeddings are bound params, never interpolated. + +The embedding dimension is mocked everywhere so the suite never loads the real +SentenceTransformer model: the live store creates ``TEST_EMBEDDING_DIM`` vectors +and the helpers build matching ones. +""" + +from __future__ import annotations + +import uuid +from unittest.mock import MagicMock + +import pytest + +import application.graphrag.store as store_module +from application.graphrag.store import GraphStore + +TEST_EMBEDDING_DIM = 8 + +_REAL_EMBEDDING_DIM = GraphStore._embedding_dim + + +@pytest.fixture(autouse=True) +def _mock_embedding_dim(monkeypatch): + monkeypatch.setattr( + GraphStore, "_embedding_dim", lambda self: TEST_EMBEDDING_DIM + ) + + +def _resolve_connection_string(): + from application.core.settings import settings + + conn = getattr(settings, "PGVECTOR_CONNECTION_STRING", None) + if not conn and getattr(settings, "POSTGRES_URI", None): + from application.core.db_uri import normalize_pgvector_connection_string + + conn = normalize_pgvector_connection_string(settings.POSTGRES_URI) + return conn + + +_GRAPH_TABLES = ( + "graph_node_chunks", + "graph_edges", + "graph_nodes", + "graph_ingest_progress", +) + + +def _drop_graph_tables(conn_string): + import psycopg + + with psycopg.connect(conn_string) as conn: + with conn.cursor() as cursor: + cursor.execute( + f"DROP TABLE IF EXISTS {', '.join(_GRAPH_TABLES)} CASCADE;" + ) + conn.commit() + + +def _live_store(): + conn = _resolve_connection_string() + if not conn: + pytest.skip("No pgvector connection string configured") + try: + _drop_graph_tables(conn) + return GraphStore(connection_string=conn) + except Exception as exc: + pytest.skip(f"pgvector DB not reachable: {exc}") + + +def _embedding(seed: float) -> list: + vec = [0.0] * TEST_EMBEDDING_DIM + vec[0] = seed + return vec + + +@pytest.mark.integration +class TestGraphStoreLive: + @pytest.fixture + def store(self): + store = _live_store() + yield store + _drop_graph_tables(store._connection_string) + + @pytest.fixture + def source_id(self): + return str(uuid.uuid4()) + + def test_ensure_tables_idempotent(self, store): + store._ensure_tables() + store._ensure_tables() + + def test_upsert_node_merges_by_normalized_name(self, store, source_id): + try: + first = store.upsert_node( + source_id=source_id, + name="Ada Lovelace", + normalized_name="ada lovelace", + type="person", + description="A mathematician.", + name_embedding=_embedding(1.0), + ) + second = store.upsert_node( + source_id=source_id, + name="Ada Lovelace", + normalized_name="ada lovelace", + type="person", + description="Wrote the first algorithm.", + ) + assert first == second + + node = store.get_node_by_normalized(source_id, "ada lovelace") + assert node is not None + assert node["id"] == first + assert node["doc_freq"] == 2 + assert "mathematician" in node["description"] + assert "first algorithm" in node["description"] + + duplicate = store.upsert_node( + source_id=source_id, + name="Ada Lovelace", + normalized_name="ada lovelace", + description="Wrote the first algorithm.", + ) + assert duplicate == first + node = store.get_node_by_normalized(source_id, "ada lovelace") + assert node["description"].count("first algorithm") == 1 + + assert store.count_nodes(source_id) == 1 + finally: + store.delete_by_source(source_id) + + def test_add_edge_and_link_chunk(self, store, source_id): + try: + a = store.upsert_node(source_id, "A", "a", "thing", "desc a") + b = store.upsert_node(source_id, "B", "b", "thing", "desc b") + store.add_edge( + source_id, a, b, "related", "a relates to b", 2.0, ["chunk-1"] + ) + store.link_node_chunk(source_id, a, "chunk-1") + store.link_node_chunk(source_id, a, "chunk-1") + store.link_node_chunk(source_id, b, "chunk-1") + + mapping = store.get_chunk_ids_for_nodes(source_id, [a, b]) + assert mapping[a] == ["chunk-1"] + assert mapping[b] == ["chunk-1"] + + store.set_node_degrees(source_id) + node_a = store.get_node_by_normalized(source_id, "a") + assert node_a["degree"] == 1 + finally: + store.delete_by_source(source_id) + + def test_self_loop_degree_agrees_across_paths(self, store, source_id): + """``add_edge``'s incremental +1 and ``set_node_degrees`` recompute must + agree on a self-loop (count it once).""" + try: + node = store.upsert_node(source_id, "Solo", "solo") + store.add_edge(source_id, node, node, "self") + + incremental = store.get_node_by_normalized(source_id, "solo")["degree"] + assert incremental == 1 + + store.set_node_degrees(source_id) + recomputed = store.get_node_by_normalized(source_id, "solo")["degree"] + assert recomputed == incremental == 1 + finally: + store.delete_by_source(source_id) + + def test_search_nodes_by_embedding(self, store, source_id): + try: + near = store.upsert_node( + source_id, "Near", "near", "thing", "d", _embedding(1.0) + ) + store.upsert_node( + source_id, "Far", "far", "thing", "d", _embedding(-1.0) + ) + results = store.search_nodes_by_embedding(source_id, _embedding(1.0), k=2) + assert len(results) == 2 + assert results[0]["id"] == near + assert results[0]["distance"] <= results[1]["distance"] + finally: + store.delete_by_source(source_id) + + def test_get_subgraph_bounded(self, store, source_id): + try: + a = store.upsert_node(source_id, "A", "a") + b = store.upsert_node(source_id, "B", "b") + c = store.upsert_node(source_id, "C", "c") + store.add_edge(source_id, a, b, "rel") + store.add_edge(source_id, b, c, "rel") + + one_hop = store.get_subgraph(source_id, [a], hops=1) + node_ids = {n["id"] for n in one_hop["nodes"]} + assert a in node_ids and b in node_ids + assert c not in node_ids + + two_hop = store.get_subgraph(source_id, [a], hops=2) + node_ids = {n["id"] for n in two_hop["nodes"]} + assert {a, b, c} <= node_ids + assert len(two_hop["edges"]) >= 2 + finally: + store.delete_by_source(source_id) + + def test_get_subgraph_frontier_truncation_is_deterministic( + self, store, source_id, monkeypatch + ): + """Bounded expansion must pick the same neighbors run-to-run so PPR (G5) + is reproducible.""" + try: + hub = store.upsert_node(source_id, "Hub", "hub") + leaves = [] + for i in range(6): + leaf = store.upsert_node(source_id, f"L{i}", f"l{i}") + store.add_edge(source_id, hub, leaf, "rel") + leaves.append(leaf) + + monkeypatch.setattr(store_module, "MAX_SUBGRAPH_NODES", 4) + + first = {n["id"] for n in store.get_subgraph(source_id, [hub])["nodes"]} + second = {n["id"] for n in store.get_subgraph(source_id, [hub])["nodes"]} + assert first == second + assert len(first) == 4 + + kept_leaves = sorted(leaves)[:3] + assert first == {hub, *kept_leaves} + finally: + store.delete_by_source(source_id) + + def test_checkpoint_pending_and_mark(self, store, source_id): + try: + all_chunks = ["c1", "c2", "c3"] + assert store.pending_chunks(source_id, all_chunks) == all_chunks + + store.mark_chunk(source_id, "c1", "done") + store.mark_chunk(source_id, "c2", "pending") + assert store.pending_chunks(source_id, all_chunks) == ["c2", "c3"] + + store.mark_chunk(source_id, "c2", "done") + assert store.pending_chunks(source_id, all_chunks) == ["c3"] + + progress = store.get_progress(source_id) + assert progress["c1"] == "done" + assert progress["c2"] == "done" + finally: + store.delete_by_source(source_id) + + def test_delete_by_source_isolation(self, store): + keep = str(uuid.uuid4()) + drop = str(uuid.uuid4()) + try: + k = store.upsert_node(keep, "K", "k") + d = store.upsert_node(drop, "D", "d") + store.add_edge(keep, k, k, "self") + store.add_edge(drop, d, d, "self") + store.link_node_chunk(keep, k, "kc") + store.link_node_chunk(drop, d, "dc") + store.mark_chunk(keep, "kc", "done") + store.mark_chunk(drop, "dc", "done") + + store.delete_by_source(drop) + + assert store.count_nodes(drop) == 0 + assert store.get_progress(drop) == {} + assert store.count_nodes(keep) == 1 + assert store.get_progress(keep) == {"kc": "done"} + finally: + store.delete_by_source(keep) + store.delete_by_source(drop) + + +@pytest.mark.unit +class TestGraphStoreParameterization: + """Asserts SQL is parameterized without touching a real DB.""" + + def _store_with_mock_conn(self): + store = GraphStore.__new__(GraphStore) + cursor = MagicMock() + cursor.fetchone.return_value = [str(uuid.uuid4())] + cursor.fetchall.return_value = [] + conn = MagicMock() + conn.cursor.return_value = cursor + store._connection = conn + store._get_connection = lambda: conn + return store, cursor + + def test_delete_by_source_binds_source_id(self): + store, cursor = self._store_with_mock_conn() + sid = str(uuid.uuid4()) + store.delete_by_source(sid) + + for call in cursor.execute.call_args_list: + sql = call.args[0] + params = call.args[1] if len(call.args) > 1 else None + assert "WHERE source_id = %s" in sql + assert sid not in sql + assert params == (sid,) + + def test_search_binds_embedding_and_source(self): + store, cursor = self._store_with_mock_conn() + sid = str(uuid.uuid4()) + embedding = _embedding(0.5) + store.search_nodes_by_embedding(sid, embedding, k=5) + + sql, params = cursor.execute.call_args.args[0], cursor.execute.call_args.args[1] + assert "%s::vector" in sql + assert "source_id = %s" in sql + assert sid not in sql + assert str(embedding) not in sql + assert params == (embedding, sid, embedding, 5) + + def test_upsert_node_binds_all_values(self): + store, cursor = self._store_with_mock_conn() + sid = str(uuid.uuid4()) + embedding = _embedding(0.1) + store.upsert_node(sid, "Name", "name", "type", "desc", embedding) + + sql, params = cursor.execute.call_args.args[0], cursor.execute.call_args.args[1] + assert "ON CONFLICT (source_id, normalized_name) DO UPDATE" in sql + assert sid not in sql + assert "name" not in [t for t in sql.split() if t == sid] + assert params[1] == sid + assert params[-1] == embedding + + +@pytest.mark.unit +class TestEmbeddingDim: + """The graph table dimension is derived from the configured model (FIX 1).""" + + def test_uses_configured_model_dimension(self, monkeypatch): + from application.vectorstore import base as base_module + + fake_embedding = MagicMock() + fake_embedding.dimension = 1536 + monkeypatch.setattr( + base_module.EmbeddingsSingleton, + "get_instance", + staticmethod(lambda *a, **k: fake_embedding), + ) + monkeypatch.setattr(GraphStore, "_embedding_dim", _REAL_EMBEDDING_DIM) + + store = GraphStore.__new__(GraphStore) + assert store._embedding_dim() == 1536 + + def test_falls_back_to_default_dimension(self, monkeypatch): + from application.vectorstore import base as base_module + + fake_embedding = object() + monkeypatch.setattr( + base_module.EmbeddingsSingleton, + "get_instance", + staticmethod(lambda *a, **k: fake_embedding), + ) + monkeypatch.setattr(GraphStore, "_embedding_dim", _REAL_EMBEDDING_DIM) + + store = GraphStore.__new__(GraphStore) + assert store._embedding_dim() == store_module.DEFAULT_NAME_EMBEDDING_DIM