mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 14:15:52 +00:00
feat(graphrag): GraphRAGRetriever (PPR local) + register + ClassicRAG fallback
Query-side: composes ClassicRAG (rephrase/budget/fallback). Per source: pgvector entity-name seed -> bounded subgraph -> networkx personalized PageRank (IDF hub down-weight) -> rank graph_node_chunks -> chunk text -> shared token budget; no LLM at query time. Citations derived from chunk metadata (matches ClassicRAG). Falls back to ClassicRAG per-source when no graph rows / unavailable / error (retrieval never breaks). get_chunk_texts queries the co-located documents table by configured table/column names (parameterized). Registered 'graphrag'. Unit G5.
This commit is contained in:
1 parent
4742aec4c6
commit
4e0c792a6d
4 files changed
+718
No files matched your search
@@ -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.
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
Reference in new issue
Block a user