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`    |
26/// | `azure/<deployment>`                                                     | `azure`      |
27pub fn infer_provider(model: &str) -> Option<&'static str> {
28    if model.is_empty() {
29        return None;
30    }
31
32    // Azure OpenAI — explicit `azure/<deployment>` prefix. The deployment name
33    // after the slash is customer-chosen (need not match an OpenAI model id), so
34    // the prefix is the only reliable signal — checked before the OpenAI
35    // gpt-*/o3-* prefixes so `azure/gpt-4o-prod` resolves to `azure`, not
36    // `openai`.
37    if azure_deployment(model).is_some() {
38        return Some("azure");
39    }
40
41    // OpenAI — gpt-*, chatgpt-*, o3, o3-*, o4-*, o5-*.
42    if model.starts_with("gpt-") || model.starts_with("chatgpt-") {
43        return Some("openai");
44    }
45    if model == "o3" || model.starts_with("o3-") {
46        return Some("openai");
47    }
48    if model.starts_with("o4-") || model.starts_with("o5-") {
49        return Some("openai");
50    }
51
52    // Anthropic
53    if model.starts_with("claude-") {
54        return Some("anthropic");
55    }
56
57    // Gemini
58    if model.starts_with("gemini-") {
59        return Some("gemini");
60    }
61
62    // Mistral family
63    if model.starts_with("mistral-")
64        || model.starts_with("mixtral-")
65        || model.starts_with("pixtral-")
66        || model.starts_with("codestral-")
67        || model.starts_with("ministral-")
68    {
69        return Some("mistral");
70    }
71
72    None
73}
74
75/// Returns `true` iff both `a` and `b` resolve to known providers AND
76/// those providers differ. `false` when either side is unknown (we don't
77/// block routes that aggregate through `openrouter` / `together` /
78/// `groq` whose model names overlap across providers).
79pub fn known_to_differ(a: &str, b: &str) -> bool {
80    match (infer_provider(a), infer_provider(b)) {
81        (Some(x), Some(y)) => x != y,
82        _ => false,
83    }
84}
85
86/// If `model` is an Azure-prefixed id (`azure/<deployment>`) with a non-empty
87/// deployment name, return the deployment (the part after `azure/`); else None.
88///
89/// Single source of truth for the Azure prefix — used by [`infer_provider`] for
90/// routing and by the Azure adapter to strip the prefix down to the bare
91/// deployment name it targets on the wire.
92pub fn azure_deployment(model: &str) -> Option<&str> {
93    model.strip_prefix("azure/").filter(|rest| !rest.is_empty())
94}
95
96/// If `model` is a local-backend-prefixed id (`ollama/…`, `vllm/…`,
97/// `lmstudio/…`) with a non-empty model name, return the backend id; else None.
98/// Single source of truth for local routing — used by the registry resolver,
99/// the same-provider exemption, and `LocalProvider`'s prefix strip.
100pub fn local_backend(model: &str) -> Option<&'static str> {
101    for id in ["ollama", "vllm", "lmstudio"] {
102        if let Some(rest) = model.strip_prefix(id).and_then(|r| r.strip_prefix('/')) {
103            if !rest.is_empty() {
104                return Some(id);
105            }
106        }
107    }
108    None
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn openai_prefixes() {
117        assert_eq!(infer_provider("gpt-4o"), Some("openai"));
118        assert_eq!(infer_provider("gpt-4o-mini"), Some("openai"));
119        assert_eq!(infer_provider("gpt-5.5"), Some("openai"));
120        assert_eq!(infer_provider("chatgpt-4o-latest"), Some("openai"));
121        assert_eq!(infer_provider("o3"), Some("openai"));
122        assert_eq!(infer_provider("o3-mini"), Some("openai"));
123        assert_eq!(infer_provider("o4-mini"), Some("openai"));
124    }
125
126    #[test]
127    fn anthropic_prefix() {
128        assert_eq!(infer_provider("claude-opus-4-7"), Some("anthropic"));
129        assert_eq!(infer_provider("claude-sonnet-4-6"), Some("anthropic"));
130        assert_eq!(infer_provider("claude-haiku-4-5"), Some("anthropic"));
131    }
132
133    #[test]
134    fn gemini_prefix() {
135        assert_eq!(infer_provider("gemini-2.5-pro"), Some("gemini"));
136        assert_eq!(infer_provider("gemini-1.5-flash"), Some("gemini"));
137    }
138
139    #[test]
140    fn mistral_family_prefixes() {
141        assert_eq!(infer_provider("mistral-large-2407"), Some("mistral"));
142        assert_eq!(infer_provider("mixtral-8x22b"), Some("mistral"));
143        assert_eq!(infer_provider("pixtral-12b"), Some("mistral"));
144        assert_eq!(infer_provider("codestral-22b"), Some("mistral"));
145        assert_eq!(infer_provider("ministral-8b"), Some("mistral"));
146    }
147
148    #[test]
149    fn azure_prefix() {
150        // `azure/<deployment>` resolves to azure, taking precedence over the
151        // OpenAI gpt-*/o3-* prefixes even when the deployment is named after an
152        // OpenAI model.
153        assert_eq!(infer_provider("azure/gpt-4o-prod"), Some("azure"));
154        assert_eq!(infer_provider("azure/gpt-4o"), Some("azure"));
155        assert_eq!(infer_provider("azure/o3"), Some("azure"));
156        assert_eq!(infer_provider("azure/my-custom-deployment"), Some("azure"));
157        // Bare prefix with no deployment name is not Azure.
158        assert_eq!(infer_provider("azure/"), None);
159    }
160
161    #[test]
162    fn azure_deployment_strips_prefix() {
163        assert_eq!(azure_deployment("azure/gpt-4o-prod"), Some("gpt-4o-prod"));
164        assert_eq!(azure_deployment("azure/o3"), Some("o3"));
165        assert_eq!(azure_deployment("azure/"), None);
166        assert_eq!(azure_deployment("gpt-4o"), None);
167        assert_eq!(azure_deployment(""), None);
168    }
169
170    #[test]
171    fn unknown_returns_none() {
172        // Aggregator-routed names overlap across providers — must NOT be
173        // assigned to a single provider.
174        assert_eq!(infer_provider("llama-3.3-70b"), None);
175        assert_eq!(infer_provider("qwen2.5-72b"), None);
176        assert_eq!(infer_provider("deepseek-r1"), None);
177        assert_eq!(infer_provider("totally-custom-model"), None);
178        assert_eq!(infer_provider(""), None);
179    }
180
181    #[test]
182    fn known_to_differ_only_blocks_known_pairs() {
183        // Same provider — not differing.
184        assert!(!known_to_differ("gpt-4o", "gpt-4o-mini"));
185        assert!(!known_to_differ("claude-sonnet-4-6", "claude-haiku-4-5"));
186        // Cross provider — differs, blocks.
187        assert!(known_to_differ("gpt-4o", "claude-sonnet-4-6"));
188        assert!(known_to_differ("claude-haiku-4-5", "gemini-2.5-pro"));
189        // Unknowns — pass through (don't block).
190        assert!(!known_to_differ("gpt-4o", "llama-3.3-70b"));
191        assert!(!known_to_differ("custom-1", "custom-2"));
192        assert!(!known_to_differ("custom-1", "gpt-4o"));
193    }
194
195    #[test]
196    fn local_backend_recognizes_prefixes() {
197        assert_eq!(local_backend("ollama/llama3.1:8b"), Some("ollama"));
198        assert_eq!(local_backend("vllm/Qwen2.5-7B"), Some("vllm"));
199        assert_eq!(local_backend("lmstudio/phi-4"), Some("lmstudio"));
200        assert_eq!(local_backend("ollama"), None);
201        assert_eq!(local_backend("ollama/"), None);
202        assert_eq!(local_backend("gpt-4o"), None);
203        assert_eq!(local_backend(""), None);
204    }
205}