yapi_core/extensions/
tool_search.rs

1//! The `tool_search` tool: a BM25 ranker over the metadata of tools that are
2//! not declared to the model (`codemode` and `deferred` exposure), which
3//! activates the matches for the next model call. Port of
4//! `extensions/tool-search` in pi `v1.0.0`.
5
6use std::collections::HashMap;
7use std::sync::{Arc, LazyLock};
8
9use futures_util::future::BoxFuture;
10use regex_lite::Regex;
11use serde_json::{Value, json};
12use tokio_util::sync::CancellationToken;
13use yapi_agent::{Tool, UpdateSink};
14use yapi_types::event::ToolResult;
15use yapi_types::message::ToolDeclaration;
16use yapi_types::rpc::SourceInfo;
17
18use super::{Extension, ToolInfo, Tools, builtin_source};
19use crate::tools::{Exposure, Namespace, RegisteredTool};
20
21/// The tool's name.
22pub const TOOL_SEARCH_TOOL_NAME: &str = "tool_search";
23/// Matches loaded when the call gives no limit.
24pub const DEFAULT_LIMIT: usize = 8;
25
26const STOP_WORDS: [&str; 21] = [
27    "a", "an", "and", "are", "as", "at", "be", "by", "for", "from", "in", "is", "it", "of", "on",
28    "or", "that", "the", "this", "to", "with",
29];
30
31/// The description, which does not list the searchable tools so it stays the
32/// same while tools are registered.
33pub const DESCRIPTION: &str = "# Tool discovery\n\nSearches over deferred tool metadata with BM25 and exposes matching tools for the next model call.\n\nSome of the tools, such as tools of MCP servers, may not have been provided to you upfront, and you should use this tool (`tool_search`) to search for the required tools. For MCP tool discovery, always use `tool_search`.";
34
35static LOWER_UPPER: LazyLock<Option<Regex>> =
36    LazyLock::new(|| Regex::new("([a-z0-9])([A-Z])").ok());
37static ACRONYM: LazyLock<Option<Regex>> = LazyLock::new(|| Regex::new("([A-Z]+)([A-Z][a-z])").ok());
38static PLURAL: LazyLock<Option<Regex>> =
39    LazyLock::new(|| Regex::new("(ches|shes|sses|xes|zes)$").ok());
40
41/// Naive singular form, so `issues` matches `issue`.
42fn stem(term: &str) -> String {
43    let length = term.chars().count();
44    if length > 4 && term.ends_with("ies") {
45        return format!("{}y", &term[..term.len() - 3]);
46    }
47    if length > 4 && PLURAL.as_ref().is_some_and(|plural| plural.is_match(term)) {
48        return term[..term.len() - 2].to_owned();
49    }
50    if length > 3 && term.ends_with('s') && !term.ends_with("ss") {
51        return term[..term.len() - 1].to_owned();
52    }
53    term.to_owned()
54}
55
56/// Lowercase terms split at camelCase boundaries and non-alphanumerics,
57/// without stop words.
58pub fn tokenize(text: &str) -> Vec<String> {
59    let mut text = text.to_owned();
60    if let Some(regex) = LOWER_UPPER.as_ref() {
61        text = regex.replace_all(&text, "$1 $2").into_owned();
62    }
63    if let Some(regex) = ACRONYM.as_ref() {
64        text = regex.replace_all(&text, "$1 $2").into_owned();
65    }
66    text.to_lowercase()
67        .split(|c: char| !(c.is_ascii_lowercase() || c.is_ascii_digit()))
68        .filter(|term| !term.is_empty() && !STOP_WORDS.contains(term))
69        .map(stem)
70        .collect()
71}
72
73/// Schema descriptions and property names, recursively.
74fn schema_text(schema: &Value, parts: &mut Vec<String>) {
75    let Some(object) = schema.as_object() else {
76        return;
77    };
78    if let Some(description) = object.get("description").and_then(Value::as_str) {
79        parts.push(description.to_owned());
80    }
81    if let Some(properties) = object.get("properties").and_then(Value::as_object) {
82        for (name, property) in properties {
83            parts.push(name.clone());
84            schema_text(property, parts);
85        }
86    }
87    if let Some(items) = object.get("items") {
88        schema_text(items, parts);
89    }
90    for key in ["anyOf", "oneOf", "allOf"] {
91        for variant in object
92            .get(key)
93            .and_then(Value::as_array)
94            .into_iter()
95            .flatten()
96        {
97            schema_text(variant, parts);
98        }
99    }
100}
101
102/// A tool's search text: its name, the name with `_` as spaces, the
103/// description, schema text, and its namespace.
104pub fn document(tool: &ToolInfo) -> String {
105    let mut parts = vec![
106        tool.name.clone(),
107        tool.name.replace('_', " "),
108        tool.description.clone(),
109    ];
110    schema_text(&tool.parameters, &mut parts);
111    if let Some(Namespace {
112        name,
113        description,
114        instructions,
115    }) = &tool.namespace
116    {
117        parts.push(name.clone());
118        parts.push(description.clone().unwrap_or_default());
119        parts.push(instructions.clone().unwrap_or_default());
120    }
121    parts
122        .into_iter()
123        .filter(|part| !part.trim().is_empty())
124        .collect::<Vec<_>>()
125        .join(" ")
126}
127
128/// Okapi BM25 (k1 1.2, b 0.75) over `documents`: the best `limit` names with a
129/// positive score. Ties keep document order.
130pub fn rank(query: &str, documents: &[(String, String)], limit: usize) -> Vec<(String, f64)> {
131    let mut terms = Vec::new();
132    for term in tokenize(query) {
133        if !terms.contains(&term) {
134            terms.push(term);
135        }
136    }
137    if terms.is_empty() || documents.is_empty() || limit == 0 {
138        return Vec::new();
139    }
140    let (k1, b) = (1.2, 0.75);
141    let counts: Vec<HashMap<String, usize>> = documents
142        .iter()
143        .map(|(_, text)| {
144            let mut counts = HashMap::new();
145            for term in tokenize(text) {
146                *counts.entry(term).or_insert(0) += 1;
147            }
148            counts
149        })
150        .collect();
151    let lengths: Vec<f64> = counts
152        .iter()
153        .map(|counts| counts.values().sum::<usize>() as f64)
154        .collect();
155    let average = lengths.iter().sum::<f64>() / documents.len() as f64;
156    let average = if average == 0.0 { 1.0 } else { average };
157    let total = documents.len() as f64;
158    let idf: HashMap<&str, f64> = terms
159        .iter()
160        .map(|term| {
161            let frequency = counts
162                .iter()
163                .filter(|counts| counts.contains_key(term))
164                .count() as f64;
165            (
166                term.as_str(),
167                (1.0 + (total - frequency + 0.5) / (frequency + 0.5)).ln(),
168            )
169        })
170        .collect();
171    let mut matches: Vec<(String, f64)> = Vec::new();
172    for (index, (name, _)) in documents.iter().enumerate() {
173        let mut score = 0.0;
174        for term in &terms {
175            let Some(&count) = counts[index].get(term) else {
176                continue;
177            };
178            let count = count as f64;
179            let norm = k1 * (1.0 - b + b * lengths[index] / average);
180            score += idf.get(term.as_str()).copied().unwrap_or(0.0) * (count * (k1 + 1.0))
181                / (count + norm);
182        }
183        if score > 0.0 {
184            matches.push((name.clone(), score));
185        }
186    }
187    matches.sort_by(|a, b| b.1.total_cmp(&a.1));
188    matches.truncate(limit);
189    matches
190}
191
192/// The `tool_search` tool over a session's tools.
193pub struct ToolSearch {
194    declaration: ToolDeclaration,
195    tools: Tools,
196}
197
198impl ToolSearch {
199    /// The tool, searching `tools`.
200    pub fn new(tools: Tools) -> ToolSearch {
201        ToolSearch {
202            declaration: ToolDeclaration {
203                name: TOOL_SEARCH_TOOL_NAME.into(),
204                description: DESCRIPTION.into(),
205                parameters: json!({
206                    "type": "object",
207                    "required": ["query"],
208                    "properties": {
209                        "query": {"type": "string", "description": "Search query for deferred tools."},
210                        "limit": {"type": "number", "description": format!("Maximum number of tools to return. Defaults to {DEFAULT_LIMIT}.")},
211                    },
212                }),
213                constrained_sampling: None,
214            },
215            tools,
216        }
217    }
218
219    /// Ranks the searchable tools that are not active and activates the matches.
220    fn search_and_load(&self, query: &str, limit: usize) -> Vec<(String, String)> {
221        let active = self.tools.active();
222        let candidates: Vec<ToolInfo> = self
223            .tools
224            .all()
225            .into_iter()
226            .filter(|tool| {
227                matches!(tool.exposure, Exposure::Codemode | Exposure::Deferred)
228                    && !active.contains(&tool.name)
229            })
230            .collect();
231        let documents: Vec<(String, String)> = candidates
232            .iter()
233            .map(|tool| (tool.name.clone(), document(tool)))
234            .collect();
235        let matches = rank(query, &documents, limit);
236        if !matches.is_empty() {
237            let mut next = active;
238            next.extend(matches.iter().map(|(name, _)| name.clone()));
239            self.tools.set_active(next);
240        }
241        matches
242            .into_iter()
243            .map(|(name, _)| {
244                let description = candidates
245                    .iter()
246                    .find(|tool| tool.name == name)
247                    .map(|tool| tool.description.clone())
248                    .unwrap_or_default();
249                (name, description)
250            })
251            .collect()
252    }
253}
254
255impl Tool for ToolSearch {
256    fn declaration(&self) -> &ToolDeclaration {
257        &self.declaration
258    }
259
260    fn execute(
261        &self,
262        _call_id: String,
263        args: Value,
264        _cancel: CancellationToken,
265        _updates: UpdateSink,
266    ) -> BoxFuture<'_, Result<ToolResult, String>> {
267        Box::pin(async move {
268            let query = args["query"].as_str().unwrap_or_default();
269            if query.trim().is_empty() {
270                return Err("query must not be empty".into());
271            }
272            let limit = match args.get("limit").and_then(Value::as_f64) {
273                None => DEFAULT_LIMIT,
274                Some(limit) if limit.fract() == 0.0 && limit > 0.0 => limit as usize,
275                Some(_) => return Err("limit must be a positive integer".into()),
276            };
277            let loaded = self.search_and_load(query, limit);
278            let text = if loaded.is_empty() {
279                "No matching tools found.".to_owned()
280            } else {
281                let lines: Vec<String> = loaded
282                    .iter()
283                    .map(|(name, description)| {
284                        let first = description.trim().lines().next().unwrap_or_default();
285                        format!("- {name}: {first}")
286                    })
287                    .collect();
288                format!(
289                    "Loaded {} tool{}. They are available from your next call:\n{}",
290                    loaded.len(),
291                    if loaded.len() == 1 { "" } else { "s" },
292                    lines.join("\n")
293                )
294            };
295            let names: Vec<&str> = loaded.iter().map(|(name, _)| name.as_str()).collect();
296            Ok(ToolResult {
297                content: vec![yapi_types::message::ContentBlock::text(text)],
298                details: Some(json!({ "loaded": names })),
299                ..ToolResult::default()
300            })
301        })
302    }
303}
304
305/// The built-in extension that registers `tool_search`, inactive until named
306/// by `--tools`, `defaultTools` or another extension.
307pub struct ToolSearchExtension;
308
309impl Extension for ToolSearchExtension {
310    fn source(&self) -> SourceInfo {
311        builtin_source("tool-search")
312    }
313
314    fn load(&self, tools: &Tools) {
315        tools.register(RegisteredTool {
316            tool: Arc::new(ToolSearch::new(tools.clone())),
317            snippet: Some("Search for tools that are not loaded yet and load the matches".into()),
318            guidelines: Vec::new(),
319            exposure: Exposure::ModelOnly,
320            namespace: None,
321            default_active: false,
322            annotations: None,
323            source: Some(builtin_source("tool-search")),
324        });
325    }
326}
327
328/// Whether `tool` is this crate's `tool_search`.
329pub fn is_tool_search(tool: &ToolInfo) -> bool {
330    tool.name == TOOL_SEARCH_TOOL_NAME && tool.exposure == Exposure::ModelOnly
331}
332
333#[cfg(test)]
334mod tests {
335    use super::*;
336
337    #[test]
338    fn tokenizes_like_pi() {
339        assert_eq!(
340            tokenize("searchIssues in GitHubAPI for the HTTPServer"),
341            ["search", "issue", "git", "hub", "api", "http", "server"]
342        );
343        assert_eq!(tokenize("watches boxes classes"), ["watch", "box", "class"]);
344        assert!(tokenize("the and of").is_empty());
345    }
346
347    #[test]
348    fn ranks_matching_tools_first() {
349        let documents = vec![
350            (
351                "mcp__gh__issues".to_owned(),
352                "mcp__gh__issues mcp gh issues List issues of a repository".to_owned(),
353            ),
354            (
355                "mcp__fs__read".to_owned(),
356                "mcp__fs__read mcp fs read Read a file".to_owned(),
357            ),
358        ];
359        let ranked = rank("find repository issues", &documents, 8);
360        assert_eq!(ranked.len(), 1);
361        assert_eq!(ranked[0].0, "mcp__gh__issues");
362        assert!(rank("", &documents, 8).is_empty());
363    }
364}
yapi_core/extensions/tool_search.rs

yapi_core/extensions/
tool_search.rs
