From 4b36fe27328fde2ec3c906912b79ae8324568028 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Wed, 23 Jul 2025 07:55:55 -0700 Subject: [PATCH 1/9] Fix finance advosor branch functionalities --- .../src/advisor/interactive.rs | 88 +++++++++++++++++-- examples/financial_advisor/src/advisor/mod.rs | 24 +++++ examples/financial_advisor/src/memory/mod.rs | 55 ++++++++++++ 3 files changed, 158 insertions(+), 9 deletions(-) diff --git a/examples/financial_advisor/src/advisor/interactive.rs b/examples/financial_advisor/src/advisor/interactive.rs index a4dbcd4..58110e2 100644 --- a/examples/financial_advisor/src/advisor/interactive.rs +++ b/examples/financial_advisor/src/advisor/interactive.rs @@ -46,7 +46,10 @@ impl<'a> InteractiveSession<'a> { }; loop { - print!("\n{} ", "đŸĻ>".blue().bold()); + let actual_branch = self.advisor.get_actual_current_branch(); + let branch_display = format!(" [{}]", actual_branch.cyan()); + + print!("\n{}{} ", "đŸĻ".blue(), branch_display); io::stdout().flush()?; let mut input = String::new(); @@ -105,7 +108,9 @@ impl<'a> InteractiveSession<'a> { println!(" {} - Show audit trail", "audit".cyan()); println!(" {} - Test injection attack", "test-inject ".cyan()); println!(" {} - Show memory tree visualization", "visualize".cyan()); - println!(" {} - Create memory branch", "branch ".cyan()); + println!(" {} - Create and switch to memory branch", "branch ".cyan()); + println!(" {} - Switch to existing branch", "switch ".cyan()); + println!(" {} - Show branch information", "branch-info".cyan()); println!(" {} - Show this help", "help".cyan()); println!(" {} - Exit", "exit".cyan()); println!(); @@ -217,6 +222,26 @@ impl<'a> InteractiveSession<'a> { self.create_branch(branch_name).await?; } + "switch" | "sw" => { + if parts.len() < 2 { + println!("{} Usage: switch ", "❓".yellow()); + return Ok(true); + } + + let branch_name = parts[1]; + self.switch_branch(branch_name).await?; + } + + "branch-info" | "bi" => { + self.show_branch_info(); + } + + "debug-branch" | "db" => { + println!("DEBUG: Testing git branch reading..."); + let actual = self.advisor.get_actual_current_branch(); + println!("Result: '{}'", actual); + } + "exit" | "quit" | "q" => { return Ok(false); } @@ -261,7 +286,9 @@ impl<'a> InteractiveSession<'a> { println!(" {} - Show audit trail", "audit".cyan()); println!(" {} - Test injection attack", "test-inject ".cyan()); println!(" {} - Show memory tree visualization", "visualize".cyan()); - println!(" {} - Create memory branch", "branch ".cyan()); + println!(" {} - Create and switch to memory branch", "branch ".cyan()); + println!(" {} - Switch to existing branch", "switch ".cyan()); + println!(" {} - Show branch information", "branch-info".cyan()); println!(" {} - Show this help", "help".cyan()); println!(" {} - Exit", "exit".cyan()); println!(); @@ -768,13 +795,56 @@ impl<'a> InteractiveSession<'a> { async fn create_branch(&mut self, name: &str) -> Result<()> { println!("{} Creating memory branch: {}", "đŸŒŋ".green(), name.bold()); - // In real implementation, use the memory store - println!("{} Branch '{}' created successfully", "✅".green(), name); - println!( - "{} You can now safely test scenarios without affecting main memory", - "💡".yellow() - ); + // Create the branch using the memory store + match self.advisor.create_and_switch_branch(name).await { + Ok(_) => { + println!("{} Branch '{}' created successfully", "✅".green(), name); + println!("{} Switched to branch '{}'", "🔀".blue(), name); + println!( + "{} You can now safely test scenarios without affecting main memory", + "💡".yellow() + ); + } + Err(e) => { + println!("{} Failed to create/switch branch: {}", "❌".red(), e); + return Err(e); + } + } Ok(()) } + + async fn switch_branch(&mut self, name: &str) -> Result<()> { + println!("{} Switching to branch: {}", "🔀".blue(), name.bold()); + + match self.advisor.switch_to_branch(name).await { + Ok(_) => { + println!("{} Switched to branch '{}'", "✅".green(), name); + } + Err(e) => { + println!("{} Failed to switch to branch: {}", "❌".red(), e); + return Err(e); + } + } + + Ok(()) + } + + fn show_branch_info(&self) { + println!("{}", "đŸŒŋ Branch Information".green().bold()); + println!("{}", "━".repeat(25).dimmed()); + + let cached_branch = self.advisor.current_branch(); + let actual_branch = self.advisor.get_actual_current_branch(); + + println!("{}: {}", "Cached branch".cyan(), cached_branch); + println!("{}: {}", "Actual git branch".cyan(), actual_branch); + + if cached_branch != actual_branch { + println!("{} Branch mismatch detected!", "⚠ī¸".yellow()); + println!("This might indicate external git operations or sync issues."); + } else { + println!("{} Branches are in sync", "✅".green()); + } + } } diff --git a/examples/financial_advisor/src/advisor/mod.rs b/examples/financial_advisor/src/advisor/mod.rs index 52c1431..f1cd737 100644 --- a/examples/financial_advisor/src/advisor/mod.rs +++ b/examples/financial_advisor/src/advisor/mod.rs @@ -309,6 +309,30 @@ impl FinancialAdvisor { self.memory_store.load_client_profile().await } + pub async fn create_and_switch_branch(&mut self, name: &str) -> Result<()> { + // Create the branch + self.memory_store.create_branch(name).await?; + + // Switch to the newly created branch + self.memory_store.checkout(name).await?; + + Ok(()) + } + + pub fn current_branch(&self) -> &str { + self.memory_store.current_branch() + } + + pub fn get_actual_current_branch(&self) -> String { + self.memory_store.get_actual_current_branch() + } + + pub async fn switch_to_branch(&mut self, name: &str) -> Result<()> { + // Just switch to the branch (no creation) + self.memory_store.checkout(name).await?; + Ok(()) + } + async fn generate_ai_reasoning( &self, symbol: &str, diff --git a/examples/financial_advisor/src/memory/mod.rs b/examples/financial_advisor/src/memory/mod.rs index 58c0bad..134fbb8 100644 --- a/examples/financial_advisor/src/memory/mod.rs +++ b/examples/financial_advisor/src/memory/mod.rs @@ -219,6 +219,16 @@ impl MemoryStore { memory: &ValidatedMemory, ) -> Result { let path = Path::new(&self.store_path); + + // Debug: Check for branch mismatch (external git operations) + let cached_branch = self.current_branch(); + let actual_branch = self.get_actual_current_branch(); + + if cached_branch != actual_branch { + println!("DEBUG: ⚠ī¸ Branch mismatch: cached='{}', actual='{}' (external git operation?)", + cached_branch, actual_branch); + } + let storage = if path.join(PROLLY_CONFIG_FILE).exists() { ProllyStorage::<32>::open(path)? } else { @@ -388,6 +398,9 @@ impl MemoryStore { .create_branch(name) .map_err(|e| anyhow::anyhow!("Failed to create branch '{}': {:?}", name, e))?; + // Important: We need to ensure any ProllyStorage instances created after this + // will see the branch. The versioned_store has created the branch but hasn't switched to it. + if self.audit_enabled { self.log_audit(&format!("Created branch: {name}"), MemoryType::System, name) .await?; @@ -528,6 +541,43 @@ impl MemoryStore { self.versioned_store.current_branch() } + /// Get the actual current branch from git HEAD (not cached) + pub fn get_actual_current_branch(&self) -> String { + let store_path = Path::new(&self.store_path); + + // Try multiple possible git directory locations + // Git repository is typically in the parent directory of the store path + let possible_git_dirs = vec![ + store_path.parent().unwrap().join(".git"), // /tmp/advisor/.git + store_path.join(".git"), // /tmp/advisor/data7/.git + std::env::current_dir().unwrap().join(".git"), // current working directory + ]; + + for git_dir in possible_git_dirs { + let head_file = git_dir.join("HEAD"); + + if head_file.exists() { + // Read the HEAD file + if let Ok(head_content) = std::fs::read_to_string(&head_file) { + let head_content = head_content.trim(); + + // Check if HEAD points to a branch (ref: refs/heads/branch_name) + if let Some(branch_ref) = head_content.strip_prefix("ref: refs/heads/") { + return branch_ref.to_string(); + } + + // If HEAD contains a commit hash (detached HEAD), show first 8 chars + if head_content.len() >= 8 && head_content.chars().all(|c| c.is_ascii_hexdigit()) { + return format!("detached@{}", &head_content[..8]); + } + } + } + } + + // Fallback to cached branch name if git read fails + self.versioned_store.current_branch().to_string() + } + /// List all branches pub fn list_branches(&self) -> Result> { self.versioned_store @@ -542,10 +592,14 @@ impl MemoryStore { /// Checkout branch or commit pub async fn checkout(&mut self, branch_or_commit: &str) -> Result<()> { + println!("DEBUG: Checking out from branch '{}' to '{}'", self.current_branch(), branch_or_commit); + self.versioned_store .checkout(branch_or_commit) .map_err(|e| anyhow::anyhow!("Failed to checkout '{}': {:?}", branch_or_commit, e))?; + println!("DEBUG: Successfully checked out to branch '{}'", self.current_branch()); + if self.audit_enabled { self.log_audit( &format!("Checked out: {branch_or_commit}"), @@ -1478,4 +1532,5 @@ impl MemoryStore { hasher.update(content.as_bytes()); hasher.finalize().into() } + } From f2929b7d77368234f5d98af357c97d71bb39c61f Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Wed, 23 Jul 2025 08:31:54 -0700 Subject: [PATCH 2/9] misc fixes --- .../src/advisor/interactive.rs | 106 +++++++++++++++--- examples/financial_advisor/src/advisor/mod.rs | 9 ++ 2 files changed, 99 insertions(+), 16 deletions(-) diff --git a/examples/financial_advisor/src/advisor/interactive.rs b/examples/financial_advisor/src/advisor/interactive.rs index 58110e2..c93fc5f 100644 --- a/examples/financial_advisor/src/advisor/interactive.rs +++ b/examples/financial_advisor/src/advisor/interactive.rs @@ -36,11 +36,25 @@ impl<'a> InteractiveSession<'a> { time_horizon: "5-10 years".to_string(), restrictions: vec![], }; - println!( - "{} Created new client profile: {}", - "🆕".blue(), - default_profile.id - ); + + // Save the default profile so it persists + match self.advisor.store_client_profile(&default_profile).await { + Ok(_) => { + println!( + "{} Created and saved new client profile: {}", + "🆕".blue(), + default_profile.id + ); + } + Err(e) => { + println!( + "{} Created client profile but failed to save: {}", + "⚠ī¸".yellow(), + e + ); + } + } + default_profile } }; @@ -110,7 +124,7 @@ impl<'a> InteractiveSession<'a> { println!(" {} - Show memory tree visualization", "visualize".cyan()); println!(" {} - Create and switch to memory branch", "branch ".cyan()); println!(" {} - Switch to existing branch", "switch ".cyan()); - println!(" {} - Show branch information", "branch-info".cyan()); + println!(" {} - List all branches", "list-branches".cyan()); println!(" {} - Show this help", "help".cyan()); println!(" {} - Exit", "exit".cyan()); println!(); @@ -242,6 +256,10 @@ impl<'a> InteractiveSession<'a> { println!("Result: '{}'", actual); } + "list-branches" | "lb" => { + self.list_branches(); + } + "exit" | "quit" | "q" => { return Ok(false); } @@ -288,7 +306,7 @@ impl<'a> InteractiveSession<'a> { println!(" {} - Show memory tree visualization", "visualize".cyan()); println!(" {} - Create and switch to memory branch", "branch ".cyan()); println!(" {} - Switch to existing branch", "switch ".cyan()); - println!(" {} - Show branch information", "branch-info".cyan()); + println!(" {} - List all branches", "list-branches".cyan()); println!(" {} - Show this help", "help".cyan()); println!(" {} - Exit", "exit".cyan()); println!(); @@ -793,6 +811,13 @@ impl<'a> InteractiveSession<'a> { } async fn create_branch(&mut self, name: &str) -> Result<()> { + // Check if branch already exists + if self.advisor.branch_exists(name) { + println!("{} Branch '{}' already exists!", "⚠ī¸".yellow(), name.bold()); + println!("{} Use 'switch {}' to switch to the existing branch", "💡".blue(), name); + return Ok(()); + } + println!("{} Creating memory branch: {}", "đŸŒŋ".green(), name.bold()); // Create the branch using the memory store @@ -815,6 +840,13 @@ impl<'a> InteractiveSession<'a> { } async fn switch_branch(&mut self, name: &str) -> Result<()> { + // Check if branch exists before trying to switch + if !self.advisor.branch_exists(name) { + println!("{} Branch '{}' does not exist!", "❌".red(), name.bold()); + println!("{} Use 'branch {}' to create a new branch", "💡".blue(), name); + return Ok(()); + } + println!("{} Switching to branch: {}", "🔀".blue(), name.bold()); match self.advisor.switch_to_branch(name).await { @@ -831,20 +863,62 @@ impl<'a> InteractiveSession<'a> { } fn show_branch_info(&self) { - println!("{}", "đŸŒŋ Branch Information".green().bold()); - println!("{}", "━".repeat(25).dimmed()); + // Show all branches like git branch command + match self.advisor.memory_store.list_branches() { + Ok(branches) => { + let current_branch = self.advisor.get_actual_current_branch(); + + if branches.is_empty() { + println!("No branches found"); + } else { + for branch in branches { + if branch == current_branch { + println!("* {}", branch.green()); + } else { + println!(" {}", branch); + } + } + } + } + Err(e) => { + println!("error: Failed to list branches: {}", e); + } + } + // Show sync status if there's a mismatch let cached_branch = self.advisor.current_branch(); let actual_branch = self.advisor.get_actual_current_branch(); - println!("{}: {}", "Cached branch".cyan(), cached_branch); - println!("{}: {}", "Actual git branch".cyan(), actual_branch); - if cached_branch != actual_branch { - println!("{} Branch mismatch detected!", "⚠ī¸".yellow()); - println!("This might indicate external git operations or sync issues."); - } else { - println!("{} Branches are in sync", "✅".green()); + println!(); + println!("{} Branch mismatch: cached='{}', actual='{}'", + "warning:".yellow(), cached_branch, actual_branch); + } + } + + fn list_branches(&self) { + println!("{}", "đŸŒŗ Available Branches".green().bold()); + println!("{}", "━".repeat(25).dimmed()); + + match self.advisor.memory_store.list_branches() { + Ok(branches) => { + let current_branch = self.advisor.get_actual_current_branch(); + + if branches.is_empty() { + println!("{} No branches found", "ℹī¸".blue()); + } else { + for branch in branches { + if branch == current_branch { + println!(" {} {} (current)", "●".green(), branch.bold()); + } else { + println!(" {} {}", "○".dimmed(), branch); + } + } + } + } + Err(e) => { + println!("{} Failed to list branches: {}", "❌".red(), e); + } } } } diff --git a/examples/financial_advisor/src/advisor/mod.rs b/examples/financial_advisor/src/advisor/mod.rs index f1cd737..866f082 100644 --- a/examples/financial_advisor/src/advisor/mod.rs +++ b/examples/financial_advisor/src/advisor/mod.rs @@ -333,6 +333,15 @@ impl FinancialAdvisor { Ok(()) } + pub fn branch_exists(&self, name: &str) -> bool { + // Check if branch exists by listing all branches + if let Ok(branches) = self.memory_store.list_branches() { + branches.contains(&name.to_string()) + } else { + false + } + } + async fn generate_ai_reasoning( &self, symbol: &str, From 8019987dc564be8374ea6cbf3357f72463d6a67e Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Wed, 23 Jul 2025 08:35:43 -0700 Subject: [PATCH 3/9] fix fmt and clippy issues --- .../src/advisor/interactive.rs | 52 +++++++++++++------ examples/financial_advisor/src/advisor/mod.rs | 4 +- examples/financial_advisor/src/memory/mod.rs | 45 +++++++++------- 3 files changed, 63 insertions(+), 38 deletions(-) diff --git a/examples/financial_advisor/src/advisor/interactive.rs b/examples/financial_advisor/src/advisor/interactive.rs index c93fc5f..5fa73de 100644 --- a/examples/financial_advisor/src/advisor/interactive.rs +++ b/examples/financial_advisor/src/advisor/interactive.rs @@ -36,7 +36,7 @@ impl<'a> InteractiveSession<'a> { time_horizon: "5-10 years".to_string(), restrictions: vec![], }; - + // Save the default profile so it persists match self.advisor.store_client_profile(&default_profile).await { Ok(_) => { @@ -54,7 +54,7 @@ impl<'a> InteractiveSession<'a> { ); } } - + default_profile } }; @@ -62,7 +62,7 @@ impl<'a> InteractiveSession<'a> { loop { let actual_branch = self.advisor.get_actual_current_branch(); let branch_display = format!(" [{}]", actual_branch.cyan()); - + print!("\n{}{} ", "đŸĻ".blue(), branch_display); io::stdout().flush()?; @@ -122,7 +122,10 @@ impl<'a> InteractiveSession<'a> { println!(" {} - Show audit trail", "audit".cyan()); println!(" {} - Test injection attack", "test-inject ".cyan()); println!(" {} - Show memory tree visualization", "visualize".cyan()); - println!(" {} - Create and switch to memory branch", "branch ".cyan()); + println!( + " {} - Create and switch to memory branch", + "branch ".cyan() + ); println!(" {} - Switch to existing branch", "switch ".cyan()); println!(" {} - List all branches", "list-branches".cyan()); println!(" {} - Show this help", "help".cyan()); @@ -253,7 +256,7 @@ impl<'a> InteractiveSession<'a> { "debug-branch" | "db" => { println!("DEBUG: Testing git branch reading..."); let actual = self.advisor.get_actual_current_branch(); - println!("Result: '{}'", actual); + println!("Result: '{actual}'"); } "list-branches" | "lb" => { @@ -304,7 +307,10 @@ impl<'a> InteractiveSession<'a> { println!(" {} - Show audit trail", "audit".cyan()); println!(" {} - Test injection attack", "test-inject ".cyan()); println!(" {} - Show memory tree visualization", "visualize".cyan()); - println!(" {} - Create and switch to memory branch", "branch ".cyan()); + println!( + " {} - Create and switch to memory branch", + "branch ".cyan() + ); println!(" {} - Switch to existing branch", "switch ".cyan()); println!(" {} - List all branches", "list-branches".cyan()); println!(" {} - Show this help", "help".cyan()); @@ -814,7 +820,11 @@ impl<'a> InteractiveSession<'a> { // Check if branch already exists if self.advisor.branch_exists(name) { println!("{} Branch '{}' already exists!", "⚠ī¸".yellow(), name.bold()); - println!("{} Use 'switch {}' to switch to the existing branch", "💡".blue(), name); + println!( + "{} Use 'switch {}' to switch to the existing branch", + "💡".blue(), + name + ); return Ok(()); } @@ -843,7 +853,11 @@ impl<'a> InteractiveSession<'a> { // Check if branch exists before trying to switch if !self.advisor.branch_exists(name) { println!("{} Branch '{}' does not exist!", "❌".red(), name.bold()); - println!("{} Use 'branch {}' to create a new branch", "💡".blue(), name); + println!( + "{} Use 'branch {}' to create a new branch", + "💡".blue(), + name + ); return Ok(()); } @@ -867,7 +881,7 @@ impl<'a> InteractiveSession<'a> { match self.advisor.memory_store.list_branches() { Ok(branches) => { let current_branch = self.advisor.get_actual_current_branch(); - + if branches.is_empty() { println!("No branches found"); } else { @@ -875,35 +889,39 @@ impl<'a> InteractiveSession<'a> { if branch == current_branch { println!("* {}", branch.green()); } else { - println!(" {}", branch); + println!(" {branch}"); } } } } Err(e) => { - println!("error: Failed to list branches: {}", e); + println!("error: Failed to list branches: {e}"); } } - + // Show sync status if there's a mismatch let cached_branch = self.advisor.current_branch(); let actual_branch = self.advisor.get_actual_current_branch(); - + if cached_branch != actual_branch { println!(); - println!("{} Branch mismatch: cached='{}', actual='{}'", - "warning:".yellow(), cached_branch, actual_branch); + println!( + "{} Branch mismatch: cached='{}', actual='{}'", + "warning:".yellow(), + cached_branch, + actual_branch + ); } } fn list_branches(&self) { println!("{}", "đŸŒŗ Available Branches".green().bold()); println!("{}", "━".repeat(25).dimmed()); - + match self.advisor.memory_store.list_branches() { Ok(branches) => { let current_branch = self.advisor.get_actual_current_branch(); - + if branches.is_empty() { println!("{} No branches found", "ℹī¸".blue()); } else { diff --git a/examples/financial_advisor/src/advisor/mod.rs b/examples/financial_advisor/src/advisor/mod.rs index 866f082..5547cfc 100644 --- a/examples/financial_advisor/src/advisor/mod.rs +++ b/examples/financial_advisor/src/advisor/mod.rs @@ -312,10 +312,10 @@ impl FinancialAdvisor { pub async fn create_and_switch_branch(&mut self, name: &str) -> Result<()> { // Create the branch self.memory_store.create_branch(name).await?; - + // Switch to the newly created branch self.memory_store.checkout(name).await?; - + Ok(()) } diff --git a/examples/financial_advisor/src/memory/mod.rs b/examples/financial_advisor/src/memory/mod.rs index 134fbb8..1b3888a 100644 --- a/examples/financial_advisor/src/memory/mod.rs +++ b/examples/financial_advisor/src/memory/mod.rs @@ -219,16 +219,15 @@ impl MemoryStore { memory: &ValidatedMemory, ) -> Result { let path = Path::new(&self.store_path); - + // Debug: Check for branch mismatch (external git operations) let cached_branch = self.current_branch(); let actual_branch = self.get_actual_current_branch(); - + if cached_branch != actual_branch { - println!("DEBUG: ⚠ī¸ Branch mismatch: cached='{}', actual='{}' (external git operation?)", - cached_branch, actual_branch); + println!("DEBUG: ⚠ī¸ Branch mismatch: cached='{cached_branch}', actual='{actual_branch}' (external git operation?)"); } - + let storage = if path.join(PROLLY_CONFIG_FILE).exists() { ProllyStorage::<32>::open(path)? } else { @@ -400,7 +399,7 @@ impl MemoryStore { // Important: We need to ensure any ProllyStorage instances created after this // will see the branch. The versioned_store has created the branch but hasn't switched to it. - + if self.audit_enabled { self.log_audit(&format!("Created branch: {name}"), MemoryType::System, name) .await?; @@ -544,36 +543,38 @@ impl MemoryStore { /// Get the actual current branch from git HEAD (not cached) pub fn get_actual_current_branch(&self) -> String { let store_path = Path::new(&self.store_path); - + // Try multiple possible git directory locations // Git repository is typically in the parent directory of the store path let possible_git_dirs = vec![ - store_path.parent().unwrap().join(".git"), // /tmp/advisor/.git - store_path.join(".git"), // /tmp/advisor/data7/.git + store_path.parent().unwrap().join(".git"), // /tmp/advisor/.git + store_path.join(".git"), // /tmp/advisor/data7/.git std::env::current_dir().unwrap().join(".git"), // current working directory ]; - + for git_dir in possible_git_dirs { let head_file = git_dir.join("HEAD"); - + if head_file.exists() { // Read the HEAD file if let Ok(head_content) = std::fs::read_to_string(&head_file) { let head_content = head_content.trim(); - + // Check if HEAD points to a branch (ref: refs/heads/branch_name) if let Some(branch_ref) = head_content.strip_prefix("ref: refs/heads/") { return branch_ref.to_string(); } - + // If HEAD contains a commit hash (detached HEAD), show first 8 chars - if head_content.len() >= 8 && head_content.chars().all(|c| c.is_ascii_hexdigit()) { + if head_content.len() >= 8 + && head_content.chars().all(|c| c.is_ascii_hexdigit()) + { return format!("detached@{}", &head_content[..8]); } } } } - + // Fallback to cached branch name if git read fails self.versioned_store.current_branch().to_string() } @@ -592,13 +593,20 @@ impl MemoryStore { /// Checkout branch or commit pub async fn checkout(&mut self, branch_or_commit: &str) -> Result<()> { - println!("DEBUG: Checking out from branch '{}' to '{}'", self.current_branch(), branch_or_commit); - + println!( + "DEBUG: Checking out from branch '{}' to '{}'", + self.current_branch(), + branch_or_commit + ); + self.versioned_store .checkout(branch_or_commit) .map_err(|e| anyhow::anyhow!("Failed to checkout '{}': {:?}", branch_or_commit, e))?; - println!("DEBUG: Successfully checked out to branch '{}'", self.current_branch()); + println!( + "DEBUG: Successfully checked out to branch '{}'", + self.current_branch() + ); if self.audit_enabled { self.log_audit( @@ -1532,5 +1540,4 @@ impl MemoryStore { hasher.update(content.as_bytes()); hasher.finalize().into() } - } From d054d30159846632e24945e626ef2dce27d159f1 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Wed, 23 Jul 2025 08:42:26 -0700 Subject: [PATCH 4/9] edit readme --- examples/financial_advisor/README.md | 94 +++++++++++++++++++++++----- 1 file changed, 80 insertions(+), 14 deletions(-) diff --git a/examples/financial_advisor/README.md b/examples/financial_advisor/README.md index c6e1f27..45123c1 100644 --- a/examples/financial_advisor/README.md +++ b/examples/financial_advisor/README.md @@ -9,9 +9,11 @@ A demonstration of an AI-powered financial advisory system using ProllyTree for - 🔒 **Security Monitoring**: Detects and prevents injection attacks and anomalies - 📚 **Versioned Memory**: Uses ProllyTree to maintain git-like versioned storage of all data - 🕐 **Temporal Queries**: Query recommendations and data as they existed at any point in time -- đŸŒŋ **Branch Management**: Create memory branches for different scenarios or clients +- đŸŒŋ **Smart Branch Management**: Git-style branch operations with validation and external tool sync +- đŸĻ **Real-time UI**: Live branch display that updates with external git operations - 📝 **Audit Trail**: Complete audit logs for compliance and debugging - đŸŽ¯ **Risk-Aware**: Adapts recommendations based on client risk tolerance +- 👤 **Persistent Profiles**: Client profiles automatically saved per branch ## Prerequisites @@ -73,8 +75,13 @@ Once the advisor is running, you can use these commands: - `memory` - Show memory system status and statistics - `audit` - Show complete audit trail +#### Branch Management +- `branch ` - Create and switch to a new memory branch +- `switch ` - Switch to an existing branch +- `list-branches` - Show all available branches with visual indicators +- `branch-info` - List branches in git-style format (like `git branch`) + #### Advanced Features -- `branch ` - Create a new memory branch - `visualize` - Show memory tree visualization - `test-inject ` - Test security monitoring (try malicious inputs) @@ -85,24 +92,45 @@ Once the advisor is running, you can use these commands: ### Example Session ```bash -đŸĻ> recommend AAPL +đŸĻ [main] recommend AAPL 📊 Recommendation Generated Symbol: AAPL Action: BUY Confidence: 52.0% Reasoning: Analysis of AAPL at $177.89 with P/E ratio 28.4... -đŸĻ> risk aggressive +đŸĻ [main] risk aggressive ✅ Risk tolerance set to: Aggressive -đŸĻ> recommend AAPL +đŸĻ [main] recommend AAPL 📊 Recommendation Generated Symbol: AAPL Action: BUY Confidence: 60.0% (Notice higher confidence for aggressive risk tolerance) -đŸĻ> history +đŸĻ [main] branch test-strategy +đŸŒŋ Creating memory branch: test-strategy +✅ Branch 'test-strategy' created successfully +🔀 Switched to branch 'test-strategy' + +đŸĻ [test-strategy] recommend MSFT +📊 Recommendation Generated +Symbol: MSFT +Action: BUY +Confidence: 58.0% + +đŸĻ [test-strategy] list-branches +đŸŒŗ Available Branches +━━━━━━━━━━━━━━━━━━━━━━━━━ + ○ main + ● test-strategy (current) + +đŸĻ [test-strategy] switch main +🔀 Switching to branch: main +✅ Switched to branch 'main' + +đŸĻ [main] history 📜 Recent Recommendations 📊 Recommendation #1 Symbol: AAPL @@ -115,7 +143,7 @@ Confidence: 60.0% Confidence: 52.0% ... -đŸĻ> memory +đŸĻ [main] memory 🧠 Memory Status ✅ Memory validation: ACTIVE 🛡ī¸ Security monitoring: ENABLED @@ -168,18 +196,56 @@ Options: ### Branch Management -Create branches for different scenarios: +Create and manage branches for different scenarios: ```bash -đŸĻ> branch conservative-strategy -✅ Created branch: conservative-strategy +# Create and switch to a new branch +đŸĻ [main] branch conservative-strategy +đŸŒŋ Creating memory branch: conservative-strategy +✅ Branch 'conservative-strategy' created successfully +🔀 Switched to branch 'conservative-strategy' + +đŸĻ [conservative-strategy] risk conservative +✅ Risk tolerance set to: Conservative -đŸĻ> risk conservative -đŸĻ> recommend MSFT +đŸĻ [conservative-strategy] recommend MSFT # Generate recommendations for conservative strategy -đŸĻ> history --branch main -# Compare with main branch recommendations +# List all available branches +đŸĻ [conservative-strategy] list-branches +đŸŒŗ Available Branches +━━━━━━━━━━━━━━━━━━━━━━━━━ + ○ main + ● conservative-strategy (current) + +# Switch back to main branch +đŸĻ [conservative-strategy] switch main +🔀 Switching to branch: main +✅ Switched to branch 'main' + +đŸĻ [main] history --branch conservative-strategy +# Compare recommendations from different branch + +# Git-style branch listing +đŸĻ [main] branch-info +* main + conservative-strategy +``` + +#### Branch Validation + +The system prevents common branching mistakes: + +```bash +# Try to create existing branch +đŸĻ [main] branch main +⚠ī¸ Branch 'main' already exists! +💡 Use 'switch main' to switch to the existing branch + +# Try to switch to non-existent branch +đŸĻ [main] switch nonexistent +❌ Branch 'nonexistent' does not exist! +💡 Use 'branch nonexistent' to create a new branch ``` ### Temporal Analysis From af724551068027ff9d23b461e4abf5ccbbd47cea Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Wed, 23 Jul 2025 08:54:38 -0700 Subject: [PATCH 5/9] add notes to recommend command --- .../src/advisor/interactive.rs | 22 +++-- examples/financial_advisor/src/advisor/mod.rs | 24 ++++- examples/financial_advisor/src/memory/mod.rs | 89 +++++++++++++++++++ 3 files changed, 125 insertions(+), 10 deletions(-) diff --git a/examples/financial_advisor/src/advisor/interactive.rs b/examples/financial_advisor/src/advisor/interactive.rs index 5fa73de..2b4fd8c 100644 --- a/examples/financial_advisor/src/advisor/interactive.rs +++ b/examples/financial_advisor/src/advisor/interactive.rs @@ -102,7 +102,7 @@ impl<'a> InteractiveSession<'a> { println!("{}", "Available commands:".yellow()); println!( " {} - Get recommendation for a stock symbol", - "recommend ".cyan() + "recommend [notes]".cyan() ); println!(" {} - Show client profile", "profile".cyan()); println!( @@ -147,12 +147,17 @@ impl<'a> InteractiveSession<'a> { "recommend" | "r" => { if parts.len() < 2 { - println!("{} Usage: recommend ", "❓".yellow()); + println!("{} Usage: recommend [notes]", "❓".yellow()); return Ok(true); } let symbol = parts[1].to_uppercase(); - self.handle_recommendation(&symbol, client).await?; + let notes = if parts.len() > 2 { + Some(parts[2..].join(" ")) + } else { + None + }; + self.handle_recommendation(&symbol, client, notes).await?; } "profile" | "p" => { @@ -287,7 +292,7 @@ impl<'a> InteractiveSession<'a> { println!("{}", "Available commands:".yellow()); println!( " {} - Get recommendation for a stock symbol", - "recommend ".cyan() + "recommend [notes]".cyan() ); println!(" {} - Show client profile", "profile".cyan()); println!( @@ -343,14 +348,19 @@ impl<'a> InteractiveSession<'a> { println!(); } - async fn handle_recommendation(&mut self, symbol: &str, client: &ClientProfile) -> Result<()> { + async fn handle_recommendation( + &mut self, + symbol: &str, + client: &ClientProfile, + notes: Option, + ) -> Result<()> { println!( "{} Generating recommendation for {}...", "🔍".yellow(), symbol ); - match self.advisor.get_recommendation(symbol, client).await { + match self.advisor.get_recommendation(symbol, client, notes).await { Ok(recommendation) => { println!(); println!("{}", "📊 Recommendation Generated".green().bold()); diff --git a/examples/financial_advisor/src/advisor/mod.rs b/examples/financial_advisor/src/advisor/mod.rs index 5547cfc..b4ee7c1 100644 --- a/examples/financial_advisor/src/advisor/mod.rs +++ b/examples/financial_advisor/src/advisor/mod.rs @@ -106,6 +106,7 @@ impl FinancialAdvisor { &mut self, symbol: &str, client_profile: &ClientProfile, + notes: Option, ) -> Result { if self.verbose { println!("🔍 Fetching market data for {symbol}..."); @@ -148,7 +149,7 @@ impl FinancialAdvisor { ); // Step 5: Store recommendation with audit trail - self.store_recommendation(&recommendation).await?; + self.store_recommendation(&recommendation, notes).await?; // Keep in session memory for quick access self.session_recommendations.push(recommendation.clone()); @@ -219,7 +220,11 @@ impl FinancialAdvisor { Ok(()) } - async fn store_recommendation(&mut self, recommendation: &Recommendation) -> Result<()> { + async fn store_recommendation( + &mut self, + recommendation: &Recommendation, + notes: Option, + ) -> Result<()> { let memory = ValidatedMemory { id: recommendation.id.clone(), content: serde_json::to_string(recommendation)?, @@ -232,9 +237,19 @@ impl FinancialAdvisor { cross_references: vec![], }; - // Store with full audit trail + // Create custom commit message based on notes + let commit_message = if let Some(user_notes) = notes { + format!( + "Finance Advisor: recommend {} ({})", + recommendation.symbol, user_notes + ) + } else { + format!("Finance Advisor: recommend {}", recommendation.symbol) + }; + + // Store with full audit trail and custom commit message self.memory_store - .store_with_audit( + .store_with_audit_and_commit( MemoryType::Recommendation, &memory, &format!( @@ -242,6 +257,7 @@ impl FinancialAdvisor { recommendation.recommendation_type.as_str(), recommendation.symbol ), + &commit_message, ) .await?; diff --git a/examples/financial_advisor/src/memory/mod.rs b/examples/financial_advisor/src/memory/mod.rs index 1b3888a..e03228d 100644 --- a/examples/financial_advisor/src/memory/mod.rs +++ b/examples/financial_advisor/src/memory/mod.rs @@ -150,6 +150,21 @@ impl MemoryStore { ) .await?; + Self::ensure_table_exists( + glue, + "memories", + r#"CREATE TABLE memories ( + id TEXT PRIMARY KEY, + content TEXT, + timestamp INTEGER, + validation_hash TEXT, + sources TEXT, + confidence FLOAT, + cross_references TEXT + )"#, + ) + .await?; + Self::ensure_table_exists( glue, "audit_log", @@ -346,6 +361,62 @@ impl MemoryStore { Ok(version) } + pub async fn store_with_commit( + &mut self, + memory: &ValidatedMemory, + commit_message: &str, + ) -> Result { + let path = Path::new(&self.store_path); + let storage = if path.join(PROLLY_CONFIG_FILE).exists() { + ProllyStorage::<32>::open(path)? + } else { + ProllyStorage::<32>::init(path)? + }; + let mut glue = Glue::new(storage); + + // Ensure schema exists + Self::init_schema(&mut glue).await?; + + // Store memory in the memories table + let memory_sql = format!( + r#"INSERT INTO memories + (id, content, timestamp, validation_hash, sources, confidence, cross_references) + VALUES ('{}', '{}', {}, '{}', '{}', {}, '{}')"#, + memory.id, + memory.content.replace('\'', "''"), + memory.timestamp.timestamp(), + hex::encode(memory.validation_hash), + memory.sources.join(","), + memory.confidence, + memory.cross_references.join(",") + ); + + glue.execute(&memory_sql).await?; + + // Store cross-references + for reference in &memory.cross_references { + // Remove existing cross-reference to avoid duplicates + let delete_sql = format!( + "DELETE FROM cross_references WHERE source_id = '{}' AND target_id = '{}'", + memory.id, reference + ); + let _ = glue.execute(&delete_sql).await; // Ignore if record doesn't exist + + let sql = format!( + r#"INSERT INTO cross_references + (source_id, target_id, reference_type, confidence) + VALUES ('{}', '{}', 'validation', {})"#, + memory.id, reference, memory.confidence + ); + glue.execute(&sql).await?; + } + + let version = memory.clone().id; + glue.storage.commit_with_message(commit_message).await?; + + Ok(version) + } + pub async fn store_with_audit( &mut self, memory_type: MemoryType, @@ -363,6 +434,24 @@ impl MemoryStore { Ok(version) } + pub async fn store_with_audit_and_commit( + &mut self, + memory_type: MemoryType, + memory: &ValidatedMemory, + action: &str, + commit_message: &str, + ) -> Result { + // Store the memory with custom commit message + let version = self.store_with_commit(memory, commit_message).await?; + + // Log audit entry + if self.audit_enabled { + self.log_audit(action, memory_type, &memory.id).await?; + } + + Ok(version) + } + pub async fn query_related(&self, content: &str, limit: usize) -> Result> { let path = Path::new(&self.store_path); let storage = if path.join(PROLLY_CONFIG_FILE).exists() { From 1326db933d7359e1eea1ba6c2ca4458b6ed175e1 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Wed, 23 Jul 2025 09:03:07 -0700 Subject: [PATCH 6/9] remove debug printout --- examples/financial_advisor/src/memory/mod.rs | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/examples/financial_advisor/src/memory/mod.rs b/examples/financial_advisor/src/memory/mod.rs index e03228d..12df0ef 100644 --- a/examples/financial_advisor/src/memory/mod.rs +++ b/examples/financial_advisor/src/memory/mod.rs @@ -682,21 +682,10 @@ impl MemoryStore { /// Checkout branch or commit pub async fn checkout(&mut self, branch_or_commit: &str) -> Result<()> { - println!( - "DEBUG: Checking out from branch '{}' to '{}'", - self.current_branch(), - branch_or_commit - ); - self.versioned_store .checkout(branch_or_commit) .map_err(|e| anyhow::anyhow!("Failed to checkout '{}': {:?}", branch_or_commit, e))?; - println!( - "DEBUG: Successfully checked out to branch '{}'", - self.current_branch() - ); - if self.audit_enabled { self.log_audit( &format!("Checked out: {branch_or_commit}"), From 0b6b4ea1944c354de7f44fa4511ca75fcdfebfa2 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Wed, 23 Jul 2025 09:42:03 -0700 Subject: [PATCH 7/9] add test-inject to memories table --- .../src/advisor/interactive.rs | 36 ++++++++++--- examples/financial_advisor/src/advisor/mod.rs | 51 +++++++++++++++++++ examples/financial_advisor/src/main.rs | 3 ++ .../financial_advisor/src/memory/types.rs | 2 + 4 files changed, 86 insertions(+), 6 deletions(-) diff --git a/examples/financial_advisor/src/advisor/interactive.rs b/examples/financial_advisor/src/advisor/interactive.rs index 2b4fd8c..3c410c5 100644 --- a/examples/financial_advisor/src/advisor/interactive.rs +++ b/examples/financial_advisor/src/advisor/interactive.rs @@ -120,7 +120,7 @@ impl<'a> InteractiveSession<'a> { ); println!(" {} - Show memory validation status", "memory".cyan()); println!(" {} - Show audit trail", "audit".cyan()); - println!(" {} - Test injection attack", "test-inject ".cyan()); + println!(" {} - Test injection attack", "test-inject [-- notes]".cyan()); println!(" {} - Show memory tree visualization", "visualize".cyan()); println!( " {} - Create and switch to memory branch", @@ -222,12 +222,25 @@ impl<'a> InteractiveSession<'a> { "test-inject" | "inject" => { if parts.len() < 2 { - println!("{} Usage: test-inject ", "❓".yellow()); + println!("{} Usage: test-inject [-- notes]", "❓".yellow()); return Ok(true); } - let payload = parts[1..].join(" "); - self.test_injection_attack(&payload).await?; + // Find if there's a "--" separator for notes + let separator_pos = parts.iter().position(|&p| p == "--"); + let (payload, notes) = if let Some(pos) = separator_pos { + let payload = parts[1..pos].join(" "); + let notes = if pos + 1 < parts.len() { + Some(parts[pos + 1..].join(" ")) + } else { + None + }; + (payload, notes) + } else { + (parts[1..].join(" "), None) + }; + + self.test_injection_attack(&payload, notes).await?; } "visualize" | "vis" => { @@ -310,7 +323,7 @@ impl<'a> InteractiveSession<'a> { ); println!(" {} - Show memory validation status", "memory".cyan()); println!(" {} - Show audit trail", "audit".cyan()); - println!(" {} - Test injection attack", "test-inject ".cyan()); + println!(" {} - Test injection attack", "test-inject [-- notes]".cyan()); println!(" {} - Show memory tree visualization", "visualize".cyan()); println!( " {} - Create and switch to memory branch", @@ -729,6 +742,7 @@ impl<'a> InteractiveSession<'a> { "MarketData" => "📈", "Audit" => "📋", "System" => "⚙ī¸", + "Security" => "🛡ī¸", _ => "📝", }; @@ -763,7 +777,7 @@ impl<'a> InteractiveSession<'a> { Ok(()) } - async fn test_injection_attack(&mut self, payload: &str) -> Result<()> { + async fn test_injection_attack(&mut self, payload: &str, notes: Option) -> Result<()> { println!("{}", "🚨 Testing Injection Attack".red().bold()); println!("{}", "━".repeat(30).dimmed()); println!("{}: {}", "Payload".yellow(), payload); @@ -790,6 +804,16 @@ impl<'a> InteractiveSession<'a> { println!( "This shows how ProllyTree's versioned memory prevents injection attacks!" ); + + // Store the security test result + match self.advisor.store_security_test(payload, &alert, notes).await { + Ok(_) => { + println!("{} Security test result saved to memory", "💾".green()); + } + Err(e) => { + println!("{} Failed to save security test: {}", "⚠ī¸".yellow(), e); + } + } } Err(e) => { println!("{} Error during attack simulation: {}", "❌".red(), e); diff --git a/examples/financial_advisor/src/advisor/mod.rs b/examples/financial_advisor/src/advisor/mod.rs index b4ee7c1..ca03fee 100644 --- a/examples/financial_advisor/src/advisor/mod.rs +++ b/examples/financial_advisor/src/advisor/mod.rs @@ -268,6 +268,57 @@ impl FinancialAdvisor { Ok(()) } + async fn store_security_test(&mut self, payload: &str, alert: &crate::security::SecurityAlert, notes: Option) -> Result<()> { + use crate::memory::MemoryType; + use uuid::Uuid; + + let security_test = serde_json::json!({ + "id": Uuid::new_v4().to_string(), + "payload": payload, + "alert_level": format!("{:?}", alert.level), + "alert_type": format!("{:?}", alert.alert_type), + "description": alert.description, + "confidence": alert.confidence, + "recommendations": alert.recommendations, + "timestamp": Utc::now().to_rfc3339(), + }); + + let memory = ValidatedMemory { + id: Uuid::new_v4().to_string(), + content: security_test.to_string(), + timestamp: Utc::now(), + validation_hash: self + .validator + .hash_content(&security_test.to_string()), + sources: vec!["security_monitor".to_string()], + confidence: alert.confidence, + cross_references: vec![], + }; + + // Create custom commit message based on notes + let commit_message = if let Some(user_notes) = notes { + format!("Finance Advisor: security test ({})", user_notes) + } else { + "Finance Advisor: security test".to_string() + }; + + // Store with full audit trail and custom commit message + self.memory_store + .store_with_audit_and_commit( + MemoryType::Security, + &memory, + &format!("Security test: {}", alert.description), + &commit_message, + ) + .await?; + + if self.verbose { + println!("✅ Security test stored. ID: {}", memory.id); + } + + Ok(()) + } + pub async fn run_interactive_session(&mut self) -> Result<()> { let session = InteractiveSession::new(self); session.run().await diff --git a/examples/financial_advisor/src/main.rs b/examples/financial_advisor/src/main.rs index 825736f..8a68073 100644 --- a/examples/financial_advisor/src/main.rs +++ b/examples/financial_advisor/src/main.rs @@ -415,6 +415,7 @@ async fn run_memory_command(storage: &str, git_command: GitCommand) -> Result<() financial_advisor::memory::MemoryType::Audit => "📋", financial_advisor::memory::MemoryType::ClientProfile => "👤", financial_advisor::memory::MemoryType::System => "⚙ī¸", + financial_advisor::memory::MemoryType::Security => "🛡ī¸", }; println!( @@ -638,6 +639,7 @@ async fn run_memory_command(storage: &str, git_command: GitCommand) -> Result<() financial_advisor::memory::MemoryType::Audit => "📋", financial_advisor::memory::MemoryType::ClientProfile => "👤", financial_advisor::memory::MemoryType::System => "⚙ī¸", + financial_advisor::memory::MemoryType::Security => "🛡ī¸", }; let connector = if i == 0 { " " } else { "│" }; @@ -736,6 +738,7 @@ async fn run_memory_command(storage: &str, git_command: GitCommand) -> Result<() financial_advisor::memory::MemoryType::Audit => "📋", financial_advisor::memory::MemoryType::ClientProfile => "👤", financial_advisor::memory::MemoryType::System => "⚙ī¸", + financial_advisor::memory::MemoryType::Security => "🛡ī¸", }; let bar_length = (percentage / 100.0 * 30.0) as usize; diff --git a/examples/financial_advisor/src/memory/types.rs b/examples/financial_advisor/src/memory/types.rs index 3ee3386..b124c45 100644 --- a/examples/financial_advisor/src/memory/types.rs +++ b/examples/financial_advisor/src/memory/types.rs @@ -11,6 +11,7 @@ pub enum MemoryType { ClientProfile, Audit, System, + Security, } impl std::fmt::Display for MemoryType { @@ -21,6 +22,7 @@ impl std::fmt::Display for MemoryType { MemoryType::ClientProfile => write!(f, "ClientProfile"), MemoryType::Audit => write!(f, "Audit"), MemoryType::System => write!(f, "System"), + MemoryType::Security => write!(f, "Security"), } } } From ca22616e537f18ca9be4fd6b1d50579983947c33 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Wed, 23 Jul 2025 10:51:59 -0700 Subject: [PATCH 8/9] fix history command --- examples/financial_advisor/src/advisor/mod.rs | 5 +- examples/financial_advisor/src/memory/mod.rs | 63 +++++++++++++++++++ 2 files changed, 66 insertions(+), 2 deletions(-) diff --git a/examples/financial_advisor/src/advisor/mod.rs b/examples/financial_advisor/src/advisor/mod.rs index ca03fee..6b8f104 100644 --- a/examples/financial_advisor/src/advisor/mod.rs +++ b/examples/financial_advisor/src/advisor/mod.rs @@ -247,9 +247,10 @@ impl FinancialAdvisor { format!("Finance Advisor: recommend {}", recommendation.symbol) }; - // Store with full audit trail and custom commit message + // Store with full audit trail and custom commit message using typed storage self.memory_store - .store_with_audit_and_commit( + .store_typed_with_audit_and_commit( + recommendation, MemoryType::Recommendation, &memory, &format!( diff --git a/examples/financial_advisor/src/memory/mod.rs b/examples/financial_advisor/src/memory/mod.rs index 12df0ef..ed4261e 100644 --- a/examples/financial_advisor/src/memory/mod.rs +++ b/examples/financial_advisor/src/memory/mod.rs @@ -417,6 +417,50 @@ impl MemoryStore { Ok(version) } + pub async fn store_typed_with_commit( + &mut self, + item: &T, + memory: &ValidatedMemory, + commit_message: &str, + ) -> Result { + let path = Path::new(&self.store_path); + let storage = if path.join(PROLLY_CONFIG_FILE).exists() { + ProllyStorage::<32>::open(path)? + } else { + ProllyStorage::<32>::init(path)? + }; + let mut glue = Glue::new(storage); + + // Ensure schema exists + Self::init_schema(&mut glue).await?; + + // Use the item's store method to handle storage + item.store_to_db(&mut glue, memory).await?; + + // Store cross-references + for reference in &memory.cross_references { + // First try to delete if exists, then insert (GlueSQL doesn't support UPSERT) + let delete_sql = format!( + "DELETE FROM cross_references WHERE source_id = '{}' AND target_id = '{}'", + memory.id, reference + ); + let _ = glue.execute(&delete_sql).await; // Ignore if record doesn't exist + + let sql = format!( + r#"INSERT INTO cross_references + (source_id, target_id, reference_type, confidence) + VALUES ('{}', '{}', 'validation', {})"#, + memory.id, reference, memory.confidence + ); + glue.execute(&sql).await?; + } + + let version = memory.clone().id; + glue.storage.commit_with_message(commit_message).await?; + + Ok(version) + } + pub async fn store_with_audit( &mut self, memory_type: MemoryType, @@ -452,6 +496,25 @@ impl MemoryStore { Ok(version) } + pub async fn store_typed_with_audit_and_commit( + &mut self, + item: &T, + memory_type: MemoryType, + memory: &ValidatedMemory, + action: &str, + commit_message: &str, + ) -> Result { + // Store using typed storage with custom commit message + let version = self.store_typed_with_commit(item, memory, commit_message).await?; + + // Log audit entry + if self.audit_enabled { + self.log_audit(action, memory_type, &memory.id).await?; + } + + Ok(version) + } + pub async fn query_related(&self, content: &str, limit: usize) -> Result> { let path = Path::new(&self.store_path); let storage = if path.join(PROLLY_CONFIG_FILE).exists() { From 0090cf4457318041293b7fc392b530478a6274ac Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Wed, 23 Jul 2025 10:53:37 -0700 Subject: [PATCH 9/9] fix fmt --- .../src/advisor/interactive.rs | 23 +++++++++++++++---- examples/financial_advisor/src/advisor/mod.rs | 15 +++++++----- examples/financial_advisor/src/memory/mod.rs | 4 +++- 3 files changed, 30 insertions(+), 12 deletions(-) diff --git a/examples/financial_advisor/src/advisor/interactive.rs b/examples/financial_advisor/src/advisor/interactive.rs index 3c410c5..bbda019 100644 --- a/examples/financial_advisor/src/advisor/interactive.rs +++ b/examples/financial_advisor/src/advisor/interactive.rs @@ -120,7 +120,10 @@ impl<'a> InteractiveSession<'a> { ); println!(" {} - Show memory validation status", "memory".cyan()); println!(" {} - Show audit trail", "audit".cyan()); - println!(" {} - Test injection attack", "test-inject [-- notes]".cyan()); + println!( + " {} - Test injection attack", + "test-inject [-- notes]".cyan() + ); println!(" {} - Show memory tree visualization", "visualize".cyan()); println!( " {} - Create and switch to memory branch", @@ -222,7 +225,10 @@ impl<'a> InteractiveSession<'a> { "test-inject" | "inject" => { if parts.len() < 2 { - println!("{} Usage: test-inject [-- notes]", "❓".yellow()); + println!( + "{} Usage: test-inject [-- notes]", + "❓".yellow() + ); return Ok(true); } @@ -239,7 +245,7 @@ impl<'a> InteractiveSession<'a> { } else { (parts[1..].join(" "), None) }; - + self.test_injection_attack(&payload, notes).await?; } @@ -323,7 +329,10 @@ impl<'a> InteractiveSession<'a> { ); println!(" {} - Show memory validation status", "memory".cyan()); println!(" {} - Show audit trail", "audit".cyan()); - println!(" {} - Test injection attack", "test-inject [-- notes]".cyan()); + println!( + " {} - Test injection attack", + "test-inject [-- notes]".cyan() + ); println!(" {} - Show memory tree visualization", "visualize".cyan()); println!( " {} - Create and switch to memory branch", @@ -806,7 +815,11 @@ impl<'a> InteractiveSession<'a> { ); // Store the security test result - match self.advisor.store_security_test(payload, &alert, notes).await { + match self + .advisor + .store_security_test(payload, &alert, notes) + .await + { Ok(_) => { println!("{} Security test result saved to memory", "💾".green()); } diff --git a/examples/financial_advisor/src/advisor/mod.rs b/examples/financial_advisor/src/advisor/mod.rs index 6b8f104..a8c1e98 100644 --- a/examples/financial_advisor/src/advisor/mod.rs +++ b/examples/financial_advisor/src/advisor/mod.rs @@ -269,10 +269,15 @@ impl FinancialAdvisor { Ok(()) } - async fn store_security_test(&mut self, payload: &str, alert: &crate::security::SecurityAlert, notes: Option) -> Result<()> { + async fn store_security_test( + &mut self, + payload: &str, + alert: &crate::security::SecurityAlert, + notes: Option, + ) -> Result<()> { use crate::memory::MemoryType; use uuid::Uuid; - + let security_test = serde_json::json!({ "id": Uuid::new_v4().to_string(), "payload": payload, @@ -288,9 +293,7 @@ impl FinancialAdvisor { id: Uuid::new_v4().to_string(), content: security_test.to_string(), timestamp: Utc::now(), - validation_hash: self - .validator - .hash_content(&security_test.to_string()), + validation_hash: self.validator.hash_content(&security_test.to_string()), sources: vec!["security_monitor".to_string()], confidence: alert.confidence, cross_references: vec![], @@ -298,7 +301,7 @@ impl FinancialAdvisor { // Create custom commit message based on notes let commit_message = if let Some(user_notes) = notes { - format!("Finance Advisor: security test ({})", user_notes) + format!("Finance Advisor: security test ({user_notes})") } else { "Finance Advisor: security test".to_string() }; diff --git a/examples/financial_advisor/src/memory/mod.rs b/examples/financial_advisor/src/memory/mod.rs index ed4261e..67b3435 100644 --- a/examples/financial_advisor/src/memory/mod.rs +++ b/examples/financial_advisor/src/memory/mod.rs @@ -505,7 +505,9 @@ impl MemoryStore { commit_message: &str, ) -> Result { // Store using typed storage with custom commit message - let version = self.store_typed_with_commit(item, memory, commit_message).await?; + let version = self + .store_typed_with_commit(item, memory, commit_message) + .await?; // Log audit entry if self.audit_enabled {