1use async_trait::async_trait;
37use serde::{Deserialize, Serialize};
38use serde_json::Value;
39
40#[async_trait]
46pub trait ToolEmbedder: Send + Sync {
47 async fn embed(&self, text: &str) -> Option<Vec<f32>>;
49}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
54#[serde(rename_all = "lowercase")]
55pub enum ToolKind {
56 Mcp,
57 Builtin,
58 Composio,
59 App,
60 #[serde(rename = "core-api")]
64 CoreApi,
65}
66
67impl ToolKind {
68 pub fn parse_filter(s: &str) -> Option<ToolKind> {
71 match s.trim().to_ascii_lowercase().as_str() {
72 "mcp" => Some(ToolKind::Mcp),
73 "builtin" => Some(ToolKind::Builtin),
74 "composio" => Some(ToolKind::Composio),
75 "app" => Some(ToolKind::App),
76 "core-api" | "core_api" | "coreapi" => Some(ToolKind::CoreApi),
79 _ => None, }
81 }
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct ToolDescriptor {
87 pub id: String,
89 pub name: String,
90 #[serde(default)]
92 pub description: String,
93 pub kind: ToolKind,
94 #[serde(default)]
95 pub arg_names: Vec<String>,
96 #[serde(default)]
97 pub arg_descriptions: Vec<String>,
98 #[serde(default, skip_serializing_if = "Option::is_none")]
99 pub score: Option<f32>,
100 #[serde(default, skip_serializing_if = "Option::is_none")]
102 pub meta: Option<Value>,
103 #[serde(default)]
105 pub widget_accessible: bool,
106 #[serde(default, skip_serializing_if = "Option::is_none")]
108 pub output_template: Option<String>,
109}
110
111impl ToolDescriptor {
112 pub fn matches_allowlist(&self, allowlist: &[String]) -> bool {
120 if self.kind == ToolKind::Composio {
121 return allowlist.iter().any(|e| e == &self.id);
122 }
123 let (server, name) = self
124 .id
125 .split_once("__")
126 .map_or((self.id.as_str(), self.name.as_str()), |(s, t)| (s, t));
127 allowlist
128 .iter()
129 .any(|e| e == &self.id || e == name || e == server)
130 }
131}
132
133#[derive(Debug, Clone, Serialize, Deserialize)]
135pub struct DescribedTool {
136 pub id: String,
137 pub name: String,
138 #[serde(default)]
139 pub description: String,
140 pub kind: ToolKind,
141 pub args: Vec<DescribedArg>,
142 #[serde(default)]
145 pub shallow: bool,
146 #[serde(default, skip_serializing_if = "Option::is_none")]
147 pub parameters: Option<Value>,
148}
149
150#[derive(Debug, Clone, Serialize, Deserialize)]
152pub struct DescribedArg {
153 pub name: String,
154 pub r#type: String,
155 #[serde(default)]
156 pub description: String,
157 pub required: bool,
158}
159
160pub fn arg_summary(schema: Option<&Value>) -> (Vec<String>, Vec<String>) {
164 let mut names = Vec::new();
165 let mut descs = Vec::new();
166 if let Some(props) = schema
167 .and_then(|s| s.get("properties"))
168 .and_then(Value::as_object)
169 {
170 for (name, def) in props {
171 names.push(name.clone());
172 descs.push(
173 def.get("description")
174 .and_then(Value::as_str)
175 .unwrap_or_default()
176 .to_string(),
177 );
178 }
179 }
180 (names, descs)
181}
182
183pub fn described_args(schema: Option<&Value>) -> Vec<DescribedArg> {
185 let Some(schema) = schema else {
186 return Vec::new();
187 };
188 let required: Vec<String> = schema
189 .get("required")
190 .and_then(Value::as_array)
191 .map(|a| {
192 a.iter()
193 .filter_map(Value::as_str)
194 .map(str::to_string)
195 .collect()
196 })
197 .unwrap_or_default();
198 let Some(props) = schema.get("properties").and_then(Value::as_object) else {
199 return Vec::new();
200 };
201 props
202 .iter()
203 .map(|(name, def)| DescribedArg {
204 name: name.clone(),
205 r#type: def
206 .get("type")
207 .and_then(Value::as_str)
208 .unwrap_or("string")
209 .to_string(),
210 description: def
211 .get("description")
212 .and_then(Value::as_str)
213 .unwrap_or_default()
214 .to_string(),
215 required: required.iter().any(|r| r == name),
216 })
217 .collect()
218}
219
220pub const RANKER_PREF_KEY: &str = "tools.active_ranker";
224
225#[derive(Debug, Clone, Copy, PartialEq, Eq)]
229pub enum ToolRanker {
230 Bm25,
232 Semantic,
236}
237
238impl ToolRanker {
239 pub fn from_pref(s: Option<&str>) -> ToolRanker {
241 match s.map(|v| v.trim().to_ascii_lowercase()).as_deref() {
242 Some("semantic") => ToolRanker::Semantic,
243 _ => ToolRanker::Bm25,
244 }
245 }
246
247 pub async fn rank(
256 self,
257 query: &str,
258 mut items: Vec<ToolDescriptor>,
259 limit: usize,
260 embedder: Option<&dyn ToolEmbedder>,
261 ) -> Vec<ToolDescriptor> {
262 let scored = match (self, embedder) {
263 (ToolRanker::Semantic, Some(embedder)) => {
264 semantic_score(query, &mut items, embedder).await
265 }
266 _ => false,
267 };
268 if !scored {
269 bm25_score(query, &mut items);
271 }
272 items.sort_by(|a, b| {
273 b.score
274 .unwrap_or(0.0)
275 .partial_cmp(&a.score.unwrap_or(0.0))
276 .unwrap_or(std::cmp::Ordering::Equal)
277 });
278 items.truncate(limit);
279 items
280 }
281}
282
283fn cosine(a: &[f32], b: &[f32]) -> f32 {
285 if a.len() != b.len() {
286 return 0.0;
287 }
288 let mut dot = 0.0_f32;
289 let mut na = 0.0_f32;
290 let mut nb = 0.0_f32;
291 for (x, y) in a.iter().zip(b.iter()) {
292 dot += x * y;
293 na += x * x;
294 nb += y * y;
295 }
296 let denom = na.sqrt() * nb.sqrt();
297 if denom > f32::EPSILON {
298 dot / denom
299 } else {
300 0.0
301 }
302}
303
304async fn semantic_score(
309 query: &str,
310 items: &mut [ToolDescriptor],
311 embedder: &dyn ToolEmbedder,
312) -> bool {
313 if query.trim().is_empty() || items.is_empty() {
314 return false;
315 }
316 let Some(q_vec) = embedder.embed(query).await else {
317 return false;
319 };
320 for d in items.iter_mut() {
321 let score = match embedder.embed(&doc_text(d)).await {
322 Some(doc_vec) => cosine(&q_vec, &doc_vec),
323 None => 0.0,
324 };
325 d.score = Some(score);
326 }
327 true
328}
329
330fn tokenize(s: &str) -> Vec<String> {
332 s.split(|c: char| !c.is_alphanumeric())
333 .filter(|t| !t.is_empty())
334 .map(|t| t.to_ascii_lowercase())
335 .collect()
336}
337
338fn doc_text(d: &ToolDescriptor) -> String {
340 let mut s = format!("{} {} {}", d.id, d.name, d.description);
341 for a in &d.arg_names {
342 s.push(' ');
343 s.push_str(a);
344 }
345 s
346}
347
348fn bm25_score(query: &str, items: &mut [ToolDescriptor]) {
351 const K1: f32 = 1.5;
352 const B: f32 = 0.75;
353 let q_terms = tokenize(query);
354 if q_terms.is_empty() {
355 for d in items.iter_mut() {
356 d.score = Some(0.0);
357 }
358 return;
359 }
360
361 let docs: Vec<Vec<String>> = items.iter().map(|d| tokenize(&doc_text(d))).collect();
362 let n = docs.len().max(1) as f32;
363 let avg_dl = docs.iter().map(|d| d.len() as f32).sum::<f32>() / n;
364 let avg_dl = if avg_dl == 0.0 { 1.0 } else { avg_dl };
365
366 let q_lower = query.trim().to_ascii_lowercase();
367
368 for (i, d) in items.iter_mut().enumerate() {
369 let doc = &docs[i];
370 let dl = doc.len() as f32;
371 let mut score = 0.0_f32;
372 for term in &q_terms {
373 let tf = doc.iter().filter(|w| *w == term).count() as f32;
374 if tf == 0.0 {
375 continue;
376 }
377 let df = docs.iter().filter(|dd| dd.contains(term)).count() as f32;
379 let idf = (((n - df + 0.5) / (df + 0.5)) + 1.0).ln();
380 let denom = tf + K1 * (1.0 - B + B * dl / avg_dl);
381 score += idf * (tf * (K1 + 1.0)) / denom;
382 }
383 if d.id.eq_ignore_ascii_case(&q_lower) || d.name.eq_ignore_ascii_case(&q_lower) {
385 score += 1000.0;
386 }
387 d.score = Some(score);
388 }
389}
390
391pub async fn run_search(
402 query: &str,
403 builtin_candidates: Vec<ToolDescriptor>,
404 composio_candidates: Vec<ToolDescriptor>,
405 kind: Option<ToolKind>,
406 limit: usize,
407 ranker: ToolRanker,
408 embedder: Option<&dyn ToolEmbedder>,
409) -> Vec<ToolDescriptor> {
410 let mut candidates: Vec<ToolDescriptor> = builtin_candidates
411 .into_iter()
412 .filter(|d| kind.is_none() || kind == Some(d.kind))
413 .collect();
414 candidates.extend(composio_candidates);
415 ranker.rank(query, candidates, limit, embedder).await
416}
417
418pub fn describe_composio(id: &str) -> DescribedTool {
422 let slug = id.strip_prefix("composio__").unwrap_or(id);
423 DescribedTool {
424 id: id.to_string(),
425 name: slug.to_string(),
426 description: String::new(),
427 kind: ToolKind::Composio,
428 args: vec![DescribedArg {
429 name: "arguments".to_string(),
430 r#type: "object".to_string(),
431 description: "Action-specific parameters for this Composio action.".to_string(),
432 required: false,
433 }],
434 shallow: true,
435 parameters: None,
436 }
437}
438
439pub fn describe_from_parts(
444 id: &str,
445 name: &str,
446 description: &str,
447 kind: ToolKind,
448 input_schema: Option<&Value>,
449) -> DescribedTool {
450 DescribedTool {
451 id: id.to_string(),
452 name: name.to_string(),
453 description: description.to_string(),
454 kind,
455 args: described_args(input_schema),
456 shallow: input_schema.is_none(),
457 parameters: input_schema.cloned(),
458 }
459}
460
461#[cfg(test)]
462mod tests {
463 use super::*;
464
465 fn desc(id: &str, name: &str, description: &str, kind: ToolKind) -> ToolDescriptor {
466 ToolDescriptor {
467 id: id.to_string(),
468 name: name.to_string(),
469 description: description.to_string(),
470 kind,
471 arg_names: Vec::new(),
472 arg_descriptions: Vec::new(),
473 score: None,
474 meta: None,
475 widget_accessible: false,
476 output_template: None,
477 }
478 }
479
480 #[test]
481 fn kind_serializes_lowercase() {
482 assert_eq!(serde_json::to_string(&ToolKind::Mcp).unwrap(), "\"mcp\"");
483 assert_eq!(
484 serde_json::to_string(&ToolKind::Builtin).unwrap(),
485 "\"builtin\""
486 );
487 assert_eq!(
488 serde_json::to_string(&ToolKind::Composio).unwrap(),
489 "\"composio\""
490 );
491 assert_eq!(serde_json::to_string(&ToolKind::App).unwrap(), "\"app\"");
492 assert_eq!(
494 serde_json::to_string(&ToolKind::CoreApi).unwrap(),
495 "\"core-api\""
496 );
497 }
498
499 #[test]
500 fn parse_filter_maps_any_to_none() {
501 assert_eq!(ToolKind::parse_filter("any"), None);
502 assert_eq!(ToolKind::parse_filter("nonsense"), None);
503 assert_eq!(ToolKind::parse_filter("mcp"), Some(ToolKind::Mcp));
504 assert_eq!(ToolKind::parse_filter("COMPOSIO"), Some(ToolKind::Composio));
505 assert_eq!(ToolKind::parse_filter("core-api"), Some(ToolKind::CoreApi));
507 assert_eq!(ToolKind::parse_filter("core_api"), Some(ToolKind::CoreApi));
508 assert_eq!(ToolKind::parse_filter("CoreApi"), Some(ToolKind::CoreApi));
509 }
510
511 #[test]
512 fn matches_allowlist_matches_id_name_or_server() {
513 let d = desc("spider__crawl", "crawl", "crawl a site", ToolKind::Mcp);
514 assert!(d.matches_allowlist(&["spider__crawl".to_string()])); assert!(d.matches_allowlist(&["crawl".to_string()])); assert!(d.matches_allowlist(&["spider".to_string()])); assert!(!d.matches_allowlist(&["other".to_string()]));
518 let c = desc("composio__slack", "Slack", "", ToolKind::Composio);
520 assert!(c.matches_allowlist(&["composio__slack".to_string()]));
521 assert!(!c.matches_allowlist(&["Slack".to_string()]));
522 }
523
524 #[tokio::test]
525 async fn bm25_ranks_exact_match_first() {
526 let items = vec![
527 desc("foo__search", "search", "search the web", ToolKind::Mcp),
528 desc(
529 "foo__send",
530 "send_message",
531 "send a search-related message",
532 ToolKind::Mcp,
533 ),
534 desc("foo__noise", "noise", "totally unrelated", ToolKind::Mcp),
535 ];
536 let ranked = ToolRanker::Bm25.rank("search", items, 8, None).await;
537 assert_eq!(ranked[0].name, "search", "exact name match ranks first");
538 assert!(ranked.iter().all(|d| d.score.is_some()));
539 assert_eq!(ranked.last().unwrap().name, "noise");
541 }
542
543 #[tokio::test]
544 async fn ranker_selectable_from_pref() {
545 assert_eq!(ToolRanker::from_pref(None), ToolRanker::Bm25);
546 assert_eq!(ToolRanker::from_pref(Some("bm25")), ToolRanker::Bm25);
547 assert_eq!(
548 ToolRanker::from_pref(Some("semantic")),
549 ToolRanker::Semantic
550 );
551 let items = vec![
554 desc("foo__search", "search", "find things", ToolKind::Mcp),
555 desc("foo__x", "x", "nothing", ToolKind::Mcp),
556 ];
557 let ranked = ToolRanker::Bm25.rank("search", items, 8, None).await;
558 assert_eq!(ranked[0].name, "search");
559 }
560
561 #[test]
562 fn described_args_extracts_required_flag() {
563 let schema = serde_json::json!({
564 "type": "object",
565 "properties": {
566 "url": { "type": "string", "description": "page url" },
567 "depth": { "type": "integer" }
568 },
569 "required": ["url"]
570 });
571 let mut args = described_args(Some(&schema));
572 args.sort_by(|a, b| a.name.cmp(&b.name));
573 assert_eq!(args.len(), 2);
574 let url = args.iter().find(|a| a.name == "url").unwrap();
575 assert_eq!(url.r#type, "string");
576 assert_eq!(url.description, "page url");
577 assert!(url.required);
578 let depth = args.iter().find(|a| a.name == "depth").unwrap();
579 assert_eq!(depth.r#type, "integer");
580 assert!(!depth.required);
581 }
582
583 #[test]
584 fn describe_composio_id_is_shallow() {
585 let d = describe_composio("composio__GITHUB_CREATE_ISSUE");
586 assert!(d.shallow);
587 assert_eq!(d.kind, ToolKind::Composio);
588 assert_eq!(d.name, "GITHUB_CREATE_ISSUE");
589 assert_eq!(d.args.len(), 1);
590 assert_eq!(d.args[0].name, "arguments");
591 assert_eq!(d.args[0].r#type, "object");
592 }
593
594 #[test]
595 fn describe_from_parts_shapes_schema_and_shallow_flag() {
596 let schema = serde_json::json!({
597 "type": "object",
598 "properties": { "url": { "type": "string" } },
599 "required": ["url"]
600 });
601 let d = describe_from_parts("spider__crawl", "crawl", "", ToolKind::Builtin, Some(&schema));
602 assert!(!d.shallow);
603 assert_eq!(d.kind, ToolKind::Builtin);
604 assert_eq!(d.args.len(), 1);
605 assert_eq!(d.parameters.as_ref(), Some(&schema));
606 let bare = describe_from_parts("foo__bar", "bar", "", ToolKind::Mcp, None);
608 assert!(bare.shallow);
609 assert!(bare.args.is_empty());
610 }
611
612 #[tokio::test]
613 async fn run_search_filters_builtins_by_kind_but_appends_composio() {
614 let builtins = vec![
617 desc("foo__search", "search", "search the web", ToolKind::Mcp),
618 desc("bar__do", "do", "do a thing", ToolKind::Builtin),
619 ];
620 let composio = vec![desc("composio__slack", "Slack", "send", ToolKind::Composio)];
621 let out = run_search(
622 "search",
623 builtins,
624 composio,
625 Some(ToolKind::Composio),
626 25,
627 ToolRanker::Bm25,
628 None,
629 )
630 .await;
631 assert!(out.iter().all(|d| d.kind == ToolKind::Composio));
632 assert!(out.iter().any(|d| d.id == "composio__slack"));
633
634 let builtins = vec![desc("foo__search", "search", "the web", ToolKind::Mcp)];
637 let out = run_search("search", builtins, Vec::new(), None, 25, ToolRanker::Bm25, None).await;
638 assert_eq!(out.len(), 1);
639 assert!(out.iter().all(|d| d.kind != ToolKind::Composio));
640 }
641}