wehub-resource-sync
fed8b2eed7
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
419 lines
17 KiB
Python
419 lines
17 KiB
Python
"""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_seed_distance_over_one_is_clamped(
|
|
self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed
|
|
):
|
|
# One seed at cosine distance > 1 (negative similarity) => raw weight
|
|
# 1 - 1.5 < 0. Paired with a positive seed the personalization sums to
|
|
# ~0, which makes networkx pagerank raise ZeroDivisionError. Clamping
|
|
# each weight to >= 0 keeps the personalization a valid distribution.
|
|
nodes = [{"id": "n1", "doc_freq": 1}, {"id": "n2", "doc_freq": 1}]
|
|
edges = [{"src_node_id": "n1", "dst_node_id": "n2", "weight": 1.0}]
|
|
node_chunks = {"n1": ["c1"], "n2": ["c2"]}
|
|
chunk_texts = {"c1": "a", "c2": "b"}
|
|
seed_rows = [
|
|
{"id": "n1", "distance": 0.5},
|
|
{"id": "n2", "distance": 1.5},
|
|
]
|
|
store = _store_with_graph(nodes, edges, node_chunks, chunk_texts, seed_rows)
|
|
mock_store_cls.return_value = store
|
|
|
|
rag = _make_retriever(chunks=2)
|
|
# Call the PPR path directly: _get_data would swallow a raise and fall
|
|
# back to ClassicRAG, hiding the regression.
|
|
docs = rag._graph_docs_for_source(store, "src1")
|
|
|
|
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_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()
|