Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 21 additions & 7 deletions crates/goose/src/agents/router_tool_selector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use rmcp::model::Tool;

use anyhow::{Context, Result};
use async_trait::async_trait;
use serde::Serialize;
use serde_json::Value;
use std::collections::HashMap;
use std::collections::VecDeque;
Expand All @@ -14,8 +15,15 @@ use tokio::sync::RwLock;
use crate::agents::tool_vectordb::ToolVectorDB;
use crate::conversation::message::Message;
use crate::model::ModelConfig;
use crate::prompt_template::render_global_file;
use crate::providers::{self, base::Provider};

#[derive(Serialize)]
struct ToolSelectorContext {
tools: String,
query: String,
}

#[derive(Debug, Clone, PartialEq)]
pub enum RouterToolSelectionStrategy {
Vector,
Expand Down Expand Up @@ -282,15 +290,21 @@ impl RouterToolSelector for LLMToolSelector {
};

if let Some(tools) = relevant_tools {
// Use LLM to search through tools
let prompt = format!(
"Given the following tools:\n{}\n\nFind the most relevant tools for the query: {}\n\nReturn the tools in this exact format for each tool:\nTool: <tool_name>\nDescription: <tool_description>\nSchema: <tool_schema>",
tools, query
);
let system_message = Message::user().with_text("You are a tool selection assistant. Your task is to find the most relevant tools based on the user's query.");
// 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| {
ToolError::ExecutionError(format!("Failed to render prompt template: {}", e))
})?;

let user_message = Message::user().with_text(&user_prompt);
let response = self
.llm_provider
.complete(&prompt, &[system_message], &[])
.complete("", &[user_message], &[])
.await
.map_err(|e| ToolError::ExecutionError(format!("Failed to search tools: {}", e)))?;

Expand Down
16 changes: 11 additions & 5 deletions crates/goose/src/permission/permission_judge.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use crate::config::permission::PermissionLevel;
use crate::config::PermissionManager;
use crate::conversation::message::{Message, MessageContent, ToolRequest};
use crate::conversation::Conversation;
use crate::prompt_template::render_global_file;
use crate::providers::base::Provider;
use chrono::Utc;
use indoc::indoc;
Expand All @@ -13,6 +14,11 @@ use serde_json::Value;
use std::collections::HashSet;
use std::sync::Arc;

#[derive(Serialize)]
struct PermissionJudgeContext {
// Empty struct for now since the current template doesn't need variables
}

/// Creates the tool definition for checking read-only permissions.
fn create_read_only_tool() -> Tool {
Tool::new(
Expand Down Expand Up @@ -133,12 +139,12 @@ pub async fn detect_read_only_tools(
let tool = create_read_only_tool();
let check_messages = create_check_messages(tool_requests);

let context = PermissionJudgeContext {};
let system_prompt = render_global_file("permission_judge.md", &context)
.unwrap_or_else(|_| "You are a good analyst and can detect operations whether they have read-only operations.".to_string());

let res = provider
.complete(
"You are a good analyst and can detect operations whether they have read-only operations.",
check_messages.messages(),
&[tool.clone()],
)
.complete(&system_prompt, check_messages.messages(), &[tool.clone()])
.await;

// Process the response and return an empty vector if the response is invalid
Expand Down
1 change: 1 addition & 0 deletions crates/goose/src/prompts/permission_judge.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
You are a good analyst and can detect operations whether they have read-only operations.
11 changes: 11 additions & 0 deletions crates/goose/src/prompts/router_tool_selector.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
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>
Loading