mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-29 04:26:20 +00:00
KaifAhmad1andClaude Sonnet 4.6
0f0800f109
- Add TemporalGraphRetriever to context_retriever.py (no new file per project convention)
- Drop-in wrapper for ContextRetriever; filters related_entities/related_relationships
via reconstruct_at_time(); at_time=None is a true passthrough
- Returns new RetrievedContext objects (no in-place mutation)
- Graceful ImportError if temporal modules unavailable
- Add at_time + header_template to ContextRetriever._generate_reasoned_response()
and query_with_reasoning()
- Temporal header prepended to LLM context block only when at_time is set
- Naive datetimes normalised to UTC before formatting
- Header built with str.replace (not .format) to prevent format-string injection
- Add TemporalQueryRewriter + TemporalQueryResult to semantica/kg/
- Regex-only (default) and LLM-assisted extraction modes
- Resolves temporal phrases via TemporalNormalizer (deterministic, zero LLM)
- Word-boundary guards on intent keywords; year fallback for noun-phrase dates
- Never calls reconstruct_at_time — extraction only
- Export TemporalGraphRetriever from semantica.context
- Export TemporalQueryRewriter, TemporalQueryResult from semantica.kg
- Add 99 tests across two new test files
- tests/context/test_temporal_retriever.py (56 tests)
- tests/kg/test_temporal_query_rewriter.py (43 tests)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
429 lines
17 KiB
Python
429 lines
17 KiB
Python
"""
|
|
Tests for TemporalGraphRetriever and temporal context header in LLM prompts.
|
|
"""
|
|
|
|
from datetime import datetime, timezone
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
import pytest
|
|
|
|
from semantica.context import ContextRetriever, RetrievedContext, TemporalGraphRetriever
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _utc(year, month=1, day=1):
|
|
return datetime(year, month, day, tzinfo=timezone.utc)
|
|
|
|
|
|
def _make_entity(eid, valid_from=None, valid_until=None, **extra):
|
|
e = {"id": eid, "name": eid}
|
|
if valid_from:
|
|
e["valid_from"] = valid_from.isoformat()
|
|
if valid_until:
|
|
e["valid_until"] = valid_until.isoformat()
|
|
e.update(extra)
|
|
return e
|
|
|
|
|
|
def _make_rel(source, target, rel_type="RELATED_TO", valid_from=None, valid_until=None):
|
|
r = {"source": source, "target": target, "type": rel_type}
|
|
if valid_from:
|
|
r["valid_from"] = valid_from.isoformat()
|
|
if valid_until:
|
|
r["valid_until"] = valid_until.isoformat()
|
|
return r
|
|
|
|
|
|
def _make_result(entities, relationships, content="test content", score=0.9):
|
|
return RetrievedContext(
|
|
content=content,
|
|
score=score,
|
|
source="graph",
|
|
related_entities=entities,
|
|
related_relationships=relationships,
|
|
)
|
|
|
|
|
|
def _base_mock(return_value=None):
|
|
base = MagicMock(spec=ContextRetriever)
|
|
base.retrieve.return_value = return_value or []
|
|
return base
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TemporalGraphRetriever — construction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestTemporalGraphRetrieverInit:
|
|
|
|
def test_default_at_time_is_none(self):
|
|
tr = TemporalGraphRetriever(_base_mock())
|
|
assert tr.at_time is None
|
|
|
|
def test_stores_base_retriever(self):
|
|
base = _base_mock()
|
|
tr = TemporalGraphRetriever(base)
|
|
assert tr.base_retriever is base
|
|
|
|
def test_custom_header_template_stored(self):
|
|
tr = TemporalGraphRetriever(_base_mock(), header_template="[{at_time}]")
|
|
assert tr.header_template == "[{at_time}]"
|
|
|
|
def test_default_header_template_contains_placeholder(self):
|
|
tr = TemporalGraphRetriever(_base_mock())
|
|
assert "{at_time}" in tr.header_template
|
|
assert "{source}" in tr.header_template
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TemporalGraphRetriever — passthrough (no at_time)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestTemporalGraphRetrieverPassthrough:
|
|
|
|
def setup_method(self):
|
|
self.entity = _make_entity("e1", _utc(2020, 1, 1), _utc(2025, 1, 1))
|
|
self.base = _base_mock([_make_result([self.entity], [])])
|
|
|
|
def test_no_at_time_returns_base_result_unchanged(self):
|
|
tr = TemporalGraphRetriever(self.base)
|
|
results = tr.retrieve("some query")
|
|
self.base.retrieve.assert_called_once_with("some query")
|
|
assert results[0].related_entities[0]["id"] == "e1"
|
|
|
|
def test_none_at_time_on_call_uses_constructor_none(self):
|
|
tr = TemporalGraphRetriever(self.base, at_time=None)
|
|
results = tr.retrieve("query", at_time=None)
|
|
assert results[0].related_entities[0]["id"] == "e1"
|
|
|
|
def test_empty_base_result_passthrough(self):
|
|
base = _base_mock([])
|
|
tr = TemporalGraphRetriever(base)
|
|
assert tr.retrieve("q") == []
|
|
|
|
def test_kwargs_forwarded_to_base_retriever(self):
|
|
tr = TemporalGraphRetriever(self.base)
|
|
tr.retrieve("q", max_results=3, min_relevance_score=0.5)
|
|
self.base.retrieve.assert_called_once_with("q", max_results=3, min_relevance_score=0.5)
|
|
|
|
def test_base_retriever_called_exactly_once(self):
|
|
tr = TemporalGraphRetriever(self.base)
|
|
tr.retrieve("q")
|
|
assert self.base.retrieve.call_count == 1
|
|
|
|
def test_result_object_identity_unchanged_on_passthrough(self):
|
|
result = _make_result([self.entity], [])
|
|
base = _base_mock([result])
|
|
tr = TemporalGraphRetriever(base)
|
|
returned = tr.retrieve("q")
|
|
assert returned[0] is result
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TemporalGraphRetriever — temporal filtering
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestTemporalGraphRetrieverFiltering:
|
|
|
|
def setup_method(self):
|
|
self.base = _base_mock()
|
|
|
|
def _set(self, entities, relationships):
|
|
self.base.retrieve.return_value = [_make_result(entities, relationships)]
|
|
|
|
# Entity validity
|
|
|
|
def test_entity_expired_before_at_time_excluded(self):
|
|
e_current = _make_entity("current", _utc(2020, 1, 1), _utc(2025, 1, 1))
|
|
e_expired = _make_entity("expired", _utc(2018, 1, 1), _utc(2020, 6, 1))
|
|
self._set([e_current, e_expired], [])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
ids = {e["id"] for e in results[0].related_entities}
|
|
assert "current" in ids
|
|
assert "expired" not in ids
|
|
|
|
def test_entity_not_yet_started_at_at_time_excluded(self):
|
|
e_future = _make_entity("future", _utc(2025, 1, 1))
|
|
e_current = _make_entity("current", _utc(2020, 1, 1))
|
|
self._set([e_future, e_current], [])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
ids = {e["id"] for e in results[0].related_entities}
|
|
assert "current" in ids
|
|
assert "future" not in ids
|
|
|
|
def test_entity_with_no_temporal_bounds_always_included(self):
|
|
e_timeless = _make_entity("timeless") # no valid_from / valid_until
|
|
self._set([e_timeless], [])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert len(results[0].related_entities) == 1
|
|
|
|
def test_entity_valid_on_boundary_date_included(self):
|
|
# valid_from == at_time — boundary should be inclusive
|
|
at = _utc(2022, 6, 1)
|
|
e = _make_entity("boundary", valid_from=at)
|
|
self._set([e], [])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=at).retrieve("q")
|
|
assert len(results[0].related_entities) == 1
|
|
|
|
def test_all_entities_expired_leaves_empty_list(self):
|
|
entities = [
|
|
_make_entity("a", _utc(2010, 1, 1), _utc(2015, 1, 1)),
|
|
_make_entity("b", _utc(2012, 1, 1), _utc(2014, 1, 1)),
|
|
]
|
|
self._set(entities, [])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert results[0].related_entities == []
|
|
|
|
# Relationship filtering
|
|
|
|
def test_dangling_relationship_removed_when_target_expired(self):
|
|
e_a = _make_entity("A", _utc(2020, 1, 1))
|
|
e_b = _make_entity("B", _utc(2022, 1, 1), _utc(2022, 6, 1))
|
|
rel = _make_rel("A", "B")
|
|
self._set([e_a, e_b], [rel])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert results[0].related_relationships == []
|
|
|
|
def test_dangling_relationship_removed_when_source_expired(self):
|
|
e_a = _make_entity("A", _utc(2020, 1, 1), _utc(2021, 1, 1))
|
|
e_b = _make_entity("B", _utc(2020, 1, 1))
|
|
rel = _make_rel("A", "B")
|
|
self._set([e_a, e_b], [rel])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert results[0].related_relationships == []
|
|
|
|
def test_valid_relationship_kept(self):
|
|
e_a = _make_entity("A", _utc(2020, 1, 1))
|
|
e_b = _make_entity("B", _utc(2020, 1, 1))
|
|
rel = _make_rel("A", "B", valid_from=_utc(2020, 1, 1))
|
|
self._set([e_a, e_b], [rel])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert len(results[0].related_relationships) == 1
|
|
|
|
def test_relationship_with_expired_valid_until_removed(self):
|
|
e_a = _make_entity("A", _utc(2020, 1, 1))
|
|
e_b = _make_entity("B", _utc(2020, 1, 1))
|
|
rel = _make_rel("A", "B", valid_from=_utc(2020, 1, 1), valid_until=_utc(2021, 1, 1))
|
|
self._set([e_a, e_b], [rel])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert results[0].related_relationships == []
|
|
|
|
def test_multiple_relationship_types_filtered_independently(self):
|
|
e_a = _make_entity("A", _utc(2020, 1, 1))
|
|
e_b = _make_entity("B", _utc(2020, 1, 1))
|
|
e_c = _make_entity("C", _utc(2020, 1, 1), _utc(2021, 1, 1)) # expires
|
|
rel_ab = _make_rel("A", "B", rel_type="USES")
|
|
rel_ac = _make_rel("A", "C", rel_type="OWNS")
|
|
self._set([e_a, e_b, e_c], [rel_ab, rel_ac])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
rels = results[0].related_relationships
|
|
types = {r["type"] for r in rels}
|
|
assert "USES" in types
|
|
assert "OWNS" not in types
|
|
|
|
# at_time precedence
|
|
|
|
def test_call_site_at_time_overrides_constructor(self):
|
|
e = _make_entity("e", _utc(2018, 1, 1), _utc(2021, 1, 1))
|
|
self._set([e], [])
|
|
|
|
tr = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1))
|
|
results = tr.retrieve("q", at_time=_utc(2020, 6, 1))
|
|
assert len(results[0].related_entities) == 1
|
|
|
|
def test_string_at_time_parsed_correctly(self):
|
|
e_current = _make_entity("current", _utc(2020, 1, 1))
|
|
e_expired = _make_entity("expired", _utc(2018, 1, 1), _utc(2020, 6, 1))
|
|
self._set([e_current, e_expired], [])
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time="2023-01-01").retrieve("q")
|
|
ids = {e["id"] for e in results[0].related_entities}
|
|
assert "current" in ids
|
|
assert "expired" not in ids
|
|
|
|
def test_datetime_at_time_accepted_directly(self):
|
|
e = _make_entity("e", _utc(2020, 1, 1))
|
|
self._set([e], [])
|
|
|
|
results = TemporalGraphRetriever(
|
|
self.base, at_time=_utc(2023, 1, 1)
|
|
).retrieve("q")
|
|
assert len(results[0].related_entities) == 1
|
|
|
|
# Multiple results
|
|
|
|
def test_all_results_filtered_independently(self):
|
|
e_valid = _make_entity("valid", _utc(2020, 1, 1))
|
|
e_old = _make_entity("old", _utc(2010, 1, 1), _utc(2015, 1, 1))
|
|
r1 = _make_result([e_valid], [])
|
|
r2 = _make_result([e_old], [])
|
|
self.base.retrieve.return_value = [r1, r2]
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert len(results[0].related_entities) == 1
|
|
assert len(results[1].related_entities) == 0
|
|
|
|
def test_result_scores_preserved_after_filtering(self):
|
|
e = _make_entity("e", _utc(2020, 1, 1))
|
|
r = _make_result([e], [], score=0.77)
|
|
self.base.retrieve.return_value = [r]
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert results[0].score == pytest.approx(0.77)
|
|
|
|
def test_result_content_preserved_after_filtering(self):
|
|
e = _make_entity("e", _utc(2020, 1, 1))
|
|
r = _make_result([e], [], content="important drug interaction fact")
|
|
self.base.retrieve.return_value = [r]
|
|
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert results[0].content == "important drug interaction fact"
|
|
|
|
def test_empty_entities_and_relationships_stays_empty(self):
|
|
self.base.retrieve.return_value = [_make_result([], [])]
|
|
results = TemporalGraphRetriever(self.base, at_time=_utc(2023, 1, 1)).retrieve("q")
|
|
assert results[0].related_entities == []
|
|
assert results[0].related_relationships == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Temporal context header — _generate_reasoned_response
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestTemporalContextHeader:
|
|
|
|
def setup_method(self):
|
|
self.retriever = ContextRetriever()
|
|
self.mock_llm = MagicMock()
|
|
self.mock_llm.generate.return_value = "LLM answer"
|
|
|
|
def _call(self, at_time=None, header_template=None, contexts=None):
|
|
if contexts is None:
|
|
contexts = [RetrievedContext(content="fact A", score=0.9, source="graph")]
|
|
kwargs = {}
|
|
if header_template is not None:
|
|
kwargs["header_template"] = header_template
|
|
return self.retriever._generate_reasoned_response(
|
|
"test query", contexts, [], self.mock_llm, at_time=at_time, **kwargs
|
|
)
|
|
|
|
def _last_prompt(self):
|
|
return self.mock_llm.generate.call_args[0][0]
|
|
|
|
def test_no_header_without_at_time(self):
|
|
self._call()
|
|
assert "Graph context valid as of" not in self._last_prompt()
|
|
|
|
def test_header_present_when_at_time_datetime(self):
|
|
self._call(at_time=_utc(2023, 6, 1))
|
|
prompt = self._last_prompt()
|
|
assert "2023-06-01" in prompt
|
|
assert "Graph context valid as of" in prompt
|
|
|
|
def test_header_present_when_at_time_string(self):
|
|
self._call(at_time="2022-03-15")
|
|
assert "2022-03-15" in self._last_prompt()
|
|
|
|
def test_header_appears_before_retrieved_context(self):
|
|
self._call(at_time=_utc(2023, 6, 1))
|
|
prompt = self._last_prompt()
|
|
assert prompt.find("Graph context valid as of") < prompt.find("Retrieved Context:")
|
|
|
|
def test_header_contains_source_label(self):
|
|
self._call(at_time=_utc(2023, 6, 1))
|
|
assert "KnowledgeGraph snapshot" in self._last_prompt()
|
|
|
|
def test_header_template_configurable(self):
|
|
self._call(
|
|
at_time=_utc(2023, 6, 1),
|
|
header_template="[Snapshot: {at_time} from {source}]",
|
|
)
|
|
prompt = self._last_prompt()
|
|
assert "[Snapshot:" in prompt
|
|
assert "2023-06-01" in prompt
|
|
|
|
def test_custom_template_source_placeholder_filled(self):
|
|
self._call(
|
|
at_time=_utc(2023, 1, 1),
|
|
header_template="Source={source}",
|
|
)
|
|
assert "Source=KnowledgeGraph snapshot" in self._last_prompt()
|
|
|
|
def test_prompt_contains_user_question(self):
|
|
self.retriever._generate_reasoned_response(
|
|
"How many suppliers?", [], [], self.mock_llm
|
|
)
|
|
assert "How many suppliers?" in self._last_prompt()
|
|
|
|
def test_prompt_contains_retrieved_context_content(self):
|
|
ctx = RetrievedContext(content="DrugA interaction warning", score=0.9, source="graph")
|
|
self.retriever._generate_reasoned_response(
|
|
"q", [ctx], [], self.mock_llm
|
|
)
|
|
assert "DrugA interaction warning" in self._last_prompt()
|
|
|
|
def test_no_at_time_prompt_identical_to_baseline(self):
|
|
ctx = RetrievedContext(content="fact", score=0.8, source="graph")
|
|
self.retriever._generate_reasoned_response("q", [ctx], [], self.mock_llm)
|
|
prompt_without = self._last_prompt()
|
|
|
|
self.retriever._generate_reasoned_response("q", [ctx], [], self.mock_llm, at_time=None)
|
|
prompt_with_none = self._last_prompt()
|
|
|
|
assert prompt_without == prompt_with_none
|
|
|
|
def test_llm_generate_called_once_per_call(self):
|
|
self._call(at_time=_utc(2023, 1, 1))
|
|
assert self.mock_llm.generate.call_count == 1
|
|
|
|
def test_query_with_reasoning_threads_at_time(self):
|
|
retriever = ContextRetriever()
|
|
retriever.retrieve = MagicMock(return_value=[])
|
|
|
|
with patch.object(
|
|
retriever, "_generate_reasoned_response", wraps=retriever._generate_reasoned_response
|
|
) as mock_gen:
|
|
mock_gen.return_value = "answer"
|
|
retriever.query_with_reasoning(
|
|
"test query", self.mock_llm, at_time=_utc(2023, 1, 1)
|
|
)
|
|
assert mock_gen.call_args.kwargs.get("at_time") == _utc(2023, 1, 1)
|
|
|
|
def test_query_with_reasoning_threads_header_template(self):
|
|
retriever = ContextRetriever()
|
|
retriever.retrieve = MagicMock(return_value=[])
|
|
custom = "[{at_time}|{source}]"
|
|
|
|
with patch.object(
|
|
retriever, "_generate_reasoned_response", wraps=retriever._generate_reasoned_response
|
|
) as mock_gen:
|
|
mock_gen.return_value = "answer"
|
|
retriever.query_with_reasoning(
|
|
"test query", self.mock_llm,
|
|
at_time=_utc(2023, 1, 1),
|
|
header_template=custom,
|
|
)
|
|
assert mock_gen.call_args.kwargs.get("header_template") == custom
|
|
|
|
def test_query_with_reasoning_no_at_time_no_header(self):
|
|
retriever = ContextRetriever()
|
|
retriever.retrieve = MagicMock(return_value=[
|
|
RetrievedContext(content="fact", score=0.9, source="graph")
|
|
])
|
|
retriever.query_with_reasoning("q", self.mock_llm)
|
|
prompt = self.mock_llm.generate.call_args[0][0]
|
|
assert "Graph context valid as of" not in prompt
|