diff --git a/crates/chat-cli/src/cli/chat/cli/knowledge.rs b/crates/chat-cli/src/cli/chat/cli/knowledge.rs index 987eda059..cccc9d8be 100644 --- a/crates/chat-cli/src/cli/chat/cli/knowledge.rs +++ b/crates/chat-cli/src/cli/chat/cli/knowledge.rs @@ -1,4 +1,5 @@ use std::io::Write; +use std::path::Path; use clap::Subcommand; use crossterm::queue; @@ -12,6 +13,8 @@ use semantic_search_client::{ OperationStatus, SystemStatus, }; +use serde::Deserialize; +use tracing::error; use crate::cli::chat::tools::sanitize_path_tool_arg; use crate::cli::chat::{ @@ -54,6 +57,13 @@ enum OperationResult { Error(String), } +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Settings { + #[serde(default)] + pub(crate) base_dir: String, +} + impl KnowledgeSubcommand { pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result { if !Self::is_feature_enabled(os) { @@ -90,7 +100,27 @@ impl KnowledgeSubcommand { } } + fn get_knowledge_base_dir(session: &mut ChatSession) -> Option + use<>> { + let agent = session.conversation.agents.get_active(); + if agent.is_none() { + return None; + } + + agent.unwrap().tools_settings.get("knowledge")?.get("base_dir"); + match agent.unwrap().tools_settings.get("knowledge") { + Some(settings) => match serde_json::from_value::(settings.clone()) { + Ok(settings) => Some(settings.base_dir), + Err(e) => { + error!("Failed to deserialize tool settings for execute_bash: {:?}", e); + None + }, + }, + None => None, + } + } + async fn execute_operation(&self, os: &Os, session: &mut ChatSession) -> OperationResult { + let knowledge_base_dir = Self::get_knowledge_base_dir(session); match self { KnowledgeSubcommand::Show => { match Self::handle_show(session).await { @@ -98,17 +128,20 @@ impl KnowledgeSubcommand { Err(e) => OperationResult::Error(format!("Failed to show contexts: {}", e)), } }, - KnowledgeSubcommand::Add { path } => Self::handle_add(os, path).await, - KnowledgeSubcommand::Remove { path } => Self::handle_remove(os, path).await, - KnowledgeSubcommand::Update { path } => Self::handle_update(os, path).await, + KnowledgeSubcommand::Add { path } => Self::handle_add(os, path, knowledge_base_dir).await, + KnowledgeSubcommand::Remove { path } => Self::handle_remove(os, path, knowledge_base_dir).await, + KnowledgeSubcommand::Update { path } => Self::handle_update(os, path, knowledge_base_dir).await, KnowledgeSubcommand::Clear => Self::handle_clear(session).await, - KnowledgeSubcommand::Status => Self::handle_status().await, - KnowledgeSubcommand::Cancel { operation_id } => Self::handle_cancel(operation_id.as_deref()).await, + KnowledgeSubcommand::Status => Self::handle_status(knowledge_base_dir).await, + KnowledgeSubcommand::Cancel { operation_id } => { + Self::handle_cancel(operation_id.as_deref(), knowledge_base_dir).await + }, } } async fn handle_show(session: &mut ChatSession) -> Result<(), std::io::Error> { - let async_knowledge_store = KnowledgeStore::get_async_instance().await; + let knowledge_base_dir = Self::get_knowledge_base_dir(session); + let async_knowledge_store = KnowledgeStore::get_async_instance(knowledge_base_dir).await; let store = async_knowledge_store.lock().await; // Use the async get_all method which is concurrent with indexing @@ -210,10 +243,10 @@ impl KnowledgeSubcommand { } /// Handle add operation - async fn handle_add(os: &Os, path: &str) -> OperationResult { + async fn handle_add(os: &Os, path: &str, knowledge_base_dir: Option>) -> OperationResult { match Self::validate_and_sanitize_path(os, path) { Ok(sanitized_path) => { - let async_knowledge_store = KnowledgeStore::get_async_instance().await; + let async_knowledge_store = KnowledgeStore::get_async_instance(knowledge_base_dir).await; let mut store = async_knowledge_store.lock().await; // Use the async add method which is fire-and-forget @@ -227,9 +260,9 @@ impl KnowledgeSubcommand { } /// Handle remove operation - async fn handle_remove(os: &Os, path: &str) -> OperationResult { + async fn handle_remove(os: &Os, path: &str, knowledge_base_dir: Option>) -> OperationResult { let sanitized_path = sanitize_path_tool_arg(os, path); - let async_knowledge_store = KnowledgeStore::get_async_instance().await; + let async_knowledge_store = KnowledgeStore::get_async_instance(knowledge_base_dir).await; let mut store = async_knowledge_store.lock().await; // Try path first, then name @@ -243,10 +276,10 @@ impl KnowledgeSubcommand { } /// Handle update operation - async fn handle_update(os: &Os, path: &str) -> OperationResult { + async fn handle_update(os: &Os, path: &str, knowledge_base_dir: Option>) -> OperationResult { match Self::validate_and_sanitize_path(os, path) { Ok(sanitized_path) => { - let async_knowledge_store = KnowledgeStore::get_async_instance().await; + let async_knowledge_store = KnowledgeStore::get_async_instance(knowledge_base_dir).await; let mut store = async_knowledge_store.lock().await; match store.update_by_path(&sanitized_path).await { @@ -278,7 +311,8 @@ impl KnowledgeSubcommand { return OperationResult::Info("Clear operation cancelled".to_string()); } - let async_knowledge_store = KnowledgeStore::get_async_instance().await; + let knowledge_base_dir = Self::get_knowledge_base_dir(session); + let async_knowledge_store = KnowledgeStore::get_async_instance(knowledge_base_dir).await; let mut store = async_knowledge_store.lock().await; // First, cancel any pending operations @@ -308,8 +342,8 @@ impl KnowledgeSubcommand { } /// Handle status operation - async fn handle_status() -> OperationResult { - let async_knowledge_store = KnowledgeStore::get_async_instance().await; + async fn handle_status(knowledge_base_dir: Option>) -> OperationResult { + let async_knowledge_store = KnowledgeStore::get_async_instance(knowledge_base_dir).await; let store = async_knowledge_store.lock().await; match store.get_status_data().await { @@ -416,8 +450,11 @@ impl KnowledgeSubcommand { } /// Handle cancel operation - async fn handle_cancel(operation_id: Option<&str>) -> OperationResult { - let async_knowledge_store = KnowledgeStore::get_async_instance().await; + async fn handle_cancel( + operation_id: Option<&str>, + knowledge_base_dir: Option>, + ) -> OperationResult { + let async_knowledge_store = KnowledgeStore::get_async_instance(knowledge_base_dir).await; let mut store = async_knowledge_store.lock().await; match store.cancel_operation(operation_id).await { diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index ad6f4c25f..23fb8ba88 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -1881,7 +1881,10 @@ impl ChatSession { ev.is_accepted = true; }); - let invoke_result = tool.tool.invoke(os, &mut self.stdout).await; + let invoke_result = tool + .tool + .invoke(os, &mut self.stdout, self.conversation.agents.get_active()) + .await; if self.spinner.is_some() { queue!( diff --git a/crates/chat-cli/src/cli/chat/tools/knowledge.rs b/crates/chat-cli/src/cli/chat/tools/knowledge.rs index dca258fb8..712f73c8f 100644 --- a/crates/chat-cli/src/cli/chat/tools/knowledge.rs +++ b/crates/chat-cli/src/cli/chat/tools/knowledge.rs @@ -1,4 +1,5 @@ use std::io::Write; +use std::path::Path; use crossterm::queue; use crossterm::style::{ @@ -7,12 +8,17 @@ use crossterm::style::{ }; use eyre::Result; use serde::Deserialize; -use tracing::warn; +use tracing::{ + error, + warn, +}; use super::{ InvokeOutput, OutputKind, }; +use crate::cli::agent::Agent; +use crate::cli::chat::cli::knowledge::Settings; use crate::database::settings::Setting; use crate::os::Os; use crate::util::knowledge_store::KnowledgeStore; @@ -305,9 +311,9 @@ impl Knowledge { Ok(()) } - pub async fn invoke(&self, os: &Os, _updates: &mut impl Write) -> Result { + pub async fn invoke(&self, os: &Os, _updates: &mut impl Write, agent: Option<&Agent>) -> Result { // Get the async knowledge store singleton - let async_knowledge_store = KnowledgeStore::get_async_instance().await; + let async_knowledge_store = KnowledgeStore::get_async_instance(Self::get_knowledge_base_dir(agent)).await; let mut store = async_knowledge_store.lock().await; let result = match self { @@ -542,4 +548,22 @@ impl Knowledge { ) } } + + fn get_knowledge_base_dir(agent: Option<&Agent>) -> Option + use<>> { + if agent.is_none() { + return None; + } + + agent.unwrap().tools_settings.get("knowledge")?.get("base_dir"); + match agent.unwrap().tools_settings.get("knowledge") { + Some(settings) => match serde_json::from_value::(settings.clone()) { + Ok(settings) => Some(settings.base_dir), + Err(e) => { + error!("Failed to deserialize tool settings for execute_bash: {:?}", e); + None + }, + }, + None => None, + } + } } diff --git a/crates/chat-cli/src/cli/chat/tools/mod.rs b/crates/chat-cli/src/cli/chat/tools/mod.rs index b640fd4f1..1eda60259 100644 --- a/crates/chat-cli/src/cli/chat/tools/mod.rs +++ b/crates/chat-cli/src/cli/chat/tools/mod.rs @@ -107,7 +107,7 @@ impl Tool { } /// Invokes the tool asynchronously - pub async fn invoke(&self, os: &Os, stdout: &mut impl Write) -> Result { + pub async fn invoke(&self, os: &Os, stdout: &mut impl Write, agent: Option<&Agent>) -> Result { match self { Tool::FsRead(fs_read) => fs_read.invoke(os, stdout).await, Tool::FsWrite(fs_write) => fs_write.invoke(os, stdout).await, @@ -115,7 +115,7 @@ impl Tool { Tool::UseAws(use_aws) => use_aws.invoke(os, stdout).await, Tool::Custom(custom_tool) => custom_tool.invoke(os, stdout).await, Tool::GhIssue(gh_issue) => gh_issue.invoke(os, stdout).await, - Tool::Knowledge(knowledge) => knowledge.invoke(os, stdout).await, + Tool::Knowledge(knowledge) => knowledge.invoke(os, stdout, agent).await, Tool::Thinking(think) => think.invoke(stdout).await, } } diff --git a/crates/chat-cli/src/util/knowledge_store.rs b/crates/chat-cli/src/util/knowledge_store.rs index 23ef1345e..bcdeec466 100644 --- a/crates/chat-cli/src/util/knowledge_store.rs +++ b/crates/chat-cli/src/util/knowledge_store.rs @@ -1,3 +1,4 @@ +use std::path::Path; use std::sync::{ Arc, LazyLock as Lazy, @@ -32,13 +33,13 @@ pub struct KnowledgeStore { impl KnowledgeStore { /// Get singleton instance - pub async fn get_async_instance() -> Arc> { + pub async fn get_async_instance(path: Option>) -> Arc> { static ASYNC_INSTANCE: Lazy>>> = Lazy::new(tokio::sync::OnceCell::new); if cfg!(test) { Arc::new(Mutex::new( - KnowledgeStore::new() + KnowledgeStore::new(path) .await .expect("Failed to create test async knowledge store"), )) @@ -46,7 +47,7 @@ impl KnowledgeStore { ASYNC_INSTANCE .get_or_init(|| async { Arc::new(Mutex::new( - KnowledgeStore::new() + KnowledgeStore::new(path) .await .expect("Failed to create async knowledge store"), )) @@ -56,12 +57,21 @@ impl KnowledgeStore { } } - pub async fn new() -> Result { - let client = AsyncSemanticSearchClient::new_with_default_dir() - .await - .map_err(|e| eyre::eyre!("Failed to create client: {}", e))?; - - Ok(Self { client }) + pub async fn new(path: Option>) -> Result { + match path { + Some(path) => { + let client = AsyncSemanticSearchClient::new(path) + .await + .map_err(|e| eyre::eyre!("Failed to create client: {}", e))?; + Ok(Self { client }) + }, + None => { + let client = AsyncSemanticSearchClient::new_with_default_dir() + .await + .map_err(|e| eyre::eyre!("Failed to create client: {}", e))?; + Ok(Self { client }) + }, + } } /// Add context - delegates to async client