From 1f25557b3a9a2058c79b4136aaa2df4f042aeb11 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Thu, 24 Jul 2025 15:04:03 -0700 Subject: [PATCH 1/6] Refactor VersionedKvStore --- examples/financial_advisor/src/memory/mod.rs | 8 +- src/bin/git-prolly.rs | 28 +- src/git/mod.rs | 2 +- src/git/operations.rs | 18 +- src/git/types.rs | 27 + src/git/versioned_store.rs | 534 ++++++++++++++----- src/sql/glue_storage.rs | 12 +- 7 files changed, 459 insertions(+), 170 deletions(-) diff --git a/examples/financial_advisor/src/memory/mod.rs b/examples/financial_advisor/src/memory/mod.rs index 67b3435..ef10f26 100644 --- a/examples/financial_advisor/src/memory/mod.rs +++ b/examples/financial_advisor/src/memory/mod.rs @@ -5,7 +5,7 @@ use anyhow::Result; use chrono::{DateTime, Utc}; use gluesql_core::prelude::{Glue, Payload}; use gluesql_core::store::Transaction; -use prollytree::git::{GitKvError, VersionedKvStore}; +use prollytree::git::{GitKvError, GitVersionedKvStore}; use prollytree::sql::ProllyStorage; use std::any::Any; use std::path::Path; @@ -68,7 +68,7 @@ pub trait Storable: serde::Serialize + serde::de::DeserializeOwned + Clone { /// Core memory store with versioning capabilities pub struct MemoryStore { store_path: String, - versioned_store: VersionedKvStore<32>, + versioned_store: GitVersionedKvStore<32>, audit_enabled: bool, } @@ -91,10 +91,10 @@ impl MemoryStore { // Initialize VersionedKvStore in dataset subdirectory // Check if prolly tree config exists to determine if we should init or open let versioned_store = if current_dir.join(PROLLY_CONFIG_FILE).exists() { - VersionedKvStore::<32>::open(¤t_dir) + GitVersionedKvStore::<32>::open(¤t_dir) .map_err(|e| anyhow::anyhow!("Failed to open versioned store: {:?}", e))? } else { - VersionedKvStore::<32>::init(¤t_dir) + GitVersionedKvStore::<32>::init(¤t_dir) .map_err(|e| anyhow::anyhow!("Failed to init versioned store: {:?}", e))? }; diff --git a/src/bin/git-prolly.rs b/src/bin/git-prolly.rs index e8449fe..7e9aedb 100644 --- a/src/bin/git-prolly.rs +++ b/src/bin/git-prolly.rs @@ -15,7 +15,7 @@ limitations under the License. use clap::{Parser, Subcommand}; #[cfg(feature = "sql")] use gluesql_core::{executor::Payload, prelude::Glue}; -use prollytree::git::{DiffOperation, GitOperations, MergeResult, VersionedKvStore}; +use prollytree::git::{DiffOperation, GitOperations, GitVersionedKvStore, MergeResult}; #[cfg(feature = "sql")] use prollytree::sql::ProllyStorage; use prollytree::tree::Tree; @@ -218,7 +218,7 @@ fn handle_init(path: Option) -> Result<(), Box> println!("Initializing ProllyTree KV store in {target_path:?}..."); - let _store = VersionedKvStore::<32>::init(&target_path)?; + let _store = GitVersionedKvStore::<32>::init(&target_path)?; println!("โœ“ Initialized empty ProllyTree KV store"); println!("โœ“ Git repository initialized"); @@ -229,7 +229,7 @@ fn handle_init(path: Option) -> Result<(), Box> fn handle_set(key: String, value: String) -> Result<(), Box> { let current_dir = env::current_dir()?; - let mut store = VersionedKvStore::<32>::open(¤t_dir)?; + let mut store = GitVersionedKvStore::<32>::open(¤t_dir)?; store.insert(key.as_bytes().to_vec(), value.as_bytes().to_vec())?; @@ -241,7 +241,7 @@ fn handle_set(key: String, value: String) -> Result<(), Box Result<(), Box> { let current_dir = env::current_dir()?; - let store = VersionedKvStore::<32>::open(¤t_dir)?; + let store = GitVersionedKvStore::<32>::open(¤t_dir)?; match store.get(key.as_bytes()) { Some(value) => { @@ -258,7 +258,7 @@ fn handle_get(key: String) -> Result<(), Box> { fn handle_delete(key: String) -> Result<(), Box> { let current_dir = env::current_dir()?; - let mut store = VersionedKvStore::<32>::open(¤t_dir)?; + let mut store = GitVersionedKvStore::<32>::open(¤t_dir)?; if store.delete(key.as_bytes())? { println!("โœ“ Staged deletion: {key}"); @@ -273,7 +273,7 @@ fn handle_delete(key: String) -> Result<(), Box> { fn handle_list(show_values: bool, show_graph: bool) -> Result<(), Box> { let current_dir = env::current_dir()?; - let mut store = VersionedKvStore::<32>::open(¤t_dir)?; + let mut store = GitVersionedKvStore::<32>::open(¤t_dir)?; if show_graph { // Show the prolly tree structure @@ -311,7 +311,7 @@ fn handle_list(show_values: bool, show_graph: bool) -> Result<(), Box Result<(), Box> { let current_dir = env::current_dir()?; - let store = VersionedKvStore::<32>::open(¤t_dir)?; + let store = GitVersionedKvStore::<32>::open(¤t_dir)?; let status = store.status(); let current_branch = store.current_branch(); @@ -340,7 +340,7 @@ fn handle_status() -> Result<(), Box> { fn handle_commit(message: String) -> Result<(), Box> { let current_dir = env::current_dir()?; - let mut store = VersionedKvStore::<32>::open(¤t_dir)?; + let mut store = GitVersionedKvStore::<32>::open(¤t_dir)?; let status = store.status(); if status.is_empty() { @@ -376,7 +376,7 @@ fn handle_diff( _keys: Option, ) -> Result<(), Box> { let current_dir = env::current_dir()?; - let store = VersionedKvStore::<32>::open(¤t_dir)?; + let store = GitVersionedKvStore::<32>::open(¤t_dir)?; let ops = GitOperations::new(store); let diffs = ops.diff(&from, &to)?; @@ -483,7 +483,7 @@ fn handle_diff( fn handle_show(commit: Option, keys_only: bool) -> Result<(), Box> { let current_dir = env::current_dir()?; - let store = VersionedKvStore::<32>::open(¤t_dir)?; + let store = GitVersionedKvStore::<32>::open(¤t_dir)?; let ops = GitOperations::new(store); let commit_ref = commit.unwrap_or_else(|| "HEAD".to_string()); @@ -537,7 +537,7 @@ fn handle_merge( _strategy: Option, ) -> Result<(), Box> { let current_dir = env::current_dir()?; - let store = VersionedKvStore::<32>::open(¤t_dir)?; + let store = GitVersionedKvStore::<32>::open(¤t_dir)?; let mut ops = GitOperations::new(store); println!("Merging branch '{branch}'..."); @@ -572,7 +572,7 @@ fn handle_merge( fn handle_stats(commit: Option) -> Result<(), Box> { let current_dir = env::current_dir()?; - let store = VersionedKvStore::<32>::open(¤t_dir)?; + let store = GitVersionedKvStore::<32>::open(¤t_dir)?; let target = commit.unwrap_or_else(|| "HEAD".to_string()); @@ -626,14 +626,14 @@ fn handle_clear(confirm: bool, keep_history: bool) -> Result<(), Box::open(¤t_dir)?; + let mut store = GitVersionedKvStore::<32>::open(¤t_dir)?; // Clear staging area first println!(" โ†ณ Clearing staging changes..."); let status = store.status(); if !status.is_empty() { // Reset staging area by recreating the store - store = VersionedKvStore::<32>::open(¤t_dir)?; + store = GitVersionedKvStore::<32>::open(¤t_dir)?; println!(" โœ“ Cleared {} staged changes", status.len()); } else { println!(" โœ“ No staged changes to clear"); diff --git a/src/git/mod.rs b/src/git/mod.rs index df5e101..f56f791 100644 --- a/src/git/mod.rs +++ b/src/git/mod.rs @@ -24,4 +24,4 @@ pub use types::{ CommitDetails, CommitInfo, DiffOperation, GitKvError, KvConflict, KvDiff, KvStorageMetadata, MergeResult, }; -pub use versioned_store::VersionedKvStore; +pub use versioned_store::{GitVersionedKvStore, VersionedKvStore}; diff --git a/src/git/operations.rs b/src/git/operations.rs index e7582bb..00a903a 100644 --- a/src/git/operations.rs +++ b/src/git/operations.rs @@ -13,17 +13,17 @@ limitations under the License. */ use crate::git::types::*; -use crate::git::versioned_store::VersionedKvStore; +use crate::git::versioned_store::GitVersionedKvStore; use gix::prelude::*; use std::collections::HashMap; /// Git operations for versioned KV store pub struct GitOperations { - store: VersionedKvStore, + store: GitVersionedKvStore, } impl GitOperations { - pub fn new(store: VersionedKvStore) -> Self { + pub fn new(store: GitVersionedKvStore) -> Self { GitOperations { store } } @@ -315,7 +315,7 @@ impl GitOperations { .map_err(|e| GitKvError::GitObjectError(format!("Failed to get current dir: {e}")))?; // Create a temporary clone of the versioned store - let mut temp_store = VersionedKvStore::::open(¤t_dir)?; + let mut temp_store = GitVersionedKvStore::::open(¤t_dir)?; // Save current state let original_branch = temp_store.current_branch().to_string(); @@ -335,7 +335,7 @@ impl GitOperations { /// Temporarily checkout a commit and extract its KV state fn checkout_commit_temporarily( &self, - store: &mut VersionedKvStore, + store: &mut GitVersionedKvStore, commit_id: &gix::ObjectId, ) -> Result, Vec>, GitKvError> { // Update the store to point to the specific commit @@ -386,7 +386,7 @@ impl GitOperations { // Try to open the store at the temp location let dataset_dir = temp_dir.join("dataset"); let result = if dataset_dir.exists() { - match VersionedKvStore::::open(&dataset_dir) { + match GitVersionedKvStore::::open(&dataset_dir) { Ok(temp_store) => self.get_current_kv_state_from_store(&temp_store), Err(_) => { // If we can't open the store, return empty state @@ -421,7 +421,7 @@ impl GitOperations { /// Get current KV state from a specific store fn get_current_kv_state_from_store( &self, - store: &VersionedKvStore, + store: &GitVersionedKvStore, ) -> Result, Vec>, GitKvError> { let mut state = HashMap::new(); @@ -463,7 +463,7 @@ mod tests { // Create subdirectory for dataset let dataset_dir = temp_dir.path().join("dataset"); std::fs::create_dir_all(&dataset_dir).unwrap(); - let store = VersionedKvStore::<32>::init(&dataset_dir).unwrap(); + let store = GitVersionedKvStore::<32>::init(&dataset_dir).unwrap(); let _ops = GitOperations::new(store); } @@ -475,7 +475,7 @@ mod tests { // Create subdirectory for dataset let dataset_dir = temp_dir.path().join("dataset"); std::fs::create_dir_all(&dataset_dir).unwrap(); - let store = VersionedKvStore::<32>::init(&dataset_dir).unwrap(); + let store = GitVersionedKvStore::<32>::init(&dataset_dir).unwrap(); let ops = GitOperations::new(store); // Test HEAD parsing diff --git a/src/git/types.rs b/src/git/types.rs index 3e5cbac..dc88fec 100644 --- a/src/git/types.rs +++ b/src/git/types.rs @@ -111,6 +111,33 @@ pub struct KvStorageMetadata { pub last_commit: Option, } +/// Storage backend types supported by VersionedKvStore +#[derive(Debug, Clone, PartialEq, Default)] +pub enum StorageBackend { + /// In-memory storage (volatile, fastest) + InMemory, + /// File-based storage (persistent, simple) + File, + /// RocksDB storage (persistent, high-performance) + #[cfg(feature = "rocksdb_storage")] + RocksDB, + /// Git object storage (development only, default) + #[default] + Git, +} + +impl std::fmt::Display for StorageBackend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + StorageBackend::InMemory => write!(f, "InMemory"), + StorageBackend::File => write!(f, "File"), + #[cfg(feature = "rocksdb_storage")] + StorageBackend::RocksDB => write!(f, "RocksDB"), + StorageBackend::Git => write!(f, "Git"), + } + } +} + impl fmt::Display for DiffOperation { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { diff --git a/src/git/versioned_store.rs b/src/git/versioned_store.rs index 0e39984..ffee824 100644 --- a/src/git/versioned_store.rs +++ b/src/git/versioned_store.rs @@ -15,24 +15,42 @@ limitations under the License. use crate::config::TreeConfig; use crate::git::storage::GitNodeStorage; use crate::git::types::*; +use crate::storage::{FileNodeStorage, InMemoryNodeStorage, NodeStorage}; use crate::tree::{ProllyTree, Tree}; use gix::prelude::*; use std::collections::HashMap; use std::path::Path; -/// A versioned key-value store backed by Git and ProllyTree +#[cfg(feature = "rocksdb_storage")] +use crate::storage::RocksDBNodeStorage; + +/// A versioned key-value store backed by Git and ProllyTree with configurable storage /// /// This combines the efficient tree operations of ProllyTree with Git's /// version control capabilities, providing a full-featured versioned /// key-value store with branching, merging, and history. -pub struct VersionedKvStore { - tree: ProllyTree>, +pub struct VersionedKvStore> { + tree: ProllyTree, git_repo: gix::Repository, staging_area: HashMap, Option>>, // None = deleted current_branch: String, + storage_backend: StorageBackend, } -impl VersionedKvStore { +/// Type alias for backward compatibility (Git storage) +pub type GitVersionedKvStore = VersionedKvStore>; + +/// Type alias for InMemory storage +pub type InMemoryVersionedKvStore = VersionedKvStore>; + +/// Type alias for File storage +pub type FileVersionedKvStore = VersionedKvStore>; + +/// Type alias for RocksDB storage +#[cfg(feature = "rocksdb_storage")] +pub type RocksDBVersionedKvStore = VersionedKvStore>; + +impl> VersionedKvStore { /// Find the git repository root by walking up the directory tree fn find_git_root>(start_path: P) -> Option { let mut current = start_path.as_ref().to_path_buf(); @@ -66,104 +84,6 @@ impl VersionedKvStore { } } - /// Initialize a new versioned KV store at the given path - pub fn init>(path: P) -> Result { - let path = path.as_ref(); - - // Reject if trying to initialize in git root directory - if Self::is_in_git_root(path)? { - return Err(GitKvError::GitObjectError( - "Cannot initialize git-prolly in git root directory. Please run from a subdirectory to create a dataset.".to_string() - )); - } - - // Find the git repository - let git_root = Self::find_git_root(path).ok_or_else(|| { - GitKvError::GitObjectError( - "Not inside a git repository. Please run from within a git repository.".to_string(), - ) - })?; - - // Open the existing git repository instead of initializing a new one - let git_repo = gix::open(&git_root).map_err(|e| GitKvError::GitOpenError(Box::new(e)))?; - - // Create GitNodeStorage with the current directory as dataset directory - let storage = GitNodeStorage::new(git_repo.clone(), path.to_path_buf())?; - - // Create ProllyTree with default config - let config: TreeConfig = TreeConfig::default(); - let tree = ProllyTree::new(storage, config); - - let mut store = VersionedKvStore { - tree, - git_repo, - staging_area: HashMap::new(), - current_branch: "main".to_string(), - }; - - // Save initial configuration - let _ = store.tree.save_config(); - - // Create initial commit (which will include prolly metadata files) - store.commit("Initial commit")?; - - Ok(store) - } - - /// Open an existing versioned KV store - pub fn open>(path: P) -> Result { - let path = path.as_ref(); - - // Reject if trying to open in git root directory - if Self::is_in_git_root(path)? { - return Err(GitKvError::GitObjectError( - "Cannot run git-prolly in git root directory. Please run from a subdirectory containing a dataset.".to_string() - )); - } - - // Find the git repository - let git_root = Self::find_git_root(path).ok_or_else(|| { - GitKvError::GitObjectError( - "Not inside a git repository. Please run from within a git repository.".to_string(), - ) - })?; - - // Open existing Git repository - let git_repo = gix::open(&git_root).map_err(|e| GitKvError::GitOpenError(Box::new(e)))?; - - // Create GitNodeStorage with the current directory as dataset directory - let storage = GitNodeStorage::new(git_repo.clone(), path.to_path_buf())?; - - // Load tree configuration from storage - let config: TreeConfig = ProllyTree::load_config(&storage).unwrap_or_default(); - - // Try to load existing tree from storage, or create new one - let tree = ProllyTree::load_from_storage(storage.clone(), config.clone()) - .unwrap_or_else(|| ProllyTree::new(storage, config)); - - // Get current branch - let current_branch = git_repo - .head_ref() - .map_err(|e| GitKvError::GitObjectError(format!("Failed to get head ref: {e}")))? - .map(|r| r.name().shorten().to_string()) - .unwrap_or_else(|| "main".to_string()); - - let mut store = VersionedKvStore { - tree, - git_repo, - staging_area: HashMap::new(), - current_branch, - }; - - // Load staging area from file if it exists - store.load_staging_area()?; - - // Reload the tree from the current HEAD - store.reload_tree_from_head()?; - - Ok(store) - } - /// Insert a key-value pair (stages the change) pub fn insert(&mut self, key: Vec, value: Vec) -> Result<(), GitKvError> { self.staging_area.insert(key, Some(value)); @@ -332,8 +252,7 @@ impl VersionedKvStore { std::fs::write(&head_file, head_content) .map_err(|e| GitKvError::GitObjectError(format!("Failed to update HEAD: {e}")))?; - // Reload tree state from the current HEAD (same as current branch) - self.reload_tree_from_head()?; + // Note: Tree reload is handled in Git-specific implementation Ok(()) } @@ -369,8 +288,7 @@ impl VersionedKvStore { } } - // Reload tree state from the new HEAD - self.reload_tree_from_head()?; + // Note: Tree reload is handled in Git-specific implementation Ok(()) } @@ -380,16 +298,6 @@ impl VersionedKvStore { &self.current_branch } - /// Get a reference to the underlying ProllyTree - pub fn tree(&self) -> &ProllyTree> { - &self.tree - } - - /// Get a mutable reference to the underlying ProllyTree - pub fn tree_mut(&mut self) -> &mut ProllyTree> { - &mut self.tree - } - /// List all branches in the repository pub fn list_branches(&self) -> Result, GitKvError> { let mut branches = Vec::new(); @@ -635,22 +543,6 @@ impl VersionedKvStore { Ok(()) } - /// Reload the ProllyTree from the current HEAD - fn reload_tree_from_head(&mut self) -> Result<(), GitKvError> { - // Since we're no longer storing prolly_tree_root in the Git tree, - // we need to reload the tree state from the GitNodeStorage - - // Load tree configuration from storage - let config: TreeConfig = ProllyTree::load_config(&self.tree.storage).unwrap_or_default(); - - // Try to load existing tree from storage, or create new one - let storage = self.tree.storage.clone(); - self.tree = ProllyTree::load_from_storage(storage.clone(), config.clone()) - .unwrap_or_else(|| ProllyTree::new(storage, config)); - - Ok(()) - } - /// Save the staging area to a file fn save_staging_area(&self) -> Result<(), GitKvError> { let staging_file = self.get_staging_file_path()?; @@ -709,6 +601,376 @@ impl VersionedKvStore { } } +// Storage-specific implementations +impl VersionedKvStore> { + /// Initialize a new versioned KV store with Git storage (default) + pub fn init>(path: P) -> Result { + let path = path.as_ref(); + + // Reject if trying to initialize in git root directory + if Self::is_in_git_root(path)? { + return Err(GitKvError::GitObjectError( + "Cannot initialize git-prolly in git root directory. Please run from a subdirectory to create a dataset.".to_string() + )); + } + + // Find the git repository + let git_root = Self::find_git_root(path).ok_or_else(|| { + GitKvError::GitObjectError( + "Not inside a git repository. Please run from within a git repository.".to_string(), + ) + })?; + + // Open the existing git repository + let git_repo = gix::open(&git_root).map_err(|e| GitKvError::GitOpenError(Box::new(e)))?; + + // Create GitNodeStorage + let storage = GitNodeStorage::new(git_repo.clone(), path.to_path_buf())?; + + // Create ProllyTree with default config + let config: TreeConfig = TreeConfig::default(); + let tree = ProllyTree::new(storage, config); + + let mut store = VersionedKvStore { + tree, + git_repo, + staging_area: HashMap::new(), + current_branch: "main".to_string(), + storage_backend: StorageBackend::Git, + }; + + // Save initial configuration + let _ = store.tree.save_config(); + + // Create initial commit (which will include prolly metadata files) + store.commit("Initial commit")?; + + Ok(store) + } + + /// Open an existing versioned KV store with Git storage (default) + pub fn open>(path: P) -> Result { + let path = path.as_ref(); + + // Reject if trying to open in git root directory + if Self::is_in_git_root(path)? { + return Err(GitKvError::GitObjectError( + "Cannot run git-prolly in git root directory. Please run from a subdirectory containing a dataset.".to_string() + )); + } + + // Find the git repository + let git_root = Self::find_git_root(path).ok_or_else(|| { + GitKvError::GitObjectError( + "Not inside a git repository. Please run from within a git repository.".to_string(), + ) + })?; + + // Open existing Git repository + let git_repo = gix::open(&git_root).map_err(|e| GitKvError::GitOpenError(Box::new(e)))?; + + // Create GitNodeStorage + let storage = GitNodeStorage::new(git_repo.clone(), path.to_path_buf())?; + + // Load tree configuration from storage + let config: TreeConfig = ProllyTree::load_config(&storage).unwrap_or_default(); + + // Try to load existing tree from storage, or create new one + let tree = ProllyTree::load_from_storage(storage.clone(), config.clone()) + .unwrap_or_else(|| ProllyTree::new(storage, config)); + + // Get current branch + let current_branch = git_repo + .head_ref() + .map_err(|e| GitKvError::GitObjectError(format!("Failed to get head ref: {e}")))? + .map(|r| r.name().shorten().to_string()) + .unwrap_or_else(|| "main".to_string()); + + let mut store = VersionedKvStore { + tree, + git_repo, + staging_area: HashMap::new(), + current_branch, + storage_backend: StorageBackend::Git, + }; + + // Load staging area from file if it exists + store.load_staging_area()?; + + // Reload the tree from the current HEAD + store.reload_tree_from_head()?; + + Ok(store) + } + + /// Get a reference to the underlying ProllyTree + pub fn tree(&self) -> &ProllyTree> { + &self.tree + } + + /// Get a mutable reference to the underlying ProllyTree + pub fn tree_mut(&mut self) -> &mut ProllyTree> { + &mut self.tree + } + + /// Reload the ProllyTree from the current HEAD (Git-specific) + fn reload_tree_from_head(&mut self) -> Result<(), GitKvError> { + // Since we're no longer storing prolly_tree_root in the Git tree, + // we need to reload the tree state from the GitNodeStorage + + // Load tree configuration from storage + let config: TreeConfig = ProllyTree::load_config(&self.tree.storage).unwrap_or_default(); + + // Try to load existing tree from storage, or create new one + let storage = self.tree.storage.clone(); + self.tree = ProllyTree::load_from_storage(storage.clone(), config.clone()) + .unwrap_or_else(|| ProllyTree::new(storage, config)); + + Ok(()) + } +} + +impl VersionedKvStore> { + /// Initialize a new versioned KV store with InMemory storage + pub fn init>(path: P) -> Result { + let path = path.as_ref(); + + // Find the git repository + let git_root = Self::find_git_root(path).ok_or_else(|| { + GitKvError::GitObjectError( + "Not inside a git repository. Please run from within a git repository.".to_string(), + ) + })?; + + // Open the existing git repository + let git_repo = gix::open(&git_root).map_err(|e| GitKvError::GitOpenError(Box::new(e)))?; + + // Create InMemoryNodeStorage + let storage = InMemoryNodeStorage::::new(); + + // Create ProllyTree with default config + let config: TreeConfig = TreeConfig::default(); + let tree = ProllyTree::new(storage, config); + + let mut store = VersionedKvStore { + tree, + git_repo, + staging_area: HashMap::new(), + current_branch: "main".to_string(), + storage_backend: StorageBackend::InMemory, + }; + + // Note: InMemory storage doesn't persist config + // Create initial commit + store.commit("Initial commit")?; + + Ok(store) + } + + /// Open an existing versioned KV store with InMemory storage + /// Note: InMemory storage is volatile, so this creates a new empty store + pub fn open>(path: P) -> Result { + // For InMemory storage, "opening" is the same as initializing + // since data is not persistent + Self::init(path) + } +} + +impl VersionedKvStore> { + /// Initialize a new versioned KV store with File storage + pub fn init>(path: P) -> Result { + let path = path.as_ref(); + + // Find the git repository + let git_root = Self::find_git_root(path).ok_or_else(|| { + GitKvError::GitObjectError( + "Not inside a git repository. Please run from within a git repository.".to_string(), + ) + })?; + + // Open the existing git repository + let git_repo = gix::open(&git_root).map_err(|e| GitKvError::GitOpenError(Box::new(e)))?; + + // Create FileNodeStorage with a subdirectory for file storage + let file_storage_path = path.join("file_storage"); + let storage = FileNodeStorage::::new(file_storage_path); + + // Create ProllyTree with default config + let config: TreeConfig = TreeConfig::default(); + let tree = ProllyTree::new(storage, config); + + let mut store = VersionedKvStore { + tree, + git_repo, + staging_area: HashMap::new(), + current_branch: "main".to_string(), + storage_backend: StorageBackend::File, + }; + + // Save initial configuration + let _ = store.tree.save_config(); + + // Create initial commit + store.commit("Initial commit")?; + + Ok(store) + } + + /// Open an existing versioned KV store with File storage + pub fn open>(path: P) -> Result { + let path = path.as_ref(); + + // Find the git repository + let git_root = Self::find_git_root(path).ok_or_else(|| { + GitKvError::GitObjectError( + "Not inside a git repository. Please run from within a git repository.".to_string(), + ) + })?; + + // Open existing Git repository + let git_repo = gix::open(&git_root).map_err(|e| GitKvError::GitOpenError(Box::new(e)))?; + + // Create FileNodeStorage with a subdirectory for file storage + let file_storage_path = path.join("file_storage"); + let storage = FileNodeStorage::::new(file_storage_path.clone()); + + // Load tree configuration from storage + let config: TreeConfig = ProllyTree::load_config(&storage).unwrap_or_default(); + + // Try to load existing tree from storage, or create new one + let tree = + if let Some(existing_tree) = ProllyTree::load_from_storage(storage, config.clone()) { + existing_tree + } else { + // Create new storage instance since the original was consumed + let new_storage = FileNodeStorage::::new(file_storage_path); + ProllyTree::new(new_storage, config) + }; + + // Get current branch + let current_branch = git_repo + .head_ref() + .map_err(|e| GitKvError::GitObjectError(format!("Failed to get head ref: {e}")))? + .map(|r| r.name().shorten().to_string()) + .unwrap_or_else(|| "main".to_string()); + + let mut store = VersionedKvStore { + tree, + git_repo, + staging_area: HashMap::new(), + current_branch, + storage_backend: StorageBackend::File, + }; + + // Load staging area from file if it exists + store.load_staging_area()?; + + // Note: File storage data is loaded directly, no need to reload from HEAD + + Ok(store) + } +} + +#[cfg(feature = "rocksdb_storage")] +impl VersionedKvStore> { + /// Initialize a new versioned KV store with RocksDB storage + pub fn init>(path: P) -> Result { + let path = path.as_ref(); + + // Find the git repository + let git_root = Self::find_git_root(path).ok_or_else(|| { + GitKvError::GitObjectError( + "Not inside a git repository. Please run from within a git repository.".to_string(), + ) + })?; + + // Open the existing git repository + let git_repo = gix::open(&git_root).map_err(|e| GitKvError::GitOpenError(Box::new(e)))?; + + // Create RocksDBNodeStorage with a subdirectory for RocksDB + let rocksdb_path = path.join("rocksdb"); + let storage = RocksDBNodeStorage::::new(rocksdb_path) + .map_err(|e| GitKvError::GitObjectError(format!("RocksDB creation failed: {e}")))?; + + // Create ProllyTree with default config + let config: TreeConfig = TreeConfig::default(); + let tree = ProllyTree::new(storage, config); + + let mut store = VersionedKvStore { + tree, + git_repo, + staging_area: HashMap::new(), + current_branch: "main".to_string(), + storage_backend: StorageBackend::RocksDB, + }; + + // Save initial configuration + let _ = store.tree.save_config(); + + // Create initial commit + store.commit("Initial commit")?; + + Ok(store) + } + + /// Open an existing versioned KV store with RocksDB storage + pub fn open>(path: P) -> Result { + let path = path.as_ref(); + + // Find the git repository + let git_root = Self::find_git_root(path).ok_or_else(|| { + GitKvError::GitObjectError( + "Not inside a git repository. Please run from within a git repository.".to_string(), + ) + })?; + + // Open existing Git repository + let git_repo = gix::open(&git_root).map_err(|e| GitKvError::GitOpenError(Box::new(e)))?; + + // Create RocksDBNodeStorage with a subdirectory for RocksDB + let rocksdb_path = path.join("rocksdb"); + let storage = RocksDBNodeStorage::::new(rocksdb_path) + .map_err(|e| GitKvError::GitObjectError(format!("RocksDB creation failed: {e}")))?; + + // Load tree configuration from storage + let config: TreeConfig = ProllyTree::load_config(&storage).unwrap_or_default(); + + // Try to load existing tree from storage, or create new one + let tree = ProllyTree::load_from_storage(storage.clone(), config.clone()) + .unwrap_or_else(|| ProllyTree::new(storage, config)); + + // Get current branch + let current_branch = git_repo + .head_ref() + .map_err(|e| GitKvError::GitObjectError(format!("Failed to get head ref: {e}")))? + .map(|r| r.name().shorten().to_string()) + .unwrap_or_else(|| "main".to_string()); + + let mut store = VersionedKvStore { + tree, + git_repo, + staging_area: HashMap::new(), + current_branch, + storage_backend: StorageBackend::RocksDB, + }; + + // Load staging area from file if it exists + store.load_staging_area()?; + + // Note: RocksDB storage data is loaded directly, no need to reload from HEAD + + Ok(store) + } +} + +// Generic implementations for all storage types +impl> VersionedKvStore { + /// Get the current storage backend type + pub fn storage_backend(&self) -> &StorageBackend { + &self.storage_backend + } +} + #[cfg(test)] mod tests { use super::*; @@ -722,7 +984,7 @@ mod tests { // Create subdirectory for dataset let dataset_dir = temp_dir.path().join("dataset"); std::fs::create_dir_all(&dataset_dir).unwrap(); - let store = VersionedKvStore::<32>::init(&dataset_dir); + let store = GitVersionedKvStore::<32>::init(&dataset_dir); assert!(store.is_ok()); } @@ -734,7 +996,7 @@ mod tests { // Create subdirectory for dataset let dataset_dir = temp_dir.path().join("dataset"); std::fs::create_dir_all(&dataset_dir).unwrap(); - let mut store = VersionedKvStore::<32>::init(&dataset_dir).unwrap(); + let mut store = GitVersionedKvStore::<32>::init(&dataset_dir).unwrap(); // Test insert and get store.insert(b"key1".to_vec(), b"value1".to_vec()).unwrap(); @@ -759,7 +1021,7 @@ mod tests { // Create subdirectory for dataset let dataset_dir = temp_dir.path().join("dataset"); std::fs::create_dir_all(&dataset_dir).unwrap(); - let mut store = VersionedKvStore::<32>::init(&dataset_dir).unwrap(); + let mut store = GitVersionedKvStore::<32>::init(&dataset_dir).unwrap(); // Stage changes store.insert(b"key1".to_vec(), b"value1".to_vec()).unwrap(); @@ -790,7 +1052,7 @@ mod tests { let dataset_dir = temp_dir.path().join("dataset"); std::fs::create_dir_all(&dataset_dir).unwrap(); - let mut store = VersionedKvStore::<32>::init(&dataset_dir).unwrap(); + let mut store = GitVersionedKvStore::<32>::init(&dataset_dir).unwrap(); // Get initial commit count let log_output = std::process::Command::new("git") diff --git a/src/sql/glue_storage.rs b/src/sql/glue_storage.rs index 26518b1..2fa2162 100644 --- a/src/sql/glue_storage.rs +++ b/src/sql/glue_storage.rs @@ -30,17 +30,17 @@ use gluesql_core::{ }, }; -use crate::git::VersionedKvStore; +use crate::git::versioned_store::GitVersionedKvStore; /// GlueSQL storage backend using ProllyTree pub struct ProllyStorage { - store: VersionedKvStore, + store: GitVersionedKvStore, schemas: HashMap, } impl ProllyStorage { /// Create a new ProllyStorage instance - pub fn new(store: VersionedKvStore) -> Self { + pub fn new(store: GitVersionedKvStore) -> Self { Self { store, schemas: HashMap::new(), @@ -52,7 +52,7 @@ impl ProllyStorage { pub fn init(path: &std::path::Path) -> Result { let dir = path.to_path_buf(); let dir_string = dir.to_string_lossy().to_string(); - let store = VersionedKvStore::init(path).map_err(|e| { + let store = GitVersionedKvStore::init(path).map_err(|e| { Error::StorageMsg(format!("Failed to initialize store: {e} from {dir_string}")) })?; Ok(Self::new(store)) @@ -61,13 +61,13 @@ impl ProllyStorage { /// Open an existing storage #[allow(clippy::result_large_err)] pub fn open(path: &std::path::Path) -> Result { - let store = VersionedKvStore::open(path) + let store = GitVersionedKvStore::open(path) .map_err(|e| Error::StorageMsg(format!("Failed to open store: {e}")))?; Ok(Self::new(store)) } // returns the underlying store - pub fn store(&self) -> &VersionedKvStore { + pub fn store(&self) -> &GitVersionedKvStore { &self.store } From 0fe6fbee3c883d149f07a88fd7bdb2aa51675826 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Thu, 24 Jul 2025 15:19:11 -0700 Subject: [PATCH 2/6] enable run ci build pipelines when pr is first openned --- .github/workflows/ci.yml | 2 +- .github/workflows/python.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4cd77f9..25e0b8a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,7 +2,7 @@ name: CI on: pull_request: - types: [ready_for_review, synchronize] + types: [opened, ready_for_review, synchronize] push: branches: - main diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index 1cef6b6..419b9fa 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -2,7 +2,7 @@ name: Build Python Package on: pull_request: - types: [ready_for_review, synchronize] + types: [opened, ready_for_review, synchronize] paths: - 'python/**' push: From 3b40198b1aa8a1fd28b9706a910f6781780a1619 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Thu, 24 Jul 2025 15:37:00 -0700 Subject: [PATCH 3/6] re-org codebase --- Cargo.toml | 30 +- .../{git_prolly_bench.rs => git_prolly.rs} | 0 benches/{sql_bench.rs => sql.rs} | 0 benches/{storage_bench.rs => storage.rs} | 0 benches/{prollytree_bench.rs => tree.rs} | 0 examples/{git_diff.rs => diff.rs} | 0 examples/{git_merge.rs => merge.rs} | 0 examples/{proof_visualization.rs => proof.rs} | 0 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 | 175 --------- .../rig_versioned_memory/src/agent/demos.rs | 265 -------------- .../rig_versioned_memory/src/agent/mod.rs | 4 - .../src/agent/versioned.rs | 175 --------- examples/rig_versioned_memory/src/lib.rs | 3 - examples/rig_versioned_memory/src/main.rs | 123 ------- .../rig_versioned_memory/src/memory/mod.rs | 6 - .../rig_versioned_memory/src/memory/schema.rs | 75 ---- .../rig_versioned_memory/src/memory/store.rs | 345 ------------------ .../rig_versioned_memory/src/memory/types.rs | 122 ------- .../rig_versioned_memory/src/utils/mod.rs | 36 -- examples/{git_sql.rs => sql.rs} | 0 examples/{rocksdb_storage.rs => storage.rs} | 0 24 files changed, 15 insertions(+), 1389 deletions(-) rename benches/{git_prolly_bench.rs => git_prolly.rs} (100%) rename benches/{sql_bench.rs => sql.rs} (100%) rename benches/{storage_bench.rs => storage.rs} (100%) rename benches/{prollytree_bench.rs => tree.rs} (100%) rename examples/{git_diff.rs => diff.rs} (100%) rename examples/{git_merge.rs => merge.rs} (100%) rename examples/{proof_visualization.rs => proof.rs} (100%) delete mode 100644 examples/rig_versioned_memory/.env.example delete mode 100644 examples/rig_versioned_memory/.gitignore delete mode 100644 examples/rig_versioned_memory/Cargo.toml delete mode 100644 examples/rig_versioned_memory/README.md delete mode 100644 examples/rig_versioned_memory/src/agent/demos.rs delete mode 100644 examples/rig_versioned_memory/src/agent/mod.rs delete mode 100644 examples/rig_versioned_memory/src/agent/versioned.rs delete mode 100644 examples/rig_versioned_memory/src/lib.rs delete mode 100644 examples/rig_versioned_memory/src/main.rs delete mode 100644 examples/rig_versioned_memory/src/memory/mod.rs delete mode 100644 examples/rig_versioned_memory/src/memory/schema.rs delete mode 100644 examples/rig_versioned_memory/src/memory/store.rs delete mode 100644 examples/rig_versioned_memory/src/memory/types.rs delete mode 100644 examples/rig_versioned_memory/src/utils/mod.rs rename examples/{git_sql.rs => sql.rs} (100%) rename examples/{rocksdb_storage.rs => storage.rs} (100%) diff --git a/Cargo.toml b/Cargo.toml index 2cb8e8f..cdaa5b0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,46 +67,46 @@ path = "src/bin/git-prolly.rs" required-features = ["git", "sql"] [[bench]] -name = "prollytree_bench" +name = "tree" harness = false [[bench]] -name = "sql_bench" +name = "sql" harness = false required-features = ["sql"] [[bench]] -name = "git_prolly_bench" +name = "git_prolly" harness = false required-features = ["git", "sql"] [[bench]] -name = "storage_bench" +name = "storage" harness = false [[example]] -name = "proof_visualization" -path = "examples/proof_visualization.rs" +name = "proof" +path = "examples/proof.rs" [[example]] -name = "git_sql" -path = "examples/git_sql.rs" +name = "sql" +path = "examples/sql.rs" required-features = ["sql"] [[example]] -name = "git_diff" -path = "examples/git_diff.rs" +name = "diff" +path = "examples/diff.rs" required-features = ["git"] [[example]] -name = "git_merge" -path = "examples/git_merge.rs" +name = "merge" +path = "examples/merge.rs" required-features = ["git"] [[example]] -name = "rocksdb_storage" -path = "examples/rocksdb_storage.rs" +name = "torage" +path = "examples/storage.rs" required-features = ["rocksdb_storage"] [workspace] -members = ["examples/rig_versioned_memory", "examples/financial_advisor"] +members = ["examples/financial_advisor"] diff --git a/benches/git_prolly_bench.rs b/benches/git_prolly.rs similarity index 100% rename from benches/git_prolly_bench.rs rename to benches/git_prolly.rs diff --git a/benches/sql_bench.rs b/benches/sql.rs similarity index 100% rename from benches/sql_bench.rs rename to benches/sql.rs diff --git a/benches/storage_bench.rs b/benches/storage.rs similarity index 100% rename from benches/storage_bench.rs rename to benches/storage.rs diff --git a/benches/prollytree_bench.rs b/benches/tree.rs similarity index 100% rename from benches/prollytree_bench.rs rename to benches/tree.rs diff --git a/examples/git_diff.rs b/examples/diff.rs similarity index 100% rename from examples/git_diff.rs rename to examples/diff.rs diff --git a/examples/git_merge.rs b/examples/merge.rs similarity index 100% rename from examples/git_merge.rs rename to examples/merge.rs diff --git a/examples/proof_visualization.rs b/examples/proof.rs similarity index 100% rename from examples/proof_visualization.rs rename to examples/proof.rs diff --git a/examples/rig_versioned_memory/.env.example b/examples/rig_versioned_memory/.env.example deleted file mode 100644 index 56ff29c..0000000 --- a/examples/rig_versioned_memory/.env.example +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 9edc8e0..0000000 --- a/examples/rig_versioned_memory/.gitignore +++ /dev/null @@ -1,17 +0,0 @@ -# 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 deleted file mode 100644 index f6e7e77..0000000 --- a/examples/rig_versioned_memory/Cargo.toml +++ /dev/null @@ -1,23 +0,0 @@ -[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 deleted file mode 100644 index 026a916..0000000 --- a/examples/rig_versioned_memory/README.md +++ /dev/null @@ -1,175 +0,0 @@ -# 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 -- 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 deleted file mode 100644 index c41a9a1..0000000 --- a/examples/rig_versioned_memory/src/agent/demos.rs +++ /dev/null @@ -1,265 +0,0 @@ -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!("{v2} (before preference change)").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(()) - } -} diff --git a/examples/rig_versioned_memory/src/agent/mod.rs b/examples/rig_versioned_memory/src/agent/mod.rs deleted file mode 100644 index e1a8ba5..0000000 --- a/examples/rig_versioned_memory/src/agent/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -pub mod demos; -pub mod versioned; - -pub use versioned::VersionedAgent; diff --git a/examples/rig_versioned_memory/src/agent/versioned.rs b/examples/rig_versioned_memory/src/agent/versioned.rs deleted file mode 100644 index d92320f..0000000 --- a/examples/rig_versioned_memory/src/agent/versioned.rs +++ /dev/null @@ -1,175 +0,0 @@ -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: {input}\nAssistant:") - } else { - format!("Context from memory:\n{context_text}\n\nUser: {input}\nAssistant:") - } - } - - 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(()) - } -} diff --git a/examples/rig_versioned_memory/src/lib.rs b/examples/rig_versioned_memory/src/lib.rs deleted file mode 100644 index 75f7c47..0000000 --- a/examples/rig_versioned_memory/src/lib.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub mod agent; -pub mod memory; -pub mod utils; diff --git a/examples/rig_versioned_memory/src/main.rs b/examples/rig_versioned_memory/src/main.rs deleted file mode 100644 index ba0f0bb..0000000 --- a/examples/rig_versioned_memory/src/main.rs +++ /dev/null @@ -1,123 +0,0 @@ -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(()) -} diff --git a/examples/rig_versioned_memory/src/memory/mod.rs b/examples/rig_versioned_memory/src/memory/mod.rs deleted file mode 100644 index 30ab0b8..0000000 --- a/examples/rig_versioned_memory/src/memory/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -pub mod schema; -pub mod store; -pub mod types; - -pub use store::VersionedMemoryStore; -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 deleted file mode 100644 index cb67156..0000000 --- a/examples/rig_versioned_memory/src/memory/schema.rs +++ /dev/null @@ -1,75 +0,0 @@ -use anyhow::Result; -use gluesql_core::prelude::Glue; -use gluesql_core::store::Transaction; -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 - - // commit the schema changes - glue.storage.commit().await?; - - Ok(()) -} diff --git a/examples/rig_versioned_memory/src/memory/store.rs b/examples/rig_versioned_memory/src/memory/store.rs deleted file mode 100644 index d1255ff..0000000 --- a/examples/rig_versioned_memory/src/memory/store.rs +++ /dev/null @@ -1,345 +0,0 @@ -use anyhow::Result; -use chrono::Utc; -use gluesql_core::prelude::{Glue, Payload}; -use gluesql_core::store::Transaction; -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("prolly_config_tree_config").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?; - glue.storage.commit().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?; - glue.storage.commit().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?; - glue.storage.commit().await?; - } - } - - // Create a version (commit to git) - let version = self - .create_version(&format!("Store {memory_type} memory")) - .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?; - glue.storage.commit().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(()) - } -} diff --git a/examples/rig_versioned_memory/src/memory/types.rs b/examples/rig_versioned_memory/src/memory/types.rs deleted file mode 100644 index e789445..0000000 --- a/examples/rig_versioned_memory/src/memory/types.rs +++ /dev/null @@ -1,122 +0,0 @@ -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 Default for MemoryContext { - fn default() -> Self { - Self::new() - } -} - -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 - } -} diff --git a/examples/rig_versioned_memory/src/utils/mod.rs b/examples/rig_versioned_memory/src/utils/mod.rs deleted file mode 100644 index 286c5ce..0000000 --- a/examples/rig_versioned_memory/src/utils/mod.rs +++ /dev/null @@ -1,36 +0,0 @@ -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); -} diff --git a/examples/git_sql.rs b/examples/sql.rs similarity index 100% rename from examples/git_sql.rs rename to examples/sql.rs diff --git a/examples/rocksdb_storage.rs b/examples/storage.rs similarity index 100% rename from examples/rocksdb_storage.rs rename to examples/storage.rs From e0d6291d9d6076ad049fd19ca269d5f2d15c2263 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Thu, 24 Jul 2025 15:46:47 -0700 Subject: [PATCH 4/6] change default features --- .github/workflows/ci.yml | 16 ---------------- Cargo.toml | 2 +- benches/git_prolly.rs | 30 +++++++++++++++++++++++++++--- 3 files changed, 28 insertions(+), 20 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 25e0b8a..cab9fbc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,19 +34,3 @@ 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/Cargo.toml b/Cargo.toml index cdaa5b0..8dd507a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,7 +51,7 @@ tempfile = "3.0" tokio = { version = "1.0", features = ["full"] } [features] -default = ["digest_base64", "prolly_balance_max_nodes", "git", "sql", "rocksdb_storage"] +default = ["digest_base64", "prolly_balance_max_nodes", "git", "sql"] tracing = ["dep:tracing"] digest_base64 = ["dep:base64"] prolly_balance_max_nodes = [] diff --git a/benches/git_prolly.rs b/benches/git_prolly.rs index e5e708a..44c56a9 100644 --- a/benches/git_prolly.rs +++ b/benches/git_prolly.rs @@ -21,7 +21,7 @@ use prollytree::config::TreeConfig; #[cfg(all(feature = "git", feature = "sql"))] use prollytree::git::GitNodeStorage; #[cfg(all(feature = "git", feature = "sql"))] -use prollytree::git::VersionedKvStore; +use prollytree::git::GitVersionedKvStore; #[cfg(all(feature = "git", feature = "sql"))] use prollytree::sql::ProllyStorage; #[cfg(all(feature = "git", feature = "sql"))] @@ -60,7 +60,19 @@ fn bench_git_versioned_commits(c: &mut Criterion) { b.iter_batched( || { let temp_dir = TempDir::new().unwrap(); - let store = VersionedKvStore::<32>::init(temp_dir.path()).unwrap(); + + // Initialize git repository first + std::process::Command::new("git") + .args(&["init"]) + .current_dir(temp_dir.path()) + .output() + .unwrap(); + + // Create dataset subdirectory + let dataset_dir = temp_dir.path().join("dataset"); + std::fs::create_dir_all(&dataset_dir).unwrap(); + + let store = GitVersionedKvStore::<32>::init(&dataset_dir).unwrap(); (store, temp_dir, generate_versioned_data(5, size)) }, |(mut store, _temp_dir, data)| { @@ -190,7 +202,19 @@ fn bench_git_branch_operations(c: &mut Criterion) { b.iter_batched( || { let temp_dir = TempDir::new().unwrap(); - let mut store = VersionedKvStore::<32>::init(temp_dir.path()).unwrap(); + + // Initialize git repository first + std::process::Command::new("git") + .args(&["init"]) + .current_dir(temp_dir.path()) + .output() + .unwrap(); + + // Create dataset subdirectory + let dataset_dir = temp_dir.path().join("dataset"); + std::fs::create_dir_all(&dataset_dir).unwrap(); + + let mut store = GitVersionedKvStore::<32>::init(&dataset_dir).unwrap(); // Initialize with some data for i in 0..size { From fc24ee0384589561417d9d3479bc89188ef0319a Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Thu, 24 Jul 2025 15:53:30 -0700 Subject: [PATCH 5/6] fix fmt --- benches/git_prolly.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/benches/git_prolly.rs b/benches/git_prolly.rs index 44c56a9..21c1d8f 100644 --- a/benches/git_prolly.rs +++ b/benches/git_prolly.rs @@ -60,18 +60,18 @@ fn bench_git_versioned_commits(c: &mut Criterion) { b.iter_batched( || { let temp_dir = TempDir::new().unwrap(); - + // Initialize git repository first std::process::Command::new("git") .args(&["init"]) .current_dir(temp_dir.path()) .output() .unwrap(); - + // Create dataset subdirectory let dataset_dir = temp_dir.path().join("dataset"); std::fs::create_dir_all(&dataset_dir).unwrap(); - + let store = GitVersionedKvStore::<32>::init(&dataset_dir).unwrap(); (store, temp_dir, generate_versioned_data(5, size)) }, @@ -202,18 +202,18 @@ fn bench_git_branch_operations(c: &mut Criterion) { b.iter_batched( || { let temp_dir = TempDir::new().unwrap(); - + // Initialize git repository first std::process::Command::new("git") .args(&["init"]) .current_dir(temp_dir.path()) .output() .unwrap(); - + // Create dataset subdirectory let dataset_dir = temp_dir.path().join("dataset"); std::fs::create_dir_all(&dataset_dir).unwrap(); - + let mut store = GitVersionedKvStore::<32>::init(&dataset_dir).unwrap(); // Initialize with some data From 6514cfd365b38e492c2c280520c013de800fcb84 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Thu, 24 Jul 2025 16:13:44 -0700 Subject: [PATCH 6/6] update dep bot --- .github/dependabot.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/dependabot.yml b/.github/dependabot.yml index e00b2db..d29a12c 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -8,4 +8,4 @@ updates: - package-ecosystem: "cargo" directory: "/" # Location of package manifests schedule: - interval: "daily" + interval: "monthly"