@bastani/atomic 0.5.11 → 0.5.12-1

package/.agents/skills/memory-systems/scripts/memory_store.py

@@ -0,0 +1,616 @@
+"""Memory System Implementation.
+
+Provides composable building blocks for agent memory: vector stores with
+metadata indexing, property graphs for entity relationships, and temporal
+knowledge graphs for facts that change over time.
+
+Use when:
+    - Building a memory persistence layer for an agent that must retain
+      knowledge across sessions.
+    - Prototyping memory architectures before committing to a production
+      framework (Mem0, Zep/Graphiti, Letta, Cognee).
+    - Combining semantic search with graph-based entity retrieval in a
+      single integrated system.
+
+Typical usage::
+
+    from memory_store import IntegratedMemorySystem
+    mem = IntegratedMemorySystem()
+    mem.start_session("session-001")
+    mem.store_fact("Alice prefers dark mode", entity="Alice")
+    results = mem.retrieve_memories("theme preference")
+"""
+
+import hashlib
+import json
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+__all__ = [
+    "VectorStore",
+    "PropertyGraph",
+    "TemporalKnowledgeGraph",
+    "IntegratedMemorySystem",
+]
+
+
+class VectorStore:
+    """Simple vector store with metadata indexing.
+
+    Use when: the agent needs semantic similarity search over stored facts
+    with optional entity and temporal filtering.
+    """
+
+    def __init__(self, dimension: int = 768) -> None:
+        self.dimension: int = dimension
+        self.vectors: List[np.ndarray] = []
+        self.metadata: List[Dict[str, Any]] = []
+        self.entity_index: Dict[str, List[int]] = {}
+        self.time_index: Dict[str, List[int]] = {}
+
+    def add(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> int:
+        """Add document to store.
+
+        Use when: persisting a new fact or observation that the agent should
+        be able to retrieve later via semantic search.
+        """
+        metadata = metadata or {}
+        embedding: np.ndarray = self._embed(text)
+        index: int = len(self.vectors)
+
+        self.vectors.append(embedding)
+        self.metadata.append(metadata)
+
+        # Index by entity
+        if "entity" in metadata:
+            entity: str = metadata["entity"]
+            if entity not in self.entity_index:
+                self.entity_index[entity] = []
+            self.entity_index[entity].append(index)
+
+        # Index by time
+        if "valid_from" in metadata:
+            time_key: str = self._time_key(metadata["valid_from"])
+            if time_key not in self.time_index:
+                self.time_index[time_key] = []
+            self.time_index[time_key].append(index)
+
+        return index
+
+    def search(
+        self,
+        query: str,
+        limit: int = 5,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> List[Dict[str, Any]]:
+        """Search for similar documents.
+
+        Use when: retrieving memories relevant to a query, optionally
+        narrowed by metadata filters (entity, session, time range).
+        """
+        query_embedding: np.ndarray = self._embed(query)
+
+        scores: List[tuple[int, float]] = []
+        for i, vec in enumerate(self.vectors):
+            score: float = float(
+                np.dot(query_embedding, vec)
+                / (np.linalg.norm(query_embedding) * np.linalg.norm(vec) + 1e-8)
+            )
+
+            # Apply filters
+            if filters and not self._matches_filters(self.metadata[i], filters):
+                score = -1.0
+
+            scores.append((i, score))
+
+        scores.sort(key=lambda x: x[1], reverse=True)
+
+        results: List[Dict[str, Any]] = []
+        for idx, score in scores[:limit]:
+            if score > 0:
+                results.append(
+                    {
+                        "index": idx,
+                        "score": score,
+                        "text": self.metadata[idx].get("text", ""),
+                        "metadata": self.metadata[idx],
+                    }
+                )
+
+        return results
+
+    def search_by_entity(
+        self, entity: str, query: str = "", limit: int = 5
+    ) -> List[Dict[str, Any]]:
+        """Search within specific entity.
+
+        Use when: the agent needs all memories associated with a known
+        entity, optionally ranked by relevance to a query.
+        """
+        indices: List[int] = self.entity_index.get(entity, [])
+
+        if not indices:
+            return []
+
+        if query:
+            query_embedding: np.ndarray = self._embed(query)
+            scored: List[tuple[int, float, Dict[str, Any]]] = []
+            for i in indices:
+                vec: np.ndarray = self.vectors[i]
+                score: float = float(
+                    np.dot(query_embedding, vec)
+                    / (np.linalg.norm(query_embedding) * np.linalg.norm(vec) + 1e-8)
+                )
+                scored.append((i, score, self.metadata[i]))
+
+            scored.sort(key=lambda x: x[1], reverse=True)
+            return [
+                {"index": i, "score": s, "metadata": m}
+                for i, s, m in scored[:limit]
+            ]
+        else:
+            return [
+                {"index": i, "score": 1.0, "metadata": self.metadata[i]}
+                for i in indices[:limit]
+            ]
+
+    def _embed(self, text: str) -> np.ndarray:
+        """Generate embedding for text.
+
+        In production, replace with an actual embedding model. This
+        deterministic stub uses the text hash as a random seed so that
+        identical texts always produce identical vectors. Uses a local
+        RNG to avoid corrupting global numpy random state.
+        """
+        rng = np.random.default_rng(hash(text) % (2**32))
+        return rng.standard_normal(self.dimension)
+
+    def _time_key(self, timestamp: Any) -> str:
+        """Create time key for indexing."""
+        if isinstance(timestamp, datetime):
+            return timestamp.strftime("%Y-%m")
+        return str(timestamp)
+
+    def _matches_filters(self, metadata: Dict[str, Any], filters: Dict[str, Any]) -> bool:
+        """Check if metadata matches filters."""
+        for key, value in filters.items():
+            if key not in metadata:
+                return False
+            if isinstance(value, list):
+                if metadata[key] not in value:
+                    return False
+            elif metadata[key] != value:
+                return False
+        return True
+
+
+class PropertyGraph:
+    """Simple property graph storage.
+
+    Use when: the agent needs to maintain entity relationships and
+    traverse connections between nodes (e.g., "find all projects
+    associated with this user").
+    """
+
+    def __init__(self) -> None:
+        self.nodes: Dict[str, Dict[str, Any]] = {}
+        self.edges: Dict[str, Dict[str, Any]] = {}
+        self.entity_registry: Dict[str, str] = {}  # name -> node_id
+        self.node_index: Dict[str, List[str]] = {}  # label -> node_ids
+        self.edge_index: Dict[str, List[str]] = {}  # type -> edge_ids
+
+    def get_or_create_node(
+        self, name: str, label: str = "Entity", properties: Optional[Dict[str, Any]] = None
+    ) -> str:
+        """Get existing node by name, or create a new one.
+
+        Use when: storing an entity that may already exist. The entity
+        registry ensures identity is maintained across interactions
+        ("John Doe" always maps to the same node).
+        """
+        if name in self.entity_registry:
+            node_id: str = self.entity_registry[name]
+            if properties:
+                self.nodes[node_id]["properties"].update(properties)
+            return node_id
+        node_id = self.create_node(label, {**(properties or {}), "name": name})
+        self.entity_registry[name] = node_id
+        return node_id
+
+    def create_node(self, label: str, properties: Optional[Dict[str, Any]] = None) -> str:
+        """Create node with label and properties.
+
+        Use when: adding a new entity to the graph that does not need
+        identity deduplication (prefer get_or_create_node otherwise).
+        """
+        node_id: str = hashlib.md5(f"{label}{datetime.now().isoformat()}".encode()).hexdigest()[:16]
+
+        self.nodes[node_id] = {
+            "id": node_id,
+            "label": label,
+            "properties": properties or {},
+            "created_at": datetime.now().isoformat(),
+        }
+
+        if label not in self.node_index:
+            self.node_index[label] = []
+        self.node_index[label].append(node_id)
+
+        return node_id
+
+    def create_relationship(
+        self,
+        source_id: str,
+        rel_type: str,
+        target_id: str,
+        properties: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        """Create directed relationship between nodes.
+
+        Use when: recording a connection between two entities (e.g.,
+        WORKS_AT, LIVES_IN, DEPENDS_ON).
+        """
+        if source_id not in self.nodes:
+            raise ValueError(f"Unknown source node: {source_id}")
+        if target_id not in self.nodes:
+            raise ValueError(f"Unknown target node: {target_id}")
+
+        edge_id: str = hashlib.md5(
+            f"{source_id}{rel_type}{target_id}{datetime.now().isoformat()}".encode()
+        ).hexdigest()[:16]
+
+        self.edges[edge_id] = {
+            "id": edge_id,
+            "source": source_id,
+            "target": target_id,
+            "type": rel_type,
+            "properties": properties or {},
+            "created_at": datetime.now().isoformat(),
+        }
+
+        if rel_type not in self.edge_index:
+            self.edge_index[rel_type] = []
+        self.edge_index[rel_type].append(edge_id)
+
+        return edge_id
+
+    def query(self, pattern: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """Query graph with simple pattern matching.
+
+        Use when: finding relationships that match a structural pattern
+        (e.g., all WORKS_AT edges from Person nodes).
+        """
+        results: List[Dict[str, Any]] = []
+
+        # Match by edge type
+        if "type" in pattern:
+            edge_ids: List[str] = self.edge_index.get(pattern["type"], [])
+            for eid in edge_ids:
+                edge: Dict[str, Any] = self.edges[eid]
+                source: Dict[str, Any] = self.nodes.get(edge["source"], {})
+                target: Dict[str, Any] = self.nodes.get(edge["target"], {})
+
+                # Match source label
+                if "source_label" in pattern:
+                    if source.get("label") != pattern["source_label"]:
+                        continue
+
+                # Match target label
+                if "target_label" in pattern:
+                    if target.get("label") != pattern["target_label"]:
+                        continue
+
+                results.append({"source": source, "edge": edge, "target": target})
+
+        return results
+
+    def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
+        """Get node by ID."""
+        return self.nodes.get(node_id)
+
+    def get_relationships(
+        self, node_id: str, direction: str = "both"
+    ) -> List[Dict[str, Any]]:
+        """Get relationships for a node.
+
+        Use when: retrieving all connections for a given entity to build
+        a complete entity context.
+        """
+        relationships: List[Dict[str, Any]] = []
+
+        for edge in self.edges.values():
+            if direction in ["outgoing", "both"] and edge["source"] == node_id:
+                relationships.append(
+                    {
+                        "edge": edge,
+                        "target": self.nodes.get(edge["target"]),
+                        "direction": "outgoing",
+                    }
+                )
+            if direction in ["incoming", "both"] and edge["target"] == node_id:
+                relationships.append(
+                    {
+                        "edge": edge,
+                        "source": self.nodes.get(edge["source"]),
+                        "direction": "incoming",
+                    }
+                )
+
+        return relationships
+
+
+class TemporalKnowledgeGraph(PropertyGraph):
+    """Property graph with temporal validity for facts.
+
+    Use when: the agent must track facts that change over time and
+    answer time-scoped queries (e.g., "where did the user live in
+    March 2024?").
+    """
+
+    def create_temporal_relationship(
+        self,
+        source_id: str,
+        rel_type: str,
+        target_id: str,
+        valid_from: datetime,
+        valid_until: Optional[datetime] = None,
+        properties: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        """Create relationship with temporal validity.
+
+        Use when: recording a fact that has a known start time and
+        may expire (e.g., employment, address, subscription status).
+        """
+        edge_id: str = super().create_relationship(
+            source_id, rel_type, target_id, properties
+        )
+
+        # Add temporal properties
+        self.edges[edge_id]["valid_from"] = valid_from.isoformat()
+        self.edges[edge_id]["valid_until"] = (
+            valid_until.isoformat() if valid_until else None
+        )
+
+        return edge_id
+
+    def query_at_time(
+        self, query: Dict[str, Any], query_time: datetime
+    ) -> List[Dict[str, Any]]:
+        """Query graph state at specific time.
+
+        Use when: answering point-in-time questions about entities
+        (e.g., "what was true on date X?").
+        """
+        results: List[Dict[str, Any]] = []
+
+        # Get base query results
+        base_results: List[Dict[str, Any]] = self.query(query)
+
+        for result in base_results:
+            edge: Dict[str, Any] = result["edge"]
+            valid_from: datetime = datetime.fromisoformat(
+                edge.get("valid_from", "1970-01-01")
+            )
+            valid_until: Optional[str] = edge.get("valid_until")
+
+            # Check temporal validity
+            if valid_from <= query_time:
+                if valid_until is None or datetime.fromisoformat(valid_until) > query_time:
+                    results.append(
+                        {
+                            **result,
+                            "valid_from": valid_from,
+                            "valid_until": valid_until,
+                        }
+                    )
+
+        return results
+
+    def query_time_range(
+        self,
+        query: Dict[str, Any],
+        start_time: datetime,
+        end_time: datetime,
+    ) -> List[Dict[str, Any]]:
+        """Query facts valid during time range.
+
+        Use when: retrieving all facts that overlap with a given time
+        window (e.g., "what changed between January and June?").
+        """
+        results: List[Dict[str, Any]] = []
+
+        base_results: List[Dict[str, Any]] = self.query(query)
+
+        for result in base_results:
+            edge: Dict[str, Any] = result["edge"]
+            valid_from: datetime = datetime.fromisoformat(
+                edge.get("valid_from", "1970-01-01")
+            )
+            valid_until: Optional[str] = edge.get("valid_until")
+
+            # Check if overlaps with query range
+            until_dt: datetime = (
+                datetime.fromisoformat(valid_until) if valid_until else datetime.max
+            )
+
+            if until_dt >= start_time and valid_from <= end_time:
+                results.append(
+                    {
+                        **result,
+                        "valid_from": valid_from,
+                        "valid_until": valid_until,
+                    }
+                )
+
+        return results
+
+
+# ---------------------------------------------------------------------------
+# Memory System Integration
+# ---------------------------------------------------------------------------
+
+
+class IntegratedMemorySystem:
+    """Integrated memory system combining vector store and graph.
+
+    Use when: the agent needs both semantic search over facts and
+    graph-based entity relationship traversal in a single unified
+    interface. This class composes VectorStore and TemporalKnowledgeGraph,
+    enriching vector search results with graph context.
+    """
+
+    def __init__(self) -> None:
+        self.vector_store: VectorStore = VectorStore()
+        self.graph: TemporalKnowledgeGraph = TemporalKnowledgeGraph()
+        self.session_id: str = ""
+
+    def start_session(self, session_id: str) -> None:
+        """Start a new memory session.
+
+        Use when: beginning a new conversation or task that should
+        scope its memories to a distinct session identifier.
+        """
+        self.session_id = session_id
+
+    def store_fact(
+        self,
+        fact: str,
+        entity: str,
+        timestamp: Optional[datetime] = None,
+        relationships: Optional[List[Dict[str, Any]]] = None,
+    ) -> None:
+        """Store a fact with entity and relationships.
+
+        Use when: the agent observes a new piece of information that
+        should be persisted for future retrieval. Stores in both the
+        vector store (for semantic search) and the graph (for entity
+        traversal).
+        """
+        # Store in vector store
+        self.vector_store.add(
+            fact,
+            {
+                "text": fact,
+                "entity": entity,
+                "valid_from": (timestamp or datetime.now()).isoformat(),
+                "session_id": self.session_id,
+            },
+        )
+
+        # Get or create entity node (uses registry for identity)
+        entity_node_id: str = self.graph.get_or_create_node(entity)
+
+        # Create relationships
+        if relationships:
+            for rel in relationships:
+                target_node_id: str = self.graph.get_or_create_node(rel["target"])
+                self.graph.create_relationship(
+                    entity_node_id,
+                    rel["type"],
+                    target_node_id,
+                    properties=rel.get("properties", {}),
+                )
+
+    def retrieve_memories(
+        self,
+        query: str,
+        entity_filter: Optional[str] = None,
+        time_filter: Optional[Dict[str, Any]] = None,
+        limit: int = 5,
+    ) -> List[Dict[str, Any]]:
+        """Retrieve memories matching query.
+
+        Use when: the agent needs to recall previously stored facts,
+        optionally filtered by entity or time. Results are enriched
+        with graph relationships for each matched entity.
+        """
+        # Vector search
+        filters: Dict[str, Any] = {"session_id": self.session_id}
+        if entity_filter:
+            filters["entity"] = entity_filter
+
+        results: List[Dict[str, Any]] = self.vector_store.search(
+            query, limit=limit, filters=filters
+        )
+
+        # Enrich with graph relationships
+        for result in results:
+            entity: Optional[str] = result["metadata"].get("entity")
+            if entity:
+                node_id: Optional[str] = self.graph.entity_registry.get(entity)
+                if node_id:
+                    result["relationships"] = self.graph.get_relationships(node_id)
+
+        return results
+
+    def retrieve_entity_context(self, entity: str) -> Dict[str, Any]:
+        """Retrieve complete context for an entity.
+
+        Use when: the agent needs a full picture of a single entity
+        including its properties, all relationships, and associated
+        vector memories.
+        """
+        node_id: Optional[str] = self.graph.entity_registry.get(entity)
+
+        # Get entity node
+        entity_node: Optional[Dict[str, Any]] = (
+            self.graph.get_node(node_id) if node_id else None
+        )
+
+        # Get relationships
+        relationships: List[Dict[str, Any]] = (
+            self.graph.get_relationships(node_id) if node_id else []
+        )
+
+        # Get vector memories
+        memories: List[Dict[str, Any]] = self.vector_store.search_by_entity(
+            entity, limit=10
+        )
+
+        return {
+            "entity": entity_node,
+            "relationships": relationships,
+            "memories": memories,
+        }
+
+    def consolidate(self) -> None:
+        """Consolidate memories and remove outdated information.
+
+        Use when: memory count exceeds a threshold, retrieval quality
+        degrades, or on a scheduled interval. In production, implement:
+        - Merge related facts into summaries
+        - Update validity periods on stale entries
+        - Archive obsolete facts (invalidate, do not discard)
+        """
+        pass
+
+
+if __name__ == "__main__":
+    # Quick smoke test demonstrating the integrated memory system.
+    mem = IntegratedMemorySystem()
+    mem.start_session("demo-session")
+
+    # Store facts with entity relationships
+    mem.store_fact(
+        "Alice prefers dark mode",
+        entity="Alice",
+        relationships=[{"target": "dark mode", "type": "PREFERS"}],
+    )
+    mem.store_fact(
+        "Alice works at Acme Corp",
+        entity="Alice",
+        relationships=[{"target": "Acme Corp", "type": "WORKS_AT"}],
+    )
+
+    # Semantic retrieval
+    results = mem.retrieve_memories("theme preference")
+    print(f"Search results: {len(results)} memories found")
+    for r in results:
+        print(f"  score={r['score']:.3f}  text={r['text']}")
+
+    # Entity context
+    context = mem.retrieve_entity_context("Alice")
+    print(f"\nAlice context: {len(context['relationships'])} relationships, "
+          f"{len(context['memories'])} memories")