From c15ea9bc592136c3b9f8cfdbc3fae4031357ae23 Mon Sep 17 00:00:00 2001 From: Divyam Talwar Date: Mon, 24 Aug 2026 01:31:06 +0530 Subject: [PATCH] feat(graph): add bitemporal fact search --- README.md | 4 + docs/temporal-graph-facts.md | 36 ++++++++ jitmind/__init__.py | 2 + jitmind/graph/graph_store.py | 106 +++++++++++++++++++++- jitmind/retriever/graph_retriever.py | 43 +++++++++ jitmind/schemas/__init__.py | 2 + jitmind/schemas/graph_fact.py | 27 ++++++ tests/test_temporal_graph_facts.py | 126 +++++++++++++++++++++++++++ 8 files changed, 343 insertions(+), 3 deletions(-) create mode 100644 docs/temporal-graph-facts.md create mode 100644 jitmind/schemas/graph_fact.py create mode 100644 tests/test_temporal_graph_facts.py diff --git a/README.md b/README.md index 6a64f47..6cd7445 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,10 @@ JITMIND directly addresses these failure modes with: - A hybrid retrieval stack (BM25 + dense + index + graph) with RRF, dynamic weighting, tier-aware boosting, and optional reranking. - A looped research process (plan -> search -> integrate -> reflect) instead of one-shot generation. - Graph memory with entity semantics, provenance edges, and Personalized PageRank for associative recall. +- Bi-temporal graph-fact search with immutable observations and source-memory/page evidence. + +See [`docs/temporal-graph-facts.md`](docs/temporal-graph-facts.md) for valid-time +and observation-time query examples. This README is intentionally long and deep. It is meant to make a new contributor productive quickly, and to make an architecture reviewer confident that the system has real technical substance. diff --git a/docs/temporal-graph-facts.md b/docs/temporal-graph-facts.md new file mode 100644 index 0000000..183ebc4 --- /dev/null +++ b/docs/temporal-graph-facts.md @@ -0,0 +1,36 @@ +# Temporal graph fact search + +JITMIND stores each extracted relation as an immutable observation keyed by +its source memory and triplet. Re-ingesting the same observation is idempotent; +observing the same triplet from a later memory creates another version instead +of overwriting history. + +```python +facts = graph_store.query_facts( + ["Alice", "Acme"], + relation_types=["WORKS_AT"], + valid_at="2026-02-03T00:00:00Z", + observed_at="2026-02-04T00:00:00Z", + namespace=("tenant", "red"), +) + +hits = graph_retriever.search_temporal( + ["Alice"], + valid_at="2026-02-03T00:00:00Z", + relation_types=["WORKS_AT"], +) +``` + +`valid_at` gates event time and uses a start-inclusive, end-exclusive +interval. `observed_at` gates ingestion time. Each result carries its source +memory, page, and full interval so an answer can cite why the fact was visible. + +The interface is an original JITMIND implementation informed by Graphiti's +public model of episodes plus facts with `valid_at` / `invalid_at`, including +its fact search filters. No Graphiti code is incorporated. + +Reference: + +- +- + diff --git a/jitmind/__init__.py b/jitmind/__init__.py index f4bbe27..821bcd9 100644 --- a/jitmind/__init__.py +++ b/jitmind/__init__.py @@ -68,6 +68,7 @@ TTLMemoryState, TTLMemoryEntry, TTLPageStore, + GraphFact, ) try: from jitmind.graph import GraphMemoryStore, GraphOntology @@ -154,6 +155,7 @@ "TTLMemoryState", "TTLMemoryEntry", "TTLPageStore", + "GraphFact", "GraphMemoryStore", "GraphOntology", "Document", diff --git a/jitmind/graph/graph_store.py b/jitmind/graph/graph_store.py index 1940a35..a5bb92f 100644 --- a/jitmind/graph/graph_store.py +++ b/jitmind/graph/graph_store.py @@ -1,8 +1,9 @@ # -*- coding: utf-8 -*- from __future__ import annotations -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Sequence from contextlib import contextmanager +from datetime import datetime, timezone import hashlib try: @@ -11,6 +12,7 @@ GraphDatabase = None # type: ignore from jitmind.graph.ontology import GraphOntology +from jitmind.schemas.graph_fact import GraphFact class GraphMemoryStore: @@ -182,6 +184,9 @@ def add_entities_relations( t_type = entity_type_map.get(tail, "Entity") h_key = f"{h_type}:{head}" t_key = f"{t_type}:{tail}" + relation_id = self._relation_id( + memory_id, h_key, rel_type, t_key + ) try: session.run( """ @@ -189,8 +194,9 @@ def add_entities_relations( SET h.name = $head, h.type = $h_type MERGE (t:Entity {key: $t_key}) SET t.name = $tail, t.type = $t_type - MERGE (h)-[r:RELATION {type: $rel_type}]->(t) - SET r.source_memory_id = $mid, + MERGE (h)-[r:RELATION {id: $relation_id}]->(t) + SET r.type = $rel_type, + r.source_memory_id = $mid, r.t_observed = $t_observed, r.t_valid = $t_valid, r.t_invalid = $t_invalid @@ -206,12 +212,100 @@ def add_entities_relations( t_observed=t_observed or r.get("t_observed"), t_valid=t_valid or r.get("t_valid"), t_invalid=t_invalid or r.get("t_invalid"), + relation_id=relation_id, ) except Exception as e: print(f"[WARN] Failed to add relation {head}->{tail}: {e}") except Exception as e: print(f"[WARN] Failed to add entities/relations for memory {memory_id}: {e}") + def query_facts( + self, + entity_names: Sequence[str] | None = None, + *, + relation_types: Sequence[str] | None = None, + valid_at: str | datetime | None = None, + observed_at: str | datetime | None = None, + namespace: Sequence[str] | None = None, + limit: int = 20, + ) -> List[GraphFact]: + """Return fact observations through valid- and transaction-time gates. + + ``valid_at`` asks what was true at an event time. ``observed_at`` asks + what the graph could have known by an ingestion time. Supplying both + provides a bi-temporal snapshot rather than a present-day projection. + """ + + if limit < 1 or limit > 1000: + raise ValueError("limit must be between 1 and 1000") + names = sorted( + {name.strip() for name in entity_names or () if isinstance(name, str) and name.strip()} + ) + types = sorted( + { + relation.strip().replace(" ", "_").upper() + for relation in relation_types or () + if isinstance(relation, str) and relation.strip() + } + ) + params = { + "names": names, + "types": types, + "valid_at": self._normalize_query_time(valid_at, "valid_at"), + "observed_at": self._normalize_query_time(observed_at, "observed_at"), + "namespace": list(namespace) if namespace is not None else None, + "limit": limit, + } + query = """ + MATCH (h:Entity)-[r:RELATION]->(t:Entity) + OPTIONAL MATCH (m:Memory {id: r.source_memory_id}) + WHERE (size($names) = 0 OR h.name IN $names OR t.name IN $names) + AND (size($types) = 0 OR r.type IN $types) + AND ($namespace IS NULL OR m.namespace = $namespace) + AND ( + $valid_at IS NULL OR ( + datetime(coalesce(r.t_valid, r.t_observed, m.t_valid, m.t_observed, m.t_created)) + <= datetime($valid_at) + AND (r.t_invalid IS NULL OR datetime($valid_at) < datetime(r.t_invalid)) + ) + ) + AND ( + $observed_at IS NULL OR + datetime(coalesce(r.t_observed, m.t_observed, m.t_created)) + <= datetime($observed_at) + ) + RETURN r.id AS id, h.name AS head, r.type AS relation, t.name AS tail, + r.source_memory_id AS source_memory_id, + m.source_page_id AS source_page_id, + r.t_observed AS t_observed, r.t_valid AS t_valid, + r.t_invalid AS t_invalid + ORDER BY coalesce(r.t_valid, r.t_observed) DESC, r.id ASC + LIMIT $limit + """ + try: + with self._get_session() as session: + rows = session.run(query, **params) + return [GraphFact(**dict(row)) for row in rows] + except Exception as exc: + print(f"[WARN] Failed to query temporal graph facts: {exc}") + return [] + + def _normalize_query_time( + self, value: str | datetime | None, field_name: str + ) -> Optional[str]: + if value is None: + return None + if isinstance(value, datetime): + parsed = value + else: + try: + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field_name} must be an ISO-8601 timestamp") from exc + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc).isoformat() + # ---- 3-tier graph architecture ---- def add_episode( self, @@ -818,3 +912,9 @@ def _node_key(self, node) -> tuple[Optional[str], str]: def _semantic_id(self, fact: str) -> str: return hashlib.sha256(fact.lower().encode("utf-8")).hexdigest() + + def _relation_id( + self, memory_id: str, head_key: str, relation: str, tail_key: str + ) -> str: + payload = "\x1f".join((memory_id, head_key, relation, tail_key)) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() diff --git a/jitmind/retriever/graph_retriever.py b/jitmind/retriever/graph_retriever.py index 742916c..34b3ad4 100644 --- a/jitmind/retriever/graph_retriever.py +++ b/jitmind/retriever/graph_retriever.py @@ -55,3 +55,46 @@ def search(self, query_list: List[str], top_k: int = 10) -> List[List[Hit]]: ) results_all.append(hits) return results_all + + def search_temporal( + self, + query_list: List[str], + *, + valid_at: str | None = None, + observed_at: str | None = None, + relation_types: List[str] | None = None, + top_k: int = 10, + ) -> List[List[Hit]]: + """Search versioned graph facts and preserve their temporal evidence.""" + + if not self.graph_store: + return [[] for _ in query_list] + results_all: List[List[Hit]] = [] + for query in query_list: + names = [name.strip() for name in query.split(",") if name.strip()] + facts = self.graph_store.query_facts( + names, + relation_types=relation_types, + valid_at=valid_at, + observed_at=observed_at, + limit=top_k, + ) + results_all.append( + [ + Hit( + page_id=fact.source_page_id, + snippet=fact.statement, + source="graph_fact", + meta={ + "fact_id": fact.id, + "memory_id": fact.source_memory_id, + "relation": fact.relation, + "t_observed": fact.t_observed, + "t_valid": fact.t_valid, + "t_invalid": fact.t_invalid, + }, + ) + for fact in facts + ] + ) + return results_all diff --git a/jitmind/schemas/__init__.py b/jitmind/schemas/__init__.py index 042ecba..ccd9e27 100644 --- a/jitmind/schemas/__init__.py +++ b/jitmind/schemas/__init__.py @@ -14,6 +14,7 @@ from .result import Result, EnoughDecision, ReflectionDecision, ResearchOutput, GenerateRequests from .memory_ops import MemoryOperationDecision, ExtractedEntity, ExtractedRelation from .self_rag import SelfRAGDecision +from .graph_fact import GraphFact # ============================= # Model rebuilding for forward references @@ -41,5 +42,6 @@ "ToolResult", "Tool", "ToolRegistry", "Result", "EnoughDecision", "ReflectionDecision", "ResearchOutput", "GenerateRequests", "MemoryOperationDecision", "ExtractedEntity", "ExtractedRelation", "SelfRAGDecision", + "GraphFact", "PLANNING_SCHEMA", "INTEGRATE_SCHEMA", "INFO_CHECK_SCHEMA", "GENERATE_REQUESTS_SCHEMA", "MEMORY_OP_SCHEMA", "SELF_RAG_SCHEMA", ] diff --git a/jitmind/schemas/graph_fact.py b/jitmind/schemas/graph_fact.py new file mode 100644 index 0000000..223bb42 --- /dev/null +++ b/jitmind/schemas/graph_fact.py @@ -0,0 +1,27 @@ +"""Typed evidence returned by temporal knowledge-graph queries.""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + +from pydantic import BaseModel, Field + + +class GraphFact(BaseModel): + id: str = Field(..., description="Stable observation id") + head: str + relation: str + tail: str + source_memory_id: Optional[str] = None + source_page_id: Optional[str] = None + t_observed: Optional[str] = None + t_valid: Optional[str] = None + t_invalid: Optional[str] = None + meta: Dict[str, Any] = Field(default_factory=dict) + + @property + def statement(self) -> str: + return f"{self.head} {self.relation} {self.tail}" + + +__all__ = ["GraphFact"] diff --git a/tests/test_temporal_graph_facts.py b/tests/test_temporal_graph_facts.py new file mode 100644 index 0000000..b81a9b1 --- /dev/null +++ b/tests/test_temporal_graph_facts.py @@ -0,0 +1,126 @@ +from datetime import datetime, timezone + +import pytest + +from jitmind.graph.graph_store import GraphMemoryStore +from jitmind.retriever.graph_retriever import GraphRetriever +from jitmind.schemas import GraphFact + + +class _Result(list): + pass + + +class _Session: + def __init__(self, rows): + self.rows = rows + self.query = None + self.params = None + + def run(self, query, **params): + self.query = query + self.params = params + return _Result(self.rows) + + def close(self): + return None + + +class _Driver: + def __init__(self, session): + self._session = session + + def session(self, database): + assert database == "neo4j" + return self._session + + +def _store(rows): + session = _Session(rows) + store = GraphMemoryStore.__new__(GraphMemoryStore) + store._database = "neo4j" + store._driver = _Driver(session) + return store, session + + +def test_query_facts_builds_a_bitemporal_snapshot() -> None: + rows = [ + { + "id": "fact-1", + "head": "Alice", + "relation": "WORKS_AT", + "tail": "Acme", + "source_memory_id": "memory-1", + "source_page_id": "7", + "t_observed": "2026-02-02T00:00:00+00:00", + "t_valid": "2026-02-01T00:00:00+00:00", + "t_invalid": None, + } + ] + store, session = _store(rows) + + facts = store.query_facts( + ["Alice"], + relation_types=["works at"], + valid_at="2026-02-03T00:00:00Z", + observed_at=datetime(2026, 2, 4, tzinfo=timezone.utc), + namespace=("tenant", "red"), + ) + + assert [fact.statement for fact in facts] == ["Alice WORKS_AT Acme"] + assert session.params == { + "names": ["Alice"], + "types": ["WORKS_AT"], + "valid_at": "2026-02-03T00:00:00+00:00", + "observed_at": "2026-02-04T00:00:00+00:00", + "namespace": ["tenant", "red"], + "limit": 20, + } + assert "r.t_invalid" in session.query + assert "m.namespace" in session.query + + +def test_each_relation_observation_has_a_distinct_stable_id() -> None: + store, _ = _store([]) + first = store._relation_id("memory-1", "Person:Alice", "WORKS_AT", "Org:Acme") + replay = store._relation_id("memory-1", "Person:Alice", "WORKS_AT", "Org:Acme") + later = store._relation_id("memory-2", "Person:Alice", "WORKS_AT", "Org:Acme") + + assert first == replay + assert first != later + + +def test_invalid_temporal_query_fails_before_hitting_neo4j() -> None: + store, session = _store([]) + with pytest.raises(ValueError, match="valid_at"): + store.query_facts(valid_at="not-a-time") + assert session.query is None + + +def test_graph_retriever_preserves_fact_provenance() -> None: + class FakeStore: + def query_facts(self, names, **kwargs): + assert names == ["Alice"] + assert kwargs["valid_at"] == "2026-02-03T00:00:00Z" + return [ + GraphFact( + id="fact-1", + head="Alice", + relation="WORKS_AT", + tail="Acme", + source_memory_id="memory-1", + source_page_id="7", + t_valid="2026-02-01T00:00:00Z", + ) + ] + + retriever = GraphRetriever({"graph_store": FakeStore()}) + hits = retriever.search_temporal( + ["Alice"], valid_at="2026-02-03T00:00:00Z" + )[0] + + assert hits[0].page_id == "7" + assert hits[0].snippet == "Alice WORKS_AT Acme" + assert hits[0].meta["fact_id"] == "fact-1" + assert hits[0].meta["memory_id"] == "memory-1" +