296 lines
8.9 KiB
Rust
296 lines
8.9 KiB
Rust
//! The merged agent catalog: lookup, aliases, groups and search.
|
|
|
|
use crate::config::AgentDef;
|
|
use std::collections::{BTreeMap, HashMap};
|
|
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct Catalog {
|
|
agents: Vec<AgentDef>,
|
|
by_name: HashMap<String, usize>,
|
|
aliases: BTreeMap<String, String>,
|
|
groups: BTreeMap<String, Vec<String>>,
|
|
}
|
|
|
|
impl Catalog {
|
|
pub fn new(
|
|
agents: Vec<AgentDef>,
|
|
aliases: BTreeMap<String, String>,
|
|
groups: BTreeMap<String, Vec<String>>,
|
|
) -> Self {
|
|
let mut by_name = HashMap::new();
|
|
for (i, a) in agents.iter().enumerate() {
|
|
by_name.insert(a.name.clone(), i);
|
|
}
|
|
Self {
|
|
agents,
|
|
by_name,
|
|
aliases,
|
|
groups,
|
|
}
|
|
}
|
|
|
|
pub fn agents(&self) -> &[AgentDef] {
|
|
&self.agents
|
|
}
|
|
|
|
pub fn len(&self) -> usize {
|
|
self.agents.len()
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.agents.is_empty()
|
|
}
|
|
|
|
pub fn get(&self, name: &str) -> Option<&AgentDef> {
|
|
self.by_name.get(name).map(|&i| &self.agents[i])
|
|
}
|
|
|
|
/// Resolve an agent name or an alias to the agent definition.
|
|
pub fn resolve(&self, name_or_alias: &str) -> Option<&AgentDef> {
|
|
if let Some(target) = self.aliases.get(name_or_alias) {
|
|
return self.get(target);
|
|
}
|
|
self.get(name_or_alias)
|
|
}
|
|
|
|
pub fn alias_target(&self, alias: &str) -> Option<&str> {
|
|
self.aliases.get(alias).map(|s| s.as_str())
|
|
}
|
|
|
|
/// Members of a group, in definition order (missing names are skipped).
|
|
pub fn group_members(&self, group: &str) -> Vec<&AgentDef> {
|
|
let mut out = Vec::new();
|
|
if let Some(names) = self.groups.get(group) {
|
|
for n in names {
|
|
if let Some(a) = self.get(n) {
|
|
out.push(a);
|
|
}
|
|
}
|
|
}
|
|
out
|
|
}
|
|
|
|
pub fn has_group(&self, group: &str) -> bool {
|
|
self.groups.contains_key(group)
|
|
}
|
|
|
|
/// Is this token a group selector ("group:<name>")?
|
|
pub fn parse_group_selector(s: &str) -> Option<&str> {
|
|
s.strip_prefix("group:")
|
|
}
|
|
|
|
/// Case-insensitive search over name, display name, description,
|
|
/// category, tags and note — typo-tolerant, ranked by relevance.
|
|
pub fn search(&self, keyword: &str, category: Option<&str>) -> Vec<&AgentDef> {
|
|
let kw = keyword.to_lowercase();
|
|
let mut scored: Vec<(u32, &AgentDef)> = self
|
|
.agents
|
|
.iter()
|
|
.filter(|a| !a.hidden)
|
|
.filter(|a| category.map_or(true, |c| a.category.as_deref() == Some(c)))
|
|
.filter_map(|a| {
|
|
let name = a.name.to_lowercase();
|
|
let title = a.title().to_lowercase();
|
|
let rest = format!(
|
|
"{} {} {}",
|
|
a.description.as_deref().unwrap_or("").to_lowercase(),
|
|
a.category.as_deref().unwrap_or("").to_lowercase(),
|
|
a.tags.join(" ").to_lowercase()
|
|
);
|
|
let score = score_keyword(&kw, &name)
|
|
.saturating_mul(3)
|
|
.saturating_add(score_keyword(&kw, &title).saturating_mul(2))
|
|
.saturating_add(score_keyword(&kw, &rest));
|
|
(score > 0).then_some((score, a))
|
|
})
|
|
.collect();
|
|
scored.sort_by(|(sa, a), (sb, b)| sb.cmp(sa).then(a.name.cmp(&b.name)));
|
|
scored.into_iter().map(|(_, a)| a).collect()
|
|
}
|
|
|
|
/// "Did you mean" candidates for a keyword with no results: closest
|
|
/// agent names and display names by edit distance.
|
|
pub fn suggest(&self, keyword: &str) -> Vec<String> {
|
|
let kw = keyword.to_lowercase();
|
|
let mut out: Vec<(usize, String)> = Vec::new();
|
|
for a in self.agents.iter().filter(|a| !a.hidden) {
|
|
for field in [&a.name, &a.title().to_string()] {
|
|
let f = field.to_lowercase();
|
|
if f.is_empty() {
|
|
continue;
|
|
}
|
|
let d = levenshtein(&kw, &f);
|
|
if d <= 3 && !out.iter().any(|(_, n)| n == &a.name) {
|
|
out.push((d, a.name.clone()));
|
|
}
|
|
}
|
|
}
|
|
out.sort_by(|(da, na), (db, nb)| da.cmp(db).then(na.cmp(nb)));
|
|
out.truncate(4);
|
|
out.into_iter().map(|(_, n)| n).collect()
|
|
}
|
|
}
|
|
|
|
/// Score a keyword against one haystack field: 100 for a full hit,
|
|
/// then subsequence and typo bonuses; 0 = no relation at all.
|
|
fn score_keyword(keyword: &str, hay: &str) -> u32 {
|
|
if keyword.is_empty() {
|
|
return 0;
|
|
}
|
|
let tokens: Vec<&str> = keyword.split_whitespace().collect();
|
|
let mut total = 0u32;
|
|
for token in &tokens {
|
|
if token.is_empty() {
|
|
continue;
|
|
}
|
|
let s = score_token(token, hay);
|
|
if s == 0 {
|
|
return 0;
|
|
}
|
|
total += s;
|
|
}
|
|
total
|
|
}
|
|
|
|
fn score_token(token: &str, hay: &str) -> u32 {
|
|
if hay.contains(token) {
|
|
return 100;
|
|
}
|
|
if subsequence(token, hay) {
|
|
return 40;
|
|
}
|
|
// Typo tolerance: compare the token with every word of the field.
|
|
let mut best = usize::MAX;
|
|
for word in hay.split(|c: char| !c.is_alphanumeric()) {
|
|
if word.is_empty() {
|
|
continue;
|
|
}
|
|
best = best.min(levenshtein(token, word));
|
|
}
|
|
match best {
|
|
0..=1 => 25,
|
|
2 => 15,
|
|
3 => 5,
|
|
_ => 0,
|
|
}
|
|
}
|
|
|
|
/// Do the characters of needle appear in order inside hay?
|
|
fn subsequence(needle: &str, hay: &str) -> bool {
|
|
let mut it = hay.chars();
|
|
needle.chars().all(|c| it.any(|h| h == c))
|
|
}
|
|
|
|
/// Classic Levenshtein distance on characters.
|
|
fn levenshtein(a: &str, b: &str) -> usize {
|
|
let a: Vec<char> = a.chars().collect();
|
|
let b: Vec<char> = b.chars().collect();
|
|
let mut prev: Vec<usize> = (0..=b.len()).collect();
|
|
let mut cur = vec![0usize; b.len() + 1];
|
|
for (i, ca) in a.iter().enumerate() {
|
|
cur[0] = i + 1;
|
|
for (j, cb) in b.iter().enumerate() {
|
|
let cost = if ca == cb { 0 } else { 1 };
|
|
cur[j + 1] = (prev[j] + cost).min(prev[j + 1] + 1).min(cur[j] + 1);
|
|
}
|
|
std::mem::swap(&mut prev, &mut cur);
|
|
}
|
|
prev[b.len()]
|
|
}
|
|
|
|
impl Catalog {
|
|
/// Sorted list of all categories present in the catalog.
|
|
pub fn categories(&self) -> Vec<String> {
|
|
let mut cats: Vec<String> = self
|
|
.agents
|
|
.iter()
|
|
.filter_map(|a| a.category.clone())
|
|
.collect();
|
|
cats.sort();
|
|
cats.dedup();
|
|
cats
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::config::Config;
|
|
|
|
fn catalog() -> Catalog {
|
|
let cfg: Config = serde_yaml::from_str(crate::config::DEFAULT_CONFIG).unwrap();
|
|
Catalog::new(cfg.agents, cfg.aliases, cfg.groups)
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_aliases() {
|
|
let cat = catalog();
|
|
let via_alias = cat.resolve("cc").map(|a| a.name.clone());
|
|
assert_eq!(via_alias.as_deref(), Some("claude-code"));
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_groups() {
|
|
let cat = catalog();
|
|
assert!(cat.has_group("dev"));
|
|
let members = cat.group_members("dev");
|
|
assert!(!members.is_empty());
|
|
assert!(members.iter().any(|a| a.name == "claude-code"));
|
|
}
|
|
|
|
#[test]
|
|
fn searches() {
|
|
let cat = catalog();
|
|
let hits = cat.search("anthropic", None);
|
|
assert!(hits.iter().any(|a| a.name == "claude-code"));
|
|
let cats = cat.categories();
|
|
assert!(cats.contains(&"coding-agent".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn fuzzy_search_tolerates_missing_separator() {
|
|
let cat = catalog();
|
|
let hits = cat.search("claude cod", None);
|
|
assert!(
|
|
hits.iter().any(|a| a.name == "claude-code"),
|
|
"expected claude-code in {:?}",
|
|
hits.iter().map(|a| &a.name).collect::<Vec<_>>()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn fuzzy_search_tolerates_typos() {
|
|
let cat = catalog();
|
|
let hits = cat.search("cludecode", None);
|
|
assert!(hits.iter().any(|a| a.name == "claude-code"));
|
|
}
|
|
|
|
#[test]
|
|
fn exact_matches_rank_first() {
|
|
let cat = catalog();
|
|
let hits = cat.search("claude", None);
|
|
assert!(!hits.is_empty());
|
|
assert_eq!(
|
|
hits[0].name, "claude-code",
|
|
"exact name match must rank first: {:?}",
|
|
hits.iter().map(|a| &a.name).collect::<Vec<_>>()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn suggest_proposes_close_names() {
|
|
let cat = catalog();
|
|
let sug = cat.suggest("aiderr");
|
|
assert!(sug.contains(&"aider".to_string()), "suggestions: {sug:?}");
|
|
assert!(sug.len() <= 4);
|
|
}
|
|
|
|
#[test]
|
|
fn levenshtein_distance() {
|
|
assert_eq!(levenshtein("aider", "aider"), 0);
|
|
assert_eq!(levenshtein("aider", "aiderr"), 1);
|
|
assert_eq!(levenshtein("abc", "axc"), 1);
|
|
assert_eq!(levenshtein("abc", "xyz"), 3);
|
|
}
|
|
}
|