Spaces:
Sleeping
Sleeping
File size: 3,819 Bytes
da21448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | from techdocs_graphrag.answering import NO_CONTEXT_ANSWER, GraphRAGAssistant
from techdocs_graphrag.extraction import GraphExtractor
from techdocs_graphrag.indexing import IndexingPipeline
from techdocs_graphrag.ingestion import ParagraphSplitter, SourceDocument
from techdocs_graphrag.models import Chunk, GraphFact, RetrievalMode
from techdocs_graphrag.retrieval import HybridRetriever
from techdocs_graphrag.visualization import facts_to_dot
from .conftest import FakeKnowledgeGraph, FakeLanguageModel, retrieved_chunk
EXTRACTION_PAYLOAD = """
{
"entities": [{"name": "StateGraph", "type": "Component"}],
"relations": [{"source": "StateGraph", "relation": "uses", "target": "Channel", "evidence": "x"}]
}
"""
def build_retriever(llm: FakeLanguageModel, store: FakeKnowledgeGraph) -> HybridRetriever:
return HybridRetriever(llm=llm, store=store, top_k=5, max_hops=2, max_facts=10)
def test_graph_mode_seeds_expansion_with_lexical_then_mentioned_entities(
fake_llm: FakeLanguageModel,
) -> None:
store = FakeKnowledgeGraph(
chunks=[
retrieved_chunk("a.md", ["checkpointer", "stategraph"]),
retrieved_chunk("b.md", ["checkpointer"]),
],
lexical_matches=["stategraph"],
)
context = build_retriever(fake_llm, store).retrieve("question", RetrievalMode.GRAPH)
assert store.expanded_with == ["stategraph", "checkpointer"]
assert context.seed_entities == ["stategraph", "checkpointer"]
def test_vector_mode_never_touches_the_graph(fake_llm: FakeLanguageModel) -> None:
store = FakeKnowledgeGraph(chunks=[retrieved_chunk("a.md", ["stategraph"])])
context = build_retriever(fake_llm, store).retrieve("question", RetrievalMode.VECTOR)
assert store.expanded_with is None
assert context.facts == []
def test_assistant_does_not_call_the_llm_without_context(fake_llm: FakeLanguageModel) -> None:
assistant = GraphRAGAssistant(retriever=build_retriever(fake_llm, FakeKnowledgeGraph()), llm=fake_llm)
answer = assistant.ask("question")
assert answer.text == NO_CONTEXT_ANSWER
assert fake_llm.completions == []
def test_assistant_sends_graph_facts_to_the_llm(fake_llm: FakeLanguageModel) -> None:
store = FakeKnowledgeGraph(
chunks=[retrieved_chunk("a.md", ["stategraph"])],
facts=[GraphFact(source="StateGraph", relation="USES", target="Channel", documents=["a.md"])],
)
assistant = GraphRAGAssistant(retriever=build_retriever(fake_llm, store), llm=fake_llm)
answer = assistant.ask("Comment fonctionne StateGraph ?")
_, user_prompt = fake_llm.completions[0]
assert "StateGraph --USES--> Channel" in user_prompt
assert answer.context.sources == ["a.md"]
def test_extractor_returns_empty_graph_on_invalid_json() -> None:
chunk = Chunk(id="1", document="a.md", position=0, text="texte")
assert GraphExtractor(FakeLanguageModel("pas du json")).extract(chunk).entities == []
def test_indexing_pipeline_reports_extracted_graph() -> None:
llm = FakeLanguageModel(EXTRACTION_PAYLOAD)
store = FakeKnowledgeGraph()
pipeline = IndexingPipeline(
splitter=ParagraphSplitter(chunk_size=50, chunk_overlap=0),
llm=llm,
extractor=GraphExtractor(llm),
store=store,
workers=2,
)
progress: list[int] = []
report = pipeline.run(
[SourceDocument("a.md", "Premier paragraphe.\n\nSecond paragraphe assez long pour couper.")],
on_progress=lambda done, _: progress.append(done),
)
assert (report.chunks, report.entities, report.relations) == (2, 4, 2)
assert progress == [1, 2]
assert store.indexes_awaited
def test_dot_output_escapes_quotes() -> None:
dot = facts_to_dot([GraphFact(source='Say "hi"', relation="USES", target="B")])
assert '"Say \\"hi\\""' in dot
|