From f05bdb19d53949d09afa37d708c3db593a2dcb36 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Wed, 18 Jun 2025 11:25:40 -0700 Subject: [PATCH 1/8] fix disable --- crates/goose/src/agents/agent.rs | 59 +++++++++++++++----------------- ui/desktop/openapi.json | 2 +- 2 files changed, 29 insertions(+), 32 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 9b7a8394e1ba..214e704c1adc 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -304,6 +304,31 @@ impl Agent { ) -> (String, Result, ToolError>) { let mut extension_manager = self.extension_manager.lock().await; + let selector = self.router_tool_selector.lock().await.clone(); + if ToolRouterIndexManager::is_tool_router_enabled(&selector) { + if let Some(selector) = selector { + let selector_action = if action == "disable" { "remove" } else { "add" }; + let extension_manager = self.extension_manager.lock().await; + let selector = Arc::new(selector); + if let Err(e) = ToolRouterIndexManager::update_extension_tools( + &selector, + &extension_manager, + &extension_name, + selector_action, + ) + .await + { + return ( + request_id, + Err(ToolError::ExecutionError(format!( + "Failed to update vector index: {}", + e + ))), + ); + } + } + } + if action == "disable" { let result = extension_manager .remove_extension(&extension_name) @@ -351,34 +376,6 @@ impl Agent { }) .map_err(|e| ToolError::ExecutionError(e.to_string())); - // Update vector index if operation was successful and vector routing is enabled - if result.is_ok() { - let selector = self.router_tool_selector.lock().await.clone(); - if ToolRouterIndexManager::is_tool_router_enabled(&selector) { - if let Some(selector) = selector { - let vector_action = if action == "disable" { "remove" } else { "add" }; - let extension_manager = self.extension_manager.lock().await; - let selector = Arc::new(selector); - if let Err(e) = ToolRouterIndexManager::update_extension_tools( - &selector, - &extension_manager, - &extension_name, - vector_action, - ) - .await - { - return ( - request_id, - Err(ToolError::ExecutionError(format!( - "Failed to update vector index: {}", - e - ))), - ); - } - } - } - } - (request_id, result) } @@ -503,9 +500,6 @@ impl Agent { } pub async fn remove_extension(&self, name: &str) -> Result<()> { - let mut extension_manager = self.extension_manager.lock().await; - extension_manager.remove_extension(name).await?; - // If vector tool selection is enabled, remove tools from the index let selector = self.router_tool_selector.lock().await.clone(); if ToolRouterIndexManager::is_tool_router_enabled(&selector) { @@ -521,6 +515,9 @@ impl Agent { } } + let mut extension_manager = self.extension_manager.lock().await; + extension_manager.remove_extension(name).await?; + Ok(()) } diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index cae672126d9d..256e61e4ae02 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -10,7 +10,7 @@ "license": { "name": "Apache-2.0" }, - "version": "1.0.28" + "version": "1.0.29" }, "paths": { "/agent/tools": { From b317321cb51dca45ca05dc4d479c527f1fd1d0c9 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Wed, 18 Jun 2025 11:58:07 -0700 Subject: [PATCH 2/8] fix --- crates/goose-server/src/routes/agent.rs | 33 ++++++++++ crates/goose/src/agents/agent.rs | 61 +++++++++++++++++++ .../ToolSelectionStrategySection.tsx | 12 ++++ 3 files changed, 106 insertions(+) diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index 354c22c639bf..8d6611c991da 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/crates/goose-server/src/routes/agent.rs @@ -217,6 +217,35 @@ async fn update_agent_provider( Ok(StatusCode::OK) } +#[utoipa::path( + post, + path = "/agent/update_tool_selection_strategy", + responses( + (status = 200, description = "Tool selection strategy updated successfully", body = String), + (status = 500, description = "Internal server error") + ) +)] +async fn update_tool_selection_strategy( + State(state): State>, + headers: HeaderMap, +) -> Result, StatusCode> { + verify_secret_key(&headers, &state)?; + + let agent = state + .get_agent() + .await + .map_err(|_| StatusCode::PRECONDITION_FAILED)?; + + agent.update_tool_selection_strategy().await.map_err(|e| { + tracing::error!("Failed to update tool selection strategy: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + Ok(Json( + "Tool selection strategy updated successfully".to_string(), + )) +} + pub fn routes(state: Arc) -> Router { Router::new() .route("/agent/versions", get(get_versions)) @@ -224,5 +253,9 @@ pub fn routes(state: Arc) -> Router { .route("/agent/prompt", post(extend_prompt)) .route("/agent/tools", get(get_tools)) .route("/agent/update_provider", post(update_agent_provider)) + .route( + "/agent/update_tool_selection_strategy", + post(update_tool_selection_strategy), + ) .with_state(state) } diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 214e704c1adc..0785d0475cd9 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -810,6 +810,67 @@ impl Agent { Ok(()) } + /// Update the tool selection strategy and re-index all tools + pub async fn update_tool_selection_strategy(&self) -> Result<()> { + let provider = self.provider().await?; + let extension_manager = self.extension_manager.lock().await; + + // Create new selector with current strategy + let config = Config::global(); + let router_tool_selection_strategy = config + .get_param("GOOSE_ROUTER_TOOL_SELECTION_STRATEGY") + .unwrap_or_else(|_| "default".to_string()); + + let strategy = match router_tool_selection_strategy.to_lowercase().as_str() { + "vector" => Some(RouterToolSelectionStrategy::Vector), + "llm" => Some(RouterToolSelectionStrategy::Llm), + _ => None, + }; + + let selector = match strategy { + Some(RouterToolSelectionStrategy::Vector) => { + let table_name = generate_table_id(); + let selector = create_tool_selector(strategy, provider.clone(), Some(table_name)) + .await + .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; + Arc::new(selector) + } + Some(RouterToolSelectionStrategy::Llm) => { + let selector = create_tool_selector(strategy, provider.clone(), None) + .await + .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; + Arc::new(selector) + } + None => return Ok(()), + }; + + // First index platform tools + ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await?; + + // Then index all currently enabled extensions + let enabled_extensions = extension_manager.list_extensions().await?; + for extension_name in enabled_extensions { + if let Err(e) = ToolRouterIndexManager::update_extension_tools( + &selector, + &extension_manager, + &extension_name, + "add", + ) + .await + { + tracing::error!( + "Failed to index tools for extension {}: {}", + extension_name, + e + ); + } + } + + // Update the selector + *self.router_tool_selector.lock().await = Some(selector.clone()); + Ok(()) + } + async fn update_router_tool_selector(&self, provider: Arc) -> Result<()> { let config = Config::global(); let router_tool_selection_strategy = config diff --git a/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx b/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx index 3851805af1e0..233ebe1d29e1 100644 --- a/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx +++ b/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx @@ -35,6 +35,18 @@ export const ToolSelectionStrategySection = ({ try { await upsert('GOOSE_ROUTER_TOOL_SELECTION_STRATEGY', newStrategy, false); setCurrentStrategy(newStrategy); + + // Call backend to update tool selection strategy and re-index tools + const response = await fetch('/api/agent/update_tool_selection_strategy', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + }); + + if (!response.ok) { + throw new Error('Failed to update tool selection strategy in backend'); + } } catch (error) { console.error('Error updating tool selection strategy:', error); throw new Error(`Failed to store new tool selection strategy: ${newStrategy}`); From 4c7b0bfa7fc827fb92aebac0897d7372163e4234 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Wed, 18 Jun 2025 12:10:31 -0700 Subject: [PATCH 3/8] update --- crates/goose-server/src/routes/agent.rs | 24 +++++- crates/goose/src/agents/agent.rs | 82 ++++++++++++------- .../ToolSelectionStrategySection.tsx | 52 ++++++++---- 3 files changed, 109 insertions(+), 49 deletions(-) diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index 8d6611c991da..502a86e81fc2 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/crates/goose-server/src/routes/agent.rs @@ -67,6 +67,11 @@ pub struct GetToolsQuery { extension_name: Option, } +#[derive(Serialize)] +struct ErrorResponse { + error: String, +} + async fn get_versions() -> Json { let versions = ["goose".to_string()]; let default_version = "goose".to_string(); @@ -228,17 +233,28 @@ async fn update_agent_provider( async fn update_tool_selection_strategy( State(state): State>, headers: HeaderMap, -) -> Result, StatusCode> { - verify_secret_key(&headers, &state)?; +) -> Result, Json> { + verify_secret_key(&headers, &state).map_err(|_| { + Json(ErrorResponse { + error: "Unauthorized - Invalid or missing API key".to_string(), + }) + })?; let agent = state .get_agent() .await - .map_err(|_| StatusCode::PRECONDITION_FAILED)?; + .map_err(|e| { + tracing::error!("Failed to get agent: {}", e); + Json(ErrorResponse { + error: format!("Failed to get agent: {}", e), + }) + })?; agent.update_tool_selection_strategy().await.map_err(|e| { tracing::error!("Failed to update tool selection strategy: {}", e); - StatusCode::INTERNAL_SERVER_ERROR + Json(ErrorResponse { + error: format!("Failed to update tool selection strategy: {}", e), + }) })?; Ok(Json( diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 0785d0475cd9..8dd48513e449 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -812,7 +812,9 @@ impl Agent { /// Update the tool selection strategy and re-index all tools pub async fn update_tool_selection_strategy(&self) -> Result<()> { - let provider = self.provider().await?; + let provider = self.provider().await.map_err(|e| { + anyhow!("Failed to get provider: {}", e) + })?; let extension_manager = self.extension_manager.lock().await; // Create new selector with current strategy @@ -827,47 +829,65 @@ impl Agent { _ => None, }; + // If no strategy is selected, clear the selector and return + if strategy.is_none() { + *self.router_tool_selector.lock().await = None; + return Ok(()); + } + + // Try to create the selector, but if it fails, fall back to no selector let selector = match strategy { Some(RouterToolSelectionStrategy::Vector) => { - let table_name = generate_table_id(); - let selector = create_tool_selector(strategy, provider.clone(), Some(table_name)) - .await - .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; - Arc::new(selector) + match create_tool_selector(strategy, provider.clone(), Some(generate_table_id())).await { + Ok(selector) => Some(Arc::new(selector)), + Err(e) => { + tracing::error!("Failed to create vector tool selector: {}", e); + None + } + } } Some(RouterToolSelectionStrategy::Llm) => { - let selector = create_tool_selector(strategy, provider.clone(), None) - .await - .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; - Arc::new(selector) + match create_tool_selector(strategy, provider.clone(), None).await { + Ok(selector) => Some(Arc::new(selector)), + Err(e) => { + tracing::error!("Failed to create LLM tool selector: {}", e); + None + } + } } - None => return Ok(()), + None => None, }; - // First index platform tools - ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await?; + // If we have a selector, try to index tools + if let Some(selector) = &selector { + // First index platform tools + if let Err(e) = ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await { + tracing::error!("Failed to index platform tools: {}", e); + } - // Then index all currently enabled extensions - let enabled_extensions = extension_manager.list_extensions().await?; - for extension_name in enabled_extensions { - if let Err(e) = ToolRouterIndexManager::update_extension_tools( - &selector, - &extension_manager, - &extension_name, - "add", - ) - .await - { - tracing::error!( - "Failed to index tools for extension {}: {}", - extension_name, - e - ); + // Then index all currently enabled extensions + if let Ok(enabled_extensions) = extension_manager.list_extensions().await { + for extension_name in enabled_extensions { + if let Err(e) = ToolRouterIndexManager::update_extension_tools( + &selector, + &extension_manager, + &extension_name, + "add", + ) + .await + { + tracing::error!( + "Failed to index tools for extension {}: {}", + extension_name, + e + ); + } + } } } - // Update the selector - *self.router_tool_selector.lock().await = Some(selector.clone()); + // Update the selector (even if it's None) + *self.router_tool_selector.lock().await = selector; Ok(()) } diff --git a/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx b/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx index 233ebe1d29e1..30e7795c2105 100644 --- a/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx +++ b/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx @@ -29,27 +29,45 @@ export const ToolSelectionStrategySection = ({ setView: _setView, }: ToolSelectionStrategySectionProps) => { const [currentStrategy, setCurrentStrategy] = useState('default'); + const [error, setError] = useState(null); const { read, upsert } = useConfig(); const handleStrategyChange = async (newStrategy: string) => { + setError(null); // Clear any previous errors try { - await upsert('GOOSE_ROUTER_TOOL_SELECTION_STRATEGY', newStrategy, false); - setCurrentStrategy(newStrategy); - - // Call backend to update tool selection strategy and re-index tools - const response = await fetch('/api/agent/update_tool_selection_strategy', { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - }); - - if (!response.ok) { - throw new Error('Failed to update tool selection strategy in backend'); + // First update the configuration + try { + await upsert('GOOSE_ROUTER_TOOL_SELECTION_STRATEGY', newStrategy, false); + } catch (error) { + console.error('Error updating configuration:', error); + setError(`Failed to update configuration: ${error}`); + return; + } + + // Then update the backend + try { + const response = await fetch('/api/agent/update_tool_selection_strategy', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + throw new Error(errorData.error || 'Unknown error from backend'); + } + } catch (error) { + console.error('Error updating backend:', error); + setError(`Failed to update backend: ${error}`); + return; } + + // If both succeeded, update the UI + setCurrentStrategy(newStrategy); } catch (error) { console.error('Error updating tool selection strategy:', error); - throw new Error(`Failed to store new tool selection strategy: ${newStrategy}`); + setError(`Failed to update tool selection strategy: ${error}`); } }; @@ -61,6 +79,7 @@ export const ToolSelectionStrategySection = ({ } } catch (error) { console.error('Error fetching current tool selection strategy:', error); + setError(`Failed to fetch current strategy: ${error}`); } }, [read]); @@ -78,6 +97,11 @@ export const ToolSelectionStrategySection = ({ Configure how Goose selects tools for your requests. Recommended when many extensions are enabled. Available only with Claude models served on Databricks for now.

+ {error && ( +
+ {error} +
+ )}
{all_tool_selection_strategies.map((strategy) => (
From ba0343a5cf69667df4fe31d7702cb3271ec0bc51 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Wed, 18 Jun 2025 13:26:20 -0700 Subject: [PATCH 4/8] dedupe index_tools --- crates/goose/src/agents/agent.rs | 82 +++++++------------ .../goose/src/agents/router_tool_selector.rs | 39 +++++++-- 2 files changed, 61 insertions(+), 60 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 8dd48513e449..0785d0475cd9 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -812,9 +812,7 @@ impl Agent { /// Update the tool selection strategy and re-index all tools pub async fn update_tool_selection_strategy(&self) -> Result<()> { - let provider = self.provider().await.map_err(|e| { - anyhow!("Failed to get provider: {}", e) - })?; + let provider = self.provider().await?; let extension_manager = self.extension_manager.lock().await; // Create new selector with current strategy @@ -829,65 +827,47 @@ impl Agent { _ => None, }; - // If no strategy is selected, clear the selector and return - if strategy.is_none() { - *self.router_tool_selector.lock().await = None; - return Ok(()); - } - - // Try to create the selector, but if it fails, fall back to no selector let selector = match strategy { Some(RouterToolSelectionStrategy::Vector) => { - match create_tool_selector(strategy, provider.clone(), Some(generate_table_id())).await { - Ok(selector) => Some(Arc::new(selector)), - Err(e) => { - tracing::error!("Failed to create vector tool selector: {}", e); - None - } - } + let table_name = generate_table_id(); + let selector = create_tool_selector(strategy, provider.clone(), Some(table_name)) + .await + .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; + Arc::new(selector) } Some(RouterToolSelectionStrategy::Llm) => { - match create_tool_selector(strategy, provider.clone(), None).await { - Ok(selector) => Some(Arc::new(selector)), - Err(e) => { - tracing::error!("Failed to create LLM tool selector: {}", e); - None - } - } + let selector = create_tool_selector(strategy, provider.clone(), None) + .await + .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; + Arc::new(selector) } - None => None, + None => return Ok(()), }; - // If we have a selector, try to index tools - if let Some(selector) = &selector { - // First index platform tools - if let Err(e) = ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await { - tracing::error!("Failed to index platform tools: {}", e); - } + // First index platform tools + ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await?; - // Then index all currently enabled extensions - if let Ok(enabled_extensions) = extension_manager.list_extensions().await { - for extension_name in enabled_extensions { - if let Err(e) = ToolRouterIndexManager::update_extension_tools( - &selector, - &extension_manager, - &extension_name, - "add", - ) - .await - { - tracing::error!( - "Failed to index tools for extension {}: {}", - extension_name, - e - ); - } - } + // Then index all currently enabled extensions + let enabled_extensions = extension_manager.list_extensions().await?; + for extension_name in enabled_extensions { + if let Err(e) = ToolRouterIndexManager::update_extension_tools( + &selector, + &extension_manager, + &extension_name, + "add", + ) + .await + { + tracing::error!( + "Failed to index tools for extension {}: {}", + extension_name, + e + ); } } - // Update the selector (even if it's None) - *self.router_tool_selector.lock().await = selector; + // Update the selector + *self.router_tool_selector.lock().await = Some(selector.clone()); Ok(()) } diff --git a/crates/goose/src/agents/router_tool_selector.rs b/crates/goose/src/agents/router_tool_selector.rs index d6d1f55bc1bb..2e989839af4c 100644 --- a/crates/goose/src/agents/router_tool_selector.rs +++ b/crates/goose/src/agents/router_tool_selector.rs @@ -166,12 +166,31 @@ impl RouterToolSelector for VectorToolSelector { }) .collect(); - // Index all tools at once + // Get vector_db lock let vector_db = self.vector_db.read().await; - vector_db - .index_tools(tool_records) - .await - .map_err(|e| ToolError::ExecutionError(format!("Failed to index tools: {}", e)))?; + + // Filter out tools that already exist in the database + let mut new_tool_records = Vec::new(); + for record in tool_records { + // Check if tool exists by searching for it + let existing_tools = vector_db + .search_tools(record.vector.clone(), 1, Some(&record.extension_name)) + .await + .map_err(|e| ToolError::ExecutionError(format!("Failed to search for existing tools: {}", e)))?; + + // Only add if no exact match found + if !existing_tools.iter().any(|t| t.tool_name == record.tool_name) { + new_tool_records.push(record); + } + } + + // Only index if there are new tools to add + if !new_tool_records.is_empty() { + vector_db + .index_tools(new_tool_records) + .await + .map_err(|e| ToolError::ExecutionError(format!("Failed to index tools: {}", e)))?; + } Ok(()) } @@ -282,7 +301,7 @@ impl RouterToolSelector for LLMToolSelector { } } - async fn index_tools(&self, tools: &[Tool], _extension_name: &str) -> Result<(), ToolError> { + async fn index_tools(&self, tools: &[Tool], extension_name: &str) -> Result<(), ToolError> { let mut tool_strings = self.tool_strings.write().await; for tool in tools { @@ -294,8 +313,11 @@ impl RouterToolSelector for LLMToolSelector { .unwrap_or_else(|_| "{}".to_string()) ); - if let Some(extension_name) = tool.name.split("__").next() { - let entry = tool_strings.entry(extension_name.to_string()).or_default(); + // Use the provided extension_name instead of parsing from tool name + let entry = tool_strings.entry(extension_name.to_string()).or_default(); + + // Check if this tool already exists in the entry + if !entry.contains(&format!("Tool: {}", tool.name)) { if !entry.is_empty() { entry.push_str("\n\n"); } @@ -305,7 +327,6 @@ impl RouterToolSelector for LLMToolSelector { Ok(()) } - async fn remove_tool(&self, tool_name: &str) -> Result<(), ToolError> { let mut tool_strings = self.tool_strings.write().await; if let Some(extension_name) = tool_name.split("__").next() { From 36ec089556682cff9deff74bb9b01a5e5e85527c Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Wed, 18 Jun 2025 16:29:40 -0700 Subject: [PATCH 5/8] working version --- crates/goose/src/agents/agent.rs | 30 +++++++++++++++++-- crates/goose/src/agents/router_tools.rs | 17 +++++++++-- .../ToolSelectionStrategySection.tsx | 19 ++++++++++-- 3 files changed, 58 insertions(+), 8 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 0785d0475cd9..a79984d8a7ba 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -873,6 +873,8 @@ impl Agent { async fn update_router_tool_selector(&self, provider: Arc) -> Result<()> { let config = Config::global(); + let extension_manager = self.extension_manager.lock().await; + let router_tool_selection_strategy = config .get_param("GOOSE_ROUTER_TOOL_SELECTION_STRATEGY") .unwrap_or_else(|_| "default".to_string()); @@ -886,21 +888,43 @@ impl Agent { let selector = match strategy { Some(RouterToolSelectionStrategy::Vector) => { let table_name = generate_table_id(); - let selector = create_tool_selector(strategy, provider, Some(table_name)) + let selector = create_tool_selector(strategy, provider.clone(), Some(table_name)) .await .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; Arc::new(selector) } Some(RouterToolSelectionStrategy::Llm) => { - let selector = create_tool_selector(strategy, provider, None) + let selector = create_tool_selector(strategy, provider.clone(), None) .await .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; Arc::new(selector) } None => return Ok(()), }; - let extension_manager = self.extension_manager.lock().await; + + // First index platform tools ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await?; + + // Then index all currently enabled extensions + let enabled_extensions = extension_manager.list_extensions().await?; + for extension_name in enabled_extensions { + if let Err(e) = ToolRouterIndexManager::update_extension_tools( + &selector, + &extension_manager, + &extension_name, + "add", + ) + .await + { + tracing::error!( + "Failed to index tools for extension {}: {}", + extension_name, + e + ); + } + } + + // Update the selector *self.router_tool_selector.lock().await = Some(selector.clone()); Ok(()) } diff --git a/crates/goose/src/agents/router_tools.rs b/crates/goose/src/agents/router_tools.rs index c5c45cccd2a7..bb3b2ad0e06a 100644 --- a/crates/goose/src/agents/router_tools.rs +++ b/crates/goose/src/agents/router_tools.rs @@ -44,13 +44,26 @@ pub fn vector_search_tool() -> Tool { } pub fn vector_search_tool_prompt() -> String { - r#"# Tool Selection Instructions + format!( + r#"# Tool Selection Instructions Important: the user has opted to dynamically enable tools, so although an extension could be enabled, \ please invoke the vector search tool to actually retrieve the most relevant tools to use according to the user's messages. For example, if the user has 3 extensions enabled, but they are asking for a tool to read a pdf file, \ you would invoke the vector_search tool to find the most relevant read pdf tool. By dynamically enabling tools, you (Goose) as the agent save context window space and allow the user to dynamically retrieve the most relevant tools. - Be sure to format the query to search rather than pass in the user's messages directly."#.to_string() + Be sure to format the query to search rather than pass in the user's messages directly. + In addition to the extension names available to you, you also have platform extension tools available to you. + The platform extension contains the following tools: + - {} + - {} + - {} + - {} + "#, + PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME, + PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME, + PLATFORM_READ_RESOURCE_TOOL_NAME, + PLATFORM_LIST_RESOURCES_TOOL_NAME + ) } pub fn llm_search_tool() -> Tool { diff --git a/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx b/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx index 30e7795c2105..806b68842d89 100644 --- a/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx +++ b/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx @@ -1,6 +1,7 @@ import { useEffect, useState, useCallback } from 'react'; import { View, ViewOptions } from '../../../App'; import { useConfig } from '../../ConfigContext'; +import { getApiUrl, getSecretKey } from '../../../config'; interface ToolSelectionStrategySectionProps { setView: (view: View, viewOptions?: ViewOptions) => void; @@ -46,17 +47,29 @@ export const ToolSelectionStrategySection = ({ // Then update the backend try { - const response = await fetch('/api/agent/update_tool_selection_strategy', { + const response = await fetch(getApiUrl('/agent/update_tool_selection_strategy'), { method: 'POST', headers: { 'Content-Type': 'application/json', + 'X-Secret-Key': getSecretKey(), }, }); - + + console.log('response', response); if (!response.ok) { - const errorData = await response.json(); + const errorData = await response + .json() + .catch(() => ({ error: 'Unknown error from backend' })); throw new Error(errorData.error || 'Unknown error from backend'); } + + // Parse the success response + const data = await response + .json() + .catch(() => ({ message: 'Tool selection strategy updated successfully' })); + if (data.error) { + throw new Error(data.error); + } } catch (error) { console.error('Error updating backend:', error); setError(`Failed to update backend: ${error}`); From 7335e6c20ad5623be2c38509df8ed7f160acf33d Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Fri, 20 Jun 2025 15:13:12 -0700 Subject: [PATCH 6/8] update index when strategy changes --- crates/goose-server/src/routes/agent.rs | 10 +- crates/goose/src/agents/agent.rs | 102 ++++-------------- .../ToolSelectionStrategySection.tsx | 34 ++++-- 3 files changed, 54 insertions(+), 92 deletions(-) diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index 502a86e81fc2..e5bff8671959 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/crates/goose-server/src/routes/agent.rs @@ -224,13 +224,13 @@ async fn update_agent_provider( #[utoipa::path( post, - path = "/agent/update_tool_selection_strategy", + path = "/agent/update_router_tool_selector", responses( (status = 200, description = "Tool selection strategy updated successfully", body = String), (status = 500, description = "Internal server error") ) )] -async fn update_tool_selection_strategy( +async fn update_router_tool_selector( State(state): State>, headers: HeaderMap, ) -> Result, Json> { @@ -250,7 +250,7 @@ async fn update_tool_selection_strategy( }) })?; - agent.update_tool_selection_strategy().await.map_err(|e| { + agent.update_router_tool_selector(None, Some(true)).await.map_err(|e| { tracing::error!("Failed to update tool selection strategy: {}", e); Json(ErrorResponse { error: format!("Failed to update tool selection strategy: {}", e), @@ -270,8 +270,8 @@ pub fn routes(state: Arc) -> Router { .route("/agent/tools", get(get_tools)) .route("/agent/update_provider", post(update_agent_provider)) .route( - "/agent/update_tool_selection_strategy", - post(update_tool_selection_strategy), + "/agent/update_router_tool_selector", + post(update_router_tool_selector), ) .with_state(state) } diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index a79984d8a7ba..9ffcfc0ce4b9 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -806,74 +806,17 @@ impl Agent { /// Update the provider used by this agent pub async fn update_provider(&self, provider: Arc) -> Result<()> { *self.provider.lock().await = Some(provider.clone()); - self.update_router_tool_selector(provider).await?; + self.update_router_tool_selector(Some(provider), None).await?; Ok(()) } - /// Update the tool selection strategy and re-index all tools - pub async fn update_tool_selection_strategy(&self) -> Result<()> { - let provider = self.provider().await?; - let extension_manager = self.extension_manager.lock().await; - - // Create new selector with current strategy - let config = Config::global(); - let router_tool_selection_strategy = config - .get_param("GOOSE_ROUTER_TOOL_SELECTION_STRATEGY") - .unwrap_or_else(|_| "default".to_string()); - - let strategy = match router_tool_selection_strategy.to_lowercase().as_str() { - "vector" => Some(RouterToolSelectionStrategy::Vector), - "llm" => Some(RouterToolSelectionStrategy::Llm), - _ => None, - }; - - let selector = match strategy { - Some(RouterToolSelectionStrategy::Vector) => { - let table_name = generate_table_id(); - let selector = create_tool_selector(strategy, provider.clone(), Some(table_name)) - .await - .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; - Arc::new(selector) - } - Some(RouterToolSelectionStrategy::Llm) => { - let selector = create_tool_selector(strategy, provider.clone(), None) - .await - .map_err(|e| anyhow!("Failed to create tool selector: {}", e))?; - Arc::new(selector) - } - None => return Ok(()), - }; - - // First index platform tools - ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await?; - - // Then index all currently enabled extensions - let enabled_extensions = extension_manager.list_extensions().await?; - for extension_name in enabled_extensions { - if let Err(e) = ToolRouterIndexManager::update_extension_tools( - &selector, - &extension_manager, - &extension_name, - "add", - ) - .await - { - tracing::error!( - "Failed to index tools for extension {}: {}", - extension_name, - e - ); - } - } - - // Update the selector - *self.router_tool_selector.lock().await = Some(selector.clone()); - Ok(()) - } - - async fn update_router_tool_selector(&self, provider: Arc) -> Result<()> { + pub async fn update_router_tool_selector(&self, provider: Option>, reindex_all: Option) -> Result<()> { let config = Config::global(); let extension_manager = self.extension_manager.lock().await; + let provider = match provider { + Some(p) => p, + None => self.provider().await?, + }; let router_tool_selection_strategy = config .get_param("GOOSE_ROUTER_TOOL_SELECTION_STRATEGY") @@ -905,22 +848,23 @@ impl Agent { // First index platform tools ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await?; - // Then index all currently enabled extensions - let enabled_extensions = extension_manager.list_extensions().await?; - for extension_name in enabled_extensions { - if let Err(e) = ToolRouterIndexManager::update_extension_tools( - &selector, - &extension_manager, - &extension_name, - "add", - ) - .await - { - tracing::error!( - "Failed to index tools for extension {}: {}", - extension_name, - e - ); + if reindex_all.unwrap_or(false) { + let enabled_extensions = extension_manager.list_extensions().await?; + for extension_name in enabled_extensions { + if let Err(e) = ToolRouterIndexManager::update_extension_tools( + &selector, + &extension_manager, + &extension_name, + "add", + ) + .await + { + tracing::error!( + "Failed to index tools for extension {}: {}", + extension_name, + e + ); + } } } diff --git a/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx b/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx index 806b68842d89..674ff02e411f 100644 --- a/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx +++ b/ui/desktop/src/components/settings/tool_selection_strategy/ToolSelectionStrategySection.tsx @@ -31,10 +31,15 @@ export const ToolSelectionStrategySection = ({ }: ToolSelectionStrategySectionProps) => { const [currentStrategy, setCurrentStrategy] = useState('default'); const [error, setError] = useState(null); + const [isLoading, setIsLoading] = useState(false); const { read, upsert } = useConfig(); const handleStrategyChange = async (newStrategy: string) => { + if (isLoading) return; // Prevent multiple simultaneous requests + setError(null); // Clear any previous errors + setIsLoading(true); + try { // First update the configuration try { @@ -42,12 +47,13 @@ export const ToolSelectionStrategySection = ({ } catch (error) { console.error('Error updating configuration:', error); setError(`Failed to update configuration: ${error}`); + setIsLoading(false); return; } // Then update the backend try { - const response = await fetch(getApiUrl('/agent/update_tool_selection_strategy'), { + const response = await fetch(getApiUrl('/agent/update_router_tool_selector'), { method: 'POST', headers: { 'Content-Type': 'application/json', @@ -55,7 +61,6 @@ export const ToolSelectionStrategySection = ({ }, }); - console.log('response', response); if (!response.ok) { const errorData = await response .json() @@ -73,6 +78,7 @@ export const ToolSelectionStrategySection = ({ } catch (error) { console.error('Error updating backend:', error); setError(`Failed to update backend: ${error}`); + setIsLoading(false); return; } @@ -81,6 +87,8 @@ export const ToolSelectionStrategySection = ({ } catch (error) { console.error('Error updating tool selection strategy:', error); setError(`Failed to update tool selection strategy: ${error}`); + } finally { + setIsLoading(false); } }; @@ -115,12 +123,21 @@ export const ToolSelectionStrategySection = ({ {error}
)} + {isLoading && ( +
+
+ Updating tool selection strategy... +
+ )}
{all_tool_selection_strategies.map((strategy) => ( -
+
handleStrategyChange(strategy.key)} + className={`flex items-center justify-between text-textStandard py-2 px-4 ${!isLoading ? 'hover:bg-bgSubtle' : ''}`} + onClick={() => !isLoading && handleStrategyChange(strategy.key)} >
@@ -135,14 +152,15 @@ export const ToolSelectionStrategySection = ({ name="tool-selection-strategy" value={strategy.key} checked={currentStrategy === strategy.key} - onChange={() => handleStrategyChange(strategy.key)} + onChange={() => !isLoading && handleStrategyChange(strategy.key)} + disabled={isLoading} className="peer sr-only" />
From a81025d9e193f7a3a5c27ee1d676f49b1fd9885b Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Fri, 20 Jun 2025 15:21:21 -0700 Subject: [PATCH 7/8] don't update openapi json --- ui/desktop/openapi.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index 256e61e4ae02..cae672126d9d 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -10,7 +10,7 @@ "license": { "name": "Apache-2.0" }, - "version": "1.0.29" + "version": "1.0.28" }, "paths": { "/agent/tools": { From 0e5d97e0d1cab762d4771278860ae10bb924853e Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Fri, 20 Jun 2025 15:24:45 -0700 Subject: [PATCH 8/8] fmt --- crates/goose-server/src/routes/agent.rs | 22 +++++++++---------- crates/goose/src/agents/agent.rs | 9 ++++++-- .../goose/src/agents/router_tool_selector.rs | 9 ++++++-- 3 files changed, 25 insertions(+), 15 deletions(-) diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index e5bff8671959..7dd154acdc70 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/crates/goose-server/src/routes/agent.rs @@ -240,23 +240,23 @@ async fn update_router_tool_selector( }) })?; - let agent = state - .get_agent() + let agent = state.get_agent().await.map_err(|e| { + tracing::error!("Failed to get agent: {}", e); + Json(ErrorResponse { + error: format!("Failed to get agent: {}", e), + }) + })?; + + agent + .update_router_tool_selector(None, Some(true)) .await .map_err(|e| { - tracing::error!("Failed to get agent: {}", e); + tracing::error!("Failed to update tool selection strategy: {}", e); Json(ErrorResponse { - error: format!("Failed to get agent: {}", e), + error: format!("Failed to update tool selection strategy: {}", e), }) })?; - agent.update_router_tool_selector(None, Some(true)).await.map_err(|e| { - tracing::error!("Failed to update tool selection strategy: {}", e); - Json(ErrorResponse { - error: format!("Failed to update tool selection strategy: {}", e), - }) - })?; - Ok(Json( "Tool selection strategy updated successfully".to_string(), )) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 9ffcfc0ce4b9..c7fc3f8f1398 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -806,11 +806,16 @@ impl Agent { /// Update the provider used by this agent pub async fn update_provider(&self, provider: Arc) -> Result<()> { *self.provider.lock().await = Some(provider.clone()); - self.update_router_tool_selector(Some(provider), None).await?; + self.update_router_tool_selector(Some(provider), None) + .await?; Ok(()) } - pub async fn update_router_tool_selector(&self, provider: Option>, reindex_all: Option) -> Result<()> { + pub async fn update_router_tool_selector( + &self, + provider: Option>, + reindex_all: Option, + ) -> Result<()> { let config = Config::global(); let extension_manager = self.extension_manager.lock().await; let provider = match provider { diff --git a/crates/goose/src/agents/router_tool_selector.rs b/crates/goose/src/agents/router_tool_selector.rs index 2e989839af4c..933316bcbdd1 100644 --- a/crates/goose/src/agents/router_tool_selector.rs +++ b/crates/goose/src/agents/router_tool_selector.rs @@ -176,10 +176,15 @@ impl RouterToolSelector for VectorToolSelector { let existing_tools = vector_db .search_tools(record.vector.clone(), 1, Some(&record.extension_name)) .await - .map_err(|e| ToolError::ExecutionError(format!("Failed to search for existing tools: {}", e)))?; + .map_err(|e| { + ToolError::ExecutionError(format!("Failed to search for existing tools: {}", e)) + })?; // Only add if no exact match found - if !existing_tools.iter().any(|t| t.tool_name == record.tool_name) { + if !existing_tools + .iter() + .any(|t| t.tool_name == record.tool_name) + { new_tool_records.push(record); } }