fix(tools): raise tool search default results

This commit is contained in:
Nightt
2026-05-29 11:46:49 +08:00
committed by Hunter Bown
parent 76089db1dc
commit 4bacda64fc
3 changed files with 132 additions and 11 deletions
+1 -1
View File
@@ -2115,7 +2115,7 @@ use self::tool_catalog::{
};
#[cfg(test)]
use self::tool_catalog::{
TOOL_SEARCH_BM25_NAME, maybe_activate_requested_deferred_tool,
TOOL_SEARCH_BM25_NAME, TOOL_SEARCH_REGEX_NAME, maybe_activate_requested_deferred_tool,
preflight_requested_deferred_tool, should_default_defer_tool,
};
use self::tool_execution::emit_tool_audit;
+90
View File
@@ -2097,6 +2097,96 @@ fn tool_search_activates_discovered_deferred_tools() {
assert!(active.contains("read_file"));
}
fn tool_search_catalog_with_matches(count: usize) -> Vec<Tool> {
let mut catalog = (0..count)
.map(|idx| Tool {
tool_type: None,
name: format!("matching_tool_{idx:03}"),
description: "Matching deferred test tool".to_string(),
input_schema: json!({"type":"object","properties":{"query":{"type":"string"}}}),
allowed_callers: Some(vec!["direct".to_string()]),
defer_loading: Some(true),
input_examples: None,
strict: None,
cache_control: None,
})
.collect::<Vec<_>>();
let always_load = HashSet::new();
ensure_advanced_tooling(&mut catalog, AppMode::Agent, &always_load);
catalog
}
fn tool_search_reference_count(result: &ToolResult) -> usize {
result
.metadata
.as_ref()
.and_then(|metadata| metadata.get("tool_references"))
.and_then(|references| references.as_array())
.map_or(0, Vec::len)
}
#[test]
fn tool_search_defaults_to_twenty_results_for_regex_and_bm25() {
let catalog = tool_search_catalog_with_matches(25);
for tool_name in [TOOL_SEARCH_REGEX_NAME, TOOL_SEARCH_BM25_NAME] {
let mut active = initial_active_tools(&catalog);
let result = execute_tool_search(
tool_name,
&json!({"query":"matching"}),
&catalog,
&mut active,
)
.expect("search succeeds");
assert_eq!(tool_search_reference_count(&result), 20);
}
}
#[test]
fn tool_search_respects_and_caps_max_results() {
let catalog = tool_search_catalog_with_matches(120);
let mut active = initial_active_tools(&catalog);
let limited = execute_tool_search(
TOOL_SEARCH_BM25_NAME,
&json!({"query":"matching","max_results":7}),
&catalog,
&mut active,
)
.expect("search succeeds");
assert_eq!(tool_search_reference_count(&limited), 7);
let mut active = initial_active_tools(&catalog);
let capped = execute_tool_search(
TOOL_SEARCH_REGEX_NAME,
&json!({"query":"matching","max_results":999}),
&catalog,
&mut active,
)
.expect("search succeeds");
assert_eq!(tool_search_reference_count(&capped), 100);
}
#[test]
fn tool_search_schema_exposes_max_results_default_and_cap() {
let mut catalog = Vec::new();
let always_load = HashSet::new();
ensure_advanced_tooling(&mut catalog, AppMode::Agent, &always_load);
for tool_name in [TOOL_SEARCH_REGEX_NAME, TOOL_SEARCH_BM25_NAME] {
let tool = catalog
.iter()
.find(|tool| tool.name == tool_name)
.expect("tool search definition exists");
let schema = &tool.input_schema["properties"]["max_results"];
assert_eq!(schema["default"], 20);
assert_eq!(schema["maximum"], 100);
assert_eq!(schema["minimum"], 1);
}
}
#[tokio::test]
async fn code_execution_runs_python_and_returns_result_payload() {
let tmp = tempdir().expect("tempdir");
+41 -10
View File
@@ -12,7 +12,7 @@ use std::time::Duration;
use serde_json::{Value, json};
use crate::models::Tool;
use crate::tools::spec::{ToolError, ToolResult, required_str};
use crate::tools::spec::{ToolError, ToolResult, optional_u64, required_str};
use crate::tui::app::AppMode;
pub(super) const MULTI_TOOL_PARALLEL_NAME: &str = "multi_tool_use.parallel";
@@ -20,10 +20,12 @@ pub(super) const REQUEST_USER_INPUT_NAME: &str = "request_user_input";
pub(super) const CODE_EXECUTION_TOOL_NAME: &str = "code_execution";
const CODE_EXECUTION_TOOL_TYPE: &str = "code_execution_20250825";
pub(super) use crate::tools::js_execution::JS_EXECUTION_TOOL_NAME;
const TOOL_SEARCH_REGEX_NAME: &str = "tool_search_tool_regex";
pub(super) const TOOL_SEARCH_REGEX_NAME: &str = "tool_search_tool_regex";
const TOOL_SEARCH_REGEX_TYPE: &str = "tool_search_tool_regex_20251119";
pub(super) const TOOL_SEARCH_BM25_NAME: &str = "tool_search_tool_bm25";
const TOOL_SEARCH_BM25_TYPE: &str = "tool_search_tool_bm25_20251119";
const TOOL_SEARCH_DEFAULT_MAX_RESULTS: usize = 20;
const TOOL_SEARCH_MAX_RESULTS_LIMIT: usize = 100;
pub(super) fn is_tool_search_tool(name: &str) -> bool {
matches!(name, TOOL_SEARCH_REGEX_NAME | TOOL_SEARCH_BM25_NAME)
@@ -178,7 +180,14 @@ pub(super) fn ensure_advanced_tooling(
input_schema: json!({
"type": "object",
"properties": {
"query": { "type": "string", "description": "Regex pattern to search tool names/descriptions/schema." }
"query": { "type": "string", "description": "Regex pattern to search tool names/descriptions/schema." },
"max_results": {
"type": "integer",
"minimum": 1,
"maximum": TOOL_SEARCH_MAX_RESULTS_LIMIT,
"default": TOOL_SEARCH_DEFAULT_MAX_RESULTS,
"description": "Maximum number of matching tool references to return."
}
},
"required": ["query"]
}),
@@ -198,7 +207,14 @@ pub(super) fn ensure_advanced_tooling(
input_schema: json!({
"type": "object",
"properties": {
"query": { "type": "string", "description": "Natural language query for tool discovery." }
"query": { "type": "string", "description": "Natural language query for tool discovery." },
"max_results": {
"type": "integer",
"minimum": 1,
"maximum": TOOL_SEARCH_MAX_RESULTS_LIMIT,
"default": TOOL_SEARCH_DEFAULT_MAX_RESULTS,
"description": "Maximum number of matching tool references to return."
}
},
"required": ["query"]
}),
@@ -280,7 +296,11 @@ fn tool_search_haystack(tool: &Tool) -> String {
)
}
fn discover_tools_with_regex(catalog: &[Tool], query: &str) -> Result<Vec<String>, ToolError> {
fn discover_tools_with_regex(
catalog: &[Tool],
query: &str,
max_results: usize,
) -> Result<Vec<String>, ToolError> {
let regex = regex::Regex::new(query)
.map_err(|err| ToolError::invalid_input(format!("Invalid regex query: {err}")))?;
@@ -293,14 +313,14 @@ fn discover_tools_with_regex(catalog: &[Tool], query: &str) -> Result<Vec<String
if regex.is_match(&hay) {
matches.push(tool.name.clone());
}
if matches.len() >= 5 {
if matches.len() >= max_results {
break;
}
}
Ok(matches)
}
fn discover_tools_with_bm25_like(catalog: &[Tool], query: &str) -> Vec<String> {
fn discover_tools_with_bm25_like(catalog: &[Tool], query: &str, max_results: usize) -> Vec<String> {
let terms: Vec<String> = query
.split_whitespace()
.map(|term| term.trim().to_lowercase())
@@ -330,7 +350,11 @@ fn discover_tools_with_bm25_like(catalog: &[Tool], query: &str) -> Vec<String> {
}
}
scored.sort_by(|a, b| b.0.cmp(&a.0).then_with(|| a.1.cmp(&b.1)));
scored.into_iter().take(5).map(|(_, name)| name).collect()
scored
.into_iter()
.take(max_results)
.map(|(_, name)| name)
.collect()
}
fn edit_distance(a: &str, b: &str) -> usize {
@@ -645,10 +669,17 @@ pub(super) fn execute_tool_search(
active_tools: &mut HashSet<String>,
) -> Result<ToolResult, ToolError> {
let query = required_str(input, "query")?;
let max_results = usize::try_from(optional_u64(
input,
"max_results",
TOOL_SEARCH_DEFAULT_MAX_RESULTS as u64,
))
.unwrap_or(TOOL_SEARCH_DEFAULT_MAX_RESULTS)
.clamp(1, TOOL_SEARCH_MAX_RESULTS_LIMIT);
let discovered = if tool_name == TOOL_SEARCH_REGEX_NAME {
discover_tools_with_regex(catalog, query)?
discover_tools_with_regex(catalog, query, max_results)?
} else {
discover_tools_with_bm25_like(catalog, query)
discover_tools_with_bm25_like(catalog, query, max_results)
};
for name in &discovered {