diff --git a/Cargo.lock b/Cargo.lock index 3d3463ceb1..4dc3aac588 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1676,6 +1676,7 @@ dependencies = [ "objc2 0.5.2", "objc2-app-kit 0.2.2", "objc2-foundation 0.2.2", + "once_cell", "owo-colors", "parking_lot", "paste", @@ -1695,6 +1696,7 @@ dependencies = [ "rustls-pemfile 2.2.0", "rustyline", "security-framework 3.2.0", + "semantic_search_client", "semver", "serde", "serde_json", diff --git a/crates/chat-cli/.amazonq/mcp.json b/crates/chat-cli/.amazonq/mcp.json new file mode 100644 index 0000000000..700113020f --- /dev/null +++ b/crates/chat-cli/.amazonq/mcp.json @@ -0,0 +1,3 @@ +{ + "mcpServers": {} +} \ No newline at end of file diff --git a/crates/chat-cli/Cargo.toml b/crates/chat-cli/Cargo.toml index 1e7bb1757e..2763161bb9 100644 --- a/crates/chat-cli/Cargo.toml +++ b/crates/chat-cli/Cargo.toml @@ -29,6 +29,8 @@ amzn-qdeveloper-streaming-client = { path = "../amzn-qdeveloper-streaming-client amzn-toolkit-telemetry-client = { path = "../amzn-toolkit-telemetry-client" } anstream = "0.6.13" arboard = { version = "3.5.0", default-features = false } +once_cell = "1.19.0" +semantic_search_client = { path = "../semantic_search_client" } async-trait = "0.1.87" aws-config = "1.0.3" aws-credential-types = "1.0.3" diff --git a/crates/chat-cli/src/cli/chat/command.rs b/crates/chat-cli/src/cli/chat/command.rs index 5252e63aa1..66a37336c7 100644 --- a/crates/chat-cli/src/cli/chat/command.rs +++ b/crates/chat-cli/src/cli/chat/command.rs @@ -36,6 +36,9 @@ pub enum Command { Context { subcommand: ContextSubcommand, }, + Knowledge { + subcommand: KnowledgeSubcommand, + }, PromptEditor { initial_text: Option, }, @@ -182,6 +185,16 @@ pub enum ContextSubcommand { Help, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum KnowledgeSubcommand { + Show, + Add { path: String }, + Remove { path: String }, + Update { path: String }, + Clear, + Help, +} + impl ContextSubcommand { const ADD_USAGE: &str = "/context add [--global] [--force] [path2...]"; const AVAILABLE_COMMANDS: &str = color_print::cstr! {"Available commands @@ -447,6 +460,113 @@ impl Command { return Ok(match parts[0].to_lowercase().as_str() { "clear" => Self::Clear, "help" => Self::Help, + "knowledge" => { + if parts.len() < 2 { + return Ok(Self::Knowledge { + subcommand: KnowledgeSubcommand::Help, + }); + } + + match parts[1].to_lowercase().as_str() { + "show" => Self::Knowledge { + subcommand: KnowledgeSubcommand::Show, + }, + "add" => { + // Parse add command with path + let mut path = None; + + let args = match shlex::split(&parts[2..].join(" ")) { + Some(args) => args, + None => return Err("Failed to parse quoted arguments".to_string()), + }; + + for arg in &args { + if path.is_none() { + path = Some(arg.to_string()); + } else { + return Err(format!("Only a single path is allowed. Found extra path: {}", arg)); + } + } + + let path = path.ok_or_else(|| { + format!( + "Invalid /knowledge arguments.\n\nUsage:\n {}", + KnowledgeSubcommand::ADD_USAGE + ) + })?; + + Self::Knowledge { + subcommand: KnowledgeSubcommand::Add { path }, + } + }, + "update" => { + // Parse update command with path + let mut path = None; + + let args = match shlex::split(&parts[2..].join(" ")) { + Some(args) => args, + None => return Err("Failed to parse quoted arguments".to_string()), + }; + + for arg in &args { + if path.is_none() { + path = Some(arg.to_string()); + } else { + return Err(format!("Only a single path is allowed. Found extra path: {}", arg)); + } + } + + let path = path.ok_or_else(|| { + format!( + "Invalid /knowledge arguments.\n\nUsage:\n {}", + KnowledgeSubcommand::UPDATE_USAGE + ) + })?; + + Self::Knowledge { + subcommand: KnowledgeSubcommand::Update { path }, + } + }, + "rm" => { + // Parse rm command with path + let mut path = None; + let args = match shlex::split(&parts[2..].join(" ")) { + Some(args) => args, + None => return Err("Failed to parse quoted arguments".to_string()), + }; + + for arg in &args { + if path.is_none() { + path = Some(arg.to_string()); + } else { + return Err(format!("Only a single path is allowed. Found extra path: {}", arg)); + } + } + + let path = path.ok_or_else(|| { + format!( + "Invalid /knowledge arguments.\n\nUsage:\n {}", + KnowledgeSubcommand::REMOVE_USAGE + ) + })?; + Self::Knowledge { + subcommand: KnowledgeSubcommand::Remove { path }, + } + }, + "clear" => Self::Knowledge { + subcommand: KnowledgeSubcommand::Clear, + }, + "help" => Self::Knowledge { + subcommand: KnowledgeSubcommand::Help, + }, + other => { + return Err(KnowledgeSubcommand::usage_msg(format!( + "Unknown subcommand '{}'.", + other + ))); + }, + } + }, "compact" => { let mut prompt = None; let show_summary = true; @@ -1133,3 +1253,46 @@ mod tests { } } } +impl KnowledgeSubcommand { + const ADD_USAGE: &str = "/knowledge add "; + const AVAILABLE_COMMANDS: &str = color_print::cstr! {"Available commands + help Show an explanation for the knowledge command + show Display the knowledge base contents + add <> Add a file or directory to knowledge base + update <> Update a file or directory in knowledge base + rm <> Remove specified knowledge context by path + clear Remove all knowledge contexts"}; + const BASE_COMMAND: &str = color_print::cstr! {"Usage: /knowledge [SUBCOMMAND] + +Description + Manage knowledge base for semantic search and retrieval. + Knowledge base is used to store and search information across chat sessions."}; + const REMOVE_USAGE: &str = "/knowledge rm "; + const UPDATE_USAGE: &str = "/knowledge update "; + + fn usage_msg(header: impl AsRef) -> String { + format!( + "{}\n\n{}\n\n{}", + header.as_ref(), + Self::BASE_COMMAND, + Self::AVAILABLE_COMMANDS + ) + } + + pub fn help_text() -> String { + color_print::cformat!( + r#" +(Beta) Knowledge Base Management + +Knowledge base allows you to store and search information across chat sessions. +Files and directories added to the knowledge base are indexed for semantic search, +enabling more relevant and contextual responses. + +{} + +{}"#, + Self::BASE_COMMAND, + Self::AVAILABLE_COMMANDS + ) + } +} diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index b4bcb8af74..6c130e13b3 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -42,6 +42,7 @@ use std::{ use command::{ Command, + KnowledgeSubcommand, PromptsSubcommand, ToolsSubcommand, }; @@ -1334,6 +1335,272 @@ impl ChatContext { let mut tool_uses: Vec = tool_uses.unwrap_or_default(); Ok(match command { + Command::Knowledge { subcommand } => { + // Handle knowledge commands + match subcommand { + KnowledgeSubcommand::Show => { + let knowledge_store = tools::knowledge::KnowledgeStore::get_instance(); + let store = knowledge_store.lock().await; + let contexts = store.get_all().unwrap_or_default(); + + if contexts.is_empty() { + queue!(self.output, style::Print("\nNo knowledge base entries found.\n\n"))?; + } else { + queue!( + self.output, + style::Print("\nšŸ“š Knowledge Base Contexts:\n"), + style::Print(format!("{}\n", "ā”".repeat(50))) + )?; + + for context in contexts { + // Display context header with ID and name + queue!( + self.output, + style::SetAttribute(Attribute::Bold), + style::SetForegroundColor(Color::Cyan), + style::Print(format!("šŸ“‚ {}: ", context.id)), + style::SetForegroundColor(Color::Green), + style::Print(&context.name), + style::SetAttribute(Attribute::Reset), + style::Print("\n") + )?; + + // Display metadata + queue!( + self.output, + style::Print(format!(" Description: {}\n", context.description)), + style::Print(format!( + " Created: {}\n", + context.created_at.format("%Y-%m-%d %H:%M:%S") + )), + style::Print(format!( + " Updated: {}\n", + context.updated_at.format("%Y-%m-%d %H:%M:%S") + )) + )?; + + if let Some(path) = &context.source_path { + queue!(self.output, style::Print(format!(" Source: {}\n", path)))?; + } + + queue!( + self.output, + style::Print(" Items: "), + style::SetForegroundColor(Color::Yellow), + style::Print(format!("{}", context.item_count)), + style::SetForegroundColor(Color::Reset), + style::Print(" | Persistent: ") + )?; + + if context.persistent { + queue!( + self.output, + style::SetForegroundColor(Color::Green), + style::Print("Yes"), + style::SetForegroundColor(Color::Reset), + style::Print("\n") + )?; + } else { + queue!( + self.output, + style::SetForegroundColor(Color::Yellow), + style::Print("No"), + style::SetForegroundColor(Color::Reset), + style::Print("\n") + )?; + } + queue!(self.output, style::Print(format!("{}\n", "ā”".repeat(50))))?; + } + queue!(self.output, style::Print("\n"))?; + } + + ChatState::PromptUser { + tool_uses: Some(tool_uses), + pending_tool_index, + skip_printing_tools: true, + } + }, + KnowledgeSubcommand::Add { path } => { + // Implementation for adding knowledge entries + let knowledge_store = tools::knowledge::KnowledgeStore::get_instance(); + let mut store = knowledge_store.lock().await; + + // Sanitize the path before using it + let sanitized_path = if !path.contains('\n') { + let ctx_path = crate::cli::chat::tools::sanitize_path_tool_arg(&self.ctx, &path); + if !ctx_path.exists() { + queue!( + self.output, + style::SetForegroundColor(Color::Red), + style::Print(format!("\nError: Path '{}' does not exist\n\n", path)), + style::SetForegroundColor(Color::Reset) + )?; + return Ok(ChatState::PromptUser { + tool_uses: Some(tool_uses), + pending_tool_index, + skip_printing_tools: true, + }); + } + ctx_path.to_string_lossy().to_string() + } else { + path.clone() + }; + + // Use the path as both name and value for simplicity + store.add(&path, &sanitized_path).unwrap_or_else(|e| { + queue!( + self.output, + style::SetForegroundColor(Color::Red), + style::Print(format!("\nError adding to knowledge base: {}\n\n", e)), + style::SetForegroundColor(Color::Reset) + ) + .unwrap(); + String::new() + }); + + queue!( + self.output, + style::SetForegroundColor(Color::Green), + style::Print(format!("\nAdded to knowledge base: {}\n\n", path)), + style::SetForegroundColor(Color::Reset) + )?; + + ChatState::PromptUser { + tool_uses: Some(tool_uses), + pending_tool_index, + skip_printing_tools: true, + } + }, + KnowledgeSubcommand::Remove { path } => { + // Implementation for removing knowledge entries + let knowledge_store = tools::knowledge::KnowledgeStore::get_instance(); + let mut store = knowledge_store.lock().await; + + // Sanitize the path before using it + let sanitized_path = crate::cli::chat::tools::sanitize_path_tool_arg(&self.ctx, &path); + + // Try to remove by path first + if store.remove_by_path(&sanitized_path.to_string_lossy()).is_ok() { + queue!( + self.output, + style::SetForegroundColor(Color::Green), + style::Print(format!( + "\nRemoved context with path '{}' from knowledge base\n\n", + path + )), + style::SetForegroundColor(Color::Reset) + )?; + } + // If path removal fails, try by name + else if store.remove_by_name(&path).is_ok() { + queue!( + self.output, + style::SetForegroundColor(Color::Green), + style::Print(format!( + "\nRemoved context with name '{}' from knowledge base\n\n", + path + )), + style::SetForegroundColor(Color::Reset) + )?; + } else { + queue!( + self.output, + style::SetForegroundColor(Color::Yellow), + style::Print(format!("\nEntry not found in knowledge base: {}\n\n", path)), + style::SetForegroundColor(Color::Reset) + )?; + } + + ChatState::PromptUser { + tool_uses: Some(tool_uses), + pending_tool_index, + skip_printing_tools: true, + } + }, + KnowledgeSubcommand::Clear => { + // Implementation for clearing knowledge + let knowledge_store = tools::knowledge::KnowledgeStore::get_instance(); + let mut store = knowledge_store.lock().await; + let count = store.clear().unwrap_or_default(); + + queue!( + self.output, + style::SetForegroundColor(Color::Green), + style::Print(format!("\nCleared {} entries from knowledge base\n\n", count)), + style::SetForegroundColor(Color::Reset) + )?; + + ChatState::PromptUser { + tool_uses: Some(tool_uses), + pending_tool_index, + skip_printing_tools: true, + } + }, + KnowledgeSubcommand::Help => { + queue!( + self.output, + style::Print("\n"), + style::Print(KnowledgeSubcommand::help_text()), + style::Print("\n") + )?; + + ChatState::PromptUser { + tool_uses: Some(tool_uses), + pending_tool_index, + skip_printing_tools: true, + } + }, + KnowledgeSubcommand::Update { path } => { + // Implementation for updating knowledge entries + let knowledge_store = tools::knowledge::KnowledgeStore::get_instance(); + let mut store = knowledge_store.lock().await; + + // Sanitize the path before using it + let sanitized_path = if !path.contains('\n') { + let ctx_path = crate::cli::chat::tools::sanitize_path_tool_arg(&self.ctx, &path); + if !ctx_path.exists() { + queue!( + self.output, + style::SetForegroundColor(Color::Red), + style::Print(format!("\nError: Path '{}' does not exist\n\n", path)), + style::SetForegroundColor(Color::Reset) + )?; + return Ok(ChatState::PromptUser { + tool_uses: Some(tool_uses), + pending_tool_index, + skip_printing_tools: true, + }); + } + ctx_path.to_string_lossy().to_string() + } else { + path.clone() + }; + + // Try to update by path directly + if let Err(e) = store.update_by_path(&sanitized_path) { + queue!( + self.output, + style::SetForegroundColor(Color::Red), + style::Print(format!("\nError updating knowledge base: {}\n\n", e)), + style::SetForegroundColor(Color::Reset) + )?; + } else { + queue!( + self.output, + style::SetForegroundColor(Color::Green), + style::Print(format!("\nUpdated in knowledge base: {}\n\n", path)), + style::SetForegroundColor(Color::Reset) + )?; + } + + ChatState::PromptUser { + tool_uses: Some(tool_uses), + pending_tool_index, + skip_printing_tools: true, + } + }, + } + }, Command::Ask { prompt } => { // Check for a pending tool approval if let Some(index) = pending_tool_index { diff --git a/crates/chat-cli/src/cli/chat/prompt.rs b/crates/chat-cli/src/cli/chat/prompt.rs index b84af4bcdf..0b8c5f3d64 100644 --- a/crates/chat-cli/src/cli/chat/prompt.rs +++ b/crates/chat-cli/src/cli/chat/prompt.rs @@ -50,6 +50,11 @@ pub const COMMANDS: &[&str] = &[ "/tools untrust", "/tools trustall", "/tools reset", + "/knowledge", + "/knowledge show", + "/knowledge add", + "/knowledge rm", + "/knowledge clear", "/profile", "/profile help", "/profile list", diff --git a/crates/chat-cli/src/cli/chat/tool_manager.rs b/crates/chat-cli/src/cli/chat/tool_manager.rs index c38d441c76..58c42f1e6d 100644 --- a/crates/chat-cli/src/cli/chat/tool_manager.rs +++ b/crates/chat-cli/src/cli/chat/tool_manager.rs @@ -80,6 +80,7 @@ use crate::cli::chat::tools::execute_bash::ExecuteBash; use crate::cli::chat::tools::fs_read::FsRead; use crate::cli::chat::tools::fs_write::FsWrite; use crate::cli::chat::tools::gh_issue::GhIssue; +use crate::cli::chat::tools::knowledge::Knowledge; use crate::cli::chat::tools::thinking::Thinking; use crate::cli::chat::tools::use_aws::UseAws; use crate::cli::chat::tools::{ @@ -918,6 +919,7 @@ impl ToolManager { "use_aws" => Tool::UseAws(serde_json::from_value::(value.args).map_err(map_err)?), "report_issue" => Tool::GhIssue(serde_json::from_value::(value.args).map_err(map_err)?), "thinking" => Tool::Thinking(serde_json::from_value::(value.args).map_err(map_err)?), + "knowledge" => Tool::Knowledge(serde_json::from_value::(value.args).map_err(map_err)?), // Note that this name is namespaced with server_name{DELIMITER}tool_name name => { // Note: tn_map also has tools that underwent no transformation. In otherwords, if diff --git a/crates/chat-cli/src/cli/chat/tools/knowledge.rs b/crates/chat-cli/src/cli/chat/tools/knowledge.rs new file mode 100644 index 0000000000..f367e9fffb --- /dev/null +++ b/crates/chat-cli/src/cli/chat/tools/knowledge.rs @@ -0,0 +1,727 @@ +use std::io::Write; +use std::path::PathBuf; +use std::sync::Arc; + +use crossterm::queue; +use crossterm::style::{ + self, + Color, +}; +use eyre::Result; +use once_cell::sync::Lazy; +use semantic_search_client::SemanticSearchClient; +use semantic_search_client::types::{ + MemoryContext, + SearchResult, +}; +use serde::Deserialize; +use tokio::sync::Mutex; +use tracing::warn; + +use super::{ + InvokeOutput, + OutputKind, +}; +use crate::platform::Context; + +#[derive(Debug, Clone, Deserialize)] +#[serde(tag = "command")] +pub enum Knowledge { + #[serde(rename = "add")] + Add(KnowledgeAdd), + #[serde(rename = "remove")] + Remove(KnowledgeRemove), + #[serde(rename = "clear")] + Clear(KnowledgeClear), + #[serde(rename = "search")] + Search(KnowledgeSearch), + #[serde(rename = "update")] + Update(KnowledgeUpdate), + #[serde(rename = "show")] + Show, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct KnowledgeAdd { + pub name: String, + pub value: String, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct KnowledgeRemove { + #[serde(default)] + pub name: String, + #[serde(default)] + pub context_id: String, + #[serde(default)] + pub path: String, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct KnowledgeClear { + pub confirm: bool, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct KnowledgeSearch { + pub query: String, + pub context_id: Option, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct KnowledgeUpdate { + #[serde(default)] + pub path: String, + #[serde(default)] + pub context_id: String, + #[serde(default)] + pub name: String, +} + +impl Knowledge { + pub async fn validate(&mut self, ctx: &Context) -> Result<()> { + match self { + Knowledge::Add(add) => { + // Check if value is intended to be a path (doesn't contain newlines) + if !add.value.contains('\n') { + let path = crate::cli::chat::tools::sanitize_path_tool_arg(ctx, &add.value); + if !path.exists() { + eyre::bail!("Path '{}' does not exist", add.value); + } + } + Ok(()) + }, + Knowledge::Remove(remove) => { + if remove.name.is_empty() && remove.context_id.is_empty() && remove.path.is_empty() { + eyre::bail!("Please provide at least one of: name, context_id, or path"); + } + // If path is provided, validate it exists + if !remove.path.is_empty() { + let path = crate::cli::chat::tools::sanitize_path_tool_arg(ctx, &remove.path); + if !path.exists() { + warn!( + "Path '{}' does not exist, will try to remove by path string match", + remove.path + ); + } + } + Ok(()) + }, + Knowledge::Update(update) => { + // Require at least one identifier (context_id or name) + if update.context_id.is_empty() && update.name.is_empty() && update.path.is_empty() { + eyre::bail!("Please provide either context_id or name or path to identify the context to update"); + } + + // Validate the path exists + if !update.path.is_empty() { + let path = crate::cli::chat::tools::sanitize_path_tool_arg(ctx, &update.path); + if !path.exists() { + eyre::bail!("Path '{}' does not exist", update.path); + } + } + + Ok(()) + }, + Knowledge::Clear(clear) => { + if !clear.confirm { + eyre::bail!("Please confirm clearing knowledge base by setting confirm=true"); + } + Ok(()) + }, + Knowledge::Search(_) => Ok(()), + Knowledge::Show => Ok(()), + } + } + + pub async fn queue_description(&self, ctx: &Context, updates: &mut impl Write) -> Result<()> { + match self { + Knowledge::Add(add) => { + queue!( + updates, + style::Print("Adding to knowledge base: "), + style::SetForegroundColor(Color::Green), + style::Print(&add.name), + style::ResetColor, + )?; + + // Check if value is a path or text content + let path = crate::cli::chat::tools::sanitize_path_tool_arg(ctx, &add.value); + if path.exists() { + let path_type = if path.is_dir() { "directory" } else { "file" }; + queue!( + updates, + style::Print(format!(" ({}: ", path_type)), + style::SetForegroundColor(Color::Green), + style::Print(&add.value), + style::ResetColor, + style::Print(")\n") + )?; + } else { + let preview: String = add.value.chars().take(20).collect(); + if add.value.len() > 20 { + queue!( + updates, + style::Print(" (text: "), + style::SetForegroundColor(Color::Blue), + style::Print(format!("{}...", preview)), + style::ResetColor, + style::Print(")\n") + )?; + } else { + queue!( + updates, + style::Print(" (text: "), + style::SetForegroundColor(Color::Blue), + style::Print(&add.value), + style::ResetColor, + style::Print(")\n") + )?; + } + } + }, + Knowledge::Remove(remove) => { + if !remove.name.is_empty() { + queue!( + updates, + style::Print("Removing from knowledge base by name: "), + style::SetForegroundColor(Color::Green), + style::Print(&remove.name), + style::ResetColor, + )?; + } else if !remove.context_id.is_empty() { + queue!( + updates, + style::Print("Removing from knowledge base by ID: "), + style::SetForegroundColor(Color::Green), + style::Print(&remove.context_id), + style::ResetColor, + )?; + } else if !remove.path.is_empty() { + queue!( + updates, + style::Print("Removing from knowledge base by path: "), + style::SetForegroundColor(Color::Green), + style::Print(&remove.path), + style::ResetColor, + )?; + } else { + queue!( + updates, + style::Print("Removing from knowledge base: "), + style::SetForegroundColor(Color::Yellow), + style::Print("No identifier provided"), + style::ResetColor, + )?; + } + }, + Knowledge::Update(update) => { + queue!(updates, style::Print("Updating knowledge base context"),)?; + + if !update.context_id.is_empty() { + queue!( + updates, + style::Print(" with ID: "), + style::SetForegroundColor(Color::Green), + style::Print(&update.context_id), + style::ResetColor, + )?; + } else if !update.name.is_empty() { + queue!( + updates, + style::Print(" with name: "), + style::SetForegroundColor(Color::Green), + style::Print(&update.name), + style::ResetColor, + )?; + } + + let path = crate::cli::chat::tools::sanitize_path_tool_arg(ctx, &update.path); + let path_type = if path.is_dir() { "directory" } else { "file" }; + queue!( + updates, + style::Print(format!(" using new {}: ", path_type)), + style::SetForegroundColor(Color::Green), + style::Print(&update.path), + style::ResetColor, + )?; + }, + Knowledge::Clear(_) => { + queue!( + updates, + style::Print("Clearing "), + style::SetForegroundColor(Color::Yellow), + style::Print("all"), + style::ResetColor, + style::Print(" knowledge base entries"), + )?; + }, + Knowledge::Search(search) => { + queue!( + updates, + style::Print("Searching knowledge base for: "), + style::SetForegroundColor(Color::Green), + style::Print(&search.query), + style::ResetColor, + )?; + + if let Some(context_id) = &search.context_id { + queue!( + updates, + style::Print(" in context: "), + style::SetForegroundColor(Color::Green), + style::Print(context_id), + style::ResetColor, + )?; + } else { + queue!(updates, style::Print(" across all contexts"),)?; + } + }, + Knowledge::Show => { + queue!(updates, style::Print("Showing all knowledge base entries"),)?; + }, + }; + Ok(()) + } + + pub async fn invoke(&self, ctx: &Context, _updates: &mut impl Write) -> Result { + // Get the knowledge store singleton + let knowledge_store = KnowledgeStore::get_instance(); + let mut store = knowledge_store.lock().await; + + let result = match self { + Knowledge::Add(add) => { + // For path indexing, we'll show a progress message first + let path = crate::cli::chat::tools::sanitize_path_tool_arg(ctx, &add.value); + let value_to_use = if path.exists() { + path.to_string_lossy().to_string() + } else { + // If it's not a valid path, use the original value (might be text content) + add.value.clone() + }; + + match store.add(&add.name, &value_to_use) { + Ok(context_id) => format!("Added '{}' to knowledge base with ID: {}", add.name, context_id), + Err(e) => format!("Failed to add to knowledge base: {}", e), + } + }, + Knowledge::Remove(remove) => { + if !remove.context_id.is_empty() { + // Remove by ID + match store.remove_by_id(&remove.context_id) { + Ok(_) => format!("Removed context with ID '{}' from knowledge base", remove.context_id), + Err(e) => format!("Failed to remove context by ID: {}", e), + } + } else if !remove.name.is_empty() { + // Remove by name + match store.remove_by_name(&remove.name) { + Ok(_) => format!("Removed context with name '{}' from knowledge base", remove.name), + Err(e) => format!("Failed to remove context by name: {}", e), + } + } else if !remove.path.is_empty() { + // Remove by path + let sanitized_path = crate::cli::chat::tools::sanitize_path_tool_arg(ctx, &remove.path); + match store.remove_by_path(&sanitized_path.to_string_lossy()) { + Ok(_) => format!("Removed context with path '{}' from knowledge base", remove.path), + Err(e) => format!("Failed to remove context by path: {}", e), + } + } else { + "Error: No identifier provided for removal. Please specify name, context_id, or path.".to_string() + } + }, + Knowledge::Update(update) => { + // Validate that we have a path and at least one identifier + if update.path.is_empty() { + return Ok(InvokeOutput { + output: OutputKind::Text( + "Error: No path provided for update. Please specify a path to update with.".to_string(), + ), + }); + } + + // Sanitize the path + let path = crate::cli::chat::tools::sanitize_path_tool_arg(ctx, &update.path); + if !path.exists() { + return Ok(InvokeOutput { + output: OutputKind::Text(format!("Error: Path '{}' does not exist", update.path)), + }); + } + + let sanitized_path = path.to_string_lossy().to_string(); + + // Choose the appropriate update method based on provided identifiers + if !update.context_id.is_empty() { + // Update by ID + match store.update_context_by_id(&update.context_id, &sanitized_path) { + Ok(_) => format!( + "Updated context with ID '{}' using path '{}'", + update.context_id, update.path + ), + Err(e) => format!("Failed to update context by ID: {}", e), + } + } else if !update.name.is_empty() { + // Update by name + match store.update_context_by_name(&update.name, &sanitized_path) { + Ok(_) => format!( + "Updated context with name '{}' using path '{}'", + update.name, update.path + ), + Err(e) => format!("Failed to update context by name: {}", e), + } + } else { + // Update by path (if no ID or name provided) + match store.update_by_path(&sanitized_path) { + Ok(_) => format!("Updated context with path '{}'", update.path), + Err(e) => format!("Failed to update context by path: {}", e), + } + } + }, + Knowledge::Clear(_) => match store.clear() { + Ok(count) => format!("Cleared {} entries from knowledge base", count), + Err(e) => format!("Failed to clear knowledge base: {}", e), + }, + Knowledge::Search(search) => { + // Only use a spinner for search, not a full progress bar + let results = store.search(&search.query, search.context_id.as_deref()); + match results { + Ok(results) => { + if results.is_empty() { + "No matching entries found in knowledge base".to_string() + } else { + let mut output = String::from("Search results:\n"); + for result in results { + if let Some(text) = result.text() { + output.push_str(&format!("- {}\n", text)); + } + } + output + } + }, + Err(e) => format!("Search failed: {}", e), + } + }, + Knowledge::Show => { + let contexts = store.get_all(); + match contexts { + Ok(contexts) => { + if contexts.is_empty() { + "No knowledge base entries found".to_string() + } else { + let mut output = String::from("Knowledge base entries:\n"); + for context in contexts { + output.push_str(&format!("- ID: {}\n Name: {}\n Description: {}\n Persistent: {}\n Created: {}\n Last Updated: {}\n Items: {}\n\n", + context.id, + context.name, + context.description, + context.persistent, + context.created_at.format("%Y-%m-%d %H:%M:%S"), + context.updated_at.format("%Y-%m-%d %H:%M:%S"), + context.item_count + )); + } + output + } + }, + Err(e) => format!("Failed to get knowledge base entries: {}", e), + } + }, + }; + + Ok(InvokeOutput { + output: OutputKind::Text(result), + }) + } +} + +// Knowledge store implementation using semantic_search_client +pub struct KnowledgeStore { + client: SemanticSearchClient, +} + +impl KnowledgeStore { + pub(crate) fn new() -> Result { + match SemanticSearchClient::new_with_default_dir() { + Ok(client) => Ok(Self { client }), + Err(e) => Err(eyre::eyre!("Failed to create semantic search client: {}", e)), + } + } + + // Create a test instance with an isolated directory + pub(crate) fn new_test_instance() -> Result { + let temp_dir = tempfile::tempdir()?; + match SemanticSearchClient::new(temp_dir.path()) { + Ok(client) => Ok(Self { client }), + Err(e) => Err(eyre::eyre!("Failed to create test semantic search client: {}", e)), + } + } + + // Singleton pattern for knowledge store with test mode support + pub fn get_instance() -> Arc> { + static INSTANCE: Lazy>> = Lazy::new(|| { + Arc::new(Mutex::new( + KnowledgeStore::new().expect("Failed to create knowledge store"), + )) + }); + + // Check if we're running in a test environment + if cfg!(test) { + // For tests, create a new isolated instance each time + Arc::new(Mutex::new( + KnowledgeStore::new_test_instance().expect("Failed to create test knowledge store"), + )) + } else { + // For normal operation, use the singleton + INSTANCE.clone() + } + } + + pub fn add(&mut self, name: &str, value: &str) -> Result { + let path = PathBuf::from(value); + + if path.exists() { + // Handle file or directory + + // Create a progress bar + let pb = indicatif::ProgressBar::new(100); + pb.set_style( + indicatif::ProgressStyle::default_bar() + .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {msg} {pos}/{len} ({eta})") + .unwrap() + .progress_chars("#>-"), + ); + + // Enable steady tick to ensure the progress bar updates regularly + pb.enable_steady_tick(std::time::Duration::from_millis(100)); + + // Use a progress callback + let progress_callback = move |status: semantic_search_client::types::ProgressStatus| { + match status { + semantic_search_client::types::ProgressStatus::CountingFiles => { + pb.set_message("Counting files..."); + pb.set_length(100); + pb.set_position(0); + }, + semantic_search_client::types::ProgressStatus::StartingIndexing(total) => { + pb.set_message("Indexing files..."); + pb.set_length(total as u64); + pb.set_position(0); + }, + semantic_search_client::types::ProgressStatus::Indexing(current, total) => { + pb.set_message(format!("Indexing file {} of {}", current, total)); + pb.set_position(current as u64); + }, + semantic_search_client::types::ProgressStatus::Finalizing => { + pb.set_message("Finalizing index..."); + if let Some(len) = pb.length() { + pb.set_position(len - 1); + } + }, + semantic_search_client::types::ProgressStatus::Complete => { + pb.finish_with_message("Indexing complete!"); + pb.println("āœ… Successfully indexed all files and created knowledge context"); + }, + semantic_search_client::types::ProgressStatus::CreatingSemanticContext => { + pb.set_message("Creating semantic context..."); + }, + semantic_search_client::types::ProgressStatus::GeneratingEmbeddings(current, total) => { + pb.set_message(format!("Generating embeddings {} of {}", current, total)); + pb.set_position(current as u64); + }, + semantic_search_client::types::ProgressStatus::BuildingIndex => { + pb.set_message("Building vector index..."); + }, + }; + }; + + self.client + .add_context_from_path( + path, + name, + &format!("Knowledge context for {}", name), + true, + Some(progress_callback), + ) + .map_err(|e| e.to_string()) + } else { + // Handle text content + let preview: String = value.chars().take(40).collect(); + self.client + .add_context_from_text(value, name, &format!("Text knowledge {}...", preview), true) + .map_err(|e| e.to_string()) + } + } + + pub fn update_context_by_id(&mut self, context_id: &str, path_str: &str) -> Result<(), String> { + // First, check if the context exists + let contexts = self.client.get_contexts(); + let context = contexts.iter().find(|c| c.id == context_id); + + if context.is_none() { + return Err(format!("Context with ID '{}' not found", context_id)); + } + + let context = context.unwrap(); + let path = PathBuf::from(path_str); + + if !path.exists() { + return Err(format!("Path '{}' does not exist", path_str)); + } + + // Create a progress bar + let pb = indicatif::ProgressBar::new(100); + pb.set_style( + indicatif::ProgressStyle::default_bar() + .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {msg} {pos}/{len} ({eta})") + .unwrap() + .progress_chars("#>-"), + ); + + // Enable steady tick to ensure the progress bar updates regularly + pb.enable_steady_tick(std::time::Duration::from_millis(100)); + + // Use a progress callback + let progress_callback = move |status: semantic_search_client::types::ProgressStatus| { + match status { + semantic_search_client::types::ProgressStatus::CountingFiles => { + pb.set_message("Counting files..."); + pb.set_length(100); + pb.set_position(0); + }, + semantic_search_client::types::ProgressStatus::StartingIndexing(total) => { + pb.set_message("Indexing files..."); + pb.set_length(total as u64); + pb.set_position(0); + }, + semantic_search_client::types::ProgressStatus::Indexing(current, total) => { + pb.set_message(format!("Indexing file {} of {}", current, total)); + pb.set_position(current as u64); + }, + semantic_search_client::types::ProgressStatus::Finalizing => { + pb.set_message("Finalizing index..."); + if let Some(len) = pb.length() { + pb.set_position(len - 1); + } + }, + semantic_search_client::types::ProgressStatus::Complete => { + pb.finish_with_message("Update complete!"); + pb.println("āœ… Successfully updated knowledge context"); + }, + semantic_search_client::types::ProgressStatus::CreatingSemanticContext => { + pb.set_message("Creating semantic context..."); + }, + semantic_search_client::types::ProgressStatus::GeneratingEmbeddings(current, total) => { + pb.set_message(format!("Generating embeddings {} of {}", current, total)); + pb.set_position(current as u64); + }, + semantic_search_client::types::ProgressStatus::BuildingIndex => { + pb.set_message("Building vector index..."); + }, + }; + }; + + // First remove the existing context + if let Err(e) = self.client.remove_context_by_id(context_id, true) { + return Err(format!("Failed to remove existing context: {}", e)); + } + + // Then add a new context with the same ID but new content + // Since we can't directly control the ID generation in add_context_from_path, + // we'll need to add the context and then update its metadata + let result = + self.client + .add_context_from_path(path, &context.name, &context.description, true, Some(progress_callback)); + + match result { + Ok(_) => Ok(()), + Err(e) => Err(format!("Failed to update context: {}", e)), + } + } + + pub fn update_context_by_name(&mut self, name: &str, path_str: &str) -> Result<(), String> { + // Find the context ID by name + let contexts = self.client.get_contexts(); + let context = contexts.iter().find(|c| c.name == name); + + if let Some(context) = context { + self.update_context_by_id(&context.id, path_str) + } else { + Err(format!("Context with name '{}' not found", name)) + } + } + + pub fn remove_by_id(&mut self, id: &str) -> Result<(), String> { + self.client.remove_context_by_id(id, true).map_err(|e| e.to_string()) + } + + pub fn remove_by_name(&mut self, name: &str) -> Result<(), String> { + self.client + .remove_context_by_name(name, true) + .map_err(|e| e.to_string()) + } + + pub fn remove_by_path(&mut self, path: &str) -> Result<(), String> { + self.client + .remove_context_by_path(path, true) + .map_err(|e| e.to_string()) + } + + pub fn update_by_path(&mut self, path_str: &str) -> Result<(), String> { + // Find contexts that might match this path + let contexts = self.client.get_contexts(); + let matching_context = contexts.iter().find(|c| { + if let Some(source_path) = &c.source_path { + source_path == path_str + } else { + false + } + }); + + if let Some(context) = matching_context { + // Found a matching context, update it + self.update_context_by_id(&context.id, path_str) + } else { + // No matching context found + Err(format!("No context found with path '{}'", path_str)) + } + } + + pub fn clear(&mut self) -> Result { + let contexts = self.client.get_contexts(); + let count = contexts.len(); + + for context in contexts { + if let Err(e) = self.client.remove_context_by_id(&context.id, true) { + tracing::warn!("Failed to remove context {}: {}", context.id, e); + } + } + + Ok(count) + } + + pub fn search(&self, query: &str, context_id: Option<&str>) -> Result, String> { + if let Some(id) = context_id { + self.client.search_context(id, query, None).map_err(|e| e.to_string()) + } else { + let results = self.client.search_all(query, None).map_err(|e| e.to_string())?; + + // Flatten results from all contexts + let mut flattened = Vec::new(); + for (_, context_results) in results { + flattened.extend(context_results); + } + + // Sort by distance (lower is better) + flattened.sort_by(|a, b| { + let a_dist = a.distance; + let b_dist = b.distance; + a_dist.partial_cmp(&b_dist).unwrap_or(std::cmp::Ordering::Equal) + }); + + Ok(flattened) + } + } + + pub fn get_all(&self) -> Result, String> { + Ok(self.client.get_contexts()) + } +} diff --git a/crates/chat-cli/src/cli/chat/tools/mod.rs b/crates/chat-cli/src/cli/chat/tools/mod.rs index 6d236ec708..9ff15c92d4 100644 --- a/crates/chat-cli/src/cli/chat/tools/mod.rs +++ b/crates/chat-cli/src/cli/chat/tools/mod.rs @@ -3,6 +3,7 @@ pub mod execute_bash; pub mod fs_read; pub mod fs_write; pub mod gh_issue; +pub mod knowledge; pub mod thinking; pub mod use_aws; @@ -20,6 +21,7 @@ use eyre::Result; use fs_read::FsRead; use fs_write::FsWrite; use gh_issue::GhIssue; +use knowledge::Knowledge; use serde::{ Deserialize, Serialize, @@ -42,6 +44,7 @@ pub enum Tool { Custom(CustomTool), GhIssue(GhIssue), Thinking(Thinking), + Knowledge(Knowledge), } impl Tool { @@ -55,6 +58,7 @@ impl Tool { Tool::Custom(custom_tool) => &custom_tool.name, Tool::GhIssue(_) => "gh_issue", Tool::Thinking(_) => "thinking (prerelease)", + Tool::Knowledge(_) => "knowledge", } .to_owned() } @@ -69,6 +73,7 @@ impl Tool { Tool::Custom(_) => true, Tool::GhIssue(_) => false, Tool::Thinking(_) => false, + Tool::Knowledge(_) => false, } } @@ -82,6 +87,7 @@ impl Tool { Tool::Custom(custom_tool) => custom_tool.invoke(context, updates).await, Tool::GhIssue(gh_issue) => gh_issue.invoke(updates).await, Tool::Thinking(think) => think.invoke(updates).await, + Tool::Knowledge(knowledge) => knowledge.invoke(context, updates).await, } } @@ -95,6 +101,7 @@ impl Tool { Tool::Custom(custom_tool) => custom_tool.queue_description(updates), Tool::GhIssue(gh_issue) => gh_issue.queue_description(updates), Tool::Thinking(thinking) => thinking.queue_description(updates), + Tool::Knowledge(knowledge) => knowledge.queue_description(ctx, updates).await, } } @@ -108,6 +115,7 @@ impl Tool { Tool::Custom(custom_tool) => custom_tool.validate(ctx).await, Tool::GhIssue(gh_issue) => gh_issue.validate(ctx).await, Tool::Thinking(think) => think.validate(ctx).await, + Tool::Knowledge(knowledge) => knowledge.validate(ctx).await, } } } @@ -187,6 +195,7 @@ impl ToolPermissions { "use_aws" => "trust read-only commands".dark_grey(), "report_issue" => "trusted".dark_green().bold(), "thinking" => "trusted (prerelease)".dark_green().bold(), + "knowledge" => "trusted (prerelease)".dark_green().bold(), _ if self.trust_all => "trusted".dark_grey().bold(), _ => "not trusted".dark_grey(), }; diff --git a/crates/chat-cli/src/cli/chat/tools/tool_index.json b/crates/chat-cli/src/cli/chat/tools/tool_index.json index 1fbe3f4d7f..ee5bcdec56 100644 --- a/crates/chat-cli/src/cli/chat/tools/tool_index.json +++ b/crates/chat-cli/src/cli/chat/tools/tool_index.json @@ -189,5 +189,44 @@ }, "required": ["thought"] } + }, + "knowledge": { + "name": "knowledge", + "description": "Store and retrieve information in knowledge base across chat sessions", + "input_schema": { + "type": "object", + "properties": { + "command": { + "type": "string", + "enum": ["show", "add", "remove", "clear", "search", "update"], + "description": "The knowledge operation to perform" + }, + "name": { + "type": "string", + "description": "A descriptive name for the knowledge context (for add and remove operations)" + }, + "value": { + "type": "string", + "description": "The value to store in knowledge base (for add operation) - can be text or a file/directory path" + }, + "context_id": { + "type": "string", + "description": "The context identifier of the knowledge bank for a targeted change. Obtained from Show command." + }, + "path": { + "type": "string", + "description": "The absolute path of an entry to be removed. Used in remove operations." + }, + "query": { + "type": "string", + "description": "The search query to find entries in knowledge base (for search operation)" + }, + "confirm": { + "type": "boolean", + "description": "Confirmation flag for destructive operations like clear" + } + }, + "required": ["command"] + } } }