mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-05 12:13:55 +00:00
Alex
a83e1dc0af
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.
110 lines
3.7 KiB
Python
110 lines
3.7 KiB
Python
"""The graph retriever's default path, end to end through ``_graph_docs_for_source``.
|
|
|
|
The shipped defaults — seed from entities, walk the passages, blend with vector
|
|
search — are the configuration that measured best, so they are what most graph
|
|
sources run. This drives that whole path with a store that returns real values,
|
|
and checks each per-source option actually switches its stage off.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from docsgpt.retriever.graph_rag import GraphRAGRetriever
|
|
from docsgpt.storage.db.source_config import RetrievalConfig
|
|
|
|
TEXTS = {
|
|
"c-alder": "Alder streams audit events to Quill.",
|
|
"c-quill": "Quill is compacted every six hours.",
|
|
}
|
|
VECTOR_ONLY = "A passage only plain vector search found."
|
|
|
|
|
|
class _Store:
|
|
"""A two-entity chain: the question matches Alder, the answer is on Quill."""
|
|
|
|
def __init__(self):
|
|
self.calls: list[str] = []
|
|
|
|
def search_nodes_by_embedding(self, source_id, query_embedding, k=10):
|
|
return [{"id": "alder", "name": "Alder", "distance": 0.1}]
|
|
|
|
def get_subgraph(self, source_id, node_ids, hops=1):
|
|
return {
|
|
"nodes": [{"id": "alder", "doc_freq": 1}, {"id": "quill", "doc_freq": 1}],
|
|
"edges": [{"src_node_id": "alder", "dst_node_id": "quill", "weight": 1.0}],
|
|
}
|
|
|
|
def get_chunk_ids_for_nodes(self, source_id, node_ids):
|
|
return {"alder": ["c-alder"], "quill": ["c-quill"]}
|
|
|
|
def chunk_similarities(self, source_id, chunk_ids, query_embedding):
|
|
self.calls.append("chunk_similarities")
|
|
return {"c-alder": 0.9, "c-quill": 0.2}
|
|
|
|
def get_chunk_texts(self, source_id, chunk_ids):
|
|
return {
|
|
c: {"text": TEXTS[c], "metadata": {"title": c}}
|
|
for c in chunk_ids
|
|
if c in TEXTS
|
|
}
|
|
|
|
|
|
def _retriever(per_source=None):
|
|
"""A retriever without its constructor (which builds a ClassicRAG)."""
|
|
retriever = object.__new__(GraphRAGRetriever)
|
|
retriever.chunks = 3
|
|
retriever.base_chunks = None
|
|
retriever.doc_token_limit = 50000
|
|
retriever.vectorstores = ["src"]
|
|
retriever.per_source_retrieval = per_source or {}
|
|
retriever.vector_calls = 0
|
|
|
|
def _vector_ranking(source_id, query_embedding):
|
|
retriever.vector_calls += 1
|
|
return [(VECTOR_ONLY, {"title": "vector"})]
|
|
|
|
retriever._vector_ranking = _vector_ranking
|
|
return retriever
|
|
|
|
|
|
def _texts(docs):
|
|
return [doc["text"] for doc in docs]
|
|
|
|
|
|
class TestDefaultPath:
|
|
def test_walks_passages_and_blends_in_vector_hits(self):
|
|
store = _Store()
|
|
retriever = _retriever()
|
|
|
|
docs = retriever._graph_docs_for_source(store, "src", [0.1, 0.2])
|
|
|
|
# The answer sits one edge away from the seed: the walk reached it.
|
|
assert TEXTS["c-quill"] in _texts(docs)
|
|
# A hit only vector search found is blended in, not lost.
|
|
assert VECTOR_ONLY in _texts(docs)
|
|
assert store.calls == ["chunk_similarities"]
|
|
assert retriever.vector_calls == 1
|
|
|
|
|
|
class TestPerSourceOptions:
|
|
def test_passage_walk_can_be_switched_off(self):
|
|
store = _Store()
|
|
retriever = _retriever(
|
|
{"src": RetrievalConfig(chunks=3, graph={"passage_nodes": False})}
|
|
)
|
|
|
|
docs = retriever._graph_docs_for_source(store, "src", [0.1, 0.2])
|
|
|
|
assert "chunk_similarities" not in store.calls
|
|
assert TEXTS["c-quill"] in _texts(docs)
|
|
|
|
def test_vector_blending_can_be_switched_off(self):
|
|
store = _Store()
|
|
retriever = _retriever(
|
|
{"src": RetrievalConfig(chunks=3, graph={"blend_vector": False})}
|
|
)
|
|
|
|
docs = retriever._graph_docs_for_source(store, "src", [0.1, 0.2])
|
|
|
|
assert retriever.vector_calls == 0
|
|
assert VECTOR_ONLY not in _texts(docs)
|