Spaces:
Sleeping
Sleeping
| import pytest | |
| from techdocs_graphrag.models import ( | |
| FALLBACK_RELATION, | |
| Entity, | |
| EntityType, | |
| GraphExtraction, | |
| GraphFact, | |
| Relation, | |
| RetrievalContext, | |
| canonical_key, | |
| ) | |
| def test_canonical_key_merges_spelling_variants(name: str) -> None: | |
| assert canonical_key(name) == "stategraph" | |
| def test_relation_type_is_normalized_to_a_safe_identifier(raw: str, expected: str) -> None: | |
| relation = Relation(source="A", relation=raw, target="B") | |
| assert relation.relation == expected | |
| def test_unknown_entity_type_falls_back_to_concept() -> None: | |
| assert Entity(name="Reducer", type="Pattern").type is EntityType.CONCEPT | |
| assert Entity(name="StateGraph", type="component").type is EntityType.COMPONENT | |
| def test_extraction_adds_missing_endpoints_and_drops_self_loops() -> None: | |
| extraction = GraphExtraction.model_validate( | |
| { | |
| "entities": [{"name": "StateGraph", "type": "Component"}], | |
| "relations": [ | |
| {"source": "StateGraph", "relation": "USES", "target": "Channel"}, | |
| {"source": "StateGraph", "relation": "IS", "target": "State Graph"}, | |
| ], | |
| } | |
| ) | |
| assert {entity.key for entity in extraction.entities} == {"stategraph", "channel"} | |
| assert [relation.relation for relation in extraction.relations] == ["USES"] | |
| def test_context_sources_combine_chunks_and_facts() -> None: | |
| context = RetrievalContext( | |
| facts=[GraphFact(source="A", relation="USES", target="B", documents=["b.md", "a.md"])] | |
| ) | |
| assert context.sources == ["a.md", "b.md"] | |
| assert not context.is_empty | |