258 lines
8.0 KiB
Python
258 lines
8.0 KiB
Python
from app.services.ontology_generator import OntologyGenerator
|
|
from app.services.graph_builder import GraphBuilderService
|
|
from app.utils.ontology import (
|
|
MAX_ONTOLOGY_ATTRIBUTES,
|
|
MAX_ONTOLOGY_SOURCE_TARGETS,
|
|
normalize_ontology_attribute,
|
|
normalize_ontology_attributes,
|
|
)
|
|
from zep_cloud.external_clients.ontology import (
|
|
edge_model_to_api_schema,
|
|
entity_model_to_api_schema,
|
|
)
|
|
|
|
|
|
def test_normalize_string_attribute():
|
|
assert normalize_ontology_attribute("role") == {
|
|
"name": "role",
|
|
"type": "text",
|
|
"description": "role",
|
|
}
|
|
|
|
|
|
def test_preserve_valid_dictionary_attribute():
|
|
original = {"name": "role", "type": "text", "description": "Public role"}
|
|
assert normalize_ontology_attribute(original) == original
|
|
assert normalize_ontology_attribute(original) is not original
|
|
|
|
|
|
def test_reject_unusable_attribute_shapes():
|
|
for value in (None, 7, [], {}, {"name": None}, {"name": ""}, " "):
|
|
assert normalize_ontology_attribute(value) is None
|
|
|
|
|
|
def test_attribute_list_is_non_empty_and_capped_for_zep():
|
|
assert normalize_ontology_attributes(None) == [{
|
|
"name": "details",
|
|
"type": "text",
|
|
"description": "Additional details about this ontology type.",
|
|
}]
|
|
|
|
attributes = [None] + [f"field_{index}" for index in range(12)]
|
|
normalized = normalize_ontology_attributes(attributes)
|
|
|
|
assert len(normalized) == MAX_ONTOLOGY_ATTRIBUTES
|
|
assert [attribute["name"] for attribute in normalized] == [
|
|
f"field_{index}" for index in range(MAX_ONTOLOGY_ATTRIBUTES)
|
|
]
|
|
|
|
|
|
def test_generator_normalizes_entity_and_edge_attributes():
|
|
result = OntologyGenerator(llm_client=object())._validate_and_process({
|
|
"entity_types": [{"name": "speaker", "attributes": ["role", None]}],
|
|
"edge_types": [{"name": "quotes", "attributes": ["source_url", {}]}],
|
|
})
|
|
|
|
assert result["entity_types"][0]["attributes"] == [{
|
|
"name": "role",
|
|
"type": "text",
|
|
"description": "role",
|
|
}]
|
|
assert result["edge_types"][0]["attributes"] == [{
|
|
"name": "source_url",
|
|
"type": "text",
|
|
"description": "source_url",
|
|
}]
|
|
|
|
|
|
def test_generator_adds_a_property_to_empty_custom_types():
|
|
result = OntologyGenerator(llm_client=object())._validate_and_process({
|
|
"entity_types": [{"name": "speaker", "attributes": []}],
|
|
"edge_types": [{"name": "quotes", "attributes": []}],
|
|
})
|
|
|
|
assert result["entity_types"][0]["attributes"][0]["name"] == "details"
|
|
assert result["edge_types"][0]["attributes"][0]["name"] == "details"
|
|
|
|
|
|
def test_graph_builder_safety_net_accepts_strings_and_skips_invalid_values():
|
|
captured = {}
|
|
|
|
class GraphApi:
|
|
def set_ontology(self, **kwargs):
|
|
captured.update(kwargs)
|
|
|
|
class Client:
|
|
graph = GraphApi()
|
|
|
|
builder = object.__new__(GraphBuilderService)
|
|
builder.client = Client()
|
|
builder.set_ontology("graph-id", {
|
|
"entity_types": [{
|
|
"name": "Speaker",
|
|
"attributes": ["role", None, {"name": "summary"}],
|
|
}],
|
|
"edge_types": [],
|
|
})
|
|
|
|
speaker = captured["entities"]["Speaker"]
|
|
assert set(speaker.__annotations__) == {"role", "entity_summary"}
|
|
|
|
|
|
def test_graph_builder_emits_a_pinned_zep_sdk_compatible_schema():
|
|
captured = {}
|
|
|
|
class GraphApi:
|
|
def set_ontology(self, **kwargs):
|
|
captured.update(kwargs)
|
|
|
|
class Client:
|
|
graph = GraphApi()
|
|
|
|
builder = object.__new__(GraphBuilderService)
|
|
builder.client = Client()
|
|
builder.set_ontology("graph-id", {
|
|
"entity_types": [{
|
|
"name": "Speaker",
|
|
"attributes": ["graph_id"] + [
|
|
f"field_{index}" for index in range(10)
|
|
],
|
|
}],
|
|
"edge_types": [{
|
|
"name": "MENTIONS",
|
|
"attributes": [],
|
|
"source_targets": [{"source": "Speaker", "target": "Speaker"}],
|
|
}],
|
|
})
|
|
|
|
assert captured["graph_ids"] == ["graph-id"]
|
|
|
|
speaker = captured["entities"]["Speaker"]
|
|
entity_schema = entity_model_to_api_schema(speaker, "Speaker")
|
|
assert len(entity_schema["properties"]) == MAX_ONTOLOGY_ATTRIBUTES
|
|
assert entity_schema["properties"][0] == {
|
|
"name": "entity_graph_id",
|
|
"type": "Text",
|
|
"description": "graph_id",
|
|
}
|
|
|
|
mentions, source_targets = captured["edges"]["MENTIONS"]
|
|
edge_schema = edge_model_to_api_schema(mentions, "MENTIONS")
|
|
assert edge_schema["properties"] == [{
|
|
"name": "details",
|
|
"type": "Text",
|
|
"description": "Additional details about this ontology type.",
|
|
}]
|
|
assert source_targets[0].source == "Speaker"
|
|
assert source_targets[0].target == "Speaker"
|
|
|
|
|
|
def test_graph_builder_passes_an_empty_entity_mapping_for_edge_only_ontology():
|
|
captured = {}
|
|
|
|
class GraphApi:
|
|
def set_ontology(self, **kwargs):
|
|
captured.update(kwargs)
|
|
|
|
class Client:
|
|
graph = GraphApi()
|
|
|
|
builder = object.__new__(GraphBuilderService)
|
|
builder.client = Client()
|
|
builder.set_ontology("graph-id", {
|
|
"entity_types": [],
|
|
"edge_types": [{
|
|
"name": "RELATED_TO",
|
|
"attributes": ["reason"],
|
|
"source_targets": [{"source": "Entity", "target": "Entity"}],
|
|
}],
|
|
})
|
|
|
|
assert captured["entities"] == {}
|
|
|
|
|
|
def test_graph_builder_deduplicates_and_caps_edge_source_targets_for_zep():
|
|
captured = {}
|
|
|
|
class GraphApi:
|
|
def set_ontology(self, **kwargs):
|
|
captured.update(kwargs)
|
|
|
|
class Client:
|
|
graph = GraphApi()
|
|
|
|
source_targets = [
|
|
{"source": f"Source{index}", "target": f"Target{index}"}
|
|
for index in range(MAX_ONTOLOGY_SOURCE_TARGETS + 2)
|
|
]
|
|
source_targets.insert(1, dict(source_targets[0]))
|
|
|
|
builder = object.__new__(GraphBuilderService)
|
|
builder.client = Client()
|
|
builder.set_ontology("graph-id", {
|
|
"entity_types": [],
|
|
"edge_types": [{
|
|
"name": "RELATED_TO",
|
|
"attributes": ["reason"],
|
|
"source_targets": source_targets,
|
|
}],
|
|
})
|
|
|
|
_, normalized_targets = captured["edges"]["RELATED_TO"]
|
|
assert len(normalized_targets) == MAX_ONTOLOGY_SOURCE_TARGETS
|
|
assert [(item.source, item.target) for item in normalized_targets] == [
|
|
(f"Source{index}", f"Target{index}")
|
|
for index in range(MAX_ONTOLOGY_SOURCE_TARGETS)
|
|
]
|
|
|
|
|
|
def test_generator_ignores_invalid_entries_and_normalizes_edge_names():
|
|
source_targets = [
|
|
{"source": "speaker", "target": "news outlet"},
|
|
{"source": "speaker", "target": "news outlet"},
|
|
None,
|
|
] + [
|
|
{"source": "speaker", "target": "news outlet" if index == 0 else "Person"}
|
|
for index in range(12)
|
|
]
|
|
|
|
result = OntologyGenerator(llm_client=object())._validate_and_process({
|
|
"entity_types": ["speaker", None, 7, {"name": "news outlet"}],
|
|
"edge_types": [
|
|
"unusable edge",
|
|
None,
|
|
{"name": "worksFor", "source_targets": source_targets},
|
|
{"name": "works-for", "source_targets": []},
|
|
],
|
|
})
|
|
|
|
assert [entity["name"] for entity in result["entity_types"][:2]] == [
|
|
"Speaker",
|
|
"NewsOutlet",
|
|
]
|
|
assert [edge["name"] for edge in result["edge_types"]] == ["WORKS_FOR"]
|
|
assert result["edge_types"][0]["source_targets"] == [
|
|
{"source": "Speaker", "target": "NewsOutlet"},
|
|
{"source": "Speaker", "target": "Person"},
|
|
]
|
|
|
|
|
|
def test_generator_caps_after_discarding_invalid_edge_endpoints():
|
|
invalid_first = [
|
|
{"source": f"Removed{index}", "target": "AlsoRemoved"}
|
|
for index in range(MAX_ONTOLOGY_SOURCE_TARGETS)
|
|
]
|
|
result = OntologyGenerator(llm_client=object())._validate_and_process({
|
|
"entity_types": [{"name": "person"}, {"name": "organization"}],
|
|
"edge_types": [{
|
|
"name": "works_for",
|
|
"source_targets": invalid_first + [
|
|
{"source": "person", "target": "organization"}
|
|
],
|
|
}],
|
|
})
|
|
|
|
assert result["edge_types"][0]["source_targets"] == [
|
|
{"source": "Person", "target": "Organization"}
|
|
]
|