mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 20:15:30 +00:00
feat(graphrag): seed the walk from what entities are, and rank with passages and vector hits
Graph retrieval tied plain vector search at best and never beat it. Measured
across five corpora, the bottleneck was seeding, not the graph: the walk
started from nodes whose embeddings were computed from bare entity names, and
a whole question shares almost nothing with a name like "Quill".
Extraction now embeds each node from "name (type): description" and each
relationship as the fact it asserts ("Alder streams_to Quill: ..."), stored on
a new nullable graph_edges.fact_embedding column that ensure_vector_schema adds
in place. Entity names are canonicalised (case, punctuation, word breaks and a
cautious plural) so "VECTOR_STORE" and "vector stores" land on one node. Extraction calls run
concurrently (GRAPHRAG_EXTRACTION_WORKERS, default 8) while embedding and graph
writes stay serial on the task thread, so ordering and idempotency are
unchanged; that measured 8.4x faster with identical output.
Retrieval gains per-source options, stored under retrieval.graph and read live
at query time:
- seed_strategy: start from matching entities (default) or matching
relationships, which can reach an entity the question never names;
- passage_nodes (on): walk the source's passages alongside entities, with
PageRank damping 0.5 instead of 0.85;
- blend_vector (on): fuse the graph ranking with the source's vector ranking
by reciprocal rank.
The defaults are the measured-best configuration. Through GraphRAGRetriever,
the new seeding moved recall@4 from 0.41 to 0.68 on a multi-hop corpus and
from 0.50 to 1.00 on the docs corpus, and regressed none of the corpora
measured. Existing graphs keep name-only embeddings until rebuilt.
This commit is contained in:
1 parent
b7e7872bf7
commit
a83e1dc0af
13 files changed
+1577
-79
No files matched your search
@@ -0,0 +1,129 @@
|
||||
"""Where the graph walk starts, per the source's graph retrieval options.
|
||||
|
||||
Seeding decides more than ranking does — a walk that starts on the wrong nodes
|
||||
cannot be rescued downstream. The options are per source and live (no
|
||||
re-ingest), carried on the per-source retrieval config the Dispatcher hands the
|
||||
retriever, so both the dispatch and how the options are resolved are pinned.
|
||||
|
||||
The fallback matters most: relationship seeding reads fact embeddings written
|
||||
at ingest, and a source built before they were recorded has none. It must keep
|
||||
retrieving through entity matching rather than returning nothing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from docsgpt.retriever.graph_rag import GraphRAGRetriever
|
||||
from docsgpt.storage.db.source_config import GraphRetrievalConfig, RetrievalConfig
|
||||
|
||||
ENTITY_ROWS = [{"id": "n1", "name": "Quill", "distance": 0.2}]
|
||||
FACT_ROWS = [
|
||||
{"id": "f1", "name": "Alder", "distance": 0.1},
|
||||
{"id": "f2", "name": "Quill", "distance": 0.1},
|
||||
]
|
||||
|
||||
|
||||
class _StubStore:
|
||||
def __init__(self, fact_rows=None):
|
||||
self.fact_rows = list(FACT_ROWS) if fact_rows is None else fact_rows
|
||||
self.calls: list[str] = []
|
||||
|
||||
def seed_nodes_from_facts(self, source_id, query_embedding, fact_limit=5, limit=10):
|
||||
self.calls.append("facts")
|
||||
return self.fact_rows
|
||||
|
||||
def search_nodes_by_embedding(self, source_id, query_embedding, k=10):
|
||||
self.calls.append("entities")
|
||||
return list(ENTITY_ROWS)
|
||||
|
||||
|
||||
def _retriever(per_source=None):
|
||||
"""A retriever without its constructor, which builds a ClassicRAG."""
|
||||
retriever = object.__new__(GraphRAGRetriever)
|
||||
if per_source is not None:
|
||||
retriever.per_source_retrieval = per_source
|
||||
return retriever
|
||||
|
||||
|
||||
def _relationships_config():
|
||||
return RetrievalConfig(graph={"seed_strategy": "relationships"})
|
||||
|
||||
|
||||
class TestDefaults:
|
||||
def test_measured_best_configuration_is_the_default(self):
|
||||
options = GraphRetrievalConfig()
|
||||
|
||||
assert options.seed_strategy == "entities"
|
||||
assert options.passage_nodes is True
|
||||
assert options.blend_vector is True
|
||||
|
||||
def test_a_source_with_no_per_source_config_gets_the_defaults(self):
|
||||
assert _retriever()._graph_options("src") == GraphRetrievalConfig()
|
||||
|
||||
def test_existing_retrieval_configs_validate_without_the_new_block(self):
|
||||
"""Source configs saved before this existed carry no ``graph`` key."""
|
||||
config = RetrievalConfig.model_validate({"retriever": "graphrag"})
|
||||
|
||||
assert config.graph == GraphRetrievalConfig()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bad", [{"seed_strategy": "vector"}, {"seed_strategy": "union"}, {"damping": 0.5}]
|
||||
)
|
||||
def test_retired_and_unknown_options_are_rejected(self, bad):
|
||||
with pytest.raises(ValidationError):
|
||||
GraphRetrievalConfig.model_validate(bad)
|
||||
|
||||
|
||||
class TestEntitySeeding:
|
||||
def test_seeds_from_entities_and_never_reads_facts(self):
|
||||
store = _StubStore()
|
||||
|
||||
rows = _retriever()._seed_rows(store, "src", [0.0, 0.1])
|
||||
|
||||
assert [r["id"] for r in rows] == ["n1"]
|
||||
assert store.calls == ["entities"]
|
||||
|
||||
|
||||
class TestRelationshipSeeding:
|
||||
def test_seeds_from_facts_when_the_source_asks_for_it(self):
|
||||
store = _StubStore()
|
||||
retriever = _retriever({"src": _relationships_config()})
|
||||
|
||||
rows = retriever._seed_rows(store, "src", [0.0, 0.1])
|
||||
|
||||
assert [r["id"] for r in rows] == ["f1", "f2"]
|
||||
assert store.calls == ["facts"]
|
||||
|
||||
def test_reads_the_option_from_a_plain_dict_config_too(self):
|
||||
store = _StubStore()
|
||||
retriever = _retriever({"src": {"graph": {"seed_strategy": "relationships"}}})
|
||||
|
||||
retriever._seed_rows(store, "src", [0.0, 0.1])
|
||||
|
||||
assert store.calls == ["facts"]
|
||||
|
||||
def test_falls_back_to_entities_without_fact_embeddings(self):
|
||||
"""A source built before fact embeddings were recorded still retrieves."""
|
||||
store = _StubStore(fact_rows=[])
|
||||
retriever = _retriever({"src": _relationships_config()})
|
||||
|
||||
rows = retriever._seed_rows(store, "src", [0.0, 0.1])
|
||||
|
||||
assert [r["id"] for r in rows] == ["n1"]
|
||||
assert store.calls == ["facts", "entities"]
|
||||
|
||||
def test_options_are_per_source(self):
|
||||
store = _StubStore()
|
||||
retriever = _retriever({"other": _relationships_config()})
|
||||
|
||||
retriever._seed_rows(store, "src", [0.0, 0.1])
|
||||
|
||||
assert store.calls == ["entities"]
|
||||
|
||||
def test_a_malformed_stored_option_falls_back_to_the_defaults(self):
|
||||
"""A bad value must not take graph retrieval down with it."""
|
||||
retriever = _retriever({"src": {"graph": {"seed_strategy": "nonsense"}}})
|
||||
|
||||
assert retriever._graph_options("src") == GraphRetrievalConfig()
|
||||
Reference in new issue
Block a user