Skip to main content

tt_shared/
providers.rs

1//! Lightweight model → provider lookup.
2//!
3//! Used by cloud-side validation (PATCH `/v1/admin/routes`) and the
4//! gateway's defensive logging to detect routes that would cross provider
5//! boundaries. The runtime registry in `tt-core` is the source of truth;
6//! this helper is a cheap heuristic so callers that can't depend on
7//! the full provider registry still get the answer right for the
8//! prefixes we actually ship.
9//!
10//! Conservative semantics: returns `Some(&'static str)` only for prefixes
11//! we know for certain belong to a single provider. Ambiguous names
12//! (e.g. `llama-3-*` which could be Groq, Together, or OpenRouter) return
13//! `None` so the caller doesn't reject a legitimate use of an aggregator.
14
15/// Returns the provider id for `model` when it's a well-known
16/// single-provider prefix; otherwise `None`.
17///
18/// Prefix table (extend in lockstep with the provider crates):
19///
20/// | Prefix(es)                                                              | Provider     |
21/// |-------------------------------------------------------------------------|--------------|
22/// | `gpt-*`, `chatgpt-*`, `o3`, `o3-*`, `o4-*`, `o5-*`                       | `openai`     |
23/// | `claude-*`                                                              | `anthropic`  |
24/// | `gemini-*`                                                              | `gemini`     |
25/// | `mistral-*`, `mixtral-*`, `pixtral-*`, `codestral-*`, `ministral-*`      | `mistral`    |
26pub fn infer_provider(model: &str) -> Option<&'static str> {
27    if model.is_empty() {
28        return None;
29    }
30
31    // OpenAI — gpt-*, chatgpt-*, o3, o3-*, o4-*, o5-*.
32    if model.starts_with("gpt-") || model.starts_with("chatgpt-") {
33        return Some("openai");
34    }
35    if model == "o3" || model.starts_with("o3-") {
36        return Some("openai");
37    }
38    if model.starts_with("o4-") || model.starts_with("o5-") {
39        return Some("openai");
40    }
41
42    // Anthropic
43    if model.starts_with("claude-") {
44        return Some("anthropic");
45    }
46
47    // Gemini
48    if model.starts_with("gemini-") {
49        return Some("gemini");
50    }
51
52    // Mistral family
53    if model.starts_with("mistral-")
54        || model.starts_with("mixtral-")
55        || model.starts_with("pixtral-")
56        || model.starts_with("codestral-")
57        || model.starts_with("ministral-")
58    {
59        return Some("mistral");
60    }
61
62    None
63}
64
65/// Returns `true` iff both `a` and `b` resolve to known providers AND
66/// those providers differ. `false` when either side is unknown (we don't
67/// block routes that aggregate through `openrouter` / `together` /
68/// `groq` whose model names overlap across providers).
69pub fn known_to_differ(a: &str, b: &str) -> bool {
70    match (infer_provider(a), infer_provider(b)) {
71        (Some(x), Some(y)) => x != y,
72        _ => false,
73    }
74}
75
76/// If `model` is a local-backend-prefixed id (`ollama/…`, `vllm/…`,
77/// `lmstudio/…`) with a non-empty model name, return the backend id; else None.
78/// Single source of truth for local routing — used by the registry resolver,
79/// the same-provider exemption, and `LocalProvider`'s prefix strip.
80pub fn local_backend(model: &str) -> Option<&'static str> {
81    for id in ["ollama", "vllm", "lmstudio"] {
82        if let Some(rest) = model.strip_prefix(id).and_then(|r| r.strip_prefix('/')) {
83            if !rest.is_empty() {
84                return Some(id);
85            }
86        }
87    }
88    None
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn openai_prefixes() {
97        assert_eq!(infer_provider("gpt-4o"), Some("openai"));
98        assert_eq!(infer_provider("gpt-4o-mini"), Some("openai"));
99        assert_eq!(infer_provider("gpt-5.5"), Some("openai"));
100        assert_eq!(infer_provider("chatgpt-4o-latest"), Some("openai"));
101        assert_eq!(infer_provider("o3"), Some("openai"));
102        assert_eq!(infer_provider("o3-mini"), Some("openai"));
103        assert_eq!(infer_provider("o4-mini"), Some("openai"));
104    }
105
106    #[test]
107    fn anthropic_prefix() {
108        assert_eq!(infer_provider("claude-opus-4-7"), Some("anthropic"));
109        assert_eq!(infer_provider("claude-sonnet-4-6"), Some("anthropic"));
110        assert_eq!(infer_provider("claude-haiku-4-5"), Some("anthropic"));
111    }
112
113    #[test]
114    fn gemini_prefix() {
115        assert_eq!(infer_provider("gemini-2.5-pro"), Some("gemini"));
116        assert_eq!(infer_provider("gemini-1.5-flash"), Some("gemini"));
117    }
118
119    #[test]
120    fn mistral_family_prefixes() {
121        assert_eq!(infer_provider("mistral-large-2407"), Some("mistral"));
122        assert_eq!(infer_provider("mixtral-8x22b"), Some("mistral"));
123        assert_eq!(infer_provider("pixtral-12b"), Some("mistral"));
124        assert_eq!(infer_provider("codestral-22b"), Some("mistral"));
125        assert_eq!(infer_provider("ministral-8b"), Some("mistral"));
126    }
127
128    #[test]
129    fn unknown_returns_none() {
130        // Aggregator-routed names overlap across providers — must NOT be
131        // assigned to a single provider.
132        assert_eq!(infer_provider("llama-3.3-70b"), None);
133        assert_eq!(infer_provider("qwen2.5-72b"), None);
134        assert_eq!(infer_provider("deepseek-r1"), None);
135        assert_eq!(infer_provider("totally-custom-model"), None);
136        assert_eq!(infer_provider(""), None);
137    }
138
139    #[test]
140    fn known_to_differ_only_blocks_known_pairs() {
141        // Same provider — not differing.
142        assert!(!known_to_differ("gpt-4o", "gpt-4o-mini"));
143        assert!(!known_to_differ("claude-sonnet-4-6", "claude-haiku-4-5"));
144        // Cross provider — differs, blocks.
145        assert!(known_to_differ("gpt-4o", "claude-sonnet-4-6"));
146        assert!(known_to_differ("claude-haiku-4-5", "gemini-2.5-pro"));
147        // Unknowns — pass through (don't block).
148        assert!(!known_to_differ("gpt-4o", "llama-3.3-70b"));
149        assert!(!known_to_differ("custom-1", "custom-2"));
150        assert!(!known_to_differ("custom-1", "gpt-4o"));
151    }
152
153    #[test]
154    fn local_backend_recognizes_prefixes() {
155        assert_eq!(local_backend("ollama/llama3.1:8b"), Some("ollama"));
156        assert_eq!(local_backend("vllm/Qwen2.5-7B"), Some("vllm"));
157        assert_eq!(local_backend("lmstudio/phi-4"), Some("lmstudio"));
158        assert_eq!(local_backend("ollama"), None);
159        assert_eq!(local_backend("ollama/"), None);
160        assert_eq!(local_backend("gpt-4o"), None);
161        assert_eq!(local_backend(""), None);
162    }
163}