wehub-skill-sync
617 lines
21 KiB
Python
617 lines
21 KiB
Python
"""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")
|