Add MemoryManager utility

README.md

@@ -192,13 +192,27 @@ graph_generator=OllamaTextToGraphParser(llm_client=phi3)
 print(graph_generator) # Should print the InstructorGraphGenerator 
 
 with GraphDB(
-    database=storage_db, 
-    vector_store=vector_store, 
+    database=storage_db,
+    vector_store=vector_store,
     graph_generator=graph_generator
   ) as db:
     print(db)
 ```
 
+### Memory Manager
+
+For a lighter interface around `GraphDB`, use `MemoryManager` to store and
+recall short text events. A runnable example is available at
+`examples/memory_manager_example.py`.
+
+```python
+from personal_graph.memory import MemoryManager
+
+with MemoryManager() as memory:
+    memory.store_event("Alice bought apples", {"source": "diary"})
+    print(memory.recall("apples"))
+```
+
 ### PersonalGraph to PyG, then back to PersonalGraph
 The following is just a sketch of the planned flow. WIP.
 

examples/memory_manager_example.py

@@ -0,0 +1,15 @@
+from personal_graph.memory import MemoryManager
+
+
+def main() -> None:
+    with MemoryManager() as memory:
+        memory.store_event("Alice started a new job at Google.", {"date": "2024-07-14"})
+        memory.store_event("Bob moved to Paris.", {"date": "2024-07-12"})
+
+        events = memory.recall("Who started a new job?", k=1)
+        for event in events:
+            print(event)
+
+
+if __name__ == "__main__":
+    main()

personal_graph/__init__.py

@@ -9,6 +9,7 @@
 from personal_graph.embeddings import OpenAIEmbeddingsModel
 from personal_graph.graph import GraphDB
 from personal_graph.models import Node, Edge, EdgeInput, KnowledgeGraph
+from personal_graph.memory import MemoryManager
 
 __all__ = [
     "GraphDB",
@@ -23,4 +24,5 @@
     "LiteLLMClient",
     "OllamaEmbeddingClient",
     "OllamaClient",
+    "MemoryManager",
 ]

personal_graph/memory.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+from personal_graph import GraphDB, Node
+
+
+class MemoryManager:
+    """Lightweight interface for storing and recalling textual events."""
+
+    def __init__(self, graph_db: Optional[GraphDB] = None) -> None:
+        self.graph = graph_db or GraphDB()
+
+    def __enter__(self) -> "MemoryManager":
+        self.graph.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.graph.__exit__(exc_type, exc_value, traceback)
+
+    def store_event(self, text: str, metadata: Dict) -> None:
+        """Store an event description with associated metadata."""
+        self.graph.insert(text, metadata)
+
+    def recall(self, query: str, k: int = 5) -> List[Node]:
+        """Recall up to ``k`` events most relevant to ``query``."""
+        results = self.graph.search(query, limit=k)
+        if not results:
+            return []
+        nodes: List[Node] = []
+        for row in results:
+            nodes.append(Node(id=row[1], label=row[2], attributes=row[3]))
+        return nodes

stub/personal_graph/memory.pyi

@@ -0,0 +1,11 @@
+from personal_graph.graph import GraphDB
+from personal_graph.models import Node
+from typing import Dict, List, Optional
+
+class MemoryManager:
+    graph: GraphDB
+    def __init__(self, graph_db: GraphDB | None = ...) -> None: ...
+    def __enter__(self) -> MemoryManager: ...
+    def __exit__(self, exc_type, exc_value, traceback) -> None: ...
+    def store_event(self, text: str, metadata: Dict) -> None: ...
+    def recall(self, query: str, k: int = 5) -> List[Node]: ...

tests/test_memory.py

@@ -0,0 +1,12 @@
+from personal_graph.memory import MemoryManager
+from personal_graph.models import Node
+
+
+def test_memory_store_and_recall(graph, mock_db_connection_and_cursor):
+    memory = MemoryManager(graph)
+    memory.store_event("Alice likes Bob", {"tag": "test"})
+
+    results = memory.recall("Alice", k=1)
+    assert isinstance(results, list)
+    if results:
+        assert isinstance(results[0], Node)