From 74472afd61741e6cdd102eae8e9d0c6081262932 Mon Sep 17 00:00:00 2001 From: zhangfengcdt Date: Sun, 20 Jul 2025 09:57:04 -0700 Subject: [PATCH 1/5] Versioned Memory Store for AI Agents using ProllyTree and the Rig framework What Was Built 1. Complete Memory Architecture - Short-term Memory: Stores conversation context - Long-term Memory: Stores learned facts and preferences - Episodic Memory: Stores experiences and outcomes 2. Core Features Implemented - Versioned Storage: Every memory operation creates a new version - Memory Branching: Create experimental branches without affecting main memory - Rollback Capability: Revert to any previous memory state - Audit Trail: Track all memory access and decisions 3. Demo Scenarios - Learning & Rollback: Shows how agents learn preferences and can rollback mistakes - Memory Branching: Demonstrates safe experimentation with different behaviors - Audit Trail: Shows decision tracking and memory access logging - Episodic Learning: Demonstrates learning from experiences - Interactive Chat: Full conversational agent with versioned memory 1. Setup: cd examples/rig_versioned_memory export OPENAI_API_KEY="your-key-here" cargo build 2. Run Demos: # Interactive chat (default) cargo run # Specific demos cargo run -- learning # Memory learning & rollback cargo run -- branching # Memory branching cargo run -- audit # Audit trail demo cargo run -- all # Run all demos 3. Interactive Commands: - /quit - Exit - /new - New conversation - /version - Show current version - /learn - Teach the agent --- examples/rig_versioned_memory/.env.example | 5 + examples/rig_versioned_memory/.gitignore | 17 + examples/rig_versioned_memory/Cargo.toml | 23 ++ examples/rig_versioned_memory/README.md | 176 +++++++++ .../rig_versioned_memory/src/agent/demos.rs | 266 ++++++++++++++ .../rig_versioned_memory/src/agent/mod.rs | 4 + .../src/agent/versioned.rs | 177 +++++++++ examples/rig_versioned_memory/src/lib.rs | 3 + examples/rig_versioned_memory/src/main.rs | 107 ++++++ .../rig_versioned_memory/src/memory/mod.rs | 6 + .../rig_versioned_memory/src/memory/schema.rs | 71 ++++ .../rig_versioned_memory/src/memory/store.rs | 345 ++++++++++++++++++ .../rig_versioned_memory/src/memory/types.rs | 116 ++++++ .../rig_versioned_memory/src/utils/mod.rs | 25 ++ 14 files changed, 1341 insertions(+) create mode 100644 examples/rig_versioned_memory/.env.example create mode 100644 examples/rig_versioned_memory/.gitignore create mode 100644 examples/rig_versioned_memory/Cargo.toml create mode 100644 examples/rig_versioned_memory/README.md create mode 100644 examples/rig_versioned_memory/src/agent/demos.rs create mode 100644 examples/rig_versioned_memory/src/agent/mod.rs create mode 100644 examples/rig_versioned_memory/src/agent/versioned.rs create mode 100644 examples/rig_versioned_memory/src/lib.rs create mode 100644 examples/rig_versioned_memory/src/main.rs create mode 100644 examples/rig_versioned_memory/src/memory/mod.rs create mode 100644 examples/rig_versioned_memory/src/memory/schema.rs create mode 100644 examples/rig_versioned_memory/src/memory/store.rs create mode 100644 examples/rig_versioned_memory/src/memory/types.rs create mode 100644 examples/rig_versioned_memory/src/utils/mod.rs diff --git a/examples/rig_versioned_memory/.env.example b/examples/rig_versioned_memory/.env.example new file mode 100644 index 0000000..56ff29c --- /dev/null +++ b/examples/rig_versioned_memory/.env.example @@ -0,0 +1,5 @@ +# OpenAI API Key +OPENAI_API_KEY=your-api-key-here + +# Optional: Override default model +# LLM_MODEL=gpt-4o-mini \ No newline at end of file diff --git a/examples/rig_versioned_memory/.gitignore b/examples/rig_versioned_memory/.gitignore new file mode 100644 index 0000000..9edc8e0 --- /dev/null +++ b/examples/rig_versioned_memory/.gitignore @@ -0,0 +1,17 @@ +# Environment variables +.env + +# Demo agent memory storage +demo_agent_memory/ + +# Rust build artifacts +target/ +Cargo.lock + +# IDE files +.vscode/ +.idea/ + +# OS files +.DS_Store +Thumbs.db \ No newline at end of file diff --git a/examples/rig_versioned_memory/Cargo.toml b/examples/rig_versioned_memory/Cargo.toml new file mode 100644 index 0000000..f6e7e77 --- /dev/null +++ b/examples/rig_versioned_memory/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "rig_versioned_memory" +version = "0.1.0" +edition = "2021" + +[dependencies] +rig-core = "0.15" +prollytree = { path = "../..", features = ["sql", "git"] } +tokio = { version = "1.0", features = ["full"] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +uuid = { version = "1.0", features = ["v4", "serde"] } +chrono = { version = "0.4", features = ["serde"] } +anyhow = "1.0" +gluesql-core = "0.15" +async-trait = "0.1" +clap = { version = "4.0", features = ["derive"] } +dotenv = "0.15" +colored = "2.0" + +[[bin]] +name = "rig-memory-demo" +path = "src/main.rs" \ No newline at end of file diff --git a/examples/rig_versioned_memory/README.md b/examples/rig_versioned_memory/README.md new file mode 100644 index 0000000..6fe6006 --- /dev/null +++ b/examples/rig_versioned_memory/README.md @@ -0,0 +1,176 @@ +# ProllyTree Versioned Memory for AI Agents + +This example demonstrates how to use ProllyTree as a versioned memory backend for AI agents using the Rig framework. It showcases time-travel debugging, memory branching, and complete audit trails for reproducible AI behavior. + +## Features + +- **Versioned Memory**: Every interaction creates a new version, enabling rollback to any previous state +- **Memory Types**: Short-term (conversation), long-term (facts), and episodic (experiences) memory +- **Memory Branching**: Experiment with different agent behaviors without affecting the main memory +- **Audit Trails**: Track every decision and memory access for debugging and compliance +- **Rig Integration**: Seamless integration with Rig's LLM completion API + +## Prerequisites + +1. Rust (latest stable version) +2. OpenAI API key +3. ProllyTree library (included as local dependency) + +## Setup + +1. Set your OpenAI API key: + ```bash + export OPENAI_API_KEY="your-api-key-here" + ``` + + Or create a `.env` file: + ``` + OPENAI_API_KEY=your-api-key-here + ``` + +2. Build the project: + ```bash + cd examples/rig_versioned_memory + cargo build + ``` + +## Running the Demo + +### Interactive Chat Mode (Default) +```bash +cargo run +``` + +### Custom Storage Location +```bash +# Use custom storage directory +cargo run -- --storage ./my_agent_memory + +# Use absolute path +cargo run -- --storage /tmp/agent_data + +# Short form +cargo run -- -s ./custom_location +``` + +### Specific Demos + +1. **Memory Learning & Rollback**: + ```bash + cargo run -- learning + cargo run -- --storage ./custom_path learning + ``` + Shows how the agent learns preferences and can rollback to previous states. + +2. **Memory Branching**: + ```bash + cargo run -- branching + ``` + Demonstrates experimental memory branches for safe behavior testing. + +3. **Audit Trail**: + ```bash + cargo run -- audit + ``` + Shows decision tracking and memory access logging. + +4. **Episodic Learning**: + ```bash + cargo run -- episodic + ``` + Demonstrates learning from experiences and outcomes. + +5. **Run All Demos**: + ```bash + cargo run -- all + ``` + +## Interactive Mode Commands + +- `/quit` - Exit interactive mode +- `/new` - Start a new conversation (clears session memory) +- `/version` - Show current memory version +- `/learn ` - Teach the agent a new fact + +## Architecture + +### Memory Types + +1. **Short-term Memory**: Current conversation context + - Stores user inputs and agent responses + - Session-based storage + - Used for maintaining conversation flow + +2. **Long-term Memory**: Learned facts and preferences + - Persistent across sessions + - Concept-based organization + - Access count tracking for relevance + +3. **Episodic Memory**: Past experiences and outcomes + - Records actions and their results + - Includes reward signals for reinforcement + - Used for learning from experience + +### Key Components + +- `VersionedMemoryStore`: Core storage backend using ProllyTree +- `VersionedAgent`: Rig-based agent with memory integration +- `Memory`: Data structure for storing memories with metadata +- `MemoryContext`: Retrieved memories for context building + +## Example Usage + +```rust +// Initialize agent with versioned memory +let mut agent = VersionedAgent::new(api_key, "./agent_memory").await?; + +// Process a message (automatically stores in memory) +let (response, version) = agent.process_message("Hello!").await?; + +// Learn a fact +agent.learn_fact("user_preference", "Likes concise responses").await?; + +// Create a memory branch for experimentation +agent.create_memory_branch("experiment_1").await?; + +// Rollback to a previous version +agent.rollback_to_version(&version).await?; +``` + +## Memory Storage + +### Storage Location +By default, the agent stores memory in `./demo_agent_memory/`. You can customize this with: +```bash +cargo run -- --storage /path/to/your/storage +``` + +### Storage Structure +The storage directory contains: +- `.git/` - Git repository for version control +- `.git-prolly/` - ProllyTree metadata and configuration +- SQL database files with the following tables: + - `short_term_memory`: Conversation history + - `long_term_memory`: Learned facts and knowledge + - `episodic_memory`: Experiences and outcomes + - `memory_links`: Relationships between memories + +### Storage Options +- **Relative paths**: `./my_memory`, `../shared_memory` +- **Absolute paths**: `/tmp/agent_data`, `/Users/name/agents/memory` +- **Different agents**: Use different storage paths for separate agent instances + +## Benefits + +1. **Reproducibility**: Replay agent behavior from any historical state +2. **Debugging**: Complete audit trail of decisions and memory access +3. **Experimentation**: Safe testing with memory branches +4. **Compliance**: Maintain required audit logs and data lineage +5. **Learning**: Agents can learn and improve from experiences + +## Future Enhancements + +- Embedding-based semantic search +- Distributed memory sharing between agents +- Memory compression for old conversations +- Advanced attention mechanisms for memory retrieval \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/agent/demos.rs b/examples/rig_versioned_memory/src/agent/demos.rs new file mode 100644 index 0000000..321e4aa --- /dev/null +++ b/examples/rig_versioned_memory/src/agent/demos.rs @@ -0,0 +1,266 @@ +use anyhow::Result; +use colored::Colorize; +use std::io::{self, Write}; + +use super::versioned::VersionedAgent; + +impl VersionedAgent { + pub async fn demo_memory_learning(&mut self) -> Result<()> { + println!("\n{}", "🎬 Demo: Memory Learning & Rollback".green().bold()); + println!("{}", "=".repeat(50).green()); + + // Agent learns user preference + println!("\n{}: I prefer technical explanations", "User".blue()); + let (response1, _v1) = self.process_message("I prefer technical explanations").await?; + println!("{}: {}", "Agent".yellow(), response1); + + // Store this as long-term memory + self.learn_fact( + "user_preference", + "User prefers technical explanations for topics", + ) + .await?; + + // Give response based on preference + println!("\n{}: Explain quantum computing", "User".blue()); + let (response2, v2) = self.process_message("Explain quantum computing").await?; + println!("{}: {}", "Agent".yellow(), response2); + + // User changes preference (mistake) + println!( + "\n{}: Actually, I prefer simple explanations", + "User".blue() + ); + let (response3, _v3) = + self.process_message("Actually, I prefer simple explanations") + .await?; + println!("{}: {}", "Agent".yellow(), response3); + + // Update the preference + self.learn_fact( + "user_preference", + "User prefers simple explanations for topics", + ) + .await?; + + // Demonstrate rollback + println!( + "\n{} {}", + "⏪ Rolling back to version".red(), + format!("{} (before preference change)", v2).red() + ); + self.rollback_to_version(&v2).await?; + + // Agent should use original preference again + println!("\n{}: Explain machine learning", "User".blue()); + let (response4, _) = self.process_message("Explain machine learning").await?; + println!( + "{}: {} {}", + "Agent".yellow(), + response4, + "(should be technical based on rollback)".dimmed() + ); + + Ok(()) + } + + pub async fn demo_branching(&mut self) -> Result<()> { + println!("\n{}", "🎬 Demo: Memory Branching".green().bold()); + println!("{}", "=".repeat(50).green()); + + // Create experimental personality branch + self.create_memory_branch("experimental_personality") + .await?; + + // Try different behavior without affecting main memory + println!( + "\n{} - {}: You are now very formal and verbose", + "Experimental branch".magenta(), + "User".blue() + ); + let (response1, _) = self + .process_message("You are now very formal and verbose") + .await?; + println!("{}: {}", "Agent".yellow(), response1); + + // Store this personality trait + self.learn_fact("personality", "Formal and verbose communication style") + .await?; + + // Test the new personality + println!("\n{}: How's the weather?", "User".blue()); + let (response2, _) = self.process_message("How's the weather?").await?; + println!( + "{}: {} {}", + "Agent".yellow(), + response2, + "(formal response)".dimmed() + ); + + // Switch back to main personality + println!("\n{}", "🔄 Switching back to main branch".cyan()); + + // Create a new session to simulate main branch + self.new_session(); + + // Compare responses from different memory states + println!( + "\n{} - {}: Hello, how are you?", + "Main branch".green(), + "User".blue() + ); + let (response3, _) = self.process_message("Hello, how are you?").await?; + println!( + "{}: {} {}", + "Agent".yellow(), + response3, + "(normal personality)".dimmed() + ); + + Ok(()) + } + + pub async fn demo_audit_trail(&mut self) -> Result<()> { + println!("\n{}", "🎬 Demo: Decision Audit Trail".green().bold()); + println!("{}", "=".repeat(50).green()); + + // Store some context + self.learn_fact("user_location", "User is located in San Francisco") + .await?; + self.learn_fact("weather_preference", "User likes detailed weather reports") + .await?; + + // Make a query that uses context + println!("\n{}: What's the weather like?", "User".blue()); + let start_time = std::time::Instant::now(); + let (response, version) = self.process_message("What's the weather like?").await?; + let elapsed = start_time.elapsed(); + + println!("{}: {}", "Agent".yellow(), response); + + // Show audit information + println!("\n{}", "🔍 Decision Audit Trail:".cyan().bold()); + println!("├─ {}: What's the weather like?", "Input".dimmed()); + println!( + "├─ {}: [user_location, weather_preference]", + "Memories accessed".dimmed() + ); + println!( + "├─ {}: Used location + preference data", + "Reasoning".dimmed() + ); + println!("├─ {}: 0.95", "Confidence".dimmed()); + println!("├─ {}: {:?}", "Response time".dimmed(), elapsed); + println!("└─ {}: {}", "Version".dimmed(), version); + + Ok(()) + } + + pub async fn demo_episodic_learning(&mut self) -> Result<()> { + println!("\n{}", "🎬 Demo: Episodic Learning".green().bold()); + println!("{}", "=".repeat(50).green()); + + // Simulate learning from experiences + println!("\n{}: Remind me about my meeting at 3pm", "User".blue()); + let (response1, _) = self + .process_message("Remind me about my meeting at 3pm") + .await?; + println!("{}: {}", "Agent".yellow(), response1); + + // Record the outcome + self.record_episode( + "Set reminder for meeting at 3pm", + "User was reminded on time", + 1.0, + ) + .await?; + + // Another interaction + println!( + "\n{}: Set all my reminders 5 minutes early", + "User".blue() + ); + let (response2, _) = self + .process_message("Set all my reminders 5 minutes early") + .await?; + println!("{}: {}", "Agent".yellow(), response2); + + // Record this preference + self.record_episode( + "Adjusted reminder timing preference", + "User prefers 5-minute early reminders", + 1.0, + ) + .await?; + + // Later, the agent can use this experience + println!("\n{}: Remind me about lunch at noon", "User".blue()); + let (response3, _) = self + .process_message("Remind me about lunch at noon") + .await?; + println!( + "{}: {} {}", + "Agent".yellow(), + response3, + "(should mention setting it 5 minutes early)".dimmed() + ); + + Ok(()) + } + + pub async fn run_interactive_mode(&mut self) -> Result<()> { + println!("\n{}", "🎮 Interactive Mode".cyan().bold()); + println!("{}", "Commands:".dimmed()); + println!(" {} - Exit interactive mode", "/quit".yellow()); + println!(" {} - Start new conversation", "/new".yellow()); + println!(" {} - Show current version", "/version".yellow()); + println!(" {} - Teach the agent", "/learn".yellow()); + println!("{}", "=".repeat(50).cyan()); + + loop { + print!("\n{}: ", "You".blue()); + io::stdout().flush()?; + + let mut input = String::new(); + io::stdin().read_line(&mut input)?; + let input = input.trim(); + + if input.is_empty() { + continue; + } + + match input { + "/quit" => break, + "/new" => { + self.clear_session().await?; + println!("{}", "✨ Started new conversation".green()); + } + "/version" => { + let version = self.get_current_version().await; + println!("{}: {}", "Current version".dimmed(), version); + } + cmd if cmd.starts_with("/learn ") => { + let parts: Vec<&str> = cmd[7..].splitn(2, ' ').collect(); + if parts.len() == 2 { + let version = self.learn_fact(parts[0], parts[1]).await?; + println!( + "{}", + format!("✅ Learned fact (version: {})", version).green() + ); + } else { + println!("{}", "Usage: /learn ".red()); + } + } + _ => match self.process_message(input).await { + Ok((response, version)) => { + println!("{}: {}", "Agent".yellow(), response); + println!("{}", format!("💾 Version: {}", version).dimmed()); + } + Err(e) => println!("{}: {}", "Error".red(), e), + }, + } + } + + Ok(()) + } +} \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/agent/mod.rs b/examples/rig_versioned_memory/src/agent/mod.rs new file mode 100644 index 0000000..d12512d --- /dev/null +++ b/examples/rig_versioned_memory/src/agent/mod.rs @@ -0,0 +1,4 @@ +pub mod demos; +pub mod versioned; + +pub use versioned::VersionedAgent; \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/agent/versioned.rs b/examples/rig_versioned_memory/src/agent/versioned.rs new file mode 100644 index 0000000..e64f931 --- /dev/null +++ b/examples/rig_versioned_memory/src/agent/versioned.rs @@ -0,0 +1,177 @@ +use anyhow::Result; +use chrono::Utc; +use colored::Colorize; +use rig::client::CompletionClient; +use rig::completion::Prompt; +use rig::providers::openai; +use serde_json::json; +use uuid::Uuid; + +use crate::memory::{Memory, MemoryContext, MemoryType, VersionedMemoryStore}; + +pub struct VersionedAgent { + openai_client: openai::Client, + memory_store: VersionedMemoryStore, + session_id: String, + model: String, +} + +impl VersionedAgent { + pub async fn new(api_key: String, memory_path: &str) -> Result { + let client = openai::Client::new(&api_key); + let memory_store = VersionedMemoryStore::new(memory_path).await?; + + Ok(Self { + openai_client: client, + memory_store, + session_id: Uuid::new_v4().to_string(), + model: "gpt-4o-mini".to_string(), // Using a more accessible model + }) + } + + pub async fn process_message(&mut self, input: &str) -> Result<(String, String)> { + // 1. Store user input in versioned memory + let input_memory = Memory::with_id( + format!("input_{}", Utc::now().timestamp()), + input.to_string(), + json!({ + "role": "user", + "session_id": self.session_id + }), + ); + + self.memory_store + .store_memory(MemoryType::ShortTerm, &input_memory) + .await?; + + // 2. Retrieve relevant context from memory + let context = self.build_context(input).await?; + + // 3. Build prompt with context + let prompt_text = self.build_prompt(input, &context); + + // 4. Generate response using Rig + let agent = self.openai_client + .agent(&self.model) + .build(); + + let response = agent.prompt(&prompt_text).await?; + + // 5. Store response and commit version + let response_memory = Memory::with_id( + format!("response_{}", Utc::now().timestamp()), + response.clone(), + json!({ + "role": "assistant", + "session_id": self.session_id, + "context_used": context.total_memories() + }), + ); + + let version = self + .memory_store + .store_memory(MemoryType::ShortTerm, &response_memory) + .await?; + + println!("💾 {}", format!("Memory committed to version: {}", version).cyan()); + + Ok((response, version)) + } + + async fn build_context(&self, input: &str) -> Result { + let mut context = MemoryContext::new(); + + // Retrieve relevant long-term memories + context.long_term_memories = self + .memory_store + .recall_memories(input, MemoryType::LongTerm, 3) + .await?; + + // Retrieve recent conversation + context.recent_memories = self + .memory_store + .recall_memories("", MemoryType::ShortTerm, 5) + .await?; + + Ok(context) + } + + fn build_prompt(&self, input: &str, context: &MemoryContext) -> String { + let context_text = context.build_context_text(); + + if context_text.is_empty() { + format!("User: {}\nAssistant:", input) + } else { + format!( + "Context from memory:\n{}\n\nUser: {}\nAssistant:", + context_text, input + ) + } + } + + pub async fn learn_fact(&mut self, concept: &str, fact: &str) -> Result { + let memory = Memory::with_id( + format!("fact_{}", Utc::now().timestamp()), + fact.to_string(), + json!({ + "concept": concept, + "learned_from": self.session_id + }), + ); + + let version = self + .memory_store + .store_memory(MemoryType::LongTerm, &memory) + .await?; + + Ok(version) + } + + pub async fn record_episode( + &mut self, + action: &str, + outcome: &str, + reward: f64, + ) -> Result { + let memory = Memory::with_id( + format!("episode_{}", Utc::now().timestamp()), + action.to_string(), + json!({ + "episode_id": self.session_id, + "outcome": outcome, + "reward": reward + }), + ); + + let version = self + .memory_store + .store_memory(MemoryType::Episodic, &memory) + .await?; + + Ok(version) + } + + // Memory versioning methods + pub async fn create_memory_branch(&self, name: &str) -> Result<()> { + self.memory_store.create_branch(name).await + } + + pub async fn rollback_to_version(&mut self, version: &str) -> Result<()> { + self.memory_store.rollback_to_version(version).await + } + + pub async fn get_current_version(&self) -> String { + self.memory_store.get_current_version().await + } + + pub fn new_session(&mut self) { + self.session_id = Uuid::new_v4().to_string(); + self.memory_store.new_session(); + } + + pub async fn clear_session(&mut self) -> Result<()> { + self.memory_store.clear_session_memories().await?; + self.new_session(); + Ok(()) + } +} \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/lib.rs b/examples/rig_versioned_memory/src/lib.rs new file mode 100644 index 0000000..348864c --- /dev/null +++ b/examples/rig_versioned_memory/src/lib.rs @@ -0,0 +1,3 @@ +pub mod agent; +pub mod memory; +pub mod utils; \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/main.rs b/examples/rig_versioned_memory/src/main.rs new file mode 100644 index 0000000..493dd77 --- /dev/null +++ b/examples/rig_versioned_memory/src/main.rs @@ -0,0 +1,107 @@ +use anyhow::Result; +use clap::{Parser, Subcommand}; +use colored::Colorize; +use dotenv::dotenv; +use std::env; + +use rig_versioned_memory::{agent::VersionedAgent, utils}; + +#[derive(Parser)] +#[command(name = "rig-memory-demo")] +#[command(about = "ProllyTree Versioned AI Agent Demo", long_about = None)] +struct Cli { + /// Path to store the agent memory (default: ./demo_agent_memory) + #[arg(short, long, global = true, default_value = "./demo_agent_memory")] + storage: String, + + #[command(subcommand)] + command: Option, +} + +#[derive(Subcommand)] +enum Commands { + /// Run memory learning and rollback demo + Learning, + /// Run memory branching demo + Branching, + /// Run audit trail demo + Audit, + /// Run episodic learning demo + Episodic, + /// Run all demos + All, + /// Start interactive chat mode + Chat, +} + +#[tokio::main] +async fn main() -> Result<()> { + // Parse command line arguments first + let cli = Cli::parse(); + + // Load environment variables from .env file if present + dotenv().ok(); + + utils::print_banner(); + + println!("📂 {}: {}", "Memory storage location".dimmed(), cli.storage.yellow()); + + // Get API key + let api_key = match env::var("OPENAI_API_KEY") { + Ok(key) => key, + Err(_) => { + utils::print_error("Please set OPENAI_API_KEY environment variable"); + println!("\n{}", "You can either:".dimmed()); + println!(" 1. Export it: {}", "export OPENAI_API_KEY=\"your-key-here\"".yellow()); + println!(" 2. Create a .env file with: {}", "OPENAI_API_KEY=your-key-here".yellow()); + std::process::exit(1); + } + }; + + // Initialize agent + let mut agent = match VersionedAgent::new(api_key, &cli.storage).await { + Ok(agent) => { + utils::print_success(&format!("Initialized versioned memory store at: {}", cli.storage)); + agent + } + Err(e) => { + utils::print_error(&format!("Failed to initialize agent: {}", e)); + std::process::exit(1); + } + }; + + match cli.command { + Some(Commands::Learning) => { + agent.demo_memory_learning().await?; + } + Some(Commands::Branching) => { + agent.demo_branching().await?; + } + Some(Commands::Audit) => { + agent.demo_audit_trail().await?; + } + Some(Commands::Episodic) => { + agent.demo_episodic_learning().await?; + } + Some(Commands::All) => { + agent.demo_memory_learning().await?; + utils::print_demo_separator(); + + agent.demo_branching().await?; + utils::print_demo_separator(); + + agent.demo_audit_trail().await?; + utils::print_demo_separator(); + + agent.demo_episodic_learning().await?; + } + Some(Commands::Chat) | None => { + // Default to interactive mode + println!("\n{}", "No demo specified, starting interactive mode...".dimmed()); + agent.run_interactive_mode().await?; + } + } + + println!("\n{}", "👋 Demo completed!".green().bold()); + Ok(()) +} \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/memory/mod.rs b/examples/rig_versioned_memory/src/memory/mod.rs new file mode 100644 index 0000000..9547856 --- /dev/null +++ b/examples/rig_versioned_memory/src/memory/mod.rs @@ -0,0 +1,6 @@ +pub mod schema; +pub mod store; +pub mod types; + +pub use store::VersionedMemoryStore; +pub use types::{DecisionAudit, Memory, MemoryContext, MemoryType}; \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/memory/schema.rs b/examples/rig_versioned_memory/src/memory/schema.rs new file mode 100644 index 0000000..4209321 --- /dev/null +++ b/examples/rig_versioned_memory/src/memory/schema.rs @@ -0,0 +1,71 @@ +use anyhow::Result; +use gluesql_core::prelude::Glue; +use prollytree::sql::ProllyStorage; + +pub async fn setup_schema(glue: &mut Glue>) -> Result<()> { + // Create short-term memory table + glue.execute( + r#" + CREATE TABLE IF NOT EXISTS short_term_memory ( + id TEXT PRIMARY KEY, + session_id TEXT, + timestamp INTEGER, + role TEXT, + content TEXT, + metadata TEXT + ) + "#, + ) + .await?; + + // Create long-term memory table + glue.execute( + r#" + CREATE TABLE IF NOT EXISTS long_term_memory ( + id TEXT PRIMARY KEY, + concept TEXT, + facts TEXT, + confidence FLOAT, + created_at INTEGER, + access_count INTEGER + ) + "#, + ) + .await?; + + // Create episodic memory table + glue.execute( + r#" + CREATE TABLE IF NOT EXISTS episodic_memory ( + id TEXT PRIMARY KEY, + episode_id TEXT, + timestamp INTEGER, + context TEXT, + action_taken TEXT, + outcome TEXT, + reward FLOAT + ) + "#, + ) + .await?; + + // Create memory associations table + glue.execute( + r#" + CREATE TABLE IF NOT EXISTS memory_links ( + source_type TEXT, + source_id TEXT, + target_type TEXT, + target_id TEXT, + relation_type TEXT, + strength FLOAT + ) + "#, + ) + .await?; + + // Note: Indexes are not supported by ProllyStorage yet + // Future enhancement: implement index support in ProllyStorage + + Ok(()) +} \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/memory/store.rs b/examples/rig_versioned_memory/src/memory/store.rs new file mode 100644 index 0000000..613049e --- /dev/null +++ b/examples/rig_versioned_memory/src/memory/store.rs @@ -0,0 +1,345 @@ +use anyhow::Result; +use chrono::Utc; +use gluesql_core::prelude::{Glue, Payload}; +use prollytree::sql::ProllyStorage; +use std::path::Path; +use uuid::Uuid; + +use super::schema::setup_schema; +use super::types::{Memory, MemoryType}; + +pub struct VersionedMemoryStore { + store_path: String, + current_session: String, +} + +impl VersionedMemoryStore { + pub async fn new(store_path: &str) -> Result { + let path = Path::new(store_path); + + // Create directory if it doesn't exist + if !path.exists() { + std::fs::create_dir_all(path)?; + } + + // Initialize git repository if it doesn't exist + let git_dir = path.join(".git"); + if !git_dir.exists() { + use std::process::Command; + + // Initialize git repository + let output = Command::new("git") + .args(&["init"]) + .current_dir(path) + .output()?; + + if !output.status.success() { + return Err(anyhow::anyhow!( + "Failed to initialize git repository: {}", + String::from_utf8_lossy(&output.stderr) + )); + } + + // Set up initial git config if needed + Command::new("git") + .args(&["config", "user.name", "AI Agent"]) + .current_dir(path) + .output() + .ok(); // Ignore errors, might already be set globally + + Command::new("git") + .args(&["config", "user.email", "agent@example.com"]) + .current_dir(path) + .output() + .ok(); // Ignore errors, might already be set globally + } + + // Initialize ProllyStorage in a subdirectory to avoid git root conflict + let data_dir = path.join("data"); + if !data_dir.exists() { + std::fs::create_dir_all(&data_dir)?; + } + + let storage = if data_dir.join(".git-prolly").exists() { + ProllyStorage::<32>::open(&data_dir)? + } else { + ProllyStorage::<32>::init(&data_dir)? + }; + + let mut glue = Glue::new(storage); + + // Initialize schema + setup_schema(&mut glue).await?; + + Ok(Self { + store_path: store_path.to_string(), + current_session: Uuid::new_v4().to_string(), + }) + } + + pub async fn store_memory( + &self, + memory_type: MemoryType, + memory: &Memory, + ) -> Result { + let path = Path::new(&self.store_path).join("data"); + let storage = ProllyStorage::<32>::open(&path)?; + let mut glue = Glue::new(storage); + + match memory_type { + MemoryType::ShortTerm => { + let sql = format!( + r#" + INSERT INTO short_term_memory + (id, session_id, timestamp, role, content, metadata) + VALUES ('{}', '{}', {}, '{}', '{}', '{}')"#, + memory.id, + self.current_session, + memory.timestamp.timestamp(), + memory + .metadata + .get("role") + .and_then(|v| v.as_str()) + .unwrap_or("user"), + memory.content.replace('\'', "''"), // Escape single quotes + memory.metadata.to_string().replace('\'', "''") + ); + glue.execute(&sql).await?; + } + MemoryType::LongTerm => { + let sql = format!( + r#" + INSERT INTO long_term_memory + (id, concept, facts, confidence, created_at, access_count) + VALUES ('{}', '{}', '{}', 0.8, {}, 1)"#, + memory.id, + memory + .metadata + .get("concept") + .and_then(|v| v.as_str()) + .unwrap_or("general"), + memory.content.replace('\'', "''"), + memory.timestamp.timestamp() + ); + glue.execute(&sql).await?; + } + MemoryType::Episodic => { + let sql = format!( + r#" + INSERT INTO episodic_memory + (id, episode_id, timestamp, context, action_taken, outcome, reward) + VALUES ('{}', '{}', {}, '{}', '{}', '{}', {})"#, + memory.id, + memory + .metadata + .get("episode_id") + .and_then(|v| v.as_str()) + .unwrap_or(&self.current_session), + memory.timestamp.timestamp(), + memory.metadata.to_string().replace('\'', "''"), + memory.content.replace('\'', "''"), + memory + .metadata + .get("outcome") + .and_then(|v| v.as_str()) + .unwrap_or(""), + memory + .metadata + .get("reward") + .and_then(|v| v.as_f64()) + .unwrap_or(0.0) + ); + glue.execute(&sql).await?; + } + } + + // Create a version (commit to git) + let version = self.create_version(&format!("Store {} memory", memory_type)).await?; + Ok(version) + } + + pub async fn recall_memories( + &self, + query: &str, + memory_type: MemoryType, + limit: usize, + ) -> Result> { + let path = Path::new(&self.store_path).join("data"); + let storage = ProllyStorage::<32>::open(&path)?; + let mut glue = Glue::new(storage); + + let sql = match memory_type { + MemoryType::ShortTerm => { + if query.is_empty() { + format!( + r#" + SELECT id, content, timestamp, metadata + FROM short_term_memory + WHERE session_id = '{}' + ORDER BY timestamp DESC + LIMIT {}"#, + self.current_session, limit + ) + } else { + format!( + r#" + SELECT id, content, timestamp, metadata + FROM short_term_memory + WHERE content LIKE '%{}%' + ORDER BY timestamp DESC + LIMIT {}"#, + query.replace('\'', "''"), + limit + ) + } + } + MemoryType::LongTerm => format!( + r#" + SELECT id, facts as content, created_at as timestamp, concept + FROM long_term_memory + WHERE facts LIKE '%{}%' + ORDER BY access_count DESC + LIMIT {}"#, + query.replace('\'', "''"), + limit + ), + MemoryType::Episodic => format!( + r#" + SELECT id, action_taken as content, timestamp, context + FROM episodic_memory + WHERE action_taken LIKE '%{}%' OR context LIKE '%{}%' + ORDER BY timestamp DESC + LIMIT {}"#, + query.replace('\'', "''"), + query.replace('\'', "''"), + limit + ), + }; + + let results = glue.execute(&sql).await?; + let memories = self.parse_query_results(results, memory_type)?; + Ok(memories) + } + + fn parse_query_results( + &self, + results: Vec, + memory_type: MemoryType, + ) -> Result> { + use gluesql_core::data::Value; + let mut memories = Vec::new(); + + for payload in results { + if let Payload::Select { labels: _, rows } = payload { + for row in rows { + if row.len() >= 4 { + let id = match &row[0] { + Value::Str(s) => s.clone(), + _ => continue, + }; + + let content = match &row[1] { + Value::Str(s) => s.clone(), + _ => continue, + }; + + let timestamp = match &row[2] { + Value::I64(ts) => *ts, + Value::I32(ts) => *ts as i64, + Value::I16(ts) => *ts as i64, + _ => Utc::now().timestamp(), + }; + + let metadata = match memory_type { + MemoryType::ShortTerm => { + match &row[3] { + Value::Str(s) => serde_json::from_str(s).unwrap_or(serde_json::json!({})), + _ => serde_json::json!({}), + } + } + MemoryType::LongTerm => { + match &row[3] { + Value::Str(concept) => serde_json::json!({ "concept": concept }), + _ => serde_json::json!({ "concept": "general" }), + } + } + MemoryType::Episodic => { + match &row[3] { + Value::Str(s) => serde_json::from_str(s).unwrap_or(serde_json::json!({})), + _ => serde_json::json!({}), + } + } + }; + + let memory = Memory { + id, + content, + timestamp: chrono::DateTime::from_timestamp(timestamp, 0) + .unwrap_or_else(Utc::now), + metadata, + embedding: None, + }; + + memories.push(memory); + } + } + } + } + + Ok(memories) + } + + pub async fn create_branch(&self, name: &str) -> Result<()> { + // In a real implementation, we would create a git branch using VersionedKvStore + println!("🌿 Created memory branch: {}", name); + Ok(()) + } + + pub async fn rollback_to_version(&self, version: &str) -> Result<()> { + // In a real implementation, we would checkout a specific git commit + println!("⏪ Rolled back to version: {}", version); + Ok(()) + } + + pub async fn get_current_version(&self) -> String { + format!("v_{}", Utc::now().timestamp()) + } + + async fn create_version(&self, _message: &str) -> Result { + // In a real implementation, this would create a git commit + let version = format!("v_{}", Utc::now().timestamp()); + Ok(version) + } + + pub async fn clear_session_memories(&self) -> Result<()> { + let path = Path::new(&self.store_path).join("data"); + let storage = ProllyStorage::<32>::open(&path)?; + let mut glue = Glue::new(storage); + + glue.execute(&format!( + "DELETE FROM short_term_memory WHERE session_id = '{}'", + self.current_session + )) + .await?; + + Ok(()) + } + + pub fn new_session(&mut self) { + self.current_session = Uuid::new_v4().to_string(); + } + + pub async fn update_access_count(&self, memory_id: &str) -> Result<()> { + let path = Path::new(&self.store_path).join("data"); + let storage = ProllyStorage::<32>::open(&path)?; + let mut glue = Glue::new(storage); + + glue.execute(&format!( + "UPDATE long_term_memory SET access_count = access_count + 1 WHERE id = '{}'", + memory_id + )) + .await?; + + Ok(()) + } +} \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/memory/types.rs b/examples/rig_versioned_memory/src/memory/types.rs new file mode 100644 index 0000000..1c97ecf --- /dev/null +++ b/examples/rig_versioned_memory/src/memory/types.rs @@ -0,0 +1,116 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use std::fmt; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Memory { + pub id: String, + pub content: String, + pub timestamp: DateTime, + pub metadata: serde_json::Value, + pub embedding: Option>, +} + +impl Memory { + pub fn new(content: String, metadata: serde_json::Value) -> Self { + Self { + id: format!("mem_{}", uuid::Uuid::new_v4()), + content, + timestamp: Utc::now(), + metadata, + embedding: None, + } + } + + pub fn with_id(id: String, content: String, metadata: serde_json::Value) -> Self { + Self { + id, + content, + timestamp: Utc::now(), + metadata, + embedding: None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MemoryType { + ShortTerm, // Current conversation context + LongTerm, // Learned facts, preferences + Episodic, // Past experiences, outcomes +} + +impl MemoryType { + pub fn table_name(&self) -> &'static str { + match self { + MemoryType::ShortTerm => "short_term_memory", + MemoryType::LongTerm => "long_term_memory", + MemoryType::Episodic => "episodic_memory", + } + } +} + +impl fmt::Display for MemoryType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + MemoryType::ShortTerm => write!(f, "short_term"), + MemoryType::LongTerm => write!(f, "long_term"), + MemoryType::Episodic => write!(f, "episodic"), + } + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct DecisionAudit { + pub timestamp: DateTime, + pub input: String, + pub memories_accessed: Vec, + pub reasoning_chain: Vec, + pub decision: String, + pub confidence: f64, + pub version: String, +} + +#[derive(Debug, Clone)] +pub struct MemoryContext { + pub long_term_memories: Vec, + pub recent_memories: Vec, +} + +impl MemoryContext { + pub fn new() -> Self { + Self { + long_term_memories: Vec::new(), + recent_memories: Vec::new(), + } + } + + pub fn total_memories(&self) -> usize { + self.long_term_memories.len() + self.recent_memories.len() + } + + pub fn build_context_text(&self) -> String { + let mut context = String::new(); + + if !self.long_term_memories.is_empty() { + context.push_str("Relevant facts:\n"); + for memory in &self.long_term_memories { + context.push_str(&format!("- {}\n", memory.content)); + } + } + + if !self.recent_memories.is_empty() { + context.push_str("\nRecent conversation:\n"); + for memory in self.recent_memories.iter().rev().take(3) { + let role = memory + .metadata + .get("role") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + context.push_str(&format!("{}: {}\n", role, memory.content)); + } + } + + context + } +} \ No newline at end of file diff --git a/examples/rig_versioned_memory/src/utils/mod.rs b/examples/rig_versioned_memory/src/utils/mod.rs new file mode 100644 index 0000000..431ae8c --- /dev/null +++ b/examples/rig_versioned_memory/src/utils/mod.rs @@ -0,0 +1,25 @@ +use colored::Colorize; + +pub fn print_banner() { + println!("\n{}", "╔════════════════════════════════════════════════════╗".cyan()); + println!("{}", "║ 🤖 ProllyTree Versioned AI Agent Demo 🤖 ║".cyan().bold()); + println!("{}", "╚════════════════════════════════════════════════════╝".cyan()); + println!("\n{}", "Powered by ProllyTree + Rig Framework".dimmed()); + println!("{}\n", "=====================================".dimmed()); +} + +pub fn print_demo_separator() { + println!("\n{}", "─".repeat(60).dimmed()); +} + +pub fn print_error(msg: &str) { + eprintln!("{}: {}", "Error".red().bold(), msg); +} + +pub fn print_warning(msg: &str) { + println!("{}: {}", "Warning".yellow().bold(), msg); +} + +pub fn print_success(msg: &str) { + println!("{}: {}", "Success".green().bold(), msg); +} \ No newline at end of file From 4e37fda95befbc905b48504b3522b3e26e3a80a9 Mon Sep 17 00:00:00 2001 From: zhangfengcdt Date: Sun, 20 Jul 2025 10:09:36 -0700 Subject: [PATCH 2/5] change ci build behavior --- .github/workflows/ci.yml | 18 +++++++++++++++++- .github/workflows/python.yml | 6 +++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index aa411dd..4cd77f9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,7 +2,7 @@ name: CI on: pull_request: - types: [ready_for_review] + types: [ready_for_review, synchronize] push: branches: - main @@ -34,3 +34,19 @@ jobs: - name: docs run: cargo doc --document-private-items --no-deps + + - name: fmt rig_versioned_memory example + run: cargo fmt --all -- --check + working-directory: examples/rig_versioned_memory + + - name: build rig_versioned_memory example + run: cargo build --verbose + working-directory: examples/rig_versioned_memory + + - name: test rig_versioned_memory example + run: cargo test --verbose + working-directory: examples/rig_versioned_memory + + - name: clippy rig_versioned_memory example + run: cargo clippy + working-directory: examples/rig_versioned_memory diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index 0a0f6ab..1cef6b6 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -2,10 +2,14 @@ name: Build Python Package on: pull_request: - types: [ready_for_review] + types: [ready_for_review, synchronize] + paths: + - 'python/**' push: branches: - main + paths: + - 'python/**' jobs: build-wheels: From e31759e93029f5562fc6d33df1bb0c986d07a749 Mon Sep 17 00:00:00 2001 From: zhangfengcdt Date: Sun, 20 Jul 2025 10:34:39 -0700 Subject: [PATCH 3/5] fix fmt --- .../rig_versioned_memory/src/agent/demos.rs | 23 +++--- .../rig_versioned_memory/src/agent/mod.rs | 2 +- .../src/agent/versioned.rs | 20 ++--- examples/rig_versioned_memory/src/lib.rs | 2 +- examples/rig_versioned_memory/src/main.rs | 40 ++++++--- .../rig_versioned_memory/src/memory/mod.rs | 2 +- .../rig_versioned_memory/src/memory/schema.rs | 2 +- .../rig_versioned_memory/src/memory/store.rs | 81 +++++++++---------- .../rig_versioned_memory/src/memory/types.rs | 14 +++- .../rig_versioned_memory/src/utils/mod.rs | 19 ++++- 10 files changed, 116 insertions(+), 89 deletions(-) diff --git a/examples/rig_versioned_memory/src/agent/demos.rs b/examples/rig_versioned_memory/src/agent/demos.rs index 321e4aa..c41a9a1 100644 --- a/examples/rig_versioned_memory/src/agent/demos.rs +++ b/examples/rig_versioned_memory/src/agent/demos.rs @@ -11,7 +11,9 @@ impl VersionedAgent { // Agent learns user preference println!("\n{}: I prefer technical explanations", "User".blue()); - let (response1, _v1) = self.process_message("I prefer technical explanations").await?; + let (response1, _v1) = self + .process_message("I prefer technical explanations") + .await?; println!("{}: {}", "Agent".yellow(), response1); // Store this as long-term memory @@ -31,9 +33,9 @@ impl VersionedAgent { "\n{}: Actually, I prefer simple explanations", "User".blue() ); - let (response3, _v3) = - self.process_message("Actually, I prefer simple explanations") - .await?; + let (response3, _v3) = self + .process_message("Actually, I prefer simple explanations") + .await?; println!("{}: {}", "Agent".yellow(), response3); // Update the preference @@ -47,7 +49,7 @@ impl VersionedAgent { println!( "\n{} {}", "⏪ Rolling back to version".red(), - format!("{} (before preference change)", v2).red() + format!("{v2} (before preference change)").red() ); self.rollback_to_version(&v2).await?; @@ -176,10 +178,7 @@ impl VersionedAgent { .await?; // Another interaction - println!( - "\n{}: Set all my reminders 5 minutes early", - "User".blue() - ); + println!("\n{}: Set all my reminders 5 minutes early", "User".blue()); let (response2, _) = self .process_message("Set all my reminders 5 minutes early") .await?; @@ -245,7 +244,7 @@ impl VersionedAgent { let version = self.learn_fact(parts[0], parts[1]).await?; println!( "{}", - format!("✅ Learned fact (version: {})", version).green() + format!("✅ Learned fact (version: {version})").green() ); } else { println!("{}", "Usage: /learn ".red()); @@ -254,7 +253,7 @@ impl VersionedAgent { _ => match self.process_message(input).await { Ok((response, version)) => { println!("{}: {}", "Agent".yellow(), response); - println!("{}", format!("💾 Version: {}", version).dimmed()); + println!("{}", format!("💾 Version: {version}").dimmed()); } Err(e) => println!("{}: {}", "Error".red(), e), }, @@ -263,4 +262,4 @@ impl VersionedAgent { Ok(()) } -} \ No newline at end of file +} diff --git a/examples/rig_versioned_memory/src/agent/mod.rs b/examples/rig_versioned_memory/src/agent/mod.rs index d12512d..e1a8ba5 100644 --- a/examples/rig_versioned_memory/src/agent/mod.rs +++ b/examples/rig_versioned_memory/src/agent/mod.rs @@ -1,4 +1,4 @@ pub mod demos; pub mod versioned; -pub use versioned::VersionedAgent; \ No newline at end of file +pub use versioned::VersionedAgent; diff --git a/examples/rig_versioned_memory/src/agent/versioned.rs b/examples/rig_versioned_memory/src/agent/versioned.rs index e64f931..56f2e54 100644 --- a/examples/rig_versioned_memory/src/agent/versioned.rs +++ b/examples/rig_versioned_memory/src/agent/versioned.rs @@ -51,10 +51,8 @@ impl VersionedAgent { let prompt_text = self.build_prompt(input, &context); // 4. Generate response using Rig - let agent = self.openai_client - .agent(&self.model) - .build(); - + let agent = self.openai_client.agent(&self.model).build(); + let response = agent.prompt(&prompt_text).await?; // 5. Store response and commit version @@ -73,7 +71,10 @@ impl VersionedAgent { .store_memory(MemoryType::ShortTerm, &response_memory) .await?; - println!("💾 {}", format!("Memory committed to version: {}", version).cyan()); + println!( + "💾 {}", + format!("Memory committed to version: {version}").cyan() + ); Ok((response, version)) } @@ -98,13 +99,12 @@ impl VersionedAgent { fn build_prompt(&self, input: &str, context: &MemoryContext) -> String { let context_text = context.build_context_text(); - + if context_text.is_empty() { - format!("User: {}\nAssistant:", input) + format!("User: {input}\nAssistant:") } else { format!( - "Context from memory:\n{}\n\nUser: {}\nAssistant:", - context_text, input + "Context from memory:\n{context_text}\n\nUser: {input}\nAssistant:" ) } } @@ -174,4 +174,4 @@ impl VersionedAgent { self.new_session(); Ok(()) } -} \ No newline at end of file +} diff --git a/examples/rig_versioned_memory/src/lib.rs b/examples/rig_versioned_memory/src/lib.rs index 348864c..75f7c47 100644 --- a/examples/rig_versioned_memory/src/lib.rs +++ b/examples/rig_versioned_memory/src/lib.rs @@ -1,3 +1,3 @@ pub mod agent; pub mod memory; -pub mod utils; \ No newline at end of file +pub mod utils; diff --git a/examples/rig_versioned_memory/src/main.rs b/examples/rig_versioned_memory/src/main.rs index 493dd77..ba0f0bb 100644 --- a/examples/rig_versioned_memory/src/main.rs +++ b/examples/rig_versioned_memory/src/main.rs @@ -13,7 +13,7 @@ struct Cli { /// Path to store the agent memory (default: ./demo_agent_memory) #[arg(short, long, global = true, default_value = "./demo_agent_memory")] storage: String, - + #[command(subcommand)] command: Option, } @@ -38,13 +38,17 @@ enum Commands { async fn main() -> Result<()> { // Parse command line arguments first let cli = Cli::parse(); - + // Load environment variables from .env file if present dotenv().ok(); utils::print_banner(); - println!("📂 {}: {}", "Memory storage location".dimmed(), cli.storage.yellow()); + println!( + "📂 {}: {}", + "Memory storage location".dimmed(), + cli.storage.yellow() + ); // Get API key let api_key = match env::var("OPENAI_API_KEY") { @@ -52,8 +56,14 @@ async fn main() -> Result<()> { Err(_) => { utils::print_error("Please set OPENAI_API_KEY environment variable"); println!("\n{}", "You can either:".dimmed()); - println!(" 1. Export it: {}", "export OPENAI_API_KEY=\"your-key-here\"".yellow()); - println!(" 2. Create a .env file with: {}", "OPENAI_API_KEY=your-key-here".yellow()); + println!( + " 1. Export it: {}", + "export OPENAI_API_KEY=\"your-key-here\"".yellow() + ); + println!( + " 2. Create a .env file with: {}", + "OPENAI_API_KEY=your-key-here".yellow() + ); std::process::exit(1); } }; @@ -61,11 +71,14 @@ async fn main() -> Result<()> { // Initialize agent let mut agent = match VersionedAgent::new(api_key, &cli.storage).await { Ok(agent) => { - utils::print_success(&format!("Initialized versioned memory store at: {}", cli.storage)); + utils::print_success(&format!( + "Initialized versioned memory store at: {}", + cli.storage + )); agent } Err(e) => { - utils::print_error(&format!("Failed to initialize agent: {}", e)); + utils::print_error(&format!("Failed to initialize agent: {e}")); std::process::exit(1); } }; @@ -86,22 +99,25 @@ async fn main() -> Result<()> { Some(Commands::All) => { agent.demo_memory_learning().await?; utils::print_demo_separator(); - + agent.demo_branching().await?; utils::print_demo_separator(); - + agent.demo_audit_trail().await?; utils::print_demo_separator(); - + agent.demo_episodic_learning().await?; } Some(Commands::Chat) | None => { // Default to interactive mode - println!("\n{}", "No demo specified, starting interactive mode...".dimmed()); + println!( + "\n{}", + "No demo specified, starting interactive mode...".dimmed() + ); agent.run_interactive_mode().await?; } } println!("\n{}", "👋 Demo completed!".green().bold()); Ok(()) -} \ No newline at end of file +} diff --git a/examples/rig_versioned_memory/src/memory/mod.rs b/examples/rig_versioned_memory/src/memory/mod.rs index 9547856..30ab0b8 100644 --- a/examples/rig_versioned_memory/src/memory/mod.rs +++ b/examples/rig_versioned_memory/src/memory/mod.rs @@ -3,4 +3,4 @@ pub mod store; pub mod types; pub use store::VersionedMemoryStore; -pub use types::{DecisionAudit, Memory, MemoryContext, MemoryType}; \ No newline at end of file +pub use types::{DecisionAudit, Memory, MemoryContext, MemoryType}; diff --git a/examples/rig_versioned_memory/src/memory/schema.rs b/examples/rig_versioned_memory/src/memory/schema.rs index 4209321..34cf08c 100644 --- a/examples/rig_versioned_memory/src/memory/schema.rs +++ b/examples/rig_versioned_memory/src/memory/schema.rs @@ -68,4 +68,4 @@ pub async fn setup_schema(glue: &mut Glue>) -> Result<()> { // Future enhancement: implement index support in ProllyStorage Ok(()) -} \ No newline at end of file +} diff --git a/examples/rig_versioned_memory/src/memory/store.rs b/examples/rig_versioned_memory/src/memory/store.rs index 613049e..558861d 100644 --- a/examples/rig_versioned_memory/src/memory/store.rs +++ b/examples/rig_versioned_memory/src/memory/store.rs @@ -16,7 +16,7 @@ pub struct VersionedMemoryStore { impl VersionedMemoryStore { pub async fn new(store_path: &str) -> Result { let path = Path::new(store_path); - + // Create directory if it doesn't exist if !path.exists() { std::fs::create_dir_all(path)?; @@ -26,29 +26,29 @@ impl VersionedMemoryStore { let git_dir = path.join(".git"); if !git_dir.exists() { use std::process::Command; - + // Initialize git repository let output = Command::new("git") - .args(&["init"]) + .args(["init"]) .current_dir(path) .output()?; - + if !output.status.success() { return Err(anyhow::anyhow!( - "Failed to initialize git repository: {}", + "Failed to initialize git repository: {}", String::from_utf8_lossy(&output.stderr) )); } - + // Set up initial git config if needed Command::new("git") - .args(&["config", "user.name", "AI Agent"]) + .args(["config", "user.name", "AI Agent"]) .current_dir(path) .output() .ok(); // Ignore errors, might already be set globally - + Command::new("git") - .args(&["config", "user.email", "agent@example.com"]) + .args(["config", "user.email", "agent@example.com"]) .current_dir(path) .output() .ok(); // Ignore errors, might already be set globally @@ -59,7 +59,7 @@ impl VersionedMemoryStore { if !data_dir.exists() { std::fs::create_dir_all(&data_dir)?; } - + let storage = if data_dir.join(".git-prolly").exists() { ProllyStorage::<32>::open(&data_dir)? } else { @@ -77,11 +77,7 @@ impl VersionedMemoryStore { }) } - pub async fn store_memory( - &self, - memory_type: MemoryType, - memory: &Memory, - ) -> Result { + pub async fn store_memory(&self, memory_type: MemoryType, memory: &Memory) -> Result { let path = Path::new(&self.store_path).join("data"); let storage = ProllyStorage::<32>::open(&path)?; let mut glue = Glue::new(storage); @@ -154,7 +150,9 @@ impl VersionedMemoryStore { } // Create a version (commit to git) - let version = self.create_version(&format!("Store {} memory", memory_type)).await?; + let version = self + .create_version(&format!("Store {memory_type} memory")) + .await?; Ok(version) } @@ -237,12 +235,12 @@ impl VersionedMemoryStore { Value::Str(s) => s.clone(), _ => continue, }; - + let content = match &row[1] { Value::Str(s) => s.clone(), _ => continue, }; - + let timestamp = match &row[2] { Value::I64(ts) => *ts, Value::I32(ts) => *ts as i64, @@ -251,24 +249,22 @@ impl VersionedMemoryStore { }; let metadata = match memory_type { - MemoryType::ShortTerm => { - match &row[3] { - Value::Str(s) => serde_json::from_str(s).unwrap_or(serde_json::json!({})), - _ => serde_json::json!({}), + MemoryType::ShortTerm => match &row[3] { + Value::Str(s) => { + serde_json::from_str(s).unwrap_or(serde_json::json!({})) } - } - MemoryType::LongTerm => { - match &row[3] { - Value::Str(concept) => serde_json::json!({ "concept": concept }), - _ => serde_json::json!({ "concept": "general" }), + _ => serde_json::json!({}), + }, + MemoryType::LongTerm => match &row[3] { + Value::Str(concept) => serde_json::json!({ "concept": concept }), + _ => serde_json::json!({ "concept": "general" }), + }, + MemoryType::Episodic => match &row[3] { + Value::Str(s) => { + serde_json::from_str(s).unwrap_or(serde_json::json!({})) } - } - MemoryType::Episodic => { - match &row[3] { - Value::Str(s) => serde_json::from_str(s).unwrap_or(serde_json::json!({})), - _ => serde_json::json!({}), - } - } + _ => serde_json::json!({}), + }, }; let memory = Memory { @@ -291,13 +287,13 @@ impl VersionedMemoryStore { pub async fn create_branch(&self, name: &str) -> Result<()> { // In a real implementation, we would create a git branch using VersionedKvStore - println!("🌿 Created memory branch: {}", name); + println!("🌿 Created memory branch: {name}"); Ok(()) } pub async fn rollback_to_version(&self, version: &str) -> Result<()> { // In a real implementation, we would checkout a specific git commit - println!("⏪ Rolled back to version: {}", version); + println!("⏪ Rolled back to version: {version}"); Ok(()) } @@ -315,13 +311,13 @@ impl VersionedMemoryStore { let path = Path::new(&self.store_path).join("data"); let storage = ProllyStorage::<32>::open(&path)?; let mut glue = Glue::new(storage); - + glue.execute(&format!( "DELETE FROM short_term_memory WHERE session_id = '{}'", self.current_session )) .await?; - + Ok(()) } @@ -333,13 +329,12 @@ impl VersionedMemoryStore { let path = Path::new(&self.store_path).join("data"); let storage = ProllyStorage::<32>::open(&path)?; let mut glue = Glue::new(storage); - + glue.execute(&format!( - "UPDATE long_term_memory SET access_count = access_count + 1 WHERE id = '{}'", - memory_id + "UPDATE long_term_memory SET access_count = access_count + 1 WHERE id = '{memory_id}'" )) .await?; - + Ok(()) } -} \ No newline at end of file +} diff --git a/examples/rig_versioned_memory/src/memory/types.rs b/examples/rig_versioned_memory/src/memory/types.rs index 1c97ecf..e789445 100644 --- a/examples/rig_versioned_memory/src/memory/types.rs +++ b/examples/rig_versioned_memory/src/memory/types.rs @@ -35,9 +35,9 @@ impl Memory { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum MemoryType { - ShortTerm, // Current conversation context - LongTerm, // Learned facts, preferences - Episodic, // Past experiences, outcomes + ShortTerm, // Current conversation context + LongTerm, // Learned facts, preferences + Episodic, // Past experiences, outcomes } impl MemoryType { @@ -77,6 +77,12 @@ pub struct MemoryContext { pub recent_memories: Vec, } +impl Default for MemoryContext { + fn default() -> Self { + Self::new() + } +} + impl MemoryContext { pub fn new() -> Self { Self { @@ -113,4 +119,4 @@ impl MemoryContext { context } -} \ No newline at end of file +} diff --git a/examples/rig_versioned_memory/src/utils/mod.rs b/examples/rig_versioned_memory/src/utils/mod.rs index 431ae8c..286c5ce 100644 --- a/examples/rig_versioned_memory/src/utils/mod.rs +++ b/examples/rig_versioned_memory/src/utils/mod.rs @@ -1,9 +1,20 @@ use colored::Colorize; pub fn print_banner() { - println!("\n{}", "╔════════════════════════════════════════════════════╗".cyan()); - println!("{}", "║ 🤖 ProllyTree Versioned AI Agent Demo 🤖 ║".cyan().bold()); - println!("{}", "╚════════════════════════════════════════════════════╝".cyan()); + println!( + "\n{}", + "╔════════════════════════════════════════════════════╗".cyan() + ); + println!( + "{}", + "║ 🤖 ProllyTree Versioned AI Agent Demo 🤖 ║" + .cyan() + .bold() + ); + println!( + "{}", + "╚════════════════════════════════════════════════════╝".cyan() + ); println!("\n{}", "Powered by ProllyTree + Rig Framework".dimmed()); println!("{}\n", "=====================================".dimmed()); } @@ -22,4 +33,4 @@ pub fn print_warning(msg: &str) { pub fn print_success(msg: &str) { println!("{}: {}", "Success".green().bold(), msg); -} \ No newline at end of file +} From e7b79f9f2aebbbaf81b70406fb2e0ec9d451b657 Mon Sep 17 00:00:00 2001 From: zhangfengcdt Date: Sun, 20 Jul 2025 14:07:43 -0700 Subject: [PATCH 4/5] fix fmt issue --- examples/rig_versioned_memory/src/agent/versioned.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/examples/rig_versioned_memory/src/agent/versioned.rs b/examples/rig_versioned_memory/src/agent/versioned.rs index 56f2e54..d92320f 100644 --- a/examples/rig_versioned_memory/src/agent/versioned.rs +++ b/examples/rig_versioned_memory/src/agent/versioned.rs @@ -103,9 +103,7 @@ impl VersionedAgent { if context_text.is_empty() { format!("User: {input}\nAssistant:") } else { - format!( - "Context from memory:\n{context_text}\n\nUser: {input}\nAssistant:" - ) + format!("Context from memory:\n{context_text}\n\nUser: {input}\nAssistant:") } } From 1f7370b7e8608683dfc0bd3551c566b40954f137 Mon Sep 17 00:00:00 2001 From: zhangfengcdt Date: Sun, 20 Jul 2025 14:17:00 -0700 Subject: [PATCH 5/5] add examples to the top cargo file --- Cargo.toml | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/Cargo.toml b/Cargo.toml index fcdfa10..2883abe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -77,3 +77,25 @@ required-features = ["sql"] name = "git_prolly_bench" harness = false required-features = ["git", "sql"] + +[[example]] +name = "proof_visualization" +path = "examples/proof_visualization.rs" + +[[example]] +name = "git_sql" +path = "examples/git_sql.rs" +required-features = ["sql"] + +[[example]] +name = "git_diff" +path = "examples/git_diff.rs" +required-features = ["git"] + +[[example]] +name = "git_merge" +path = "examples/git_merge.rs" +required-features = ["git"] + +[workspace] +members = ["examples/rig_versioned_memory"]