tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

remove Tool Selection Strategy preview (#6250)

Browse Source
This commit is contained in:
Angie Jones
2025-12-23 14:09:06 -06:00
committed by GitHub
parent 8e1eb1ebb6
commit 40246432dc
32 changed files with 10 additions and 1270 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/configure.rs
-35
View File
@@ -1154,11 +1154,6 @@ pub fn remove_extension_dialog() -> anyhow::Result<()> {
pub async fn configure_settings_dialog() -> anyhow::Result<()> {
let setting_type = cliclack::select("What setting would you like to configure?")
.item("goose_mode", "goose mode", "Configure goose mode")
.item(
"goose_router_strategy",
"Router Tool Selection Strategy",
"Experimental: configure a strategy for auto selecting tools to use",
)
.item(
"tool_permission",
"Tool Permission",
@@ -1197,9 +1192,6 @@ pub async fn configure_settings_dialog() -> anyhow::Result<()> {
"goose_mode" => {
configure_goose_mode_dialog()?;
}
"goose_router_strategy" => {
configure_goose_router_strategy_dialog()?;
}
"tool_permission" => {
configure_tool_permissions_dialog().await.and(Ok(()))?;
// No need to print config file path since it's already handled.
@@ -1271,33 +1263,6 @@ pub fn configure_goose_mode_dialog() -> anyhow::Result<()> {
Ok(())
}
pub fn configure_goose_router_strategy_dialog() -> anyhow::Result<()> {
let config = Config::global();
let enable_router = cliclack::select("Would you like to enable smart tool routing?")
.item(
true,
"Enable Router",
"Use LLM-based intelligence to select tools",
)
.item(
false,
"Disable Router",
"Use the default tool selection strategy",
)
.interact()?;
config.set_param("GOOSE_ENABLE_ROUTER", enable_router)?;
let msg = if enable_router {
"Router enabled - using LLM-based intelligence for tool selection"
} else {
"Router disabled - using default tool selection"
};
cliclack::outro(msg)?;
Ok(())
}
pub fn configure_tool_output_dialog() -> anyhow::Result<()> {
let config = Config::global();
crates/goose-cli/src/session/builder.rs
-2
View File
@@ -386,7 +386,6 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
.set_context(PlatformExtensionContext {
session_id: Some(session_id.clone()),
extension_manager: Some(Arc::downgrade(&agent.extension_manager)),
tool_route_manager: Some(Arc::downgrade(&agent.tool_route_manager)),
})
.await;
@@ -425,7 +424,6 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
// Extensions need to be added after the session is created because we change directory when resuming a session
// If we get extensions_override, only run those extensions and none other
let extensions_to_run: Vec<_> = if let Some(extensions) = session_config.extensions_override {
agent.disable_router_for_recipe().await;
extensions.into_iter().collect()
} else if session_config.resume {
match SessionManager::get_session(&session_id, false).await {
crates/goose-server/src/openapi.rs
-2
View File
@@ -361,7 +361,6 @@ derive_utoipa!(Icon as IconSchema);
super::routes::agent::agent_add_extension,
super::routes::agent::agent_remove_extension,
super::routes::agent::update_agent_provider,
super::routes::agent::update_router_tool_selector,
super::routes::action_required::confirm_tool_action,
super::routes::reply::reply,
super::routes::session::list_sessions,
@@ -525,7 +524,6 @@ derive_utoipa!(Icon as IconSchema);
super::routes::agent::ReadResourceResponse,
super::routes::agent::CallToolRequest,
super::routes::agent::CallToolResponse,
super::routes::agent::UpdateRouterToolSelectorRequest,
super::routes::agent::StartAgentRequest,
super::routes::agent::ResumeAgentRequest,
super::routes::agent::UpdateFromSessionRequest,
crates/goose-server/src/routes/agent.rs
+1 -39
View File
@@ -53,11 +53,6 @@ pub struct GetToolsQuery {
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct UpdateRouterToolSelectorRequest {
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct StartAgentRequest {
working_dir: String,
@@ -288,7 +283,7 @@ async fn resume_agent(
})?;
agent
.load_provider(provider)
.update_provider(provider, &payload.session_id)
.await
.map_err(|e| ErrorResponse {
message: format!("Could not configure agent: {}", e),
@@ -501,35 +496,6 @@ async fn update_agent_provider(
Ok(())
}
#[utoipa::path(
post,
path = "/agent/update_router_tool_selector",
request_body = UpdateRouterToolSelectorRequest,
responses(
(status = 200, description = "Tool selection strategy updated successfully", body = String),
(status = 401, description = "Unauthorized - invalid secret key"),
(status = 424, description = "Agent not initialized"),
(status = 500, description = "Internal server error")
)
)]
async fn update_router_tool_selector(
State(state): State<Arc<AppState>>,
Json(payload): Json<UpdateRouterToolSelectorRequest>,
) -> Result<Json<String>, StatusCode> {
let agent = state.get_agent_for_route(payload.session_id).await?;
agent
.update_router_tool_selector(None, Some(true))
.await
.map_err(|e| {
error!("Failed to update tool selection strategy: {}", e);
StatusCode::INTERNAL_SERVER_ERROR
})?;
Ok(Json(
"Tool selection strategy updated successfully".to_string(),
))
}
#[utoipa::path(
post,
path = "/agent/add_extension",
@@ -695,10 +661,6 @@ pub fn routes(state: Arc<AppState>) -> Router {
.route("/agent/read_resource", post(read_resource))
.route("/agent/call_tool", post(call_tool))
.route("/agent/update_provider", post(update_agent_provider))
.route(
"/agent/update_router_tool_selector",
post(update_router_tool_selector),
)
.route("/agent/update_from_session", post(update_from_session))
.route("/agent/add_extension", post(agent_add_extension))
.route("/agent/remove_extension", post(agent_remove_extension))
crates/goose/src/agents/agent.rs
+1 -93
View File
@@ -12,20 +12,17 @@ use super::final_output_tool::FinalOutputTool;
use super::platform_tools;
use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
use crate::action_required_manager::ActionRequiredManager;
use crate::agents::extension::{ExtensionConfig, ExtensionError, ExtensionResult, ToolInfo};
use crate::agents::extension::{ExtensionConfig, ExtensionResult, ToolInfo};
use crate::agents::extension_manager::{get_parameter_names, ExtensionManager};
use crate::agents::extension_manager_extension::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE;
use crate::agents::final_output_tool::{FINAL_OUTPUT_CONTINUATION_MESSAGE, FINAL_OUTPUT_TOOL_NAME};
use crate::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME;
use crate::agents::prompt_manager::PromptManager;
use crate::agents::retry::{RetryManager, RetryResult};
use crate::agents::router_tools::ROUTER_LLM_SEARCH_TOOL_NAME;
use crate::agents::subagent_task_config::TaskConfig;
use crate::agents::subagent_tool::{
create_subagent_tool, handle_subagent_tool, SUBAGENT_TOOL_NAME,
};
use crate::agents::tool_route_manager::ToolRouteManager;
use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
use crate::agents::types::SessionConfig;
use crate::agents::types::{FrontendTool, SharedProvider, ToolResultReceiver};
use crate::config::{get_enabled_extensions, Config, GooseMode};
@@ -95,7 +92,6 @@ pub struct Agent {
pub(super) tool_result_tx: mpsc::Sender<(String, ToolResult<CallToolResult>)>,
pub(super) tool_result_rx: ToolResultReceiver,
pub tool_route_manager: Arc<ToolRouteManager>,
pub(super) scheduler_service: Mutex<Option<Arc<dyn SchedulerTrait>>>,
pub(super) retry_manager: RetryManager,
pub(super) tool_inspection_manager: ToolInspectionManager,
@@ -169,7 +165,6 @@ impl Agent {
confirmation_rx: Mutex::new(confirm_rx),
tool_result_tx: tool_tx,
tool_result_rx: Arc::new(Mutex::new(tool_rx)),
tool_route_manager: Arc::new(ToolRouteManager::new()),
scheduler_service: Mutex::new(None),
retry_manager: RetryManager::new(),
tool_inspection_manager: Self::create_default_tool_inspection_manager(),
@@ -369,10 +364,6 @@ impl Agent {
*scheduler_service = Some(scheduler);
}
pub async fn disable_router_for_recipe(&self) {
self.tool_route_manager.disable_router_for_recipe().await;
}
/// Get a reference count clone to the provider
pub async fn provider(&self) -> Result<Arc<dyn Provider>, anyhow::Error> {
match &*self.provider.lock().await {
@@ -518,15 +509,6 @@ impl Agent {
"Frontend tool execution required".to_string(),
None,
)))
} else if tool_call.name == ROUTER_LLM_SEARCH_TOOL_NAME {
match self
.tool_route_manager
.dispatch_route_search_tool(tool_call.arguments.unwrap_or_default())
.await
{
Ok(tool_result) => tool_result,
Err(e) => return (request_id, Err(e)),
}
} else {
// Clone the result to ensure no references to extension_manager are returned
let result = self
@@ -617,28 +599,6 @@ impl Agent {
}
}
// If LLM tool selection is functional, index the tools
if self.tool_route_manager.is_router_functional().await {
let selector = self.tool_route_manager.get_router_tool_selector().await;
if let Some(selector) = selector {
let selector = Arc::new(selector);
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
&selector,
&self.extension_manager,
&extension.name(),
"add",
)
.await
{
return Err(ExtensionError::SetupError(format!(
"Failed to index tools for extension {}: {}",
extension.name(),
e
)));
}
}
}
Ok(())
}
@@ -702,29 +662,8 @@ impl Agent {
prefixed_tools
}
pub async fn list_tools_for_router(&self) -> Vec<Tool> {
self.tool_route_manager
.list_tools_for_router(&self.extension_manager)
.await
}
pub async fn remove_extension(&self, name: &str) -> Result<()> {
self.extension_manager.remove_extension(name).await?;
// If LLM tool selection is functional, remove tools from the index
if self.tool_route_manager.is_router_functional().await {
let selector = self.tool_route_manager.get_router_tool_selector().await;
if let Some(selector) = selector {
ToolRouterIndexManager::update_extension_tools(
&selector,
&self.extension_manager,
name,
"remove",
)
.await?;
}
}
Ok(())
}
@@ -1052,10 +991,6 @@ impl Agent {
remaining_requests,
filtered_response,
} = self.categorize_tools(&response, &tools).await;
let requests_to_record: Vec<ToolRequest> = frontend_requests.iter().chain(remaining_requests.iter()).cloned().collect();
self.tool_route_manager
.record_tool_requests(&requests_to_record)
.await;
yield AgentEvent::Message(filtered_response.clone());
tokio::task::yield_now().await;
@@ -1377,9 +1312,6 @@ impl Agent {
let mut current_provider = self.provider.lock().await;
*current_provider = Some(provider.clone());
self.update_router_tool_selector(Some(provider.clone()), None)
.await?;
SessionManager::update_session(session_id)
.provider_name(provider.get_name())
.model_config(provider.get_model_config())
@@ -1388,30 +1320,6 @@ impl Agent {
.context("Failed to persist provider config to session")
}
pub async fn load_provider(&self, provider: Arc<dyn Provider>) -> Result<()> {
let mut current_provider = self.provider.lock().await;
*current_provider = Some(provider.clone());
self.update_router_tool_selector(Some(provider.clone()), None)
.await
}
pub async fn update_router_tool_selector(
&self,
provider: Option<Arc<dyn Provider>>,
reindex_all: Option<bool>,
) -> Result<()> {
let provider = match provider {
Some(p) => p,
None => self.provider().await?,
};
// Delegate to ToolRouteManager
self.tool_route_manager
.update_router_tool_selector(provider, reindex_all, &self.extension_manager)
.await
}
/// Override the system prompt with a custom template
pub async fn override_system_prompt(&self, template: String) {
let mut prompt_manager = self.prompt_manager.lock().await;
crates/goose/src/agents/code_execution_extension.rs
-2
View File
@@ -923,7 +923,6 @@ mod tests {
let context = PlatformExtensionContext {
session_id: None,
extension_manager: None,
tool_route_manager: None,
};
let client = CodeExecutionClient::new(context).unwrap();
@@ -948,7 +947,6 @@ mod tests {
let context = PlatformExtensionContext {
session_id: None,
extension_manager: None,
tool_route_manager: None,
};
let client = CodeExecutionClient::new(context).unwrap();
crates/goose/src/agents/extension.rs
-2
View File
@@ -109,8 +109,6 @@ pub struct PlatformExtensionContext {
pub session_id: Option<String>,
pub extension_manager:
Option<std::sync::Weak<crate::agents::extension_manager::ExtensionManager>>,
pub tool_route_manager:
Option<std::sync::Weak<crate::agents::tool_route_manager::ToolRouteManager>>,
}
#[derive(Debug, Clone)]
crates/goose/src/agents/extension_manager.rs
-1
View File
@@ -264,7 +264,6 @@ impl ExtensionManager {
context: Mutex::new(PlatformExtensionContext {
session_id: None,
extension_manager: None,
tool_route_manager: None,
}),
provider,
}
crates/goose/src/agents/extension_manager_extension.rs
+3 -74
View File
@@ -1,6 +1,5 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
use crate::config::get_extension_by_name;
use anyhow::Result;
use async_trait::async_trait;
@@ -163,7 +162,6 @@ impl ExtensionManagerClient {
}
}
#[allow(clippy::too_many_lines)]
async fn manage_extensions_impl(
&self,
action: ManageExtensionAction,
@@ -182,43 +180,8 @@ impl ExtensionManagerClient {
)
})?;
let tool_route_manager = self
.context
.tool_route_manager
.as_ref()
.and_then(|weak| weak.upgrade());
// Update tool router index if router is functional
if let Some(tool_route_manager) = &tool_route_manager {
if tool_route_manager.is_router_functional().await {
let selector = tool_route_manager.get_router_tool_selector().await;
if let Some(selector) = selector {
let selector_action = if action == ManageExtensionAction::Disable {
"remove"
} else {
"add"
};
let selector = Arc::new(selector);
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
&selector,
&extension_manager,
&extension_name,
selector_action,
)
.await
{
return Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to update LLM index: {}", e),
None,
));
}
}
}
}
if action == ManageExtensionAction::Disable {
let result = extension_manager
return extension_manager
.remove_extension(&extension_name)
.await
.map(|_| {
@@ -228,7 +191,6 @@ impl ExtensionManagerClient {
))]
})
.map_err(|e| ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), None));
return result;
}
let config = match get_extension_by_name(&extension_name) {
@@ -245,7 +207,7 @@ impl ExtensionManagerClient {
}
};
let result = extension_manager
extension_manager
.add_extension(config)
.await
.map(|_| {
@@ -254,40 +216,7 @@ impl ExtensionManagerClient {
extension_name
))]
})
.map_err(|e| ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), None));
// Update LLM index if operation was successful and LLM routing is functional
if result.is_ok() {
if let Some(tool_route_manager) = &tool_route_manager {
if tool_route_manager.is_router_functional().await {
let selector = tool_route_manager.get_router_tool_selector().await;
if let Some(selector) = selector {
let llm_action = if action == ManageExtensionAction::Disable {
"remove"
} else {
"add"
};
let selector = Arc::new(selector);
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
&selector,
&extension_manager,
&extension_name,
llm_action,
)
.await
{
return Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to update LLM index: {}", e),
None,
));
}
}
}
}
}
result
.map_err(|e| ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), None))
}
async fn handle_list_resources(
crates/goose/src/agents/mod.rs
-4
View File
@@ -14,8 +14,6 @@ pub mod platform_tools;
pub mod prompt_manager;
mod reply_parts;
pub mod retry;
mod router_tool_selector;
mod router_tools;
mod schedule_tool;
pub(crate) mod skills_extension;
pub mod subagent_execution_tool;
@@ -24,8 +22,6 @@ mod subagent_task_config;
pub mod subagent_tool;
pub(crate) mod todo_extension;
mod tool_execution;
mod tool_route_manager;
mod tool_router_index_manager;
pub mod types;
pub use agent::{Agent, AgentEvent};
crates/goose/src/agents/prompt_manager.rs
-13
View File
@@ -6,7 +6,6 @@ use serde_json::Value;
use std::collections::HashMap;
use crate::agents::extension::ExtensionInfo;
use crate::agents::router_tools::llm_search_tool_prompt;
use crate::hints::load_hints::{load_hint_files, AGENTS_MD_FILENAME, GOOSE_HINTS_FILENAME};
use crate::{
config::{Config, GooseMode},
@@ -33,8 +32,6 @@ impl Default for PromptManager {
#[derive(Serialize)]
struct SystemPromptContext {
extensions: Vec<ExtensionInfo>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_selection_strategy: Option<String>,
current_date_time: String,
#[serde(skip_serializing_if = "Option::is_none")]
extension_tool_limits: Option<(usize, usize)>,
@@ -52,7 +49,6 @@ pub struct SystemPromptBuilder<'a, M> {
extensions_info: Vec<ExtensionInfo>,
frontend_instructions: Option<String>,
extension_tool_count: Option<(usize, usize)>,
router_enabled: bool,
subagents_enabled: bool,
hints: Option<String>,
code_execution_mode: bool,
@@ -85,11 +81,6 @@ impl<'a> SystemPromptBuilder<'a, PromptManager> {
self
}
pub fn with_router_enabled(mut self, enabled: bool) -> Self {
self.router_enabled = enabled;
self
}
pub fn with_code_execution_mode(mut self, enabled: bool) -> Self {
self.code_execution_mode = enabled;
self
@@ -158,7 +149,6 @@ impl<'a> SystemPromptBuilder<'a, PromptManager> {
let context = SystemPromptContext {
extensions: sanitized_extensions_info,
tool_selection_strategy: self.router_enabled.then(llm_search_tool_prompt),
current_date_time: self.manager.current_date_timestamp.clone(),
extension_tool_limits,
goose_mode,
@@ -247,7 +237,6 @@ impl PromptManager {
extensions_info: vec![],
frontend_instructions: None,
extension_tool_count: None,
router_enabled: false,
subagents_enabled: false,
hints: None,
code_execution_mode: false,
@@ -369,7 +358,6 @@ mod tests {
"how to use this extension",
true,
))
.with_router_enabled(true)
.build();
assert_snapshot!(system_prompt)
@@ -391,7 +379,6 @@ mod tests {
"<instructions on how to use extension B (no resources)>",
false,
))
.with_router_enabled(true)
.with_extension_and_tool_counts(MAX_EXTENSIONS + 1, MAX_TOOLS + 1)
.build();
crates/goose/src/agents/reply_parts.rs
+4 -19
View File
@@ -113,16 +113,8 @@ impl Agent {
&self,
working_dir: &std::path::Path,
) -> Result<(Vec<Tool>, Vec<Tool>, String)> {
// Get router enabled status
let router_enabled = self.tool_route_manager.is_router_enabled().await;
// Get tools from extension manager
let mut tools = self.list_tools_for_router().await;
// If router is disabled and no tools were returned, fall back to regular tools
if !router_enabled && tools.is_empty() {
tools = self.list_tools(None).await;
}
let mut tools = self.list_tools(None).await;
// Add frontend tools
let frontend_tools = self.frontend_tools.lock().await;
@@ -139,10 +131,8 @@ impl Agent {
tools.retain(|tool| tool.name.starts_with(&code_exec_prefix));
}
if !router_enabled {
// Stable tool ordering is important for multi session prompt caching.
tools.sort_by(|a, b| a.name.cmp(&b.name));
}
// Stable tool ordering is important for multi session prompt caching.
tools.sort_by(|a, b| a.name.cmp(&b.name));
// Prepare system prompt
let extensions_info = self.extension_manager.get_extensions_info().await;
@@ -159,7 +149,6 @@ impl Agent {
.with_extensions(extensions_info.into_iter())
.with_frontend_instructions(self.frontend_instructions.lock().await.clone())
.with_extension_and_tool_counts(extension_count, tool_count)
.with_router_enabled(router_enabled)
.with_code_execution_mode(code_execution_active)
.with_hints(working_dir)
.with_enable_subagents(self.subagents_enabled().await)
@@ -437,8 +426,7 @@ mod tests {
}
#[tokio::test]
async fn prepare_tools_sorts_when_router_disabled_and_includes_frontend_and_list_tools(
) -> anyhow::Result<()> {
async fn prepare_tools_sorts_and_includes_frontend_and_list_tools() -> anyhow::Result<()> {
let agent = crate::agents::Agent::new();
let session = SessionManager::create_session(
@@ -452,9 +440,6 @@ mod tests {
let provider = std::sync::Arc::new(MockProvider { model_config });
agent.update_provider(provider, &session.id).await?;
// Disable the router to trigger sorting
agent.disable_router_for_recipe().await;
// Add unsorted frontend tools
let frontend_tools = vec![
Tool::new(
crates/goose/src/agents/router_tool_selector.rs
-180
View File
@@ -1,180 +0,0 @@
use rmcp::model::{Content, ErrorCode, ErrorData};
use rmcp::model::{JsonObject, Tool};
use anyhow::Result;
use async_trait::async_trait;
use serde::Serialize;
use std::borrow::Cow;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::conversation::message::Message;
use crate::prompt_template::render_global_file;
use crate::providers::base::Provider;
#[derive(Serialize)]
struct ToolSelectorContext {
tools: String,
query: String,
}
#[async_trait]
pub trait RouterToolSelector: Send + Sync {
async fn select_tools(&self, params: JsonObject) -> Result<Vec<Content>, ErrorData>;
async fn index_tools(&self, tools: &[Tool], extension_name: &str) -> Result<(), ErrorData>;
async fn remove_tool(&self, tool_name: &str) -> Result<(), ErrorData>;
async fn record_tool_call(&self, tool_name: &str) -> Result<(), ErrorData>;
async fn get_recent_tool_calls(&self, limit: usize) -> Result<Vec<String>, ErrorData>;
}
pub struct LLMToolSelector {
llm_provider: Arc<dyn Provider>,
tool_strings: Arc<RwLock<HashMap<String, String>>>, // extension_name -> tool_string
recent_tool_calls: Arc<RwLock<VecDeque<String>>>,
}
impl LLMToolSelector {
pub async fn new(provider: Arc<dyn Provider>) -> Result<Self> {
Ok(Self {
llm_provider: provider.clone(),
tool_strings: Arc::new(RwLock::new(HashMap::new())),
recent_tool_calls: Arc::new(RwLock::new(VecDeque::with_capacity(100))),
})
}
}
#[async_trait]
impl RouterToolSelector for LLMToolSelector {
async fn select_tools(&self, params: JsonObject) -> Result<Vec<Content>, ErrorData> {
let query = params
.get("query")
.and_then(|v| v.as_str())
.ok_or_else(|| ErrorData {
code: ErrorCode::INVALID_PARAMS,
message: Cow::from("Missing 'query' parameter"),
data: None,
})?;
let extension_name = params
.get("extension_name")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
// Get relevant tool strings based on extension_name
let tool_strings = self.tool_strings.read().await;
let relevant_tools = if let Some(ext) = &extension_name {
tool_strings.get(ext).cloned()
} else {
// If no extension specified, use all tools
Some(
tool_strings
.values()
.cloned()
.collect::<Vec<String>>()
.join("\n"),
)
};
if let Some(tools) = relevant_tools {
// Use template to generate the prompt
let context = ToolSelectorContext {
tools: tools.clone(),
query: query.to_string(),
};
let user_prompt =
render_global_file("router_tool_selector.md", &context).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to render prompt template: {}", e)),
data: None,
})?;
let user_message = Message::user().with_text(&user_prompt);
let response = self
.llm_provider
.complete("system", &[user_message], &[])
.await
.map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to search tools: {}", e)),
data: None,
})?;
// Extract just the message content from the response
let (message, _usage) = response;
let text = message.content[0].as_text().unwrap_or_default();
// Split the response into individual tool entries
let tool_entries: Vec<Content> = text
.split("\n\n")
.filter(|entry| entry.trim().starts_with("Tool:"))
.map(|entry| Content::text(entry.trim().to_string()))
.collect();
Ok(tool_entries)
} else {
Ok(vec![])
}
}
async fn index_tools(&self, tools: &[Tool], extension_name: &str) -> Result<(), ErrorData> {
let mut tool_strings = self.tool_strings.write().await;
for tool in tools {
let tool_string = format!(
"Tool: {}\nDescription: {}\nSchema: {}",
tool.name,
tool.description
.as_ref()
.map(|d| d.as_ref())
.unwrap_or_default(),
serde_json::to_string_pretty(&tool.input_schema)
.unwrap_or_else(|_| "{}".to_string())
);
// 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");
}
entry.push_str(&tool_string);
}
}
Ok(())
}
async fn remove_tool(&self, tool_name: &str) -> Result<(), ErrorData> {
let mut tool_strings = self.tool_strings.write().await;
if let Some(extension_name) = tool_name.split("__").next() {
tool_strings.remove(extension_name);
}
Ok(())
}
async fn record_tool_call(&self, tool_name: &str) -> Result<(), ErrorData> {
let mut recent_calls = self.recent_tool_calls.write().await;
if recent_calls.len() >= 100 {
recent_calls.pop_front();
}
recent_calls.push_back(tool_name.to_string());
Ok(())
}
async fn get_recent_tool_calls(&self, limit: usize) -> Result<Vec<String>, ErrorData> {
let recent_calls = self.recent_tool_calls.read().await;
Ok(recent_calls.iter().rev().take(limit).cloned().collect())
}
}
// Helper function to create a boxed tool selector
pub async fn create_tool_selector(
provider: Arc<dyn Provider>,
) -> Result<Box<dyn RouterToolSelector>> {
let selector = LLMToolSelector::new(provider).await?;
Ok(Box::new(selector))
}
crates/goose/src/agents/router_tools.rs
-65
View File
@@ -1,65 +0,0 @@
use crate::agents::extension_manager_extension::{
LIST_RESOURCES_TOOL_NAME, MANAGE_EXTENSIONS_TOOL_NAME, READ_RESOURCE_TOOL_NAME,
SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
};
use indoc::indoc;
use rmcp::model::{Tool, ToolAnnotations};
use rmcp::object;
pub const ROUTER_LLM_SEARCH_TOOL_NAME: &str = "router__llm_search";
pub fn llm_search_tool() -> Tool {
Tool::new(
ROUTER_LLM_SEARCH_TOOL_NAME.to_string(),
indoc! {r#"
Searches for relevant tools based on the user's messages.
Format a query to search for the most relevant tools based on the user's messages.
Pay attention to the keywords in the user's messages, especially the last message and potential tools they are asking for.
This tool should be invoked when the user's messages suggest they are asking for a tool to be run.
Use the extension_name parameter to filter tools by the appropriate extension.
For example, if the user is asking to list the files in the current directory, you filter for the "developer" extension.
Example: {"User": "list the files in the current directory", "Query": "list files in current directory", "Extension Name": "developer", "k": 5}
Extension name is not optional, it is required.
The returned result will be a list of tool names, descriptions, and schemas from which you, the agent can select the most relevant tool to invoke.
"#}
.to_string(),
object!({
"type": "object",
"required": ["query", "extension_name"],
"properties": {
"extension_name": {"type": "string", "description": "The name of the extension to filter tools by"},
"query": {"type": "string", "description": "The query to search for the most relevant tools based on the user's messages"},
"k": {"type": "integer", "description": "The number of tools to retrieve (defaults to 5)", "default": 5}
}
})
).annotate(ToolAnnotations {
title: Some("LLM search for relevant tools".to_string()),
read_only_hint: Some(true),
destructive_hint: Some(false),
idempotent_hint: Some(false),
open_world_hint: Some(false),
})
}
pub fn llm_search_tool_prompt() -> String {
format!(
r#"# LLM Tool Selection Instructions
Important: the user has opted to dynamically enable tools, so although an extension could be enabled, \
please invoke the llm 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 llm_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 a query packed with relevant keywords to search for the most relevant tools.
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:
- {}
- {}
- {}
- {}
"#,
SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
MANAGE_EXTENSIONS_TOOL_NAME,
READ_RESOURCE_TOOL_NAME,
LIST_RESOURCES_TOOL_NAME
)
}
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__basic.snap
-2
View File
@@ -26,8 +26,6 @@ new ones.
No extensions are defined. You should let the user know that they should add extensions.
# Response Guidelines
- Use Markdown formatting for all responses.
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__one_extension.snap
-15
View File
@@ -37,21 +37,6 @@ and platform__list_resources on this extension.
how to use this extension
# LLM Tool Selection Instructions
Important: the user has opted to dynamically enable tools, so although an extension could be enabled, \
please invoke the llm 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 llm_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 a query packed with relevant keywords to search for the most relevant tools.
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:
- search_available_extensions
- manage_extensions
- read_resource
- list_resources
# Response Guidelines
- Use Markdown formatting for all responses.
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__typical_setup.snap
-15
View File
@@ -51,21 +51,6 @@ You should only disable extensions found from the search_available_extensions to
List all the extensions available to disable in the response.
Explain that minimizing extensions helps with the recall of the correct tools to use.
# LLM Tool Selection Instructions
Important: the user has opted to dynamically enable tools, so although an extension could be enabled, \
please invoke the llm 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 llm_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 a query packed with relevant keywords to search for the most relevant tools.
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:
- search_available_extensions
- manage_extensions
- read_resource
- list_resources
# Response Guidelines
- Use Markdown formatting for all responses.
crates/goose/src/agents/tool_route_manager.rs
-185
View File
@@ -1,185 +0,0 @@
use crate::agents::extension_manager::ExtensionManager;
use crate::agents::router_tool_selector::{create_tool_selector, RouterToolSelector};
use crate::agents::router_tools::{self};
use crate::agents::tool_execution::ToolCallResult;
use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
use crate::config::Config;
use crate::conversation::message::ToolRequest;
use crate::providers::base::Provider;
use anyhow::{anyhow, Result};
use rmcp::model::{ErrorCode, ErrorData, JsonObject, Tool};
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::error;
pub struct ToolRouteManager {
router_tool_selector: Mutex<Option<Arc<Box<dyn RouterToolSelector>>>>,
router_disabled_override: Mutex<bool>,
}
impl Default for ToolRouteManager {
fn default() -> Self {
Self::new()
}
}
impl ToolRouteManager {
pub fn new() -> Self {
Self {
router_tool_selector: Mutex::new(None),
router_disabled_override: Mutex::new(false),
}
}
pub async fn disable_router_for_recipe(&self) {
*self.router_disabled_override.lock().await = true;
*self.router_tool_selector.lock().await = None;
}
pub async fn record_tool_requests(&self, requests: &[ToolRequest]) {
let selector = self.router_tool_selector.lock().await.clone();
for request in requests {
if let Ok(tool_call) = &request.tool_call {
if let Some(ref selector) = selector {
if let Err(e) = selector.record_tool_call(&tool_call.name).await {
error!("Failed to record tool call: {}", e);
}
}
}
}
}
pub async fn dispatch_route_search_tool(
&self,
arguments: JsonObject,
) -> Result<ToolCallResult, ErrorData> {
let selector = self.router_tool_selector.lock().await.clone();
match selector.as_ref() {
Some(selector) => match selector.select_tools(arguments).await {
Ok(content) => Ok(ToolCallResult::from(Ok(rmcp::model::CallToolResult {
content,
structured_content: None,
is_error: Some(false),
meta: None,
}))),
Err(e) => Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to select tools: {}", e),
None,
)),
},
None => Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
"No tool selector available".to_string(),
None,
)),
}
}
pub async fn is_router_enabled(&self) -> bool {
if *self.router_disabled_override.lock().await {
return false;
}
let config = Config::global();
if let Ok(config_value) = config.get_param::<String>("GOOSE_ENABLE_ROUTER") {
return config_value.to_lowercase() == "true";
}
// Default to false if neither is set
false
}
pub async fn update_router_tool_selector(
&self,
provider: Arc<dyn Provider>,
reindex_all: Option<bool>,
extension_manager: &ExtensionManager,
) -> Result<()> {
let enabled = self.is_router_enabled().await;
if !enabled {
return Ok(());
}
let selector = create_tool_selector(provider.clone())
.await
.map_err(|e| anyhow!("Failed to create tool selector: {}", e))?;
// Wrap selector in Arc for the index manager methods
let selector_arc = Arc::new(selector);
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_arc,
extension_manager,
&extension_name,
"add",
)
.await
{
error!(
"Failed to index tools for extension {}: {}",
extension_name, e
);
}
}
}
// Update the selector
*self.router_tool_selector.lock().await = Some(selector_arc);
Ok(())
}
pub async fn get_router_tool_selector(&self) -> Option<Arc<Box<dyn RouterToolSelector>>> {
self.router_tool_selector.lock().await.clone()
}
/// Check if the router is actually functional (enabled in config AND initialized)
pub async fn is_router_functional(&self) -> bool {
if !self.is_router_enabled().await {
return false;
}
// Check if the selector actually exists (meaning it was successfully initialized)
self.router_tool_selector.lock().await.is_some()
}
pub async fn list_tools_for_router(&self, extension_manager: &ExtensionManager) -> Vec<Tool> {
// If router is disabled or overridden, return empty
if *self.router_disabled_override.lock().await {
return vec![];
}
let mut prefixed_tools = vec![];
// If router is enabled but not functional (no provider), just return the search tool
if !self.is_router_functional().await {
return prefixed_tools;
}
prefixed_tools.push(router_tools::llm_search_tool());
// Get recent tool calls from router tool selector
let selector = self.router_tool_selector.lock().await.clone();
if let Some(selector) = selector {
if let Ok(recent_calls) = selector.get_recent_tool_calls(20).await {
// Add recent tool calls to the list, avoiding duplicates
for tool_name in recent_calls {
// Find the tool in the extension manager's tools
if let Ok(extension_tools) = extension_manager.get_prefixed_tools(None).await {
if let Some(tool) = extension_tools.iter().find(|t| t.name == tool_name) {
// Only add if not already in prefixed_tools
if !prefixed_tools.iter().any(|t| t.name == tool.name) {
prefixed_tools.push(tool.clone());
}
}
}
}
}
}
prefixed_tools
}
}
crates/goose/src/agents/tool_router_index_manager.rs
-76
View File
@@ -1,76 +0,0 @@
use anyhow::{anyhow, Result};
use std::sync::Arc;
use tracing;
use crate::agents::extension_manager::ExtensionManager;
use crate::agents::router_tool_selector::RouterToolSelector;
/// Manages tool indexing operations for the router when LLM routing is enabled
pub struct ToolRouterIndexManager;
impl ToolRouterIndexManager {
/// Updates the LLM index for tools when extensions are added or removed
pub async fn update_extension_tools(
selector: &Arc<Box<dyn RouterToolSelector>>,
extension_manager: &ExtensionManager,
extension_name: &str,
action: &str,
) -> Result<()> {
match action {
"add" => {
// Get tools for specific extension
let tools = extension_manager
.get_prefixed_tools(Some(extension_name.to_string()))
.await?;
if !tools.is_empty() {
// Index all tools at once
selector
.index_tools(&tools, extension_name)
.await
.map_err(|e| {
anyhow!(
"Failed to index tools for extension {}: {}",
extension_name,
e
)
})?;
tracing::info!(
"Indexed {} tools for extension {}",
tools.len(),
extension_name
);
}
}
"remove" => {
// Remove all tools for this extension
let tools = extension_manager
.get_prefixed_tools(Some(extension_name.to_string()))
.await?;
for tool in &tools {
selector.remove_tool(&tool.name).await.map_err(|e| {
anyhow!(
"Failed to remove tool {} for extension {}: {}",
tool.name,
extension_name,
e
)
})?;
}
tracing::info!(
"Removed {} tools for extension {}",
tools.len(),
extension_name
);
}
_ => {
return Err(anyhow!("Invalid action: {}", action));
}
}
Ok(())
}
}
crates/goose/src/execution/manager.rs
-1
View File
@@ -87,7 +87,6 @@ impl AgentManager {
.set_context(PlatformExtensionContext {
session_id: Some(session_id.clone()),
extension_manager: Some(Arc::downgrade(&agent.extension_manager)),
tool_route_manager: Some(Arc::downgrade(&agent.tool_route_manager)),
})
.await;
if let Some(provider) = &*self.default_provider.read().await {
crates/goose/src/posthog.rs
-5
View File
@@ -363,11 +363,6 @@ async fn send_session_event(installation: &InstallationData) -> Result<(), Strin
if let Ok(max_turns) = config.get_param::<i64>("GOOSE_MAX_TURNS") {
event.insert_prop("setting_max_turns", max_turns).ok();
}
if let Ok(router_enabled) = config.get_param::<bool>("GOOSE_ENABLE_ROUTER") {
event
.insert_prop("setting_router_enabled", router_enabled)
.ok();
}
if let Ok(lead_model) = config.get_param::<String>("GOOSE_LEAD_MODEL") {
event.insert_prop("setting_lead_model", lead_model).ok();
crates/goose/src/prompts/router_tool_selector.md
-11
View File
@@ -1,11 +0,0 @@
You are a tool selection assistant. Your task is to find the most relevant tools based on the user's query.
Given the following tools:
{{ tools }}
Find the most relevant tools for the query: {{ query }}
Return the tools in this exact format for each tool:
Tool: <tool_name>
Description: <tool_description>
Schema: <tool_schema>
crates/goose/src/prompts/system.md
-2
View File
@@ -58,8 +58,6 @@ Explain that minimizing extensions helps with the recall of the correct tools to
{% endwith %}
{% endif %}
{{tool_selection_strategy}}
# Response Guidelines
- Use Markdown formatting for all responses.
crates/goose/tests/agent.rs
-1
View File
@@ -475,7 +475,6 @@ mod tests {
.set_context(PlatformExtensionContext {
session_id: Some("test_session".to_string()),
extension_manager: Some(Arc::downgrade(&agent.extension_manager)),
tool_route_manager: Some(Arc::downgrade(&agent.tool_route_manager)),
})
.await;