Spaces:
Sleeping
Sleeping
File size: 1,969 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 | import pytest
from techdocs_graphrag.models import (
FALLBACK_RELATION,
Entity,
EntityType,
GraphExtraction,
GraphFact,
Relation,
RetrievalContext,
canonical_key,
)
@pytest.mark.parametrize("name", ["StateGraph", "State Graph", "state_graph", " stategraph "])
def test_canonical_key_merges_spelling_variants(name: str) -> None:
assert canonical_key(name) == "stategraph"
@pytest.mark.parametrize(
("raw", "expected"),
[
("persists with", "PERSISTS_WITH"),
("dépend-de", "DEPEND_DE"),
("MENTIONS", FALLBACK_RELATION),
("42_uses", FALLBACK_RELATION),
("", FALLBACK_RELATION),
],
)
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
|