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

feat: (tool router) llm tool selector (#2866)

Browse Source
This commit is contained in:
Wendy Tang
2025-06-13 10:29:11 -07:00
committed by GitHub
parent de5fc9b450
commit b1d2a02f66
8 changed files with 260 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+32 -18
View File
@@ -33,7 +33,7 @@ use crate::agents::prompt_manager::PromptManager;
use crate::agents::router_tool_selector::{
create_tool_selector, RouterToolSelectionStrategy, RouterToolSelector,
};
use crate::agents::router_tools::ROUTER_VECTOR_SEARCH_TOOL_NAME;
use crate::agents::router_tools::{ROUTER_LLM_SEARCH_TOOL_NAME, ROUTER_VECTOR_SEARCH_TOOL_NAME};
use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
use crate::agents::tool_vectordb::generate_table_id;
use crate::agents::types::SessionConfig;
@@ -244,7 +244,9 @@ impl Agent {
ToolCallResult::from(Err(ToolError::ExecutionError(
"Frontend tool execution required".to_string(),
)))
} else if tool_call.name == ROUTER_VECTOR_SEARCH_TOOL_NAME {
} else if tool_call.name == ROUTER_VECTOR_SEARCH_TOOL_NAME
|| tool_call.name == ROUTER_LLM_SEARCH_TOOL_NAME
{
let selector = self.router_tool_selector.lock().await.clone();
let selected_tools = match selector.as_ref() {
Some(selector) => match selector.select_tools(tool_call.arguments.clone()).await {
@@ -351,10 +353,11 @@ impl Agent {
// 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::vector_tool_router_enabled(&selector) {
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,
@@ -414,9 +417,10 @@ impl Agent {
// If vector tool selection is enabled, index the tools
let selector = self.router_tool_selector.lock().await.clone();
if ToolRouterIndexManager::vector_tool_router_enabled(&selector) {
if ToolRouterIndexManager::is_tool_router_enabled(&selector) {
if let Some(selector) = selector {
let extension_manager = self.extension_manager.lock().await;
let selector = Arc::new(selector);
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
&selector,
&extension_manager,
@@ -468,6 +472,9 @@ impl Agent {
Some(RouterToolSelectionStrategy::Vector) => {
prefixed_tools.push(router_tools::vector_search_tool());
}
Some(RouterToolSelectionStrategy::Llm) => {
prefixed_tools.push(router_tools::llm_search_tool());
}
None => {}
}
@@ -500,7 +507,7 @@ impl Agent {
// If vector tool selection is enabled, remove tools from the index
let selector = self.router_tool_selector.lock().await.clone();
if ToolRouterIndexManager::vector_tool_router_enabled(&selector) {
if ToolRouterIndexManager::is_tool_router_enabled(&selector) {
if let Some(selector) = selector {
let extension_manager = self.extension_manager.lock().await;
ToolRouterIndexManager::update_extension_tools(
@@ -793,22 +800,29 @@ impl Agent {
let strategy = match router_tool_selection_strategy.to_lowercase().as_str() {
"vector" => Some(RouterToolSelectionStrategy::Vector),
"llm" => Some(RouterToolSelectionStrategy::Llm),
_ => None,
};
if let Some(strategy) = strategy {
let table_name = generate_table_id();
let selector = create_tool_selector(Some(strategy), provider, table_name)
.await
.map_err(|e| anyhow!("Failed to create tool selector: {}", e))?;
let selector = Arc::new(selector);
*self.router_tool_selector.lock().await = Some(selector.clone());
let extension_manager = self.extension_manager.lock().await;
ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await?;
}
let selector = match strategy {
Some(RouterToolSelectionStrategy::Vector) => {
let table_name = generate_table_id();
let selector = create_tool_selector(strategy, provider, 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)
.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;
ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await?;
*self.router_tool_selector.lock().await = Some(selector.clone());
Ok(())
}