1"""Memory System Implementation.
2 
3Provides composable building blocks for agent memory: vector stores with
4metadata indexing, property graphs for entity relationships, and temporal
5knowledge graphs for facts that change over time.
6 
7Use when:
8    - Building a memory persistence layer for an agent that must retain
9      knowledge across sessions.
10    - Prototyping memory architectures before committing to a production
11      framework (Mem0, Zep/Graphiti, Letta, Cognee).
12    - Combining semantic search with graph-based entity retrieval in a
13      single integrated system.
14 
15Typical usage::
16 
17    from memory_store import IntegratedMemorySystem
18    mem = IntegratedMemorySystem()
19    mem.start_session("session-001")
20    mem.store_fact("Alice prefers dark mode", entity="Alice")
21    results = mem.retrieve_memories("theme preference")
22"""
23 
24import hashlib
25import json
26from datetime import datetime
27from typing import Any, Dict, List, Optional
28 
29import numpy as np
30 
31__all__ = [
32    "VectorStore",
33    "PropertyGraph",
34    "TemporalKnowledgeGraph",
35    "IntegratedMemorySystem",
36]
37 
38 
39class VectorStore:
40    """Simple vector store with metadata indexing.
41 
42    Use when: the agent needs semantic similarity search over stored facts
43    with optional entity and temporal filtering.
44    """
45 
46    def __init__(self, dimension: int = 768) -> None:
47        self.dimension: int = dimension
48        self.vectors: List[np.ndarray] = []
49        self.metadata: List[Dict[str, Any]] = []
50        self.entity_index: Dict[str, List[int]] = {}
51        self.time_index: Dict[str, List[int]] = {}
52 
53    def add(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> int:
54        """Add document to store.
55 
56        Use when: persisting a new fact or observation that the agent should
57        be able to retrieve later via semantic search.
58        """
59        metadata = metadata or {}
60        embedding: np.ndarray = self._embed(text)
61        index: int = len(self.vectors)
62 
63        self.vectors.append(embedding)
64        self.metadata.append(metadata)
65 
66        # Index by entity
67        if "entity" in metadata:
68            entity: str = metadata["entity"]
69            if entity not in self.entity_index:
70                self.entity_index[entity] = []
71            self.entity_index[entity].append(index)
72 
73        # Index by time
74        if "valid_from" in metadata:
75            time_key: str = self._time_key(metadata["valid_from"])
76            if time_key not in self.time_index:
77                self.time_index[time_key] = []
78            self.time_index[time_key].append(index)
79 
80        return index
81 
82    def search(
83        self,
84        query: str,
85        limit: int = 5,
86        filters: Optional[Dict[str, Any]] = None,
87    ) -> List[Dict[str, Any]]:
88        """Search for similar documents.
89 
90        Use when: retrieving memories relevant to a query, optionally
91        narrowed by metadata filters (entity, session, time range).
92        """
93        query_embedding: np.ndarray = self._embed(query)
94 
95        scores: List[tuple[int, float]] = []
96        for i, vec in enumerate(self.vectors):
97            score: float = float(
98                np.dot(query_embedding, vec)
99                / (np.linalg.norm(query_embedding) * np.linalg.norm(vec) + 1e-8)
100            )
101 
102            # Apply filters
103            if filters and not self._matches_filters(self.metadata[i], filters):
104                score = -1.0
105 
106            scores.append((i, score))
107 
108        scores.sort(key=lambda x: x[1], reverse=True)
109 
110        results: List[Dict[str, Any]] = []
111        for idx, score in scores[:limit]:
112            if score > 0:
113                results.append(
114                    {
115                        "index": idx,
116                        "score": score,
117                        "text": self.metadata[idx].get("text", ""),
118                        "metadata": self.metadata[idx],
119                    }
120                )
121 
122        return results
123 
124    def search_by_entity(
125        self, entity: str, query: str = "", limit: int = 5
126    ) -> List[Dict[str, Any]]:
127        """Search within specific entity.
128 
129        Use when: the agent needs all memories associated with a known
130        entity, optionally ranked by relevance to a query.
131        """
132        indices: List[int] = self.entity_index.get(entity, [])
133 
134        if not indices:
135            return []
136 
137        if query:
138            query_embedding: np.ndarray = self._embed(query)
139            scored: List[tuple[int, float, Dict[str, Any]]] = []
140            for i in indices:
141                vec: np.ndarray = self.vectors[i]
142                score: float = float(
143                    np.dot(query_embedding, vec)
144                    / (np.linalg.norm(query_embedding) * np.linalg.norm(vec) + 1e-8)
145                )
146                scored.append((i, score, self.metadata[i]))
147 
148            scored.sort(key=lambda x: x[1], reverse=True)
149            return [
150                {"index": i, "score": s, "metadata": m}
151                for i, s, m in scored[:limit]
152            ]
153        else:
154            return [
155                {"index": i, "score": 1.0, "metadata": self.metadata[i]}
156                for i in indices[:limit]
157            ]
158 
159    def _embed(self, text: str) -> np.ndarray:
160        """Generate embedding for text.
161 
162        In production, replace with an actual embedding model. This
163        deterministic stub uses the text hash as a random seed so that
164        identical texts always produce identical vectors. Uses a local
165        RNG to avoid corrupting global numpy random state.
166        """
167        rng = np.random.default_rng(hash(text) % (2**32))
168        return rng.standard_normal(self.dimension)
169 
170    def _time_key(self, timestamp: Any) -> str:
171        """Create time key for indexing."""
172        if isinstance(timestamp, datetime):
173            return timestamp.strftime("%Y-%m")
174        return str(timestamp)
175 
176    def _matches_filters(self, metadata: Dict[str, Any], filters: Dict[str, Any]) -> bool:
177        """Check if metadata matches filters."""
178        for key, value in filters.items():
179            if key not in metadata:
180                return False
181            if isinstance(value, list):
182                if metadata[key] not in value:
183                    return False
184            elif metadata[key] != value:
185                return False
186        return True
187 
188 
189class PropertyGraph:
190    """Simple property graph storage.
191 
192    Use when: the agent needs to maintain entity relationships and
193    traverse connections between nodes (e.g., "find all projects
194    associated with this user").
195    """
196 
197    def __init__(self) -> None:
198        self.nodes: Dict[str, Dict[str, Any]] = {}
199        self.edges: Dict[str, Dict[str, Any]] = {}
200        self.entity_registry: Dict[str, str] = {}  # name -> node_id
201        self.node_index: Dict[str, List[str]] = {}  # label -> node_ids
202        self.edge_index: Dict[str, List[str]] = {}  # type -> edge_ids
203 
204    def get_or_create_node(
205        self, name: str, label: str = "Entity", properties: Optional[Dict[str, Any]] = None
206    ) -> str:
207        """Get existing node by name, or create a new one.
208 
209        Use when: storing an entity that may already exist. The entity
210        registry ensures identity is maintained across interactions
211        ("John Doe" always maps to the same node).
212        """
213        if name in self.entity_registry:
214            node_id: str = self.entity_registry[name]
215            if properties:
216                self.nodes[node_id]["properties"].update(properties)
217            return node_id
218        node_id = self.create_node(label, {**(properties or {}), "name": name})
219        self.entity_registry[name] = node_id
220        return node_id
221 
222    def create_node(self, label: str, properties: Optional[Dict[str, Any]] = None) -> str:
223        """Create node with label and properties.
224 
225        Use when: adding a new entity to the graph that does not need
226        identity deduplication (prefer get_or_create_node otherwise).
227        """
228        node_id: str = hashlib.md5(f"{label}{datetime.now().isoformat()}".encode()).hexdigest()[:16]
229 
230        self.nodes[node_id] = {
231            "id": node_id,
232            "label": label,
233            "properties": properties or {},
234            "created_at": datetime.now().isoformat(),
235        }
236 
237        if label not in self.node_index:
238            self.node_index[label] = []
239        self.node_index[label].append(node_id)
240 
241        return node_id
242 
243    def create_relationship(
244        self,
245        source_id: str,
246        rel_type: str,
247        target_id: str,
248        properties: Optional[Dict[str, Any]] = None,
249    ) -> str:
250        """Create directed relationship between nodes.
251 
252        Use when: recording a connection between two entities (e.g.,
253        WORKS_AT, LIVES_IN, DEPENDS_ON).
254        """
255        if source_id not in self.nodes:
256            raise ValueError(f"Unknown source node: {source_id}")
257        if target_id not in self.nodes:
258            raise ValueError(f"Unknown target node: {target_id}")
259 
260        edge_id: str = hashlib.md5(
261            f"{source_id}{rel_type}{target_id}{datetime.now().isoformat()}".encode()
262        ).hexdigest()[:16]
263 
264        self.edges[edge_id] = {
265            "id": edge_id,
266            "source": source_id,
267            "target": target_id,
268            "type": rel_type,
269            "properties": properties or {},
270            "created_at": datetime.now().isoformat(),
271        }
272 
273        if rel_type not in self.edge_index:
274            self.edge_index[rel_type] = []
275        self.edge_index[rel_type].append(edge_id)
276 
277        return edge_id
278 
279    def query(self, pattern: Dict[str, Any]) -> List[Dict[str, Any]]:
280        """Query graph with simple pattern matching.
281 
282        Use when: finding relationships that match a structural pattern
283        (e.g., all WORKS_AT edges from Person nodes).
284        """
285        results: List[Dict[str, Any]] = []
286 
287        # Match by edge type
288        if "type" in pattern:
289            edge_ids: List[str] = self.edge_index.get(pattern["type"], [])
290            for eid in edge_ids:
291                edge: Dict[str, Any] = self.edges[eid]
292                source: Dict[str, Any] = self.nodes.get(edge["source"], {})
293                target: Dict[str, Any] = self.nodes.get(edge["target"], {})
294 
295                # Match source label
296                if "source_label" in pattern:
297                    if source.get("label") != pattern["source_label"]:
298                        continue
299 
300                # Match target label
301                if "target_label" in pattern:
302                    if target.get("label") != pattern["target_label"]:
303                        continue
304 
305                results.append({"source": source, "edge": edge, "target": target})
306 
307        return results
308 
309    def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
310        """Get node by ID."""
311        return self.nodes.get(node_id)
312 
313    def get_relationships(
314        self, node_id: str, direction: str = "both"
315    ) -> List[Dict[str, Any]]:
316        """Get relationships for a node.
317 
318        Use when: retrieving all connections for a given entity to build
319        a complete entity context.
320        """
321        relationships: List[Dict[str, Any]] = []
322 
323        for edge in self.edges.values():
324            if direction in ["outgoing", "both"] and edge["source"] == node_id:
325                relationships.append(
326                    {
327                        "edge": edge,
328                        "target": self.nodes.get(edge["target"]),
329                        "direction": "outgoing",
330                    }
331                )
332            if direction in ["incoming", "both"] and edge["target"] == node_id:
333                relationships.append(
334                    {
335                        "edge": edge,
336                        "source": self.nodes.get(edge["source"]),
337                        "direction": "incoming",
338                    }
339                )
340 
341        return relationships
342 
343 
344class TemporalKnowledgeGraph(PropertyGraph):
345    """Property graph with temporal validity for facts.
346 
347    Use when: the agent must track facts that change over time and
348    answer time-scoped queries (e.g., "where did the user live in
349    March 2024?").
350    """
351 
352    def create_temporal_relationship(
353        self,
354        source_id: str,
355        rel_type: str,
356        target_id: str,
357        valid_from: datetime,
358        valid_until: Optional[datetime] = None,
359        properties: Optional[Dict[str, Any]] = None,
360    ) -> str:
361        """Create relationship with temporal validity.
362 
363        Use when: recording a fact that has a known start time and
364        may expire (e.g., employment, address, subscription status).
365        """
366        edge_id: str = super().create_relationship(
367            source_id, rel_type, target_id, properties
368        )
369 
370        # Add temporal properties
371        self.edges[edge_id]["valid_from"] = valid_from.isoformat()
372        self.edges[edge_id]["valid_until"] = (
373            valid_until.isoformat() if valid_until else None
374        )
375 
376        return edge_id
377 
378    def query_at_time(
379        self, query: Dict[str, Any], query_time: datetime
380    ) -> List[Dict[str, Any]]:
381        """Query graph state at specific time.
382 
383        Use when: answering point-in-time questions about entities
384        (e.g., "what was true on date X?").
385        """
386        results: List[Dict[str, Any]] = []
387 
388        # Get base query results
389        base_results: List[Dict[str, Any]] = self.query(query)
390 
391        for result in base_results:
392            edge: Dict[str, Any] = result["edge"]
393            valid_from: datetime = datetime.fromisoformat(
394                edge.get("valid_from", "1970-01-01")
395            )
396            valid_until: Optional[str] = edge.get("valid_until")
397 
398            # Check temporal validity
399            if valid_from <= query_time:
400                if valid_until is None or datetime.fromisoformat(valid_until) > query_time:
401                    results.append(
402                        {
403                            **result,
404                            "valid_from": valid_from,
405                            "valid_until": valid_until,
406                        }
407                    )
408 
409        return results
410 
411    def query_time_range(
412        self,
413        query: Dict[str, Any],
414        start_time: datetime,
415        end_time: datetime,
416    ) -> List[Dict[str, Any]]:
417        """Query facts valid during time range.
418 
419        Use when: retrieving all facts that overlap with a given time
420        window (e.g., "what changed between January and June?").
421        """
422        results: List[Dict[str, Any]] = []
423 
424        base_results: List[Dict[str, Any]] = self.query(query)
425 
426        for result in base_results:
427            edge: Dict[str, Any] = result["edge"]
428            valid_from: datetime = datetime.fromisoformat(
429                edge.get("valid_from", "1970-01-01")
430            )
431            valid_until: Optional[str] = edge.get("valid_until")
432 
433            # Check if overlaps with query range
434            until_dt: datetime = (
435                datetime.fromisoformat(valid_until) if valid_until else datetime.max
436            )
437 
438            if until_dt >= start_time and valid_from <= end_time:
439                results.append(
440                    {
441                        **result,
442                        "valid_from": valid_from,
443                        "valid_until": valid_until,
444                    }
445                )
446 
447        return results
448 
449 
450# ---------------------------------------------------------------------------
451# Memory System Integration
452# ---------------------------------------------------------------------------
453 
454 
455class IntegratedMemorySystem:
456    """Integrated memory system combining vector store and graph.
457 
458    Use when: the agent needs both semantic search over facts and
459    graph-based entity relationship traversal in a single unified
460    interface. This class composes VectorStore and TemporalKnowledgeGraph,
461    enriching vector search results with graph context.
462    """
463 
464    def __init__(self) -> None:
465        self.vector_store: VectorStore = VectorStore()
466        self.graph: TemporalKnowledgeGraph = TemporalKnowledgeGraph()
467        self.session_id: str = ""
468 
469    def start_session(self, session_id: str) -> None:
470        """Start a new memory session.
471 
472        Use when: beginning a new conversation or task that should
473        scope its memories to a distinct session identifier.
474        """
475        self.session_id = session_id
476 
477    def store_fact(
478        self,
479        fact: str,
480        entity: str,
481        timestamp: Optional[datetime] = None,
482        relationships: Optional[List[Dict[str, Any]]] = None,
483    ) -> None:
484        """Store a fact with entity and relationships.
485 
486        Use when: the agent observes a new piece of information that
487        should be persisted for future retrieval. Stores in both the
488        vector store (for semantic search) and the graph (for entity
489        traversal).
490        """
491        # Store in vector store
492        self.vector_store.add(
493            fact,
494            {
495                "text": fact,
496                "entity": entity,
497                "valid_from": (timestamp or datetime.now()).isoformat(),
498                "session_id": self.session_id,
499            },
500        )
501 
502        # Get or create entity node (uses registry for identity)
503        entity_node_id: str = self.graph.get_or_create_node(entity)
504 
505        # Create relationships
506        if relationships:
507            for rel in relationships:
508                target_node_id: str = self.graph.get_or_create_node(rel["target"])
509                self.graph.create_relationship(
510                    entity_node_id,
511                    rel["type"],
512                    target_node_id,
513                    properties=rel.get("properties", {}),
514                )
515 
516    def retrieve_memories(
517        self,
518        query: str,
519        entity_filter: Optional[str] = None,
520        time_filter: Optional[Dict[str, Any]] = None,
521        limit: int = 5,
522    ) -> List[Dict[str, Any]]:
523        """Retrieve memories matching query.
524 
525        Use when: the agent needs to recall previously stored facts,
526        optionally filtered by entity or time. Results are enriched
527        with graph relationships for each matched entity.
528        """
529        # Vector search
530        filters: Dict[str, Any] = {"session_id": self.session_id}
531        if entity_filter:
532            filters["entity"] = entity_filter
533 
534        results: List[Dict[str, Any]] = self.vector_store.search(
535            query, limit=limit, filters=filters
536        )
537 
538        # Enrich with graph relationships
539        for result in results:
540            entity: Optional[str] = result["metadata"].get("entity")
541            if entity:
542                node_id: Optional[str] = self.graph.entity_registry.get(entity)
543                if node_id:
544                    result["relationships"] = self.graph.get_relationships(node_id)
545 
546        return results
547 
548    def retrieve_entity_context(self, entity: str) -> Dict[str, Any]:
549        """Retrieve complete context for an entity.
550 
551        Use when: the agent needs a full picture of a single entity
552        including its properties, all relationships, and associated
553        vector memories.
554        """
555        node_id: Optional[str] = self.graph.entity_registry.get(entity)
556 
557        # Get entity node
558        entity_node: Optional[Dict[str, Any]] = (
559            self.graph.get_node(node_id) if node_id else None
560        )
561 
562        # Get relationships
563        relationships: List[Dict[str, Any]] = (
564            self.graph.get_relationships(node_id) if node_id else []
565        )
566 
567        # Get vector memories
568        memories: List[Dict[str, Any]] = self.vector_store.search_by_entity(
569            entity, limit=10
570        )
571 
572        return {
573            "entity": entity_node,
574            "relationships": relationships,
575            "memories": memories,
576        }
577 
578    def consolidate(self) -> None:
579        """Consolidate memories and remove outdated information.
580 
581        Use when: memory count exceeds a threshold, retrieval quality
582        degrades, or on a scheduled interval. In production, implement:
583        - Merge related facts into summaries
584        - Update validity periods on stale entries
585        - Archive obsolete facts (invalidate, do not discard)
586        """
587        pass
588 
589 
590if __name__ == "__main__":
591    # Quick smoke test demonstrating the integrated memory system.
592    mem = IntegratedMemorySystem()
593    mem.start_session("demo-session")
594 
595    # Store facts with entity relationships
596    mem.store_fact(
597        "Alice prefers dark mode",
598        entity="Alice",
599        relationships=[{"target": "dark mode", "type": "PREFERS"}],
600    )
601    mem.store_fact(
602        "Alice works at Acme Corp",
603        entity="Alice",
604        relationships=[{"target": "Acme Corp", "type": "WORKS_AT"}],
605    )
606 
607    # Semantic retrieval
608    results = mem.retrieve_memories("theme preference")
609    print(f"Search results: {len(results)} memories found")
610    for r in results:
611        print(f"  score={r['score']:.3f}  text={r['text']}")
612 
613    # Entity context
614    context = mem.retrieve_entity_context("Alice")
615    print(f"\nAlice context: {len(context['relationships'])} relationships, "
616          f"{len(context['memories'])} memories")
617