diff --git a/application/graphrag/store.py b/application/graphrag/store.py index b1c6c456..3da835a6 100644 --- a/application/graphrag/store.py +++ b/application/graphrag/store.py @@ -24,6 +24,41 @@ DEFAULT_NAME_EMBEDDING_DIM = 768 MAX_SUBGRAPH_NODES = 500 MAX_SUBGRAPH_EDGES = 2000 +PGVECTOR_SOURCE_COLUMN = "source_id" + + +def _safe_identifier(name: str) -> str: + """Return ``name`` if it is a bare SQL identifier, else raise. + + Guards the interpolated table/column names against injection; pgvector uses + plain identifiers, so anything outside ``[A-Za-z_][A-Za-z0-9_]*`` is rejected. + """ + if not isinstance(name, str) or not name.isidentifier(): + raise ValueError(f"Unsafe SQL identifier: {name!r}") + return name + + +def _pgvector_identifiers() -> tuple[str, str, str, str]: + """Resolve ``(table, text_col, metadata_col, source_col)`` from ``PGVectorStore``. + + Reads the table and column defaults from ``PGVectorStore.__init__`` so the + graph store queries the same names a customized deployment configured. + """ + import inspect + + from application.vectorstore.pgvector import PGVectorStore + + params = inspect.signature(PGVectorStore.__init__).parameters + table = params["table_name"].default + text_col = params["text_column"].default + metadata_col = params["metadata_column"].default + return ( + _safe_identifier(table), + _safe_identifier(text_col), + _safe_identifier(metadata_col), + _safe_identifier(PGVECTOR_SOURCE_COLUMN), + ) + class GraphStore: """Stores and queries a per-source knowledge graph in the pgvector DB.""" @@ -511,6 +546,42 @@ class GraphStore: finally: cursor.close() + def get_chunk_texts( + self, + source_id: str, + chunk_ids: List[str], + ) -> Dict[str, Dict[str, Any]]: + """Map chunk ids to ``{"text": ..., "metadata": {...}}`` from the pgvector table. + + Reads the co-located documents table, deriving its name and the text, + metadata and source-id column names from the same defaults + ``PGVectorStore`` uses so a customized deployment still resolves. Chunk + ids are pgvector document ids (SERIAL) cast to text to match the + JSONB-sourced string ids without per-id round trips. + """ + if not chunk_ids: + return {} + table, text_col, metadata_col, source_col = _pgvector_identifiers() + conn = self._get_connection() + cursor = conn.cursor() + try: + cursor.execute( + f""" + SELECT id, {text_col}, {metadata_col} FROM {table} + WHERE {source_col} = %s AND id::text = ANY(%s); + """, + (source_id, [str(c) for c in chunk_ids]), + ) + return { + str(row[0]): {"text": row[1], "metadata": row[2] or {}} + for row in cursor.fetchall() + } + except Exception as e: + logging.error(f"Error getting chunk texts: {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. diff --git a/application/retriever/graph_rag.py b/application/retriever/graph_rag.py new file mode 100644 index 00000000..aeb0dc35 --- /dev/null +++ b/application/retriever/graph_rag.py @@ -0,0 +1,255 @@ +"""GraphRAG local retriever — Personalized PageRank over a per-source graph (D27/D31). + +Local mode only (v1): rephrased query -> entity-name NN seeds -> bounded 1-2-hop +fetch -> networkx Personalized PageRank (IDF-down-weighted hubs) -> chunks ranked +by landed PPR mass -> shared token budget. No LLM call at query time beyond the +(optional, reused) rephrase. + +Composes :class:`ClassicRAG` rather than subclassing: PPR doesn't fit the +``_fetch_candidates`` hook, but the composed instance supplies the rephrase, the +token-budget loop, and the per-source fallback when a source has no graph (D31). +""" + +from __future__ import annotations + +import logging +import math +import os +from typing import Any, Dict, List + +import networkx as nx + +from application.core.settings import settings +from application.graphrag import graphrag_available +from application.graphrag.store import GraphStore +from application.retriever.base import BaseRetriever +from application.retriever.classic_rag import ClassicRAG +from application.utils import num_tokens_from_string +from application.vectorstore.base import EmbeddingsSingleton + +SEED_NODES = 10 +SUBGRAPH_HOPS = 1 + + +def _idf(doc_freq: Any) -> float: + """Node-specificity weight: rarer entities (low ``doc_freq``) score higher.""" + return 1.0 / math.log(1.0 + max(int(doc_freq or 0), 0) + 1.0) + + +def _labels_from_metadata( + metadata: Dict[str, Any], text: str, source_id: str +) -> Dict[str, str]: + """Derive ``title``/``source``/``filename`` from chunk metadata as ClassicRAG does.""" + metadata = metadata or {} + title = metadata.get("title", metadata.get("post_title", text)) + if not isinstance(title, str): + title = str(title) + title = title.split("/")[-1] + + filename = ( + metadata.get("filename") + or metadata.get("file_name") + or metadata.get("source") + ) + if isinstance(filename, str): + filename = os.path.basename(filename) or filename + else: + filename = title + if not filename: + filename = title + source_path = metadata.get("source") or source_id + return {"title": title, "source": source_path, "filename": filename} + + +class GraphRAGRetriever(BaseRetriever): + """Per-source PPR retriever; falls back to ClassicRAG when a source has no graph.""" + + def __init__( + self, + source, + chat_history=None, + prompt="", + chunks=2, + doc_token_limit=50000, + model_id="docsgpt-local", + user_api_key=None, + agent_id=None, + llm_name=settings.LLM_PROVIDER, + api_key=settings.API_KEY, + decoded_token=None, + model_user_id=None, + defer_rephrase=False, + request_id=None, + ): + self._classic = ClassicRAG( + source=source, + chat_history=chat_history, + prompt=prompt, + chunks=chunks, + doc_token_limit=doc_token_limit, + model_id=model_id, + user_api_key=user_api_key, + agent_id=agent_id, + llm_name=llm_name, + api_key=api_key, + decoded_token=decoded_token, + model_user_id=model_user_id, + defer_rephrase=defer_rephrase, + request_id=request_id, + ) + self.original_question = self._classic.original_question + self.chunks = self._classic.chunks + self.doc_token_limit = doc_token_limit + self.vectorstores = self._classic.vectorstores + self.per_source_retrieval = {} + + def _embed_query(self, question: str) -> List[float]: + embedding = EmbeddingsSingleton.get_instance( + settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY + ) + return embedding.embed_query(question) + + def _ppr_scores(self, subgraph, seeds) -> Dict[str, float]: + """Run Personalized PageRank, then down-weight hub nodes by IDF. + + ``seeds`` maps seed node id -> personalization weight (seed similarity). + After PPR, each node's mass is scaled by ``1/log(2 + doc_freq)`` so a + high-degree hub contributes less than a specific entity at equal mass. + """ + graph = nx.Graph() + for node in subgraph.get("nodes", []): + graph.add_node(node["id"], doc_freq=node.get("doc_freq", 0)) + for edge in subgraph.get("edges", []): + src, dst = edge["src_node_id"], edge["dst_node_id"] + if src in graph and dst in graph: + weight = float(edge.get("weight") or 1.0) + graph.add_edge(src, dst, weight=weight) + if graph.number_of_nodes() == 0: + return {} + + personalization = {n: seeds.get(n, 0.0) for n in graph.nodes} + if not any(personalization.values()): + personalization = None + + ranks = nx.pagerank(graph, personalization=personalization, weight="weight") + return { + node: rank * _idf(graph.nodes[node].get("doc_freq", 0)) + for node, rank in ranks.items() + } + + def _rank_chunks(self, store, source_id, node_scores) -> List[str]: + """Score chunks by summed (PPR mass x IDF) of their linked nodes; top candidates. + + Over-fetches beyond ``self.chunks`` so chunks with missing text don't drop + the final count below the budget; the budget loop caps the real total. + """ + node_ids = list(node_scores.keys()) + chunk_links = store.get_chunk_ids_for_nodes(source_id, node_ids) + chunk_scores: Dict[str, float] = {} + for node_id, chunk_ids in chunk_links.items(): + node_score = node_scores.get(node_id, 0.0) + for chunk_id in chunk_ids: + chunk_scores[chunk_id] = chunk_scores.get(chunk_id, 0.0) + node_score + ranked = sorted(chunk_scores, key=lambda c: chunk_scores[c], reverse=True) + candidates = max(self.chunks * 2, self.chunks + 5) + return ranked[: max(1, candidates)] + + def _graph_docs_for_source(self, store, source_id) -> List[Dict[str, Any]]: + """Local PPR retrieval for one source (caller guarantees it has a graph).""" + question = self._classic._get_rephrased_question() + query_embedding = self._embed_query(question) + seed_rows = store.search_nodes_by_embedding( + source_id, query_embedding, k=SEED_NODES + ) + if not seed_rows: + return [] + + seed_ids = [row["id"] for row in seed_rows] + seeds = {row["id"]: 1.0 - float(row.get("distance") or 0.0) for row in seed_rows} + + subgraph = store.get_subgraph(source_id, seed_ids, hops=SUBGRAPH_HOPS) + node_scores = self._ppr_scores(subgraph, seeds) + if not node_scores: + return [] + + chunk_ids = self._rank_chunks(store, source_id, node_scores) + chunk_data = store.get_chunk_texts(source_id, chunk_ids) + + docs: List[Dict[str, Any]] = [] + token_budget = max(int(self.doc_token_limit * 0.9), 100) + cumulative_tokens = 0 + for chunk_id in chunk_ids: + if len(docs) >= self.chunks: + break + chunk = chunk_data.get(chunk_id) + text = chunk.get("text") if chunk else None + if not text: + continue + labels = _labels_from_metadata(chunk.get("metadata"), text, source_id) + doc_tokens = num_tokens_from_string(f"{labels['filename']}\n{text}") + if cumulative_tokens + doc_tokens >= token_budget: + break + docs.append({"text": text, **labels}) + cumulative_tokens += doc_tokens + return docs + + def _classic_for_source(self, source_id) -> List[Dict[str, Any]]: + """Reuse the composed ClassicRAG to retrieve one source's chunks (D31).""" + original = self._classic.vectorstores + original_overrides = self._classic.per_source_retrieval + try: + self._classic.vectorstores = [source_id] + self._classic.per_source_retrieval = { + k: v for k, v in self.per_source_retrieval.items() if k == source_id + } + return self._classic._get_data() + finally: + self._classic.vectorstores = original + self._classic.per_source_retrieval = original_overrides + + def _get_data(self) -> List[Dict[str, Any]]: + if not self.vectorstores: + return [] + + store = None + if graphrag_available(): + try: + store = GraphStore() + except Exception as e: + logging.error(f"GraphRAG store unavailable, falling back: {e}") + store = None + + all_docs: List[Dict[str, Any]] = [] + for source_id in self.vectorstores: + if not source_id: + continue + has_graph = False + if store is not None: + try: + has_graph = store.count_nodes(source_id) > 0 + except Exception as e: + logging.error(f"GraphRAG count_nodes failed for {source_id}: {e}") + has_graph = False + + if not has_graph: + all_docs.extend(self._classic_for_source(source_id)) + continue + + try: + all_docs.extend(self._graph_docs_for_source(store, source_id)) + except Exception as e: + logging.error( + f"GraphRAG retrieval failed for {source_id}, falling back: {e}", + exc_info=True, + ) + all_docs.extend(self._classic_for_source(source_id)) + return all_docs + + def search(self, query: str = "") -> List[Dict[str, Any]]: + if query: + self.original_question = query + self._classic.original_question = query + self._classic._rephrased_question = None + self._classic.question = self._classic._rephrase_query() + self._classic._rephrased_question = self._classic.question + return self._get_data() diff --git a/application/retriever/retriever_creator.py b/application/retriever/retriever_creator.py index 43d75035..7588fb41 100644 --- a/application/retriever/retriever_creator.py +++ b/application/retriever/retriever_creator.py @@ -1,4 +1,5 @@ from application.retriever.classic_rag import ClassicRAG +from application.retriever.graph_rag import GraphRAGRetriever from application.retriever.hybrid_rag import HybridRetriever @@ -7,6 +8,7 @@ class RetrieverCreator: "classic": ClassicRAG, "default": ClassicRAG, "hybrid": HybridRetriever, + "graphrag": GraphRAGRetriever, } @classmethod diff --git a/tests/retriever/test_graph_rag.py b/tests/retriever/test_graph_rag.py new file mode 100644 index 00000000..ad2a828c --- /dev/null +++ b/tests/retriever/test_graph_rag.py @@ -0,0 +1,390 @@ +"""Tests for the GraphRAG local PPR retriever. + +The GraphStore and embeddings are mocked (no DB, no model load); ``networkx`` +runs for real on small crafted graphs. The composed ClassicRAG is mocked when +exercising the fallback path. +""" + +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from application.retriever.graph_rag import GraphRAGRetriever +from application.retriever.retriever_creator import RetrieverCreator + + +@pytest.fixture +def _patch_llm_creator(mock_llm, monkeypatch): + monkeypatch.setattr( + "application.retriever.classic_rag.LLMCreator.create_llm", + Mock(return_value=mock_llm), + ) + return mock_llm + + +def _make_retriever(source=None, **overrides): + defaults = dict( + source=source or {"question": "q", "active_docs": ["src1"]}, + chat_history=None, + prompt="", + chunks=2, + doc_token_limit=50000, + model_id="test-model", + llm_name="openai", + api_key="fake", + decoded_token={"sub": "user1"}, + ) + defaults.update(overrides) + return GraphRAGRetriever(**defaults) + + +@pytest.fixture +def _patch_embed(monkeypatch): + monkeypatch.setattr( + GraphRAGRetriever, "_embed_query", lambda self, q: [0.1, 0.2, 0.3] + ) + + +# ── Fallback to ClassicRAG ──────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestGraphRAGFallback: + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_no_graph_delegates_to_classic( + self, _avail, mock_store_cls, _patch_llm_creator + ): + store = MagicMock() + store.count_nodes.return_value = 0 + mock_store_cls.return_value = store + + rag = _make_retriever() + classic_docs = [{"title": "c", "text": "classic", "source": "src1", "filename": "c"}] + rag._classic._get_data = Mock(return_value=list(classic_docs)) + + docs = rag._get_data() + + assert docs == classic_docs + store.search_nodes_by_embedding.assert_not_called() + store.get_subgraph.assert_not_called() + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=False) + def test_graphrag_unavailable_delegates_to_classic( + self, _avail, mock_store_cls, _patch_llm_creator + ): + rag = _make_retriever() + classic_docs = [{"title": "c", "text": "classic", "source": "src1", "filename": "c"}] + rag._classic._get_data = Mock(return_value=list(classic_docs)) + + docs = rag._get_data() + + assert docs == classic_docs + mock_store_cls.assert_not_called() + + +# ── Happy path: seed -> subgraph -> PPR -> rank ─────────────────────────────── + + +def _as_chunk_data(chunk_texts, metadata_by_chunk=None): + """Wrap plain ``{chunk_id: text}`` into the richer get_chunk_texts shape.""" + metadata_by_chunk = metadata_by_chunk or {} + return { + chunk_id: {"text": text, "metadata": metadata_by_chunk.get(chunk_id, {})} + for chunk_id, text in chunk_texts.items() + } + + +def _store_with_graph( + nodes, edges, node_chunks, chunk_texts, seed_rows, metadata_by_chunk=None +): + store = MagicMock() + store.count_nodes.return_value = len(nodes) + store.search_nodes_by_embedding.return_value = seed_rows + store.get_subgraph.return_value = {"nodes": nodes, "edges": edges} + store.get_chunk_ids_for_nodes.return_value = node_chunks + store.get_chunk_texts.return_value = _as_chunk_data(chunk_texts, metadata_by_chunk) + return store + + +@pytest.mark.unit +class TestGraphRAGHappyPath: + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_ppr_ranks_near_seed_higher( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + # Chain: seed(n1) - n2 - n3. Personalization on n1 biases the walk toward + # the seed neighborhood, so the far node n3 lands the least PPR mass and + # must rank below the seed and its direct neighbor. + nodes = [ + {"id": "n1", "doc_freq": 1}, + {"id": "n2", "doc_freq": 1}, + {"id": "n3", "doc_freq": 1}, + ] + edges = [ + {"src_node_id": "n1", "dst_node_id": "n2", "weight": 1.0}, + {"src_node_id": "n2", "dst_node_id": "n3", "weight": 1.0}, + ] + node_chunks = {"n1": ["c1"], "n2": ["c2"], "n3": ["c3"]} + chunk_texts = {"c1": "near", "c2": "mid", "c3": "far"} + seed_rows = [{"id": "n1", "distance": 0.0}] + store = _store_with_graph(nodes, edges, node_chunks, chunk_texts, seed_rows) + mock_store_cls.return_value = store + + rag = _make_retriever(chunks=3) + docs = rag._get_data() + + texts = [d["text"] for d in docs] + assert texts[-1] == "far" + assert texts.index("near") < texts.index("far") + assert docs[0].keys() == {"title", "text", "source", "filename"} + + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_topk_respected( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + nodes = [{"id": f"n{i}", "doc_freq": 1} for i in range(1, 5)] + edges = [ + {"src_node_id": "n1", "dst_node_id": "n2", "weight": 1.0}, + {"src_node_id": "n1", "dst_node_id": "n3", "weight": 1.0}, + {"src_node_id": "n1", "dst_node_id": "n4", "weight": 1.0}, + ] + node_chunks = {f"n{i}": [f"c{i}"] for i in range(1, 5)} + chunk_texts = {f"c{i}": f"t{i}" for i in range(1, 5)} + seed_rows = [{"id": "n1", "distance": 0.0}] + store = _store_with_graph(nodes, edges, node_chunks, chunk_texts, seed_rows) + mock_store_cls.return_value = store + + rag = _make_retriever(chunks=2) + docs = rag._get_data() + + assert len(docs) == 2 + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_token_budget_honored( + self, _avail, mock_store_cls, _patch_llm_creator, _patch_embed + ): + nodes = [{"id": f"n{i}", "doc_freq": 1} for i in range(1, 4)] + edges = [ + {"src_node_id": "n1", "dst_node_id": "n2", "weight": 1.0}, + {"src_node_id": "n2", "dst_node_id": "n3", "weight": 1.0}, + ] + node_chunks = {f"n{i}": [f"c{i}"] for i in range(1, 4)} + chunk_texts = {f"c{i}": f"t{i}" for i in range(1, 4)} + seed_rows = [{"id": "n1", "distance": 0.0}] + store = _store_with_graph(nodes, edges, node_chunks, chunk_texts, seed_rows) + mock_store_cls.return_value = store + + # Tiny budget: 0.9 * 100 = 90; each chunk costs 50 tokens → only one fits. + rag = _make_retriever(chunks=3, doc_token_limit=100) + with patch( + "application.retriever.graph_rag.num_tokens_from_string", return_value=50 + ): + docs = rag._get_data() + + assert len(docs) == 1 + + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_labels_derived_from_metadata_not_source_id( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + nodes = [{"id": "n1", "doc_freq": 1}] + edges = [] + node_chunks = {"n1": ["c1"]} + chunk_texts = {"c1": "near"} + metadata = {"c1": {"title": "My Title", "source": "/docs/report.pdf"}} + seed_rows = [{"id": "n1", "distance": 0.0}] + store = _store_with_graph( + nodes, edges, node_chunks, chunk_texts, seed_rows, metadata + ) + mock_store_cls.return_value = store + + rag = _make_retriever(chunks=1) + docs = rag._get_data() + + assert len(docs) == 1 + doc = docs[0] + assert doc["title"] == "My Title" + assert doc["filename"] == "report.pdf" + assert doc["source"] == "/docs/report.pdf" + assert "src1" not in (doc["title"], doc["filename"]) + + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_overfetch_fills_when_some_text_missing( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + # n2 ranks above n3 but its chunk text is missing; over-fetching past + # ``chunks`` lets c3 fill the gap so the result still reaches ``chunks``. + nodes = [{"id": f"n{i}", "doc_freq": 1} for i in range(1, 4)] + edges = [ + {"src_node_id": "n1", "dst_node_id": "n2", "weight": 2.0}, + {"src_node_id": "n2", "dst_node_id": "n3", "weight": 1.0}, + ] + node_chunks = {"n1": ["c1"], "n2": ["c2"], "n3": ["c3"]} + chunk_texts = {"c1": "first", "c3": "third"} # c2 missing + seed_rows = [{"id": "n1", "distance": 0.0}] + store = _store_with_graph(nodes, edges, node_chunks, chunk_texts, seed_rows) + mock_store_cls.return_value = store + + rag = _make_retriever(chunks=2) + docs = rag._get_data() + + texts = [d["text"] for d in docs] + assert len(docs) == 2 + assert texts == ["first", "third"] + + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_no_seeds_returns_empty( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + store = _store_with_graph([], [], {}, {}, []) + store.count_nodes.return_value = 5 + mock_store_cls.return_value = store + + rag = _make_retriever() + assert rag._get_data() == [] + + +# ── IDF down-weighting ──────────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestGraphRAGIdf: + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_hub_downweighted_below_specific_node( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + # Star: seed n1 links a hub node (huge doc_freq) and a specific node + # (doc_freq=1). PPR mass is symmetric across the two leaves, so only IDF + # can break the tie — the specific node must rank above the hub. + nodes = [ + {"id": "n1", "doc_freq": 1}, + {"id": "hub", "doc_freq": 100000}, + {"id": "specific", "doc_freq": 1}, + ] + edges = [ + {"src_node_id": "n1", "dst_node_id": "hub", "weight": 1.0}, + {"src_node_id": "n1", "dst_node_id": "specific", "weight": 1.0}, + ] + node_chunks = {"hub": ["c_hub"], "specific": ["c_spec"]} + chunk_texts = {"c_hub": "hub_text", "c_spec": "spec_text"} + seed_rows = [{"id": "n1", "distance": 0.0}] + store = _store_with_graph(nodes, edges, node_chunks, chunk_texts, seed_rows) + mock_store_cls.return_value = store + + rag = _make_retriever(chunks=2) + docs = rag._get_data() + texts = [d["text"] for d in docs] + + assert texts.index("spec_text") < texts.index("hub_text") + + @pytest.mark.unit + def test_idf_helper_monotonic(self): + from application.retriever.graph_rag import _idf + + assert _idf(1) > _idf(10) > _idf(1000) + + +# ── Registry resolution ────────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestGraphRAGRegistration: + def test_graphrag_resolves_via_creator(self): + assert RetrieverCreator.retrievers["graphrag"] is GraphRAGRetriever + + def test_create_retriever_builds_graphrag(self, _patch_llm_creator): + retriever = RetrieverCreator.create_retriever( + "graphrag", + source={"question": "q", "active_docs": ["src1"]}, + chunks=2, + doc_token_limit=50000, + model_id="m", + llm_name="openai", + api_key="fake", + decoded_token={"sub": "u"}, + ) + assert isinstance(retriever, GraphRAGRetriever) + + +# ── get_chunk_texts parameterization ───────────────────────────────────────── + + +@pytest.mark.unit +class TestGetChunkTexts: + def _store_with_mock_conn(self): + from application.graphrag.store import GraphStore + + store = GraphStore.__new__(GraphStore) + cursor = MagicMock() + cursor.fetchall.return_value = [ + (1, "alpha", {"filename": "a.pdf"}), + (2, "beta", None), + ] + conn = MagicMock() + conn.cursor.return_value = cursor + store._connection = conn + store._get_connection = lambda: conn + return store, cursor + + def test_returns_text_and_metadata_shape(self): + import uuid + + store, cursor = self._store_with_mock_conn() + sid = str(uuid.uuid4()) + result = store.get_chunk_texts(sid, ["1", "2"]) + + assert result == { + "1": {"text": "alpha", "metadata": {"filename": "a.pdf"}}, + "2": {"text": "beta", "metadata": {}}, + } + + def test_uses_configured_identifiers_and_binds_params(self): + import uuid + + from application.graphrag.store import _pgvector_identifiers + + table, text_col, metadata_col, source_col = _pgvector_identifiers() + store, cursor = self._store_with_mock_conn() + sid = str(uuid.uuid4()) + store.get_chunk_texts(sid, ["1", "2"]) + + sql, params = cursor.execute.call_args.args[0], cursor.execute.call_args.args[1] + assert f"FROM {table}" in sql + assert text_col in sql + assert metadata_col in sql + assert f"{source_col} = %s" in sql + assert "id::text = ANY(%s)" in sql + assert sid not in sql + assert params == (sid, ["1", "2"]) + + def test_identifiers_match_pgvector_defaults(self): + from application.graphrag.store import _pgvector_identifiers + from application.vectorstore.pgvector import PGVectorStore + import inspect + + params = inspect.signature(PGVectorStore.__init__).parameters + table, text_col, metadata_col, source_col = _pgvector_identifiers() + assert table == params["table_name"].default + assert text_col == params["text_column"].default + assert metadata_col == params["metadata_column"].default + assert source_col == "source_id" + + def test_empty_chunk_ids_short_circuits(self): + store, cursor = self._store_with_mock_conn() + assert store.get_chunk_texts("sid", []) == {} + cursor.execute.assert_not_called()