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

feat: lancedb vector tool selection (#2654)

Browse Source 
Co-authored-by: Wendy Tang <wendytang@squareup.com>
Co-authored-by: Alice Hau <ahau@squareup.com>
This commit is contained in:
Alice Hau
2025-05-28 23:23:02 -04:00
committed by GitHub
parent cf7bb08ee1
commit bf1c0d51e4
23 changed files with 3661 additions and 301 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/configure.rs
+51
View File
@@ -824,6 +824,11 @@ pub fn remove_extension_dialog() -> Result<(), Box<dyn Error>> {
pub async fn configure_settings_dialog() -> Result<(), Box<dyn Error>> {
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",
"Configure the strategy for selecting tools to use",
)
.item(
"tool_permission",
"Tool Permission",
@@ -850,6 +855,9 @@ pub async fn configure_settings_dialog() -> Result<(), Box<dyn Error>> {
"goose_mode" => {
configure_goose_mode_dialog()?;
}
"goose_router_strategy" => {
configure_goose_router_strategy_dialog()?;
}
"tool_permission" => {
configure_tool_permissions_dialog().await.and(Ok(()))?;
}
@@ -921,6 +929,49 @@ pub fn configure_goose_mode_dialog() -> Result<(), Box<dyn Error>> {
Ok(())
}
pub fn configure_goose_router_strategy_dialog() -> Result<(), Box<dyn Error>> {
let config = Config::global();
// Check if GOOSE_ROUTER_STRATEGY is set as an environment variable
if std::env::var("GOOSE_ROUTER_TOOL_SELECTION_STRATEGY").is_ok() {
let _ = cliclack::log::info("Notice: GOOSE_ROUTER_TOOL_SELECTION_STRATEGY environment variable is set. Configuration will override this.");
}
let strategy = cliclack::select("Which router strategy would you like to use?")
.item(
"vector",
"Vector Strategy",
"Use vector-based similarity to select tools",
)
.item(
"default",
"Default Strategy",
"Use the default tool selection strategy",
)
.interact()?;
match strategy {
"vector" => {
config.set_param(
"GOOSE_ROUTER_TOOL_SELECTION_STRATEGY",
Value::String("vector".to_string()),
)?;
cliclack::outro(
"Set to Vector Strategy - using vector-based similarity for tool selection",
)?;
}
"default" => {
config.set_param(
"GOOSE_ROUTER_TOOL_SELECTION_STRATEGY",
Value::String("default".to_string()),
)?;
cliclack::outro("Set to Default Strategy - using default tool selection")?;
}
_ => unreachable!(),
};
Ok(())
}
pub fn configure_tool_output_dialog() -> Result<(), Box<dyn Error>> {
let config = Config::global();
// Check if GOOSE_CLI_MIN_PRIORITY is set as an environment variable
crates/goose-cli/src/session/builder.rs
+7 -1
View File
@@ -55,7 +55,13 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
// Create the agent
let agent: Agent = Agent::new();
let new_provider = create(&provider_name, model_config).unwrap();
let _ = agent.update_provider(new_provider).await;
agent
.update_provider(new_provider)
.await
.unwrap_or_else(|e| {
output::render_error(&format!("Failed to initialize agent: {}", e));
process::exit(1);
});
// Configure tool monitoring if max_tool_repetitions is set
if let Some(max_repetitions) = session_config.max_tool_repetitions {
crates/goose-server/src/routes/extension.rs
+10 -6
View File
@@ -268,12 +268,16 @@ async fn remove_extension(
.get_agent()
.await
.map_err(|_| StatusCode::PRECONDITION_FAILED)?;
agent.remove_extension(&name).await;
Ok(Json(ExtensionResponse {
error: false,
message: None,
}))
match agent.remove_extension(&name).await {
Ok(_) => Ok(Json(ExtensionResponse {
error: false,
message: None,
})),
Err(e) => Ok(Json(ExtensionResponse {
error: true,
message: Some(format!("Failed to remove extension: {:?}", e)),
})),
}
}
/// Registers the extension management routes with the Axum router.
crates/goose/Cargo.toml
+4
View File
@@ -76,6 +76,10 @@ blake3 = "1.5"
fs2 = "0.4.3"
futures-util = "0.3.31"
# Vector database for tool selection
lancedb = "0.13"
arrow = "52.2"
[target.'cfg(target_os = "windows")'.dependencies]
winapi = { version = "0.3", features = ["wincred"] }
crates/goose/src/agents/agent.rs
+144 -31
View File
@@ -18,7 +18,7 @@ use serde_json::Value;
use tokio::sync::{mpsc, Mutex};
use tracing::{debug, error, instrument};
use crate::agents::extension::{ExtensionConfig, ExtensionResult, ToolInfo};
use crate::agents::extension::{ExtensionConfig, ExtensionError, ExtensionResult, ToolInfo};
use crate::agents::extension_manager::{get_parameter_names, ExtensionManager};
use crate::agents::platform_tools::{
PLATFORM_LIST_RESOURCES_TOOL_NAME, PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME,
@@ -29,6 +29,8 @@ use crate::agents::router_tool_selector::{
create_tool_selector, RouterToolSelectionStrategy, RouterToolSelector,
};
use crate::agents::router_tools::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;
use crate::agents::types::{FrontendTool, ToolResultReceiver};
use mcp_core::{
@@ -51,7 +53,7 @@ pub struct Agent {
pub(super) tool_result_tx: mpsc::Sender<(String, ToolResult<Vec<Content>>)>,
pub(super) tool_result_rx: ToolResultReceiver,
pub(super) tool_monitor: Mutex<Option<ToolMonitor>>,
pub(super) router_tool_selector: Mutex<Option<Box<dyn RouterToolSelector>>>,
pub(super) router_tool_selector: Mutex<Option<Arc<Box<dyn RouterToolSelector>>>>,
}
impl Agent {
@@ -60,16 +62,6 @@ impl Agent {
let (confirm_tx, confirm_rx) = mpsc::channel(32);
let (tool_tx, tool_rx) = mpsc::channel(32);
let router_tool_selection_strategy = std::env::var("GOOSE_ROUTER_TOOL_SELECTION_STRATEGY")
.ok()
.and_then(|s| {
if s.eq_ignore_ascii_case("vector") {
Some(RouterToolSelectionStrategy::Vector)
} else {
None
}
});
Self {
provider: Mutex::new(None),
extension_manager: Mutex::new(ExtensionManager::new()),
@@ -81,9 +73,7 @@ impl Agent {
tool_result_tx: tool_tx,
tool_result_rx: Arc::new(Mutex::new(tool_rx)),
tool_monitor: Mutex::new(None),
router_tool_selector: Mutex::new(Some(create_tool_selector(
router_tool_selection_strategy,
))),
router_tool_selector: Mutex::new(None),
}
}
@@ -204,8 +194,8 @@ impl Agent {
"Frontend tool execution required".to_string(),
))
} else if tool_call.name == ROUTER_VECTOR_SEARCH_TOOL_NAME {
let router_tool_selector = self.router_tool_selector.lock().await;
if let Some(selector) = router_tool_selector.as_ref() {
let selector = self.router_tool_selector.lock().await.clone();
if let Some(selector) = selector {
selector.select_tools(tool_call.arguments.clone()).await
} else {
Err(ToolError::ExecutionError(
@@ -284,6 +274,33 @@ 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::vector_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;
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)
}
@@ -317,10 +334,32 @@ impl Agent {
}
_ => {
let mut extension_manager = self.extension_manager.lock().await;
extension_manager.add_extension(extension).await?;
extension_manager.add_extension(extension.clone()).await?;
}
};
// 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 let Some(selector) = selector {
let extension_manager = self.extension_manager.lock().await;
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
&selector,
&extension_manager,
&extension.name(),
"add",
)
.await
{
return Err(ExtensionError::SetupError(format!(
"Failed to index tools for extension {}: {}",
extension.name(),
e
)));
}
}
}
Ok(())
}
@@ -350,8 +389,6 @@ impl Agent {
&self,
strategy: Option<RouterToolSelectionStrategy>,
) -> Vec<Tool> {
let extension_manager = self.extension_manager.lock().await;
let mut prefixed_tools = vec![];
match strategy {
Some(RouterToolSelectionStrategy::Vector) => {
@@ -359,22 +396,50 @@ impl Agent {
}
None => {}
}
prefixed_tools.push(platform_tools::search_available_extensions_tool());
prefixed_tools.push(platform_tools::manage_extensions_tool());
if extension_manager.supports_resources() {
prefixed_tools.push(platform_tools::read_resource_tool());
prefixed_tools.push(platform_tools::list_resources_tool());
// Get recent tool calls from router tool selector if available
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 {
let extension_manager = self.extension_manager.lock().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
}
pub async fn remove_extension(&self, name: &str) {
pub async fn remove_extension(&self, name: &str) -> Result<()> {
let mut extension_manager = self.extension_manager.lock().await;
extension_manager
.remove_extension(name)
.await
.expect("Failed to remove extension");
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::vector_tool_router_enabled(&selector) {
if let Some(selector) = selector {
let extension_manager = self.extension_manager.lock().await;
ToolRouterIndexManager::update_extension_tools(
&selector,
&extension_manager,
name,
"remove",
)
.await?;
}
}
Ok(())
}
pub async fn list_extensions(&self) -> Vec<String> {
@@ -447,6 +512,26 @@ impl Agent {
filtered_response) =
self.categorize_tool_requests(&response).await;
// Record tool calls in the router selector
let selector = self.router_tool_selector.lock().await.clone();
if let Some(selector) = selector {
// Record frontend tool calls
for request in &frontend_requests {
if let Ok(tool_call) = &request.tool_call {
if let Err(e) = selector.record_tool_call(&tool_call.name).await {
tracing::error!("Failed to record frontend tool call: {}", e);
}
}
}
// Record remaining tool calls
for request in &remaining_requests {
if let Ok(tool_call) = &request.tool_call {
if let Err(e) = selector.record_tool_call(&tool_call.name).await {
tracing::error!("Failed to record tool call: {}", e);
}
}
}
}
// Yield the assistant's response with frontend tool requests filtered out
yield filtered_response.clone();
@@ -598,7 +683,35 @@ impl Agent {
/// Update the provider used by this agent
pub async fn update_provider(&self, provider: Arc<dyn Provider>) -> Result<()> {
*self.provider.lock().await = Some(provider);
*self.provider.lock().await = Some(provider.clone());
self.update_router_tool_selector(provider).await?;
Ok(())
}
async fn update_router_tool_selector(&self, provider: Arc<dyn Provider>) -> Result<()> {
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),
_ => 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?;
}
Ok(())
}
crates/goose/src/agents/mod.rs
+2
View File
@@ -9,6 +9,8 @@ mod reply_parts;
mod router_tool_selector;
mod router_tools;
mod tool_execution;
mod tool_router_index_manager;
pub(crate) mod tool_vectordb;
mod types;
pub use agent::Agent;
crates/goose/src/agents/reply_parts.rs
+12 -10
View File
@@ -3,6 +3,7 @@ use std::collections::HashSet;
use std::sync::Arc;
use crate::agents::router_tool_selector::RouterToolSelectionStrategy;
use crate::config::Config;
use crate::message::{Message, MessageContent, ToolRequest};
use crate::providers::base::{Provider, ProviderUsage};
use crate::providers::errors::ProviderError;
@@ -19,16 +20,17 @@ impl Agent {
pub(crate) async fn prepare_tools_and_prompt(
&self,
) -> anyhow::Result<(Vec<Tool>, Vec<Tool>, String)> {
// Get tool selection strategy
let tool_selection_strategy = std::env::var("GOOSE_ROUTER_TOOL_SELECTION_STRATEGY")
.ok()
.and_then(|s| {
if s.eq_ignore_ascii_case("vector") {
Some(RouterToolSelectionStrategy::Vector)
} else {
None
}
});
// Get tool selection strategy from config
let config = Config::global();
let router_tool_selection_strategy = config
.get_param("GOOSE_ROUTER_TOOL_SELECTION_STRATEGY")
.unwrap_or_else(|_| "default".to_string());
let tool_selection_strategy = match router_tool_selection_strategy.to_lowercase().as_str() {
"vector" => Some(RouterToolSelectionStrategy::Vector),
_ => None,
};
// Get tools from extension manager
let mut tools = match tool_selection_strategy {
Some(RouterToolSelectionStrategy::Vector) => {
crates/goose/src/agents/router_tool_selector.rs
+188 -9
View File
@@ -1,8 +1,20 @@
use mcp_core::content::TextContent;
use mcp_core::tool::Tool;
use mcp_core::{Content, ToolError};
use anyhow::{Context, Result};
use async_trait::async_trait;
use serde_json::Value;
use std::collections::VecDeque;
use std::env;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::agents::tool_vectordb::ToolVectorDB;
use crate::model::ModelConfig;
use crate::providers::{self, base::Provider};
#[derive(Debug, Clone, PartialEq)]
pub enum RouterToolSelectionStrategy {
Vector,
}
@@ -10,27 +22,194 @@ pub enum RouterToolSelectionStrategy {
#[async_trait]
pub trait RouterToolSelector: Send + Sync {
async fn select_tools(&self, params: Value) -> Result<Vec<Content>, ToolError>;
async fn index_tools(&self, tools: &[Tool]) -> Result<(), ToolError>;
async fn remove_tool(&self, tool_name: &str) -> Result<(), ToolError>;
async fn record_tool_call(&self, tool_name: &str) -> Result<(), ToolError>;
async fn get_recent_tool_calls(&self, limit: usize) -> Result<Vec<String>, ToolError>;
fn selector_type(&self) -> RouterToolSelectionStrategy;
}
pub struct VectorToolSelector;
pub struct VectorToolSelector {
vector_db: Arc<RwLock<ToolVectorDB>>,
embedding_provider: Arc<dyn Provider>,
recent_tool_calls: Arc<RwLock<VecDeque<String>>>,
}
impl VectorToolSelector {
pub async fn new(provider: Arc<dyn Provider>, table_name: String) -> Result<Self> {
let vector_db = ToolVectorDB::new(Some(table_name)).await?;
let embedding_provider = if env::var("EMBEDDING_MODEL_PROVIDER").is_ok() {
// If env var is set, create a new provider for embeddings
// Get embedding model and provider from environment variables
let embedding_model = env::var("EMBEDDING_MODEL")
.unwrap_or_else(|_| "text-embedding-3-small".to_string());
let embedding_provider_name =
env::var("EMBEDDING_MODEL_PROVIDER").unwrap_or_else(|_| "openai".to_string());
// Create the provider using the factory
let model_config = ModelConfig::new(embedding_model);
providers::create(&embedding_provider_name, model_config).context(format!(
"Failed to create {} provider for embeddings. If using OpenAI, make sure OPENAI_API_KEY env var is set or that you have configured the OpenAI provider via Goose before.",
embedding_provider_name
))?
} else {
// Otherwise fall back to using the same provider instance as used for base goose model
provider.clone()
};
Ok(Self {
vector_db: Arc::new(RwLock::new(vector_db)),
embedding_provider,
recent_tool_calls: Arc::new(RwLock::new(VecDeque::with_capacity(100))),
})
}
}
#[async_trait]
impl RouterToolSelector for VectorToolSelector {
async fn select_tools(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let query = params.get("query").and_then(|v| v.as_str());
println!("query: {:?}", query);
let selected_tools = Vec::new();
// TODO: placeholder for vector tool selection
let query = params
.get("query")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("Missing 'query' parameter".to_string()))?;
let k = params.get("k").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
// Check if provider supports embeddings
if !self.embedding_provider.supports_embeddings() {
return Err(ToolError::ExecutionError(
"Embedding provider does not support embeddings".to_string(),
));
}
let embeddings = self
.embedding_provider
.create_embeddings(vec![query.to_string()])
.await
.map_err(|e| {
ToolError::ExecutionError(format!("Failed to generate query embedding: {}", e))
})?;
let query_embedding = embeddings
.into_iter()
.next()
.ok_or_else(|| ToolError::ExecutionError("No embedding returned".to_string()))?;
let vector_db = self.vector_db.read().await;
let tools = vector_db
.search_tools(query_embedding, k)
.await
.map_err(|e| ToolError::ExecutionError(format!("Failed to search tools: {}", e)))?;
let selected_tools: Vec<Content> = tools
.into_iter()
.map(|tool| {
let text = format!(
"Tool: {}\nDescription: {}\nSchema: {}",
tool.tool_name, tool.description, tool.schema
);
Content::Text(TextContent {
text,
annotations: None,
})
})
.collect();
Ok(selected_tools)
}
async fn index_tools(&self, tools: &[Tool]) -> Result<(), ToolError> {
let texts_to_embed: Vec<String> = tools
.iter()
.map(|tool| {
let schema_str = serde_json::to_string_pretty(&tool.input_schema)
.unwrap_or_else(|_| "{}".to_string());
format!("{} {} {}", tool.name, tool.description, schema_str)
})
.collect();
if !self.embedding_provider.supports_embeddings() {
return Err(ToolError::ExecutionError(
"Embedding provider does not support embeddings".to_string(),
));
}
let embeddings = self
.embedding_provider
.create_embeddings(texts_to_embed)
.await
.map_err(|e| {
ToolError::ExecutionError(format!("Failed to generate tool embeddings: {}", e))
})?;
// Create tool records
let tool_records: Vec<crate::agents::tool_vectordb::ToolRecord> = tools
.iter()
.zip(embeddings.into_iter())
.map(|(tool, vector)| {
let schema_str = serde_json::to_string_pretty(&tool.input_schema)
.unwrap_or_else(|_| "{}".to_string());
crate::agents::tool_vectordb::ToolRecord {
tool_name: tool.name.clone(),
description: tool.description.clone(),
schema: schema_str,
vector,
}
})
.collect();
// Index all tools at once
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)))?;
Ok(())
}
async fn remove_tool(&self, tool_name: &str) -> Result<(), ToolError> {
let vector_db = self.vector_db.read().await;
vector_db.remove_tool(tool_name).await.map_err(|e| {
ToolError::ExecutionError(format!("Failed to remove tool {}: {}", tool_name, e))
})?;
Ok(())
}
async fn record_tool_call(&self, tool_name: &str) -> Result<(), ToolError> {
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>, ToolError> {
let recent_calls = self.recent_tool_calls.read().await;
Ok(recent_calls.iter().rev().take(limit).cloned().collect())
}
fn selector_type(&self) -> RouterToolSelectionStrategy {
RouterToolSelectionStrategy::Vector
}
}
// Helper function to create a boxed tool selector
pub fn create_tool_selector(
pub async fn create_tool_selector(
strategy: Option<RouterToolSelectionStrategy>,
) -> Box<dyn RouterToolSelector> {
provider: Arc<dyn Provider>,
table_name: String,
) -> Result<Box<dyn RouterToolSelector>> {
match strategy {
Some(RouterToolSelectionStrategy::Vector) => Box::new(VectorToolSelector),
_ => Box::new(VectorToolSelector), // Default to VectorToolSelector
Some(RouterToolSelectionStrategy::Vector) => {
let selector = VectorToolSelector::new(provider, table_name).await?;
Ok(Box::new(selector))
}
None => {
let selector = VectorToolSelector::new(provider, table_name).await?;
Ok(Box::new(selector))
}
}
}
crates/goose/src/agents/router_tools.rs
+2 -1
View File
@@ -22,7 +22,8 @@ pub fn vector_search_tool() -> Tool {
"type": "object",
"required": ["query"],
"properties": {
"query": {"type": "string", "description": "The query to search for the most relevant tools based on the user's messages"}
"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}
}
}),
Some(ToolAnnotations {
crates/goose/src/agents/tool_router_index_manager.rs
+106
View File
@@ -0,0 +1,106 @@
use anyhow::{anyhow, Result};
use std::sync::Arc;
use tracing;
use crate::agents::extension_manager::ExtensionManager;
use crate::agents::platform_tools;
use crate::agents::router_tool_selector::{RouterToolSelectionStrategy, RouterToolSelector};
/// Manages tool indexing operations for the router when vector routing is enabled
pub struct ToolRouterIndexManager;
impl ToolRouterIndexManager {
/// Updates the vector 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).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" => {
// Get tool names for the extension to remove them
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 {}: {}", tool.name, e))?;
}
tracing::info!(
"Removed {} tools for extension {}",
tools.len(),
extension_name
);
}
_ => {
anyhow::bail!("Invalid action '{}' for tool indexing", action);
}
}
Ok(())
}
/// Indexes platform tools (search_available_extensions, manage_extensions, etc.)
pub async fn index_platform_tools(
selector: &Arc<Box<dyn RouterToolSelector>>,
extension_manager: &ExtensionManager,
) -> Result<()> {
let mut tools = Vec::new();
// Add the standard platform tools
tools.push(platform_tools::search_available_extensions_tool());
tools.push(platform_tools::manage_extensions_tool());
// Add resource tools if supported
if extension_manager.supports_resources() {
tools.push(platform_tools::read_resource_tool());
tools.push(platform_tools::list_resources_tool());
}
// Index all platform tools at once
selector
.index_tools(&tools)
.await
.map_err(|e| anyhow!("Failed to index platform tools: {}", e))?;
tracing::info!("Indexed platform tools for vector search");
Ok(())
}
/// Helper to check if vector tool router is enabled
pub fn vector_tool_router_enabled(selector: &Option<Arc<Box<dyn RouterToolSelector>>>) -> bool {
if let Some(selector) = selector {
selector.selector_type() == RouterToolSelectionStrategy::Vector
} else {
false
}
}
}
crates/goose/src/agents/tool_vectordb.rs
+440
View File
@@ -0,0 +1,440 @@
use anyhow::{Context, Result};
use arrow::array::{FixedSizeListBuilder, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use chrono::Local;
use etcetera::base_strategy::{BaseStrategy, Xdg};
use futures::TryStreamExt;
use lancedb::connect;
use lancedb::connection::Connection;
use lancedb::query::{ExecutableQuery, QueryBase};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolRecord {
pub tool_name: String,
pub description: String,
pub schema: String,
pub vector: Vec<f32>,
}
pub struct ToolVectorDB {
connection: Arc<RwLock<Connection>>,
table_name: String,
}
impl ToolVectorDB {
pub async fn new(table_name: Option<String>) -> Result<Self> {
let db_path = Self::get_db_path()?;
// Ensure the directory exists
if let Some(parent) = db_path.parent() {
tokio::fs::create_dir_all(parent)
.await
.context("Failed to create database directory")?;
}
let connection = connect(db_path.to_str().unwrap())
.execute()
.await
.context("Failed to connect to LanceDB")?;
let tool_db = Self {
connection: Arc::new(RwLock::new(connection)),
table_name: table_name.unwrap_or_else(|| "tools".to_string()),
};
// Initialize the table if it doesn't exist
tool_db.init_table().await?;
Ok(tool_db)
}
fn get_db_path() -> Result<PathBuf> {
let data_dir = Xdg::new()
.context("Failed to determine base strategy")?
.data_dir();
Ok(data_dir.join("goose").join("tool_db"))
}
async fn init_table(&self) -> Result<()> {
let connection = self.connection.read().await;
// Check if table exists
let table_names = connection
.table_names()
.execute()
.await
.context("Failed to list tables")?;
if !table_names.contains(&self.table_name) {
// Create the table schema
let schema = Arc::new(Schema::new(vec![
Field::new("tool_name", DataType::Utf8, false),
Field::new("description", DataType::Utf8, false),
Field::new("schema", DataType::Utf8, false),
Field::new(
"vector",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
1536, // OpenAI embedding dimension
),
false,
),
]));
// Create empty table
let tool_names = StringArray::from(vec![] as Vec<&str>);
let descriptions = StringArray::from(vec![] as Vec<&str>);
let schemas = StringArray::from(vec![] as Vec<&str>);
// Create empty fixed size list array for vectors
let mut vectors_builder =
FixedSizeListBuilder::new(arrow::array::Float32Builder::new(), 1536);
let vectors = vectors_builder.finish();
let batch = arrow::record_batch::RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(tool_names),
Arc::new(descriptions),
Arc::new(schemas),
Arc::new(vectors),
],
)
.context("Failed to create record batch")?;
// Create an empty table with the schema
// LanceDB will create the table from the RecordBatch
drop(connection);
let connection = self.connection.write().await;
// Use the RecordBatch directly
let reader = arrow::record_batch::RecordBatchIterator::new(
vec![Ok(batch)].into_iter(),
schema.clone(),
);
connection
.create_table(&self.table_name, Box::new(reader))
.execute()
.await
.map_err(|e| {
anyhow::anyhow!("Failed to create tools table '{}': {}", self.table_name, e)
})?;
}
Ok(())
}
#[cfg(test)]
pub async fn clear_tools(&self) -> Result<()> {
let connection = self.connection.write().await;
// Try to open the table first
match connection.open_table(&self.table_name).execute().await {
Ok(table) => {
// Delete all records instead of dropping the table
table
.delete("1=1") // This will match all records
.await
.context("Failed to delete all records")?;
}
Err(_) => {
// If table doesn't exist, that's fine - we'll create it
}
}
drop(connection);
// Ensure table exists with correct schema
self.init_table().await?;
Ok(())
}
pub async fn index_tools(&self, tools: Vec<ToolRecord>) -> Result<()> {
if tools.is_empty() {
return Ok(());
}
let tool_names: Vec<&str> = tools.iter().map(|t| t.tool_name.as_str()).collect();
let descriptions: Vec<&str> = tools.iter().map(|t| t.description.as_str()).collect();
let schemas: Vec<&str> = tools.iter().map(|t| t.schema.as_str()).collect();
let vectors_data: Vec<Option<Vec<Option<f32>>>> = tools
.iter()
.map(|t| Some(t.vector.iter().map(|&v| Some(v)).collect()))
.collect();
let schema = Arc::new(Schema::new(vec![
Field::new("tool_name", DataType::Utf8, false),
Field::new("description", DataType::Utf8, false),
Field::new("schema", DataType::Utf8, false),
Field::new(
"vector",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
1536,
),
false,
),
]));
let tool_names_array = StringArray::from(tool_names);
let descriptions_array = StringArray::from(descriptions);
let schemas_array = StringArray::from(schemas);
// Build vectors array
let mut vectors_builder =
FixedSizeListBuilder::new(arrow::array::Float32Builder::new(), 1536);
for vector_opt in vectors_data {
if let Some(vector) = vector_opt {
let values = vectors_builder.values();
for val_opt in vector {
if let Some(val) = val_opt {
values.append_value(val);
} else {
values.append_null();
}
}
vectors_builder.append(true);
} else {
vectors_builder.append(false);
}
}
let vectors_array = vectors_builder.finish();
let batch = arrow::record_batch::RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(tool_names_array),
Arc::new(descriptions_array),
Arc::new(schemas_array),
Arc::new(vectors_array),
],
)
.context("Failed to create record batch")?;
let connection = self.connection.read().await;
let table = connection
.open_table(&self.table_name)
.execute()
.await
.context("Failed to open tools table")?;
// Add batch to table using RecordBatchIterator
let reader = arrow::record_batch::RecordBatchIterator::new(
vec![Ok(batch)].into_iter(),
schema.clone(),
);
table
.add(Box::new(reader))
.execute()
.await
.context("Failed to add tools to table")?;
Ok(())
}
pub async fn search_tools(&self, query_vector: Vec<f32>, k: usize) -> Result<Vec<ToolRecord>> {
let connection = self.connection.read().await;
let table = connection
.open_table(&self.table_name)
.execute()
.await
.context("Failed to open tools table")?;
let results = table
.vector_search(query_vector)
.context("Failed to create vector search")?
.limit(k)
.execute()
.await
.context("Failed to execute vector search")?;
let batches: Vec<_> = results.try_collect().await?;
let mut tools = Vec::new();
for batch in batches {
let tool_names = batch
.column_by_name("tool_name")
.context("Missing tool_name column")?
.as_any()
.downcast_ref::<StringArray>()
.context("Invalid tool_name column type")?;
let descriptions = batch
.column_by_name("description")
.context("Missing description column")?
.as_any()
.downcast_ref::<StringArray>()
.context("Invalid description column type")?;
let schemas = batch
.column_by_name("schema")
.context("Missing schema column")?
.as_any()
.downcast_ref::<StringArray>()
.context("Invalid schema column type")?;
// Get the distance scores
let distances = batch
.column_by_name("_distance")
.context("Missing _distance column")?
.as_any()
.downcast_ref::<arrow::array::Float32Array>()
.context("Invalid _distance column type")?;
for i in 0..batch.num_rows() {
let tool_name = tool_names.value(i).to_string();
let _distance = distances.value(i);
tools.push(ToolRecord {
tool_name,
description: descriptions.value(i).to_string(),
schema: schemas.value(i).to_string(),
vector: vec![], // We don't need to return the vector
});
}
}
Ok(tools)
}
pub async fn remove_tool(&self, tool_name: &str) -> Result<()> {
let connection = self.connection.read().await;
let table = connection
.open_table(&self.table_name)
.execute()
.await
.context("Failed to open tools table")?;
// Delete records matching the tool name
table
.delete(&format!("tool_name = '{}'", tool_name))
.await
.context("Failed to delete tool")?;
Ok(())
}
}
pub fn generate_table_id() -> String {
Local::now().format("%Y%m%d_%H%M%S").to_string()
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_tool_vectordb_creation() {
let db = ToolVectorDB::new(Some("test_tools_vectordb_creation".to_string()))
.await
.unwrap();
db.clear_tools().await.unwrap();
assert_eq!(db.table_name, "test_tools_vectordb_creation");
}
#[tokio::test]
async fn test_tool_vectordb_operations() -> Result<()> {
// Create a new database instance with a unique table name
let db = ToolVectorDB::new(Some("test_tool_vectordb_operations".to_string())).await?;
// Clear any existing tools
db.clear_tools().await?;
// Create test tool records
let test_tools = vec![
ToolRecord {
tool_name: "test_tool_1".to_string(),
description: "A test tool for reading files".to_string(),
schema: r#"{"type": "object", "properties": {"path": {"type": "string"}}}"#
.to_string(),
vector: vec![0.1; 1536], // Mock embedding vector
},
ToolRecord {
tool_name: "test_tool_2".to_string(),
description: "A test tool for writing files".to_string(),
schema: r#"{"type": "object", "properties": {"path": {"type": "string"}}}"#
.to_string(),
vector: vec![0.2; 1536], // Different mock embedding vector
},
];
// Index the test tools
db.index_tools(test_tools).await?;
// Search for tools using a query vector similar to test_tool_1
let query_vector = vec![0.1; 1536];
let results = db.search_tools(query_vector, 2).await?;
// Verify results
assert_eq!(results.len(), 2, "Should find both tools");
assert_eq!(
results[0].tool_name, "test_tool_1",
"First result should be test_tool_1"
);
assert_eq!(
results[1].tool_name, "test_tool_2",
"Second result should be test_tool_2"
);
Ok(())
}
#[tokio::test]
async fn test_empty_db() -> Result<()> {
// Create a new database instance with a unique table name
let db = ToolVectorDB::new(Some("test_empty_db".to_string())).await?;
// Clear any existing tools
db.clear_tools().await?;
// Search in empty database
let query_vector = vec![0.1; 1536];
let results = db.search_tools(query_vector, 2).await?;
// Verify no results returned
assert_eq!(results.len(), 0, "Empty database should return no results");
Ok(())
}
#[tokio::test]
async fn test_tool_deletion() -> Result<()> {
// Create a new database instance with a unique table name
let db = ToolVectorDB::new(Some("test_tool_deletion".to_string())).await?;
// Clear any existing tools
db.clear_tools().await?;
// Create and index a test tool
let test_tool = ToolRecord {
tool_name: "test_tool_to_delete".to_string(),
description: "A test tool that will be deleted".to_string(),
schema: r#"{"type": "object", "properties": {"path": {"type": "string"}}}"#.to_string(),
vector: vec![0.1; 1536],
};
db.index_tools(vec![test_tool]).await?;
// Verify tool exists
let query_vector = vec![0.1; 1536];
let results = db.search_tools(query_vector.clone(), 1).await?;
assert_eq!(results.len(), 1, "Tool should exist before deletion");
// Delete the tool
db.remove_tool("test_tool_to_delete").await?;
// Verify tool is gone
let results = db.search_tools(query_vector.clone(), 1).await?;
assert_eq!(results.len(), 0, "Tool should be deleted");
Ok(())
}
}
crates/goose/src/providers/base.rs
+12
View File
@@ -183,6 +183,18 @@ pub trait Provider: Send + Sync {
async fn fetch_supported_models_async(&self) -> Result<Option<Vec<String>>, ProviderError> {
Ok(None)
}
/// Check if this provider supports embeddings
fn supports_embeddings(&self) -> bool {
false
}
/// Create embeddings if supported. Default implementation returns an error.
async fn create_embeddings(&self, _texts: Vec<String>) -> Result<Vec<Vec<f32>>, ProviderError> {
Err(ProviderError::ExecutionError(
"This provider does not support embeddings".to_string(),
))
}
}
#[cfg(test)]
crates/goose/src/providers/databricks.rs
+64 -10
View File
@@ -1,11 +1,5 @@
use anyhow::Result;
use async_trait::async_trait;
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::Duration;
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::embedding::EmbeddingCapable;
use super::errors::ProviderError;
use super::formats::databricks::{create_request, get_usage, response_to_message};
use super::oauth;
@@ -14,8 +8,16 @@ use crate::config::ConfigError;
use crate::message::Message;
use crate::model::ModelConfig;
use mcp_core::tool::Tool;
use serde_json::json;
use url::Url;
use anyhow::Result;
use async_trait::async_trait;
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::Duration;
const DEFAULT_CLIENT_ID: &str = "databricks-cli";
const DEFAULT_REDIRECT_URL: &str = "http://localhost:8020";
// "offline_access" scope is used to request an OAuth 2.0 Refresh Token
@@ -128,7 +130,6 @@ impl DatabricksProvider {
///
/// * `host` - The Databricks host URL
/// * `token` - The Databricks API token
/// * `model` - The model configuration
///
/// # Returns
///
@@ -166,7 +167,17 @@ impl DatabricksProvider {
async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
let base_url = Url::parse(&self.host)
.map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
let path = format!("serving-endpoints/{}/invocations", self.model.model_name);
// Check if this is an embedding request by looking at the payload structure
let is_embedding = payload.get("input").is_some() && payload.get("messages").is_none();
let path = if is_embedding {
// For embeddings, use the embeddings endpoint
format!("serving-endpoints/{}/invocations", "text-embedding-3-small")
} else {
// For chat completions, use the model name in the path
format!("serving-endpoints/{}/invocations", self.model.model_name)
};
let url = base_url.join(&path).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;
@@ -184,7 +195,7 @@ impl DatabricksProvider {
let payload: Option<Value> = response.json().await.ok();
match status {
StatusCode::OK => payload.ok_or_else( || ProviderError::RequestFailed("Response body is not valid JSON".to_string()) ),
StatusCode::OK => payload.ok_or_else(|| ProviderError::RequestFailed("Response body is not valid JSON".to_string())),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
Err(ProviderError::Authentication(format!("Authentication failed. Please ensure your API keys are valid and have the required permissions. \
Status: {}. Response: {:?}", status, payload)))
@@ -295,4 +306,47 @@ impl Provider for DatabricksProvider {
Ok((message, ProviderUsage::new(model, usage)))
}
fn supports_embeddings(&self) -> bool {
true
}
async fn create_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>, ProviderError> {
EmbeddingCapable::create_embeddings(self, texts)
.await
.map_err(|e| ProviderError::ExecutionError(e.to_string()))
}
}
#[async_trait]
impl EmbeddingCapable for DatabricksProvider {
async fn create_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>> {
if texts.is_empty() {
return Ok(vec![]);
}
// Create request in Databricks format for embeddings
let request = json!({
"input": texts,
});
let response = self.post(request).await?;
let embeddings = response["data"]
.as_array()
.ok_or_else(|| anyhow::anyhow!("Invalid response format: missing data array"))?
.iter()
.map(|item| {
item["embedding"]
.as_array()
.ok_or_else(|| anyhow::anyhow!("Invalid embedding format"))?
.iter()
.map(|v| v.as_f64().map(|f| f as f32))
.collect::<Option<Vec<f32>>>()
.ok_or_else(|| anyhow::anyhow!("Invalid embedding values"))
})
.collect::<Result<Vec<Vec<f32>>>>()?;
Ok(embeddings)
}
}
crates/goose/src/providers/embedding.rs
+24
View File
@@ -0,0 +1,24 @@
use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingRequest {
pub input: Vec<String>,
pub model: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingResponse {
pub data: Vec<EmbeddingData>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingData {
pub embedding: Vec<f32>,
}
#[async_trait]
pub trait EmbeddingCapable {
async fn create_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>>;
}
crates/goose/src/providers/mod.rs
+1
View File
@@ -4,6 +4,7 @@ pub mod azureauth;
pub mod base;
pub mod bedrock;
pub mod databricks;
pub mod embedding;
pub mod errors;
mod factory;
pub mod formats;
crates/goose/src/providers/openai.rs
+85 -12
View File
@@ -6,6 +6,7 @@ use std::collections::HashMap;
use std::time::Duration;
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::embedding::{EmbeddingCapable, EmbeddingRequest, EmbeddingResponse};
use super::errors::ProviderError;
use super::formats::openai::{create_request, get_usage, response_to_message};
use super::utils::{emit_debug_trace, get_model, handle_response_openai_compat, ImageFormat};
@@ -80,18 +81,8 @@ impl OpenAiProvider {
})
}
async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
let base_url = url::Url::parse(&self.host)
.map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
let url = base_url.join(&self.base_path).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;
let mut request = self
.client
.post(url)
.header("Authorization", format!("Bearer {}", self.api_key));
/// Helper function to add OpenAI-specific headers to a request
fn add_headers(&self, mut request: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
// Add organization header if present
if let Some(org) = &self.organization {
request = request.header("OpenAI-Organization", org);
@@ -102,12 +93,30 @@ impl OpenAiProvider {
request = request.header("OpenAI-Project", project);
}
// Add custom headers if present
if let Some(custom_headers) = &self.custom_headers {
for (key, value) in custom_headers {
request = request.header(key, value);
}
}
request
}
async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
let base_url = url::Url::parse(&self.host)
.map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
let url = base_url.join(&self.base_path).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;
let request = self
.client
.post(url)
.header("Authorization", format!("Bearer {}", self.api_key));
let request = self.add_headers(request);
let response = request.json(&payload).send().await?;
handle_response_openai_compat(response).await
@@ -209,6 +218,16 @@ impl Provider for OpenAiProvider {
models.sort();
Ok(Some(models))
}
fn supports_embeddings(&self) -> bool {
true
}
async fn create_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>, ProviderError> {
EmbeddingCapable::create_embeddings(self, texts)
.await
.map_err(|e| ProviderError::ExecutionError(e.to_string()))
}
}
fn parse_custom_headers(s: String) -> HashMap<String, String> {
@@ -221,3 +240,57 @@ fn parse_custom_headers(s: String) -> HashMap<String, String> {
})
.collect()
}
#[async_trait]
impl EmbeddingCapable for OpenAiProvider {
async fn create_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>> {
if texts.is_empty() {
return Ok(vec![]);
}
// Get embedding model from env var or use default
let embedding_model = std::env::var("EMBEDDING_MODEL")
.unwrap_or_else(|_| "text-embedding-3-small".to_string());
let request = EmbeddingRequest {
input: texts,
model: embedding_model,
};
// Construct embeddings endpoint URL
let base_url =
url::Url::parse(&self.host).map_err(|e| anyhow::anyhow!("Invalid base URL: {e}"))?;
let url = base_url
.join("v1/embeddings")
.map_err(|e| anyhow::anyhow!("Failed to construct embeddings URL: {e}"))?;
let req = self
.client
.post(url)
.header("Authorization", format!("Bearer {}", self.api_key))
.json(&request);
let req = self.add_headers(req);
let response = req
.send()
.await
.map_err(|e| anyhow::anyhow!("Failed to send embedding request: {e}"))?;
if !response.status().is_success() {
let error_text = response.text().await.unwrap_or_default();
return Err(anyhow::anyhow!("Embedding API error: {}", error_text));
}
let embedding_response: EmbeddingResponse = response
.json()
.await
.map_err(|e| anyhow::anyhow!("Failed to parse embedding response: {e}"))?;
Ok(embedding_response
.data
.into_iter()
.map(|d| d.embedding)
.collect())
}
}