From 4c2c86f8ca43abdd30c2c277512926a251c058b2 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Mon, 28 Jul 2025 14:01:13 -0700 Subject: [PATCH 1/3] Extended the ProllyTree Python API with two major new components MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🧠 Agent Memory System New Classes Added: - AgentMemorySystem - Comprehensive memory system for AI agents - MemoryType - Enum for different memory types (ShortTerm, Semantic, Episodic, Procedural) Key Features: - Short-term Memory: Conversation history with thread-based organization - Semantic Memory: Facts and knowledge about entities with confidence scores - Procedural Memory: Task procedures and instructions with prerequisites - Git-backed Persistence: All memories are version-controlled - Memory Optimization: Cleanup, consolidation, and archival capabilities - Checkpointing: Create memory snapshots with commit messages 🗃️ VersionedKvStore New Classes Added: - VersionedKvStore - Git-backed versioned key-value store - StorageBackend - Enum for storage types (InMemory, File, Git) Key Features: - Git Integration: Full version control with branching, commits, and history - Staging Area: Git-like workflow for changes before committing - Branch Management: Create, switch, and list branches - Commit History: Track all changes with author and timestamp info - Multi-Backend Support: Currently exposes Git backend (others available in Rust) --- python/prollytree/__init__.py | 4 +- python/prollytree/prollytree.pyi | 342 ++++++++++++++++++++ python/tests/test_agent.py | 111 +++++++ python/tests/test_versioned_kv.py | 121 +++++++ src/python.rs | 506 ++++++++++++++++++++++++++++++ 5 files changed, 1082 insertions(+), 2 deletions(-) create mode 100644 python/tests/test_agent.py create mode 100644 python/tests/test_versioned_kv.py diff --git a/python/prollytree/__init__.py b/python/prollytree/__init__.py index 66e1405..9104046 100644 --- a/python/prollytree/__init__.py +++ b/python/prollytree/__init__.py @@ -17,7 +17,7 @@ to provide efficient data access with verifiable integrity. """ -from .prollytree import ProllyTree, TreeConfig +from .prollytree import ProllyTree, TreeConfig, AgentMemorySystem, MemoryType, VersionedKvStore, StorageBackend __version__ = "0.2.1" -__all__ = ["ProllyTree", "TreeConfig"] \ No newline at end of file +__all__ = ["ProllyTree", "TreeConfig", "AgentMemorySystem", "MemoryType", "VersionedKvStore", "StorageBackend"] \ No newline at end of file diff --git a/python/prollytree/prollytree.pyi b/python/prollytree/prollytree.pyi index ff03c46..4b4d6c2 100644 --- a/python/prollytree/prollytree.pyi +++ b/python/prollytree/prollytree.pyi @@ -107,4 +107,346 @@ class ProllyTree: def save_config(self) -> None: """Save the tree configuration to storage""" + ... + +class MemoryType: + """Enum representing different types of memory in the agent system""" + ShortTerm: "MemoryType" + Semantic: "MemoryType" + Episodic: "MemoryType" + Procedural: "MemoryType" + + def __str__(self) -> str: ... + +class AgentMemorySystem: + """Comprehensive memory system for AI agents""" + + def __init__(self, path: str, agent_id: str) -> None: + """ + Initialize a new agent memory system. + + Args: + path: Directory path for memory storage + agent_id: Unique identifier for the agent + """ + ... + + @staticmethod + def open(path: str, agent_id: str) -> "AgentMemorySystem": + """ + Open an existing agent memory system. + + Args: + path: Directory path where memory is stored + agent_id: Unique identifier for the agent + """ + ... + + def store_conversation_turn( + self, + thread_id: str, + role: str, + content: str, + metadata: Optional[Dict[str, str]] = None + ) -> str: + """ + Store a conversation turn in short-term memory. + + Args: + thread_id: Conversation thread identifier + role: Role of the speaker (e.g., "user", "assistant") + content: The message content + metadata: Optional metadata dictionary + + Returns: + Unique ID of the stored memory + """ + ... + + def get_conversation_history( + self, + thread_id: str, + limit: Optional[int] = None + ) -> List[Dict[str, Union[str, float]]]: + """ + Retrieve conversation history for a thread. + + Args: + thread_id: Conversation thread identifier + limit: Maximum number of messages to retrieve + + Returns: + List of message dictionaries with id, content, and created_at fields + """ + ... + + def store_fact( + self, + entity_type: str, + entity_id: str, + facts: str, # JSON string + confidence: float, + source: str + ) -> str: + """ + Store a fact in semantic memory. + + Args: + entity_type: Type of entity (e.g., "person", "place") + entity_id: Unique identifier for the entity + facts: JSON string containing the facts + confidence: Confidence score (0.0 to 1.0) + source: Source of the information + + Returns: + Unique ID of the stored fact + """ + ... + + def get_entity_facts( + self, + entity_type: str, + entity_id: str + ) -> List[Dict[str, Union[str, float]]]: + """ + Retrieve facts about an entity. + + Args: + entity_type: Type of entity + entity_id: Unique identifier for the entity + + Returns: + List of fact dictionaries + """ + ... + + def store_procedure( + self, + category: str, + name: str, + description: str, + steps: List[str], # List of JSON strings + prerequisites: Optional[List[str]] = None, + priority: int = 1 + ) -> str: + """ + Store a procedure in procedural memory. + + Args: + category: Category of the procedure + name: Name of the procedure + description: Description of what the procedure does + steps: List of JSON strings describing each step + prerequisites: Optional list of prerequisites + priority: Priority level (default: 1) + + Returns: + Unique ID of the stored procedure + """ + ... + + def get_procedures_by_category( + self, + category: str + ) -> List[Dict[str, str]]: + """ + Retrieve procedures by category. + + Args: + category: Category to search for + + Returns: + List of procedure dictionaries + """ + ... + + def checkpoint(self, message: str) -> str: + """ + Create a memory checkpoint. + + Args: + message: Commit message for the checkpoint + + Returns: + Checkpoint ID + """ + ... + + def optimize(self) -> Dict[str, int]: + """ + Optimize the memory system by cleaning up and consolidating memories. + + Returns: + Dictionary with optimization statistics + """ + ... + +class StorageBackend: + """Enum representing different storage backend types""" + InMemory: "StorageBackend" + File: "StorageBackend" + Git: "StorageBackend" + + def __str__(self) -> str: ... + +class VersionedKvStore: + """A versioned key-value store backed by Git and ProllyTree""" + + def __init__(self, path: str) -> None: + """ + Initialize a new versioned key-value store. + + Args: + path: Directory path for the store (must be within a git repository) + """ + ... + + @staticmethod + def open(path: str) -> "VersionedKvStore": + """ + Open an existing versioned key-value store. + + Args: + path: Directory path where the store is located + """ + ... + + def insert(self, key: bytes, value: bytes) -> None: + """ + Insert a key-value pair (stages the change). + + Args: + key: The key as bytes + value: The value as bytes + """ + ... + + def get(self, key: bytes) -> Optional[bytes]: + """ + Get a value by key. + + Args: + key: The key to look up + + Returns: + The value as bytes, or None if not found + """ + ... + + def update(self, key: bytes, value: bytes) -> bool: + """ + Update an existing key-value pair (stages the change). + + Args: + key: The key to update + value: The new value + + Returns: + True if the key existed and was updated, False otherwise + """ + ... + + def delete(self, key: bytes) -> bool: + """ + Delete a key-value pair (stages the change). + + Args: + key: The key to delete + + Returns: + True if the key existed and was deleted, False otherwise + """ + ... + + def list_keys(self) -> List[bytes]: + """ + List all keys in the store (includes staged changes). + + Returns: + List of keys as bytes + """ + ... + + def status(self) -> List[Tuple[bytes, str]]: + """ + Show current staging area status. + + Returns: + List of tuples (key, status) where status is "added", "modified", or "deleted" + """ + ... + + def commit(self, message: str) -> str: + """ + Commit staged changes. + + Args: + message: Commit message + + Returns: + Commit hash as hex string + """ + ... + + def branch(self, name: str) -> None: + """ + Create a new branch. + + Args: + name: Name of the new branch + """ + ... + + def create_branch(self, name: str) -> None: + """ + Create a new branch and switch to it. + + Args: + name: Name of the new branch + """ + ... + + def checkout(self, branch_or_commit: str) -> None: + """ + Switch to a different branch or commit. + + Args: + branch_or_commit: Branch name or commit hash + """ + ... + + def current_branch(self) -> str: + """ + Get the current branch name. + + Returns: + Current branch name + """ + ... + + def list_branches(self) -> List[str]: + """ + List all branches in the repository. + + Returns: + List of branch names + """ + ... + + def log(self) -> List[Dict[str, Union[str, int]]]: + """ + Get commit history. + + Returns: + List of commit dictionaries with id, author, committer, message, and timestamp + """ + ... + + def storage_backend(self) -> StorageBackend: + """ + Get the current storage backend type. + + Returns: + Storage backend enum value + """ ... \ No newline at end of file diff --git a/python/tests/test_agent.py b/python/tests/test_agent.py new file mode 100644 index 0000000..6fdbfef --- /dev/null +++ b/python/tests/test_agent.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""Test script for ProllyTree Agent Memory System Python bindings.""" + +import json +from prollytree import AgentMemorySystem, MemoryType +import tempfile +import os + +def test_agent_memory_system(): + """Test the agent memory system functionality.""" + + # Create a temporary directory for the memory store + with tempfile.TemporaryDirectory() as tmpdir: + print(f"📁 Creating memory system in: {tmpdir}") + + # Initialize the agent memory system + memory_system = AgentMemorySystem(tmpdir, "test_agent_001") + print("✅ Agent memory system initialized") + + # Test 1: Store conversation turns + print("\n🧪 Test 1: Short-term memory (conversation)") + conv_id1 = memory_system.store_conversation_turn( + "thread_123", + "user", + "Hello, how are you?", + {"source": "chat", "session": "morning"} + ) + print(f" Stored user message: {conv_id1}") + + conv_id2 = memory_system.store_conversation_turn( + "thread_123", + "assistant", + "I'm doing well, thank you for asking! How can I help you today?" + ) + print(f" Stored assistant message: {conv_id2}") + + # Retrieve conversation history + history = memory_system.get_conversation_history("thread_123", limit=10) + print(f" Retrieved {len(history)} messages from conversation history") + for msg in history: + print(f" - {msg['created_at']}: {json.loads(msg['content'])}") + + # Test 2: Store semantic facts + print("\n🧪 Test 2: Semantic memory (facts)") + fact_id = memory_system.store_fact( + "person", + "john_doe", + json.dumps({ + "age": 30, + "occupation": "software engineer", + "location": "San Francisco" + }), + 0.95, # confidence + "user_input" + ) + print(f" Stored fact about john_doe: {fact_id}") + + # Retrieve facts + facts = memory_system.get_entity_facts("person", "john_doe") + print(f" Retrieved {len(facts)} facts about john_doe") + for fact in facts: + print(f" - Confidence: {fact['confidence']}, Source: {fact['source']}") + print(f" Facts: {fact['facts']}") + + # Test 3: Store procedures + print("\n🧪 Test 3: Procedural memory") + proc_id = memory_system.store_procedure( + "task_management", + "create_project", + "How to create a new project in the system", + [ + json.dumps({"step": 1, "action": "Define project name and description"}), + json.dumps({"step": 2, "action": "Set project timeline and milestones"}), + json.dumps({"step": 3, "action": "Assign team members and roles"}), + json.dumps({"step": 4, "action": "Initialize project repository"}) + ], + ["admin_access", "project_creation_permission"], + priority=2 + ) + print(f" Stored procedure: {proc_id}") + + # Get procedures by category + procedures = memory_system.get_procedures_by_category("task_management") + print(f" Retrieved {len(procedures)} procedures in task_management category") + for proc in procedures: + print(f" - {proc['id']}: {proc['content']}") + + # Test 4: Create checkpoint + print("\n🧪 Test 4: Memory checkpoint") + checkpoint_id = memory_system.checkpoint("Initial test data loaded") + print(f" Created checkpoint: {checkpoint_id}") + + # Test 5: Optimize memory + print("\n🧪 Test 5: Memory optimization") + optimization_report = memory_system.optimize() + print(" Optimization report:") + for key, value in optimization_report.items(): + print(f" - {key}: {value}") + + print("\n✅ All tests completed successfully!") + + # Test MemoryType enum + print("\n🧪 Test 6: MemoryType enum") + print(f" MemoryType.ShortTerm: {MemoryType.ShortTerm}") + print(f" MemoryType.Semantic: {MemoryType.Semantic}") + print(f" MemoryType.Episodic: {MemoryType.Episodic}") + print(f" MemoryType.Procedural: {MemoryType.Procedural}") + + +if __name__ == "__main__": + test_agent_memory_system() \ No newline at end of file diff --git a/python/tests/test_versioned_kv.py b/python/tests/test_versioned_kv.py new file mode 100644 index 0000000..11ec868 --- /dev/null +++ b/python/tests/test_versioned_kv.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +"""Test script for VersionedKvStore Python bindings.""" + +import tempfile +import os +import subprocess +from prollytree import VersionedKvStore, StorageBackend + +def test_versioned_kv_store(): + """Test the versioned key-value store functionality.""" + + # Create a temporary directory and initialize a git repository + with tempfile.TemporaryDirectory() as tmpdir: + print(f"📁 Creating test in: {tmpdir}") + + # Initialize git repository + subprocess.run(["git", "init"], cwd=tmpdir, check=True, capture_output=True) + subprocess.run(["git", "config", "user.name", "Test User"], cwd=tmpdir, check=True) + subprocess.run(["git", "config", "user.email", "test@example.com"], cwd=tmpdir, check=True) + + # Create a subdirectory for our dataset + dataset_dir = os.path.join(tmpdir, "dataset") + os.makedirs(dataset_dir) + os.chdir(dataset_dir) + + print("✅ Git repository initialized") + + # Test 1: Initialize VersionedKvStore + print("\n🧪 Test 1: Initialize VersionedKvStore") + store = VersionedKvStore(dataset_dir) + print(f" Storage backend: {store.storage_backend()}") + print(f" Current branch: {store.current_branch()}") + + # Test 2: Basic key-value operations + print("\n🧪 Test 2: Basic key-value operations") + store.insert(b"name", b"Alice") + store.insert(b"age", b"30") + store.insert(b"city", b"San Francisco") + + # Check values + name = store.get(b"name") + age = store.get(b"age") + city = store.get(b"city") + + print(f" name: {name}") + print(f" age: {age}") + print(f" city: {city}") + + # Test 3: List keys and status + print("\n🧪 Test 3: List keys and status") + keys = store.list_keys() + print(f" Keys: {[k.decode() for k in keys]}") + + status = store.status() + print(" Status:") + for key, status_str in status: + print(f" - {key.decode()}: {status_str}") + + # Test 4: Commit changes + print("\n🧪 Test 4: Commit changes") + commit_hash = store.commit("Add initial user data") + print(f" Commit hash: {commit_hash}") + + # Check status after commit + status = store.status() + print(f" Status after commit: {len(status)} staged changes") + + # Test 5: Update and delete operations + print("\n🧪 Test 5: Update and delete operations") + updated = store.update(b"age", b"31") + print(f" Updated age: {updated}") + + deleted = store.delete(b"city") + print(f" Deleted city: {deleted}") + + # Add new key + store.insert(b"country", b"USA") + + # Check status + status = store.status() + print(" Status after changes:") + for key, status_str in status: + print(f" - {key.decode()}: {status_str}") + + # Test 6: Branch operations + print("\n🧪 Test 6: Branch operations") + store.create_branch("feature-branch") + print(" Created and switched to feature-branch") + print(f" Current branch: {store.current_branch()}") + + # Make changes on feature branch + store.insert(b"feature", b"new-feature") + store.commit("Add feature on feature branch") + + # List all branches + branches = store.list_branches() + print(f" Available branches: {branches}") + + # Test 7: Switch back to main + print("\n🧪 Test 7: Switch back to main") + store.checkout("main") + print(f" Current branch: {store.current_branch()}") + + # Check if feature key exists (should not exist on main) + feature = store.get(b"feature") + print(f" Feature key on main: {feature}") + + # Test 8: Commit history + print("\n🧪 Test 8: Commit history") + history = store.log() + print(f" Commit history ({len(history)} commits):") + for i, commit in enumerate(history[:3]): # Show first 3 commits + print(f" {i+1}. {commit['id'][:8]} - {commit['message']}") + print(f" Author: {commit['author']}") + print(f" Timestamp: {commit['timestamp']}") + + print("\n✅ All VersionedKvStore tests completed successfully!") + + +if __name__ == "__main__": + test_versioned_kv_store() \ No newline at end of file diff --git a/src/python.rs b/src/python.rs index 42e4381..d0994f1 100644 --- a/src/python.rs +++ b/src/python.rs @@ -20,7 +20,9 @@ use std::path::PathBuf; use std::sync::{Arc, Mutex}; use crate::{ + agent::{AgentMemorySystem, MemoryType}, config::TreeConfig, + git::{GitVersionedKvStore, types::StorageBackend}, proof::Proof, storage::{FileNodeStorage, InMemoryNodeStorage}, tree::{ProllyTree, Tree}, @@ -349,9 +351,513 @@ impl PyProllyTree { } } +#[pyclass(name = "MemoryType", eq, eq_int)] +#[derive(Clone, PartialEq)] +enum PyMemoryType { + ShortTerm, + Semantic, + Episodic, + Procedural, +} + +#[pymethods] +impl PyMemoryType { + fn __str__(&self) -> &str { + match self { + PyMemoryType::ShortTerm => "ShortTerm", + PyMemoryType::Semantic => "Semantic", + PyMemoryType::Episodic => "Episodic", + PyMemoryType::Procedural => "Procedural", + } + } +} + +impl From for MemoryType { + fn from(py_type: PyMemoryType) -> Self { + match py_type { + PyMemoryType::ShortTerm => MemoryType::ShortTerm, + PyMemoryType::Semantic => MemoryType::Semantic, + PyMemoryType::Episodic => MemoryType::Episodic, + PyMemoryType::Procedural => MemoryType::Procedural, + } + } +} + +impl From for PyMemoryType { + fn from(mem_type: MemoryType) -> Self { + match mem_type { + MemoryType::ShortTerm => PyMemoryType::ShortTerm, + MemoryType::Semantic => PyMemoryType::Semantic, + MemoryType::Episodic => PyMemoryType::Episodic, + MemoryType::Procedural => PyMemoryType::Procedural, + } + } +} + +#[pyclass(name = "AgentMemorySystem")] +struct PyAgentMemorySystem { + inner: Arc>, +} + +#[pymethods] +impl PyAgentMemorySystem { + #[new] + #[pyo3(signature = (path, agent_id))] + fn new(path: String, agent_id: String) -> PyResult { + let memory_system = AgentMemorySystem::init(path, agent_id, None) + .map_err(|e| PyValueError::new_err(format!("Failed to initialize memory system: {}", e)))?; + + Ok(PyAgentMemorySystem { + inner: Arc::new(Mutex::new(memory_system)), + }) + } + + #[staticmethod] + fn open(path: String, agent_id: String) -> PyResult { + let memory_system = AgentMemorySystem::open(path, agent_id, None) + .map_err(|e| PyValueError::new_err(format!("Failed to open memory system: {}", e)))?; + + Ok(PyAgentMemorySystem { + inner: Arc::new(Mutex::new(memory_system)), + }) + } + + #[pyo3(signature = (thread_id, role, content, metadata=None))] + fn store_conversation_turn( + &self, + py: Python, + thread_id: String, + role: String, + content: String, + metadata: Option>, + ) -> PyResult { + py.allow_threads(|| { + let runtime = tokio::runtime::Runtime::new() + .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; + + let mut memory_system = self.inner.lock().unwrap(); + + // Convert HashMap to HashMap + let metadata_values = metadata.map(|m| { + m.into_iter() + .map(|(k, v)| (k, serde_json::Value::String(v))) + .collect() + }); + + runtime.block_on(async { + memory_system.short_term + .store_conversation_turn(&thread_id, &role, &content, metadata_values) + .await + .map_err(|e| PyValueError::new_err(format!("Failed to store conversation: {}", e))) + }) + }) + } + + #[pyo3(signature = (thread_id, limit=None))] + fn get_conversation_history( + &self, + py: Python, + thread_id: String, + limit: Option, + ) -> PyResult>>> { + py.allow_threads(|| { + let runtime = tokio::runtime::Runtime::new() + .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; + + let memory_system = self.inner.lock().unwrap(); + + runtime.block_on(async { + let history = memory_system.short_term + .get_conversation_history(&thread_id, limit) + .await + .map_err(|e| PyValueError::new_err(format!("Failed to get history: {}", e)))?; + + Python::with_gil(|py| { + let results: PyResult>>> = history.iter().map(|doc| { + let mut map = HashMap::new(); + map.insert("id".to_string(), doc.id.clone().into_py(py)); + map.insert("content".to_string(), doc.content.to_string().into_py(py)); + map.insert("created_at".to_string(), doc.metadata.created_at.to_rfc3339().into_py(py)); + Ok(map) + }).collect(); + results + }) + }) + }) + } + + fn store_fact( + &self, + py: Python, + entity_type: String, + entity_id: String, + facts: String, // JSON string + confidence: f64, + source: String, + ) -> PyResult { + py.allow_threads(|| { + let runtime = tokio::runtime::Runtime::new() + .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; + + let mut memory_system = self.inner.lock().unwrap(); + + let facts_value: serde_json::Value = serde_json::from_str(&facts) + .map_err(|e| PyValueError::new_err(format!("Invalid JSON: {}", e)))?; + + runtime.block_on(async { + memory_system.semantic + .store_fact(&entity_type, &entity_id, facts_value, confidence, &source) + .await + .map_err(|e| PyValueError::new_err(format!("Failed to store fact: {}", e))) + }) + }) + } + + fn get_entity_facts( + &self, + py: Python, + entity_type: String, + entity_id: String, + ) -> PyResult>>> { + py.allow_threads(|| { + let runtime = tokio::runtime::Runtime::new() + .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; + + let memory_system = self.inner.lock().unwrap(); + + runtime.block_on(async { + let facts = memory_system.semantic + .get_entity_facts(&entity_type, &entity_id) + .await + .map_err(|e| PyValueError::new_err(format!("Failed to get facts: {}", e)))?; + + Python::with_gil(|py| { + let results: PyResult>>> = facts.iter().map(|doc| { + let mut map = HashMap::new(); + map.insert("id".to_string(), doc.id.clone().into_py(py)); + map.insert("facts".to_string(), doc.content.to_string().into_py(py)); + map.insert("confidence".to_string(), doc.metadata.confidence.into_py(py)); + map.insert("source".to_string(), doc.metadata.source.clone().into_py(py)); + Ok(map) + }).collect(); + results + }) + }) + }) + } + + #[pyo3(signature = (category, name, description, steps, prerequisites=None, priority=1))] + fn store_procedure( + &self, + py: Python, + category: String, + name: String, + description: String, + steps: Vec, // JSON strings + prerequisites: Option>, + priority: u32, + ) -> PyResult { + py.allow_threads(|| { + let runtime = tokio::runtime::Runtime::new() + .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; + + let mut memory_system = self.inner.lock().unwrap(); + + let steps_values: Result, _> = steps.iter() + .map(|s| serde_json::from_str(s)) + .collect(); + let steps_values = steps_values + .map_err(|e| PyValueError::new_err(format!("Invalid JSON in steps: {}", e)))?; + + // Convert prerequisites to serde_json::Value + let conditions = prerequisites.map(|p| serde_json::Value::Array( + p.into_iter().map(serde_json::Value::String).collect() + )); + + runtime.block_on(async { + memory_system.procedural + .store_procedure(&category, &name, &description, steps_values, conditions, priority) + .await + .map_err(|e| PyValueError::new_err(format!("Failed to store procedure: {}", e))) + }) + }) + } + + fn get_procedures_by_category( + &self, + py: Python, + category: String, + ) -> PyResult>>> { + py.allow_threads(|| { + let runtime = tokio::runtime::Runtime::new() + .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; + + let memory_system = self.inner.lock().unwrap(); + + runtime.block_on(async { + let procedures = memory_system.procedural + .get_procedures_by_category(&category) + .await + .map_err(|e| PyValueError::new_err(format!("Failed to get procedures: {}", e)))?; + + Python::with_gil(|py| { + let results: PyResult>>> = procedures.iter().map(|doc| { + let mut map = HashMap::new(); + map.insert("id".to_string(), doc.id.clone().into_py(py)); + map.insert("content".to_string(), doc.content.to_string().into_py(py)); + map.insert("created_at".to_string(), doc.metadata.created_at.to_rfc3339().into_py(py)); + Ok(map) + }).collect(); + results + }) + }) + }) + } + + fn checkpoint(&self, py: Python, message: String) -> PyResult { + py.allow_threads(|| { + let runtime = tokio::runtime::Runtime::new() + .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; + + let mut memory_system = self.inner.lock().unwrap(); + + runtime.block_on(async { + memory_system.checkpoint(&message) + .await + .map_err(|e| PyValueError::new_err(format!("Failed to create checkpoint: {}", e))) + }) + }) + } + + fn optimize(&self, py: Python) -> PyResult> { + py.allow_threads(|| { + let runtime = tokio::runtime::Runtime::new() + .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; + + let mut memory_system = self.inner.lock().unwrap(); + + runtime.block_on(async { + let report = memory_system.optimize() + .await + .map_err(|e| PyValueError::new_err(format!("Failed to optimize: {}", e)))?; + + let mut result = HashMap::new(); + result.insert("expired_cleaned".to_string(), report.expired_cleaned); + result.insert("memories_consolidated".to_string(), report.memories_consolidated); + result.insert("memories_archived".to_string(), report.memories_archived); + result.insert("memories_pruned".to_string(), report.memories_pruned); + result.insert("total_processed".to_string(), report.total_processed()); + Ok(result) + }) + }) + } +} + +#[pyclass(name = "StorageBackend", eq, eq_int)] +#[derive(Clone, PartialEq)] +enum PyStorageBackend { + InMemory, + File, + Git, +} + +#[pymethods] +impl PyStorageBackend { + fn __str__(&self) -> &str { + match self { + PyStorageBackend::InMemory => "InMemory", + PyStorageBackend::File => "File", + PyStorageBackend::Git => "Git", + } + } +} + +impl From for StorageBackend { + fn from(py_backend: PyStorageBackend) -> Self { + match py_backend { + PyStorageBackend::InMemory => StorageBackend::InMemory, + PyStorageBackend::File => StorageBackend::File, + PyStorageBackend::Git => StorageBackend::Git, + } + } +} + +impl From for PyStorageBackend { + fn from(backend: StorageBackend) -> Self { + match backend { + StorageBackend::InMemory => PyStorageBackend::InMemory, + StorageBackend::File => PyStorageBackend::File, + StorageBackend::Git => PyStorageBackend::Git, + #[cfg(feature = "rocksdb_storage")] + StorageBackend::RocksDB => PyStorageBackend::Git, // Fallback to Git for RocksDB + } + } +} + +#[pyclass(name = "VersionedKvStore")] +struct PyVersionedKvStore { + inner: Arc>>, +} + +#[pymethods] +impl PyVersionedKvStore { + #[new] + fn new(path: String) -> PyResult { + let store = GitVersionedKvStore::<32>::init(path) + .map_err(|e| PyValueError::new_err(format!("Failed to initialize store: {}", e)))?; + + Ok(PyVersionedKvStore { + inner: Arc::new(Mutex::new(store)), + }) + } + + #[staticmethod] + fn open(path: String) -> PyResult { + let store = GitVersionedKvStore::<32>::open(path) + .map_err(|e| PyValueError::new_err(format!("Failed to open store: {}", e)))?; + + Ok(PyVersionedKvStore { + inner: Arc::new(Mutex::new(store)), + }) + } + + fn insert(&self, key: &Bound<'_, PyBytes>, value: &Bound<'_, PyBytes>) -> PyResult<()> { + let key_vec = key.as_bytes().to_vec(); + let value_vec = value.as_bytes().to_vec(); + + let mut store = self.inner.lock().unwrap(); + store.insert(key_vec, value_vec) + .map_err(|e| PyValueError::new_err(format!("Failed to insert: {}", e)))?; + + Ok(()) + } + + fn get(&self, py: Python, key: &Bound<'_, PyBytes>) -> PyResult>> { + let key_vec = key.as_bytes().to_vec(); + + let store = self.inner.lock().unwrap(); + match store.get(&key_vec) { + Some(value) => Ok(Some(PyBytes::new_bound(py, &value).into())), + None => Ok(None), + } + } + + fn update(&self, key: &Bound<'_, PyBytes>, value: &Bound<'_, PyBytes>) -> PyResult { + let key_vec = key.as_bytes().to_vec(); + let value_vec = value.as_bytes().to_vec(); + + let mut store = self.inner.lock().unwrap(); + store.update(key_vec, value_vec) + .map_err(|e| PyValueError::new_err(format!("Failed to update: {}", e))) + } + + fn delete(&self, key: &Bound<'_, PyBytes>) -> PyResult { + let key_vec = key.as_bytes().to_vec(); + + let mut store = self.inner.lock().unwrap(); + store.delete(&key_vec) + .map_err(|e| PyValueError::new_err(format!("Failed to delete: {}", e))) + } + + fn list_keys(&self, py: Python) -> PyResult>> { + let store = self.inner.lock().unwrap(); + let keys = store.list_keys(); + + let py_keys: Vec> = keys.iter() + .map(|key| PyBytes::new_bound(py, key).into()) + .collect(); + + Ok(py_keys) + } + + fn status(&self, py: Python) -> PyResult, String)>> { + let store = self.inner.lock().unwrap(); + let status = store.status(); + + let py_status: Vec<(Py, String)> = status.iter() + .map(|(key, status_str)| { + (PyBytes::new_bound(py, key).into(), status_str.clone()) + }) + .collect(); + + Ok(py_status) + } + + fn commit(&self, message: String) -> PyResult { + let mut store = self.inner.lock().unwrap(); + let commit_id = store.commit(&message) + .map_err(|e| PyValueError::new_err(format!("Failed to commit: {}", e)))?; + + Ok(commit_id.to_hex().to_string()) + } + + fn branch(&self, name: String) -> PyResult<()> { + let mut store = self.inner.lock().unwrap(); + store.branch(&name) + .map_err(|e| PyValueError::new_err(format!("Failed to create branch: {}", e)))?; + + Ok(()) + } + + fn create_branch(&self, name: String) -> PyResult<()> { + let mut store = self.inner.lock().unwrap(); + store.create_branch(&name) + .map_err(|e| PyValueError::new_err(format!("Failed to create and switch branch: {}", e)))?; + + Ok(()) + } + + fn checkout(&self, branch_or_commit: String) -> PyResult<()> { + let mut store = self.inner.lock().unwrap(); + store.checkout(&branch_or_commit) + .map_err(|e| PyValueError::new_err(format!("Failed to checkout: {}", e)))?; + + Ok(()) + } + + fn current_branch(&self) -> PyResult { + let store = self.inner.lock().unwrap(); + Ok(store.current_branch().to_string()) + } + + fn list_branches(&self) -> PyResult> { + let store = self.inner.lock().unwrap(); + store.list_branches() + .map_err(|e| PyValueError::new_err(format!("Failed to list branches: {}", e))) + } + + fn log(&self) -> PyResult>>> { + let store = self.inner.lock().unwrap(); + let history = store.log() + .map_err(|e| PyValueError::new_err(format!("Failed to get log: {}", e)))?; + + Python::with_gil(|py| { + let results: PyResult>>> = history.iter().map(|commit| { + let mut map = HashMap::new(); + map.insert("id".to_string(), commit.id.to_hex().to_string().into_py(py)); + map.insert("author".to_string(), commit.author.clone().into_py(py)); + map.insert("committer".to_string(), commit.committer.clone().into_py(py)); + map.insert("message".to_string(), commit.message.clone().into_py(py)); + map.insert("timestamp".to_string(), commit.timestamp.into_py(py)); + Ok(map) + }).collect(); + results + }) + } + + fn storage_backend(&self) -> PyResult { + let store = self.inner.lock().unwrap(); + Ok(store.storage_backend().clone().into()) + } +} + #[pymodule] fn prollytree(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } From 39e300d2430ba60f4e7b6bfd542375bf6ac79311 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Mon, 28 Jul 2025 14:07:17 -0700 Subject: [PATCH 2/3] revise readme file --- python/README.md | 268 ++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 220 insertions(+), 48 deletions(-) diff --git a/python/README.md b/python/README.md index 14ec0df..d55c615 100644 --- a/python/README.md +++ b/python/README.md @@ -1,8 +1,18 @@ # ProllyTree Python Bindings -This directory contains Python bindings for the ProllyTree Rust library, providing a Pythonic interface to the probabilistic tree data structure. +This directory contains Python bindings for the ProllyTree Rust library, providing a comprehensive toolkit for: -## Usage Example +- **Probabilistic Trees**: High-performance prolly trees for efficient data storage and retrieval +- **AI Agent Memory Systems**: Multi-layered memory systems for intelligent agents +- **Versioned Key-Value Storage**: Git-backed versioned storage with branching and history + +## Overview + +ProllyTree combines B-trees and Merkle trees to provide both efficient data access and verifiable integrity, making it ideal for distributed systems and applications requiring data verification. + +## Usage Examples + +### Basic ProllyTree Operations ```python from prollytree import ProllyTree, TreeConfig @@ -14,38 +24,124 @@ tree = ProllyTree(storage_type="memory") tree.insert(b"key1", b"value1") tree.insert(b"key2", b"value2") -# Batch insert +# Batch operations items = [(b"key3", b"value3"), (b"key4", b"value4")] tree.insert_batch(items) -# Find values +# Find and update values value = tree.find(b"key1") # Returns b"value1" - -# Update a value tree.update(b"key1", b"new_value1") - -# Delete keys tree.delete(b"key2") -# Get tree properties -print(f"Size: {tree.size()}") -print(f"Depth: {tree.depth()}") -print(f"Root hash: {tree.get_root_hash().hex()}") - -# Generate and verify Merkle proofs +# Tree properties and verification +print(f"Size: {tree.size()}, Depth: {tree.depth()}") proof = tree.generate_proof(b"key3") is_valid = tree.verify_proof(proof, b"key3", b"value3") -# Compare trees -tree2 = ProllyTree() -tree2.insert(b"key1", b"different_value") -diff = tree.diff(tree2) - # File-based storage config = TreeConfig(base=4, modulus=64) file_tree = ProllyTree(storage_type="file", path="/tmp/my_tree", config=config) -file_tree.insert(b"persistent_key", b"persistent_value") -file_tree.save_config() +``` + +### AI Agent Memory System + +```python +import json +from prollytree import AgentMemorySystem, MemoryType + +# Initialize agent memory system +memory = AgentMemorySystem("/path/to/memory", "agent_001") + +# Short-term memory (conversations) +memory.store_conversation_turn( + "thread_123", + "user", + "What's the weather like?", + {"session": "morning", "platform": "chat"} +) + +memory.store_conversation_turn( + "thread_123", + "assistant", + "I'd be happy to help with weather information!" +) + +# Retrieve conversation history +history = memory.get_conversation_history("thread_123", limit=10) + +# Semantic memory (facts about entities) +memory.store_fact( + "person", + "john_doe", + json.dumps({ + "name": "John Doe", + "role": "Software Engineer", + "location": "San Francisco" + }), + confidence=0.95, + source="user_profile" +) + +# Get facts about an entity +facts = memory.get_entity_facts("person", "john_doe") + +# Procedural memory (task instructions) +memory.store_procedure( + "development", + "code_review", + "How to conduct a code review", + [ + json.dumps({"step": 1, "action": "Check code style and formatting"}), + json.dumps({"step": 2, "action": "Review logic and algorithms"}), + json.dumps({"step": 3, "action": "Test edge cases and error handling"}) + ], + prerequisites=["git_access", "reviewer_permissions"], + priority=2 +) + +# Memory management +checkpoint_id = memory.checkpoint("Saved conversation and facts") +optimization_report = memory.optimize() +``` + +### Versioned Key-Value Store + +```python +from prollytree import VersionedKvStore, StorageBackend + +# Initialize in a git repository subdirectory +store = VersionedKvStore("/path/to/git/repo/dataset") + +# Basic key-value operations (staged until commit) +store.insert(b"user:1", b'{"name": "Alice", "age": 30}') +store.insert(b"user:2", b'{"name": "Bob", "age": 25}') +store.update(b"user:1", b'{"name": "Alice", "age": 31}') + +# Check staging status +status = store.status() # Shows added/modified/deleted keys +keys = store.list_keys() + +# Commit changes with message +commit_hash = store.commit("Add initial user data") + +# Branch operations +store.create_branch("feature-branch") +store.insert(b"feature:1", b"experimental_data") +store.commit("Add experimental feature") + +# Switch branches +store.checkout("main") +print(f"Current branch: {store.current_branch()}") +print(f"Available branches: {store.list_branches()}") + +# View commit history +history = store.log() +for commit in history[:5]: # Last 5 commits + print(f"{commit['id'][:8]} - {commit['message']}") + print(f" Author: {commit['author']}") + +# Storage backend info +backend = store.storage_backend() # Returns StorageBackend.Git ``` ## Publishing to PyPI @@ -73,35 +169,111 @@ Publish to production PyPI: ./publish_python.sh prod ``` -## API Reference +## Installation + +```bash +pip install prollytree +``` -### TreeConfig +## API Reference -Configuration class for ProllyTree: +### ProllyTree Classes -- `base`: Base for the rolling hash (default: 4) -- `modulus`: Modulus for the rolling hash (default: 64) +#### TreeConfig +Configuration for ProllyTree instances: +- `base`: Rolling hash base (default: 4) +- `modulus`: Rolling hash modulus (default: 64) - `min_chunk_size`: Minimum chunk size (default: 1) - `max_chunk_size`: Maximum chunk size (default: 4096) -- `pattern`: Pattern for chunk boundaries (default: 0) - -### ProllyTree - -Main tree class with the following methods: - -- `__init__(storage_type="memory", path=None, config=None)`: Create a new tree -- `insert(key: bytes, value: bytes)`: Insert a key-value pair -- `insert_batch(items: List[Tuple[bytes, bytes]])`: Batch insert -- `find(key: bytes) -> Optional[bytes]`: Find a value by key -- `update(key: bytes, value: bytes)`: Update an existing key -- `delete(key: bytes)`: Delete a key -- `delete_batch(keys: List[bytes])`: Batch delete -- `size() -> int`: Get number of key-value pairs -- `depth() -> int`: Get tree depth -- `get_root_hash() -> bytes`: Get the root hash -- `stats() -> Dict[str, int]`: Get tree statistics -- `generate_proof(key: bytes) -> bytes`: Generate a Merkle proof -- `verify_proof(proof: bytes, key: bytes, expected_value: Optional[bytes]) -> bool`: Verify a proof -- `diff(other: ProllyTree) -> Dict`: Compare two trees -- `traverse() -> str`: Get string representation of tree structure -- `save_config()`: Save tree configuration to storage +- `pattern`: Chunk boundary pattern (default: 0) + +#### ProllyTree +High-performance probabilistic tree with Merkle verification: + +**Core Operations:** +- `insert(key: bytes, value: bytes)`: Insert key-value pair +- `find(key: bytes) -> Optional[bytes]`: Find value by key +- `update(key: bytes, value: bytes)`: Update existing key +- `delete(key: bytes)`: Delete key +- `insert_batch(items)`, `delete_batch(keys)`: Batch operations + +**Properties & Verification:** +- `size() -> int`, `depth() -> int`: Tree metrics +- `get_root_hash() -> bytes`: Cryptographic root hash +- `generate_proof(key) -> bytes`: Create Merkle proof +- `verify_proof(proof, key, value) -> bool`: Verify proof +- `stats() -> Dict`: Detailed tree statistics + +### Agent Memory System + +#### AgentMemorySystem +Comprehensive memory system for AI agents with multiple memory types: + +**Initialization:** +- `AgentMemorySystem(path: str, agent_id: str)`: Create new system +- `AgentMemorySystem.open(path: str, agent_id: str)`: Open existing + +**Short-term Memory (Conversations):** +- `store_conversation_turn(thread_id, role, content, metadata=None) -> str` +- `get_conversation_history(thread_id, limit=None) -> List[Dict]` + +**Semantic Memory (Facts & Knowledge):** +- `store_fact(entity_type, entity_id, facts_json, confidence, source) -> str` +- `get_entity_facts(entity_type, entity_id) -> List[Dict]` + +**Procedural Memory (Instructions & Procedures):** +- `store_procedure(category, name, description, steps, prerequisites=None, priority=1) -> str` +- `get_procedures_by_category(category) -> List[Dict]` + +**System Operations:** +- `checkpoint(message: str) -> str`: Create memory snapshot +- `optimize() -> Dict[str, int]`: Cleanup and consolidation + +#### MemoryType +Enum for memory classification: `ShortTerm`, `Semantic`, `Episodic`, `Procedural` + +### Versioned Key-Value Store + +#### VersionedKvStore +Git-backed versioned storage with full branching support: + +**Initialization:** +- `VersionedKvStore(path: str)`: Initialize new store +- `VersionedKvStore.open(path: str)`: Open existing store + +**Basic Operations:** +- `insert(key: bytes, value: bytes)`: Stage insertion +- `get(key: bytes) -> Optional[bytes]`: Get current value +- `update(key: bytes, value: bytes) -> bool`: Stage update +- `delete(key: bytes) -> bool`: Stage deletion +- `list_keys() -> List[bytes]`: List all keys + +**Git Operations:** +- `commit(message: str) -> str`: Commit staged changes +- `status() -> List[Tuple[bytes, str]]`: Show staging status +- `branch(name: str)`: Create branch +- `create_branch(name: str)`: Create and switch to branch +- `checkout(branch_or_commit: str)`: Switch branch/commit +- `current_branch() -> str`: Get current branch +- `list_branches() -> List[str]`: List all branches +- `log() -> List[Dict]`: Commit history + +#### StorageBackend +Enum for storage types: `InMemory`, `File`, `Git` + +## Requirements + +- **VersionedKvStore**: Requires git repository (must be run within a git repo) +- **AgentMemorySystem**: Requires filesystem access for persistence +- **ProllyTree**: Works with memory, file, or custom storage backends + +## Testing + +Run the comprehensive test suite: + +```bash +cd python/tests +python test_prollytree.py # Basic ProllyTree functionality +python test_agent.py # Agent memory system +python test_versioned_kv.py # Versioned key-value store +``` From eb8bea8e3f7520aa1b1050815abdc9b9f4894633 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Mon, 28 Jul 2025 17:01:16 -0700 Subject: [PATCH 3/3] fix fmt error --- src/python.rs | 279 ++++++++++++++++++++++++++++++-------------------- 1 file changed, 168 insertions(+), 111 deletions(-) diff --git a/src/python.rs b/src/python.rs index d0994f1..97bc295 100644 --- a/src/python.rs +++ b/src/python.rs @@ -22,7 +22,7 @@ use std::sync::{Arc, Mutex}; use crate::{ agent::{AgentMemorySystem, MemoryType}, config::TreeConfig, - git::{GitVersionedKvStore, types::StorageBackend}, + git::{types::StorageBackend, GitVersionedKvStore}, proof::Proof, storage::{FileNodeStorage, InMemoryNodeStorage}, tree::{ProllyTree, Tree}, @@ -404,9 +404,10 @@ impl PyAgentMemorySystem { #[new] #[pyo3(signature = (path, agent_id))] fn new(path: String, agent_id: String) -> PyResult { - let memory_system = AgentMemorySystem::init(path, agent_id, None) - .map_err(|e| PyValueError::new_err(format!("Failed to initialize memory system: {}", e)))?; - + let memory_system = AgentMemorySystem::init(path, agent_id, None).map_err(|e| { + PyValueError::new_err(format!("Failed to initialize memory system: {}", e)) + })?; + Ok(PyAgentMemorySystem { inner: Arc::new(Mutex::new(memory_system)), }) @@ -416,7 +417,7 @@ impl PyAgentMemorySystem { fn open(path: String, agent_id: String) -> PyResult { let memory_system = AgentMemorySystem::open(path, agent_id, None) .map_err(|e| PyValueError::new_err(format!("Failed to open memory system: {}", e)))?; - + Ok(PyAgentMemorySystem { inner: Arc::new(Mutex::new(memory_system)), }) @@ -434,21 +435,24 @@ impl PyAgentMemorySystem { py.allow_threads(|| { let runtime = tokio::runtime::Runtime::new() .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; - + let mut memory_system = self.inner.lock().unwrap(); - + // Convert HashMap to HashMap let metadata_values = metadata.map(|m| { m.into_iter() .map(|(k, v)| (k, serde_json::Value::String(v))) .collect() }); - + runtime.block_on(async { - memory_system.short_term + memory_system + .short_term .store_conversation_turn(&thread_id, &role, &content, metadata_values) .await - .map_err(|e| PyValueError::new_err(format!("Failed to store conversation: {}", e))) + .map_err(|e| { + PyValueError::new_err(format!("Failed to store conversation: {}", e)) + }) }) }) } @@ -463,23 +467,30 @@ impl PyAgentMemorySystem { py.allow_threads(|| { let runtime = tokio::runtime::Runtime::new() .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; - + let memory_system = self.inner.lock().unwrap(); - + runtime.block_on(async { - let history = memory_system.short_term + let history = memory_system + .short_term .get_conversation_history(&thread_id, limit) .await .map_err(|e| PyValueError::new_err(format!("Failed to get history: {}", e)))?; - + Python::with_gil(|py| { - let results: PyResult>>> = history.iter().map(|doc| { - let mut map = HashMap::new(); - map.insert("id".to_string(), doc.id.clone().into_py(py)); - map.insert("content".to_string(), doc.content.to_string().into_py(py)); - map.insert("created_at".to_string(), doc.metadata.created_at.to_rfc3339().into_py(py)); - Ok(map) - }).collect(); + let results: PyResult>>> = history + .iter() + .map(|doc| { + let mut map = HashMap::new(); + map.insert("id".to_string(), doc.id.clone().into_py(py)); + map.insert("content".to_string(), doc.content.to_string().into_py(py)); + map.insert( + "created_at".to_string(), + doc.metadata.created_at.to_rfc3339().into_py(py), + ); + Ok(map) + }) + .collect(); results }) }) @@ -498,14 +509,15 @@ impl PyAgentMemorySystem { py.allow_threads(|| { let runtime = tokio::runtime::Runtime::new() .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; - + let mut memory_system = self.inner.lock().unwrap(); - + let facts_value: serde_json::Value = serde_json::from_str(&facts) .map_err(|e| PyValueError::new_err(format!("Invalid JSON: {}", e)))?; - + runtime.block_on(async { - memory_system.semantic + memory_system + .semantic .store_fact(&entity_type, &entity_id, facts_value, confidence, &source) .await .map_err(|e| PyValueError::new_err(format!("Failed to store fact: {}", e))) @@ -522,24 +534,34 @@ impl PyAgentMemorySystem { py.allow_threads(|| { let runtime = tokio::runtime::Runtime::new() .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; - + let memory_system = self.inner.lock().unwrap(); - + runtime.block_on(async { - let facts = memory_system.semantic + let facts = memory_system + .semantic .get_entity_facts(&entity_type, &entity_id) .await .map_err(|e| PyValueError::new_err(format!("Failed to get facts: {}", e)))?; - + Python::with_gil(|py| { - let results: PyResult>>> = facts.iter().map(|doc| { - let mut map = HashMap::new(); - map.insert("id".to_string(), doc.id.clone().into_py(py)); - map.insert("facts".to_string(), doc.content.to_string().into_py(py)); - map.insert("confidence".to_string(), doc.metadata.confidence.into_py(py)); - map.insert("source".to_string(), doc.metadata.source.clone().into_py(py)); - Ok(map) - }).collect(); + let results: PyResult>>> = facts + .iter() + .map(|doc| { + let mut map = HashMap::new(); + map.insert("id".to_string(), doc.id.clone().into_py(py)); + map.insert("facts".to_string(), doc.content.to_string().into_py(py)); + map.insert( + "confidence".to_string(), + doc.metadata.confidence.into_py(py), + ); + map.insert( + "source".to_string(), + doc.metadata.source.clone().into_py(py), + ); + Ok(map) + }) + .collect(); results }) }) @@ -560,23 +582,30 @@ impl PyAgentMemorySystem { py.allow_threads(|| { let runtime = tokio::runtime::Runtime::new() .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; - + let mut memory_system = self.inner.lock().unwrap(); - - let steps_values: Result, _> = steps.iter() - .map(|s| serde_json::from_str(s)) - .collect(); + + let steps_values: Result, _> = + steps.iter().map(|s| serde_json::from_str(s)).collect(); let steps_values = steps_values .map_err(|e| PyValueError::new_err(format!("Invalid JSON in steps: {}", e)))?; - + // Convert prerequisites to serde_json::Value - let conditions = prerequisites.map(|p| serde_json::Value::Array( - p.into_iter().map(serde_json::Value::String).collect() - )); - + let conditions = prerequisites.map(|p| { + serde_json::Value::Array(p.into_iter().map(serde_json::Value::String).collect()) + }); + runtime.block_on(async { - memory_system.procedural - .store_procedure(&category, &name, &description, steps_values, conditions, priority) + memory_system + .procedural + .store_procedure( + &category, + &name, + &description, + steps_values, + conditions, + priority, + ) .await .map_err(|e| PyValueError::new_err(format!("Failed to store procedure: {}", e))) }) @@ -591,23 +620,32 @@ impl PyAgentMemorySystem { py.allow_threads(|| { let runtime = tokio::runtime::Runtime::new() .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; - + let memory_system = self.inner.lock().unwrap(); - + runtime.block_on(async { - let procedures = memory_system.procedural + let procedures = memory_system + .procedural .get_procedures_by_category(&category) .await - .map_err(|e| PyValueError::new_err(format!("Failed to get procedures: {}", e)))?; - + .map_err(|e| { + PyValueError::new_err(format!("Failed to get procedures: {}", e)) + })?; + Python::with_gil(|py| { - let results: PyResult>>> = procedures.iter().map(|doc| { - let mut map = HashMap::new(); - map.insert("id".to_string(), doc.id.clone().into_py(py)); - map.insert("content".to_string(), doc.content.to_string().into_py(py)); - map.insert("created_at".to_string(), doc.metadata.created_at.to_rfc3339().into_py(py)); - Ok(map) - }).collect(); + let results: PyResult>>> = procedures + .iter() + .map(|doc| { + let mut map = HashMap::new(); + map.insert("id".to_string(), doc.id.clone().into_py(py)); + map.insert("content".to_string(), doc.content.to_string().into_py(py)); + map.insert( + "created_at".to_string(), + doc.metadata.created_at.to_rfc3339().into_py(py), + ); + Ok(map) + }) + .collect(); results }) }) @@ -618,13 +656,13 @@ impl PyAgentMemorySystem { py.allow_threads(|| { let runtime = tokio::runtime::Runtime::new() .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; - + let mut memory_system = self.inner.lock().unwrap(); - + runtime.block_on(async { - memory_system.checkpoint(&message) - .await - .map_err(|e| PyValueError::new_err(format!("Failed to create checkpoint: {}", e))) + memory_system.checkpoint(&message).await.map_err(|e| { + PyValueError::new_err(format!("Failed to create checkpoint: {}", e)) + }) }) }) } @@ -633,17 +671,21 @@ impl PyAgentMemorySystem { py.allow_threads(|| { let runtime = tokio::runtime::Runtime::new() .map_err(|e| PyValueError::new_err(format!("Failed to create runtime: {}", e)))?; - + let mut memory_system = self.inner.lock().unwrap(); - + runtime.block_on(async { - let report = memory_system.optimize() + let report = memory_system + .optimize() .await .map_err(|e| PyValueError::new_err(format!("Failed to optimize: {}", e)))?; - + let mut result = HashMap::new(); result.insert("expired_cleaned".to_string(), report.expired_cleaned); - result.insert("memories_consolidated".to_string(), report.memories_consolidated); + result.insert( + "memories_consolidated".to_string(), + report.memories_consolidated, + ); result.insert("memories_archived".to_string(), report.memories_archived); result.insert("memories_pruned".to_string(), report.memories_pruned); result.insert("total_processed".to_string(), report.total_processed()); @@ -705,7 +747,7 @@ impl PyVersionedKvStore { fn new(path: String) -> PyResult { let store = GitVersionedKvStore::<32>::init(path) .map_err(|e| PyValueError::new_err(format!("Failed to initialize store: {}", e)))?; - + Ok(PyVersionedKvStore { inner: Arc::new(Mutex::new(store)), }) @@ -715,7 +757,7 @@ impl PyVersionedKvStore { fn open(path: String) -> PyResult { let store = GitVersionedKvStore::<32>::open(path) .map_err(|e| PyValueError::new_err(format!("Failed to open store: {}", e)))?; - + Ok(PyVersionedKvStore { inner: Arc::new(Mutex::new(store)), }) @@ -724,17 +766,18 @@ impl PyVersionedKvStore { fn insert(&self, key: &Bound<'_, PyBytes>, value: &Bound<'_, PyBytes>) -> PyResult<()> { let key_vec = key.as_bytes().to_vec(); let value_vec = value.as_bytes().to_vec(); - + let mut store = self.inner.lock().unwrap(); - store.insert(key_vec, value_vec) + store + .insert(key_vec, value_vec) .map_err(|e| PyValueError::new_err(format!("Failed to insert: {}", e)))?; - + Ok(()) } fn get(&self, py: Python, key: &Bound<'_, PyBytes>) -> PyResult>> { let key_vec = key.as_bytes().to_vec(); - + let store = self.inner.lock().unwrap(); match store.get(&key_vec) { Some(value) => Ok(Some(PyBytes::new_bound(py, &value).into())), @@ -745,73 +788,79 @@ impl PyVersionedKvStore { fn update(&self, key: &Bound<'_, PyBytes>, value: &Bound<'_, PyBytes>) -> PyResult { let key_vec = key.as_bytes().to_vec(); let value_vec = value.as_bytes().to_vec(); - + let mut store = self.inner.lock().unwrap(); - store.update(key_vec, value_vec) + store + .update(key_vec, value_vec) .map_err(|e| PyValueError::new_err(format!("Failed to update: {}", e))) } fn delete(&self, key: &Bound<'_, PyBytes>) -> PyResult { let key_vec = key.as_bytes().to_vec(); - + let mut store = self.inner.lock().unwrap(); - store.delete(&key_vec) + store + .delete(&key_vec) .map_err(|e| PyValueError::new_err(format!("Failed to delete: {}", e))) } fn list_keys(&self, py: Python) -> PyResult>> { let store = self.inner.lock().unwrap(); let keys = store.list_keys(); - - let py_keys: Vec> = keys.iter() + + let py_keys: Vec> = keys + .iter() .map(|key| PyBytes::new_bound(py, key).into()) .collect(); - + Ok(py_keys) } fn status(&self, py: Python) -> PyResult, String)>> { let store = self.inner.lock().unwrap(); let status = store.status(); - - let py_status: Vec<(Py, String)> = status.iter() - .map(|(key, status_str)| { - (PyBytes::new_bound(py, key).into(), status_str.clone()) - }) + + let py_status: Vec<(Py, String)> = status + .iter() + .map(|(key, status_str)| (PyBytes::new_bound(py, key).into(), status_str.clone())) .collect(); - + Ok(py_status) } fn commit(&self, message: String) -> PyResult { let mut store = self.inner.lock().unwrap(); - let commit_id = store.commit(&message) + let commit_id = store + .commit(&message) .map_err(|e| PyValueError::new_err(format!("Failed to commit: {}", e)))?; - + Ok(commit_id.to_hex().to_string()) } fn branch(&self, name: String) -> PyResult<()> { let mut store = self.inner.lock().unwrap(); - store.branch(&name) + store + .branch(&name) .map_err(|e| PyValueError::new_err(format!("Failed to create branch: {}", e)))?; - + Ok(()) } fn create_branch(&self, name: String) -> PyResult<()> { let mut store = self.inner.lock().unwrap(); - store.create_branch(&name) - .map_err(|e| PyValueError::new_err(format!("Failed to create and switch branch: {}", e)))?; - + store.create_branch(&name).map_err(|e| { + PyValueError::new_err(format!("Failed to create and switch branch: {}", e)) + })?; + Ok(()) } fn checkout(&self, branch_or_commit: String) -> PyResult<()> { let mut store = self.inner.lock().unwrap(); - store.checkout(&branch_or_commit) + store + .checkout(&branch_or_commit) .map_err(|e| PyValueError::new_err(format!("Failed to checkout: {}", e)))?; - + Ok(()) } @@ -822,25 +871,33 @@ impl PyVersionedKvStore { fn list_branches(&self) -> PyResult> { let store = self.inner.lock().unwrap(); - store.list_branches() + store + .list_branches() .map_err(|e| PyValueError::new_err(format!("Failed to list branches: {}", e))) } fn log(&self) -> PyResult>>> { let store = self.inner.lock().unwrap(); - let history = store.log() + let history = store + .log() .map_err(|e| PyValueError::new_err(format!("Failed to get log: {}", e)))?; - + Python::with_gil(|py| { - let results: PyResult>>> = history.iter().map(|commit| { - let mut map = HashMap::new(); - map.insert("id".to_string(), commit.id.to_hex().to_string().into_py(py)); - map.insert("author".to_string(), commit.author.clone().into_py(py)); - map.insert("committer".to_string(), commit.committer.clone().into_py(py)); - map.insert("message".to_string(), commit.message.clone().into_py(py)); - map.insert("timestamp".to_string(), commit.timestamp.into_py(py)); - Ok(map) - }).collect(); + let results: PyResult>>> = history + .iter() + .map(|commit| { + let mut map = HashMap::new(); + map.insert("id".to_string(), commit.id.to_hex().to_string().into_py(py)); + map.insert("author".to_string(), commit.author.clone().into_py(py)); + map.insert( + "committer".to_string(), + commit.committer.clone().into_py(py), + ); + map.insert("message".to_string(), commit.message.clone().into_py(py)); + map.insert("timestamp".to_string(), commit.timestamp.into_py(py)); + Ok(map) + }) + .collect(); results }) }