edgequake_llm/discovery/
search.rs

1//! Model name search — exact, substring, and fuzzy matching.
2//!
3//! Shared by static registry search and live discovery so ranking logic stays DRY.
4
5use serde::{Deserialize, Serialize};
6
7use super::registry::{all_static_models, static_lookup};
8use super::types::{CapabilityFilter, DiscoveredModel};
9
10/// How closely a model matched the query.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
12pub enum ModelMatchKind {
13    /// Exact model ID (case-insensitive).
14    ExactId,
15    /// Exact display name (case-insensitive).
16    ExactName,
17    /// `provider/model` composite exact match.
18    ExactQualifiedId,
19    /// ID or name starts with the query.
20    Prefix,
21    /// ID, name, or tag contains the query.
22    Contains,
23    /// Token / similarity fuzzy match.
24    Fuzzy,
25}
26
27/// A model ranked by search relevance.
28#[derive(Debug, Clone)]
29pub struct ModelSearchMatch {
30    pub model: DiscoveredModel,
31    pub score: f64,
32    pub match_kind: ModelMatchKind,
33}
34
35/// Query for searching models by name or ID.
36#[derive(Debug, Clone, Default, Serialize, Deserialize)]
37pub struct ModelSearchQuery {
38    /// Search text — model ID, display name, tag, or `provider/model`.
39    pub query: String,
40    /// Restrict to one provider (canonical ID).
41    pub provider: Option<String>,
42    /// Enable fuzzy / similarity matching (default: exact + substring only).
43    pub fuzzy: bool,
44    /// Minimum score in `[0.0, 1.0]` (defaults: `0.80` exact mode, `0.55` fuzzy).
45    pub min_score: Option<f64>,
46    /// Maximum number of results (default: `20`).
47    pub limit: Option<usize>,
48    /// Minimum input context window (tokens).
49    pub min_context_length: Option<usize>,
50    /// Maximum input context window (tokens).
51    pub max_context_length: Option<usize>,
52    /// Minimum max output tokens.
53    pub min_output_tokens: Option<usize>,
54    /// Maximum max output tokens.
55    pub max_output_tokens: Option<usize>,
56    /// Optional capability constraints (AND logic with name match).
57    pub capability_filter: Option<CapabilityFilter>,
58}
59
60impl ModelSearchQuery {
61    pub fn new(query: impl Into<String>) -> Self {
62        Self {
63            query: query.into(),
64            ..Default::default()
65        }
66    }
67
68    pub fn fuzzy(mut self, enabled: bool) -> Self {
69        self.fuzzy = enabled;
70        self
71    }
72
73    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
74        self.provider = Some(provider.into());
75        self
76    }
77
78    pub fn with_min_score(mut self, score: f64) -> Self {
79        self.min_score = Some(score);
80        self
81    }
82
83    pub fn with_limit(mut self, limit: usize) -> Self {
84        self.limit = Some(limit);
85        self
86    }
87
88    pub fn with_capability_filter(mut self, filter: CapabilityFilter) -> Self {
89        self.capability_filter = Some(filter);
90        self
91    }
92
93    pub fn with_min_context_length(mut self, tokens: usize) -> Self {
94        self.min_context_length = Some(tokens);
95        self
96    }
97
98    pub fn with_max_context_length(mut self, tokens: usize) -> Self {
99        self.max_context_length = Some(tokens);
100        self
101    }
102
103    pub fn with_min_output_tokens(mut self, tokens: usize) -> Self {
104        self.min_output_tokens = Some(tokens);
105        self
106    }
107
108    pub fn with_max_output_tokens(mut self, tokens: usize) -> Self {
109        self.max_output_tokens = Some(tokens);
110        self
111    }
112
113    /// Merge query-level length constraints with an optional capability filter.
114    pub fn effective_filter(&self) -> CapabilityFilter {
115        let mut filter = self.capability_filter.clone().unwrap_or_default();
116        merge_min_usize(&mut filter.min_context_length, self.min_context_length);
117        merge_max_usize(&mut filter.max_context_length, self.max_context_length);
118        merge_min_usize(&mut filter.min_output_tokens, self.min_output_tokens);
119        merge_max_usize(&mut filter.max_output_tokens, self.max_output_tokens);
120        filter
121    }
122
123    fn effective_min_score(&self) -> f64 {
124        self.min_score
125            .unwrap_or(if self.fuzzy { 0.55 } else { 0.80 })
126    }
127
128    fn effective_limit(&self) -> usize {
129        self.limit.unwrap_or(20)
130    }
131}
132
133/// Search a model list by name/ID.
134pub fn search_models(
135    models: impl IntoIterator<Item = DiscoveredModel>,
136    query: &ModelSearchQuery,
137) -> Vec<ModelSearchMatch> {
138    let (parsed_provider, parsed_model) = parse_qualified_query(&query.query);
139    let provider_filter = query.provider.as_deref().or(parsed_provider);
140    let search_text = parsed_model.unwrap_or(&query.query);
141
142    if search_text.trim().is_empty() {
143        return Vec::new();
144    }
145
146    let mut matches = Vec::new();
147    let qualified = parsed_model.is_some();
148    for model in models {
149        if let Some(provider) = provider_filter {
150            if !model.provider.eq_ignore_ascii_case(provider) {
151                continue;
152            }
153        }
154        let filter = query.effective_filter();
155        if !filter.matches(&model) {
156            continue;
157        }
158        if let Some((score, kind)) = score_model(search_text, &model, query.fuzzy, qualified) {
159            if score >= query.effective_min_score() {
160                matches.push(ModelSearchMatch {
161                    model,
162                    score,
163                    match_kind: kind,
164                });
165            }
166        }
167    }
168
169    matches.sort_by(|a, b| {
170        b.score
171            .partial_cmp(&a.score)
172            .unwrap_or(std::cmp::Ordering::Equal)
173            .then_with(|| a.model.provider.cmp(&b.model.provider))
174            .then_with(|| a.model.id.cmp(&b.model.id))
175    });
176    matches.truncate(query.effective_limit());
177    matches
178}
179
180/// Search the offline static registry by name/ID (no network).
181pub fn search_static_models(query: &ModelSearchQuery) -> Vec<ModelSearchMatch> {
182    search_models(all_static_models(), query)
183}
184
185/// Resolve an exact model from the static registry by provider + name/ID.
186///
187/// Tries exact ID match first, then exact display name (case-insensitive).
188pub fn static_lookup_by_name(provider: &str, name_or_id: &str) -> Option<DiscoveredModel> {
189    if let Some(model) = static_lookup(provider, name_or_id) {
190        return Some(model);
191    }
192
193    let models = match provider {
194        "openai" => super::registry::openai_models(),
195        "anthropic" => super::registry::anthropic_models(),
196        "gemini" => super::registry::gemini_models(),
197        "vertexai" => super::registry::vertexai_models(),
198        "mistral" => super::registry::mistral_models(),
199        "xai" => super::registry::xai_models(),
200        "cohere" => super::registry::cohere_models(),
201        "nvidia" => super::registry::nvidia_models(),
202        _ => return None,
203    };
204
205    models
206        .into_iter()
207        .find(|m| m.name.eq_ignore_ascii_case(name_or_id))
208}
209
210fn merge_min_usize(target: &mut Option<usize>, value: Option<usize>) {
211    if let Some(value) = value {
212        *target = Some(target.map(|current| current.max(value)).unwrap_or(value));
213    }
214}
215
216fn merge_max_usize(target: &mut Option<usize>, value: Option<usize>) {
217    if let Some(value) = value {
218        *target = Some(target.map(|current| current.min(value)).unwrap_or(value));
219    }
220}
221
222fn parse_qualified_query(query: &str) -> (Option<&str>, Option<&str>) {
223    let trimmed = query.trim();
224    if let Some((provider, model)) = trimmed.split_once('/') {
225        let provider = provider.trim();
226        let model = model.trim();
227        if !provider.is_empty() && !model.is_empty() {
228            return (Some(provider), Some(model));
229        }
230    }
231    (None, None)
232}
233
234fn normalize(value: &str) -> String {
235    value
236        .to_lowercase()
237        .chars()
238        .map(|c| if c.is_alphanumeric() { c } else { ' ' })
239        .collect::<String>()
240        .split_whitespace()
241        .collect::<Vec<_>>()
242        .join(" ")
243}
244
245fn compact(value: &str) -> String {
246    normalize(value)
247        .chars()
248        .filter(|c| !c.is_whitespace())
249        .collect()
250}
251
252fn score_model(
253    query: &str,
254    model: &DiscoveredModel,
255    fuzzy: bool,
256    qualified: bool,
257) -> Option<(f64, ModelMatchKind)> {
258    let q_norm = normalize(query);
259    let q_compact = compact(query);
260    if q_norm.is_empty() {
261        return None;
262    }
263
264    let id_norm = normalize(&model.id);
265    let name_norm = normalize(&model.name);
266    let id_compact = compact(&model.id);
267    let name_compact = compact(&model.name);
268    let qualified_id = format!("{} {}", model.provider, model.id);
269    let qualified_norm = normalize(&qualified_id);
270
271    if id_norm == q_norm || id_compact == q_compact {
272        return Some((1.0, ModelMatchKind::ExactId));
273    }
274    if name_norm == q_norm || name_compact == q_compact {
275        return Some((0.98, ModelMatchKind::ExactName));
276    }
277    if qualified_norm == q_norm {
278        return Some((0.96, ModelMatchKind::ExactQualifiedId));
279    }
280    if qualified {
281        return None;
282    }
283    if id_norm.starts_with(&q_norm)
284        || name_norm.starts_with(&q_norm)
285        || id_compact.starts_with(&q_compact)
286    {
287        return Some((0.92, ModelMatchKind::Prefix));
288    }
289    if id_norm.contains(&q_norm)
290        || name_norm.contains(&q_norm)
291        || id_compact.contains(&q_compact)
292        || name_compact.contains(&q_compact)
293        || model
294            .tags
295            .iter()
296            .any(|tag| normalize(tag).contains(&q_norm))
297    {
298        return Some((0.85, ModelMatchKind::Contains));
299    }
300
301    if !fuzzy {
302        return None;
303    }
304
305    let token_score = token_overlap_score(&q_norm, &id_norm, &name_norm, &model.tags);
306    let similarity = id_norm
307        .split_whitespace()
308        .chain(name_norm.split_whitespace())
309        .map(|part| dice_coefficient(&q_compact, &compact(part)))
310        .fold(0.0_f64, f64::max)
311        .max(dice_coefficient(&q_compact, &id_compact))
312        .max(dice_coefficient(&q_compact, &name_compact))
313        .max(token_score);
314
315    if similarity >= 0.55 {
316        Some((similarity, ModelMatchKind::Fuzzy))
317    } else {
318        None
319    }
320}
321
322fn token_overlap_score(query: &str, id: &str, name: &str, tags: &[String]) -> f64 {
323    let tokens: Vec<&str> = query.split_whitespace().collect();
324    if tokens.is_empty() {
325        return 0.0;
326    }
327    let haystacks = [id, name];
328    let matched = tokens
329        .iter()
330        .filter(|token| {
331            haystacks.iter().any(|h| h.contains(*token))
332                || tags.iter().any(|tag| normalize(tag).contains(*token))
333        })
334        .count();
335    (matched as f64 / tokens.len() as f64) * 0.72
336}
337
338fn dice_coefficient(a: &str, b: &str) -> f64 {
339    if a.is_empty() || b.is_empty() {
340        return 0.0;
341    }
342    if a == b {
343        return 1.0;
344    }
345
346    let a_bigrams = bigrams(a);
347    let b_bigrams = bigrams(b);
348    if a_bigrams.is_empty() || b_bigrams.is_empty() {
349        return if a.contains(b) || b.contains(a) {
350            0.65
351        } else {
352            0.0
353        };
354    }
355
356    let overlap = a_bigrams.iter().filter(|bg| b_bigrams.contains(bg)).count();
357    (2.0 * overlap as f64) / (a_bigrams.len() + b_bigrams.len()) as f64
358}
359
360fn bigrams(value: &str) -> Vec<(char, char)> {
361    let chars: Vec<char> = value.chars().collect();
362    if chars.len() < 2 {
363        return Vec::new();
364    }
365    chars.windows(2).map(|w| (w[0], w[1])).collect()
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371    use crate::model_config::ModelCapabilities;
372
373    fn sample(id: &str, name: &str, provider: &str) -> DiscoveredModel {
374        DiscoveredModel {
375            id: id.into(),
376            name: name.into(),
377            provider: provider.into(),
378            context_length: 128_000,
379            max_output_tokens: 4096,
380            capabilities: ModelCapabilities {
381                supports_function_calling: true,
382                ..Default::default()
383            },
384            ..Default::default()
385        }
386    }
387
388    #[test]
389    fn exact_id_match() {
390        let models = vec![sample("gpt-4.1", "GPT-4.1", "openai")];
391        let query = ModelSearchQuery::new("gpt-4.1");
392        let hits = search_models(models, &query);
393        assert_eq!(hits.len(), 1);
394        assert_eq!(hits[0].match_kind, ModelMatchKind::ExactId);
395        assert!((hits[0].score - 1.0).abs() < f64::EPSILON);
396    }
397
398    #[test]
399    fn qualified_query_parses_provider() {
400        let models = vec![
401            sample("gpt-4.1", "GPT-4.1", "openai"),
402            sample("gpt-4.1", "GPT-4.1", "azure"),
403        ];
404        let query = ModelSearchQuery::new("openai/gpt-4.1");
405        let hits = search_models(models, &query);
406        assert_eq!(hits.len(), 1);
407        assert_eq!(hits[0].model.provider, "openai");
408    }
409
410    #[test]
411    fn substring_search_finds_display_name() {
412        let models = vec![sample("claude-opus-4-8", "Claude Opus 4.8", "anthropic")];
413        let query = ModelSearchQuery::new("opus 4.8");
414        let hits = search_models(models, &query);
415        assert_eq!(hits.len(), 1);
416        assert_eq!(hits[0].match_kind, ModelMatchKind::Contains);
417    }
418
419    #[test]
420    fn fuzzy_search_handles_typos() {
421        let models = vec![sample("gpt-4.1", "GPT-4.1", "openai")];
422        let query = ModelSearchQuery::new("gpt-4l").fuzzy(true);
423        let hits = search_models(models, &query);
424        assert_eq!(hits.len(), 1);
425        assert_eq!(hits[0].match_kind, ModelMatchKind::Fuzzy);
426    }
427
428    #[test]
429    fn static_search_finds_known_model() {
430        let query = ModelSearchQuery::new("claude sonnet").fuzzy(true);
431        let hits = search_static_models(&query);
432        assert!(hits.iter().any(|m| m.model.id.contains("claude-sonnet")));
433    }
434
435    #[test]
436    fn search_respects_output_length_bounds() {
437        let mut model = sample("big-out", "Big Out", "openai");
438        model.max_output_tokens = 32_768;
439        let models = vec![model];
440        let query = ModelSearchQuery::new("big-out")
441            .with_min_output_tokens(16_384)
442            .with_max_output_tokens(65_536);
443        let hits = search_models(models, &query);
444        assert_eq!(hits.len(), 1);
445
446        let mut small = sample("big-out", "Big Out", "openai");
447        small.max_output_tokens = 4096;
448        let query = ModelSearchQuery::new("big-out").with_min_output_tokens(65_536);
449        assert!(search_models(vec![small], &query).is_empty());
450    }
451
452    #[test]
453    fn search_respects_input_length_bounds() {
454        let mut model = sample("big-ctx", "Big Context", "openai");
455        model.context_length = 1_000_000;
456        let models = vec![model];
457        let query = ModelSearchQuery::new("big")
458            .with_min_context_length(500_000)
459            .with_max_context_length(2_000_000);
460        assert_eq!(search_models(models.clone(), &query).len(), 1);
461
462        let mut small = sample("big-ctx", "Big Context", "openai");
463        small.context_length = 1_000_000;
464        let query = ModelSearchQuery::new("big").with_max_context_length(500_000);
465        assert!(search_models(vec![small], &query).is_empty());
466    }
467
468    #[test]
469    fn static_lookup_by_name_falls_back_to_display_name() {
470        let model = static_lookup_by_name("openai", "GPT-4.1");
471        assert!(model.is_some());
472        assert_eq!(model.unwrap().id, "gpt-4.1");
473    }
474}
edgequake_llm/discovery/search.rs

edgequake_llm/discovery/
search.rs
