Spaces:
Sleeping
Sleeping
| 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 | |