use serde::{Deserialize, Serialize};
use crate::{ConnectorDescription, Snippet, generate_types};
const LIMIT: usize = 50;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SearchResult {
pub path: String,
pub connector: String,
pub method: String,
pub description: Option<String>,
pub types: String,
pub requires_approval: bool,
pub kind: String,
pub score: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SearchOutput {
pub results: Vec<SearchResult>,
pub total: usize,
pub truncated: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DescribeOutput {
pub path: String,
pub description: Option<String>,
pub requires_approval: bool,
pub types: String,
pub kind: String,
}
pub fn search(
query: &str,
connectors: &[ConnectorDescription],
snippets: &[Snippet],
) -> SearchOutput {
let normalized = normalize(query);
let tokens = tokenize(query);
let mut results = Vec::new();
for connector in connectors {
for tool in &connector.tools {
let path = format!("{}.{}", connector.name, tool.name);
if let Some(score) = score(
&normalized,
&tokens,
[
(&path, 12),
(&connector.name, 8),
(&tool.name, 10),
(tool.description.as_deref().unwrap_or_default(), 5),
],
) {
results.push(SearchResult {
path,
connector: connector.name.clone(),
method: tool.name.clone(),
description: tool.description.clone(),
types: {
let mut single = connector.clone();
single.instructions = None;
single.tools = vec![tool.clone()];
generate_types(&single)
},
requires_approval: tool.policy.requires_approval,
kind: "method".to_string(),
score,
});
}
}
}
for snippet in snippets {
if let Some(score) = score(
&normalized,
&tokens,
[
(&snippet.name, 12),
("snippet", 8),
(&snippet.name, 10),
(&snippet.description, 5),
],
) {
results.push(SearchResult {
path: snippet.name.clone(),
connector: "snippet".to_string(),
method: snippet.name.clone(),
description: Some(snippet.description.clone()),
types: snippet.code.clone(),
requires_approval: false,
kind: "snippet".to_string(),
score,
});
}
}
results.sort_by(|left, right| {
right
.score
.cmp(&left.score)
.then_with(|| left.path.cmp(&right.path))
});
let total = results.len();
results.truncate(LIMIT);
SearchOutput {
results,
total,
truncated: total > LIMIT,
}
}
pub fn describe(
target: &str,
connectors: &[ConnectorDescription],
snippets: &[Snippet],
) -> DescribeOutput {
if let Some(snippet) = snippets.iter().find(|snippet| snippet.name == target) {
return DescribeOutput {
path: target.to_string(),
description: Some(snippet.description.clone()),
requires_approval: false,
types: format!("{}\n\n```ts\n{}\n```", snippet.description, snippet.code),
kind: "snippet".to_string(),
};
}
let (connector_name, method_name) = target
.split_once('.')
.map(|(connector, method)| (Some(connector), method))
.unwrap_or((None, target));
if connector_name.is_none()
&& let Some(connector) = connectors.iter().find(|connector| connector.name == target)
{
return DescribeOutput {
path: target.to_string(),
description: connector.instructions.clone(),
requires_approval: false,
types: generate_types(connector),
kind: "connector".to_string(),
};
}
for connector in connectors
.iter()
.filter(|connector| connector_name.is_none_or(|name| connector.name == name))
{
if let Some(tool) = connector.tools.iter().find(|tool| tool.name == method_name) {
let mut single = connector.clone();
single.instructions = None;
single.tools = vec![tool.clone()];
return DescribeOutput {
path: format!("{}.{}", connector.name, tool.name),
description: tool.description.clone(),
requires_approval: tool.policy.requires_approval,
types: generate_types(&single),
kind: "method".to_string(),
};
}
}
DescribeOutput {
path: target.to_string(),
description: None,
requires_approval: false,
types: format!("\"{target}\" not found."),
kind: "method".to_string(),
}
}
fn score<const COUNT: usize>(
query: &str,
query_tokens: &[String],
fields: [(&str, u32); COUNT],
) -> Option<u32> {
if query.is_empty() || query_tokens.is_empty() {
return None;
}
let mut score = 0;
let mut matched = vec![false; query_tokens.len()];
let mut phrase = false;
for (value, weight) in fields {
let value = normalize(value);
let tokens = tokenize(&value);
if value == query {
score += weight * 14;
} else if value.starts_with(query) {
score += weight * 9;
} else if value.contains(query) {
score += weight * 6;
phrase = true;
}
for (index, token) in query_tokens.iter().enumerate() {
if tokens.contains(token) {
score += weight * 4;
matched[index] = true;
} else if tokens
.iter()
.any(|candidate| candidate.starts_with(token) || token.starts_with(candidate))
{
score += weight * 2;
matched[index] = true;
} else if value.contains(token) {
score += weight;
matched[index] = true;
}
}
}
let matched_count = matched.iter().filter(|matched| **matched).count();
if matched_count == 0 {
return None;
}
let coverage = matched_count as f32 / query_tokens.len() as f32;
if coverage < if query_tokens.len() <= 2 { 1.0 } else { 0.6 } && !phrase {
return None;
}
Some(
score
+ if matched_count == query_tokens.len() {
25
} else {
10
},
)
}
fn normalize(value: &str) -> String {
let mut result = String::new();
let mut previous_lower = false;
for ch in value.chars() {
if ch.is_ascii_uppercase() && previous_lower {
result.push(' ');
}
if ch.is_ascii_alphanumeric() {
result.push(ch.to_ascii_lowercase());
previous_lower = ch.is_ascii_lowercase() || ch.is_ascii_digit();
} else {
result.push(' ');
previous_lower = false;
}
}
result.split_whitespace().collect::<Vec<_>>().join(" ")
}
fn tokenize(value: &str) -> Vec<String> {
normalize(value)
.split_whitespace()
.map(str::to_string)
.collect()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use crate::{ConnectorTool, ToolAnnotations};
use super::*;
#[test]
fn ranks_exact_method_matches_first() {
let connectors = [ConnectorDescription {
name: "github".to_string(),
instructions: None,
tools: vec![ConnectorTool {
name: "listIssues".to_string(),
description: Some("List repository issues".to_string()),
input_schema: json!({"type": "object"}),
output_schema: None,
instructions: None,
examples: Vec::new(),
annotations: ToolAnnotations::default(),
policy: crate::ToolPolicy::default(),
}],
}];
let result = search("list issues", &connectors, &[]);
assert_eq!(result.results[0].path, "github.listIssues");
assert!(result.results[0].types.contains("listIssues"));
}
}