1use crate::Result;
2use crate::catalog::LlmModel;
3#[cfg(feature = "bedrock")]
4use crate::providers::bedrock::BedrockProvider;
5#[cfg(feature = "codex")]
6use crate::providers::codex::CodexProvider;
7use crate::providers::generic::{self, GenericProvider};
8use crate::{
9 LlmError, ProviderConnectionConfig, ProviderConnectionOverrides, StreamingModelProvider,
10 alloyed::AlloyedModelProvider,
11};
12use futures::future::BoxFuture;
13use std::collections::{HashMap, HashSet};
14use std::future::{Future, ready};
15use std::sync::{Arc, Mutex, PoisonError};
16use tokio::sync::OnceCell;
17
18#[doc = include_str!("docs/parser.md")]
19pub struct ModelProviderParser {
20 factories: HashMap<String, CreateProviderFn>,
21}
22
23impl Default for ModelProviderParser {
24 fn default() -> Self {
26 let mut parser = Self { factories: HashMap::new() };
27
28 for config in generic::BUILT_INS {
29 parser = parser.with_generic_provider(config);
30 }
31
32 #[cfg(feature = "bedrock")]
33 {
34 parser = parser.register(
35 "bedrock",
36 |connection| async move { Ok(BedrockProvider::new(connection).await) },
37 BedrockProvider::with_model,
38 );
39 }
40
41 parser
42 }
43}
44
45impl ModelProviderParser {
46 #[cfg(feature = "codex")]
47 pub fn with_codex_provider(self, codex: CodexProvider) -> Self {
48 self.register_sync(
49 "codex",
50 move |connection| Ok(codex.clone().with_connection(connection)),
51 CodexProvider::with_model,
52 )
53 }
54
55 pub fn with_generic_provider(self, config: &'static generic::ProviderConfig) -> Self {
56 self.register_sync(
57 config.provider.parser_name(),
58 move |connection| GenericProvider::from_env_with_connection(config, connection),
59 GenericProvider::with_model,
60 )
61 }
62
63 pub async fn parse(
74 &self,
75 models_str: &str,
76 connections: &ProviderConnectionOverrides,
77 ) -> Result<(Box<dyn StreamingModelProvider>, LlmModel)> {
78 let provider_model_pairs: Vec<&str> = models_str.split(',').map(str::trim).collect();
79 if provider_model_pairs.is_empty() {
80 return Err(LlmError::EmptyModelSpec);
81 }
82
83 let bedrock_has_inference_profile_arn = connections.config_for("bedrock").inference_profile_arn.is_some();
84 let mut seen_bedrock = false;
85 let mut seen_request_model_providers = HashSet::new();
86 let mut providers = Vec::new();
87 let mut first_identity: Option<LlmModel> = None;
88
89 for pair in provider_model_pairs {
90 let (provider_name, model) = pair.split_once(':').unwrap_or((pair, ""));
91
92 if provider_name == "bedrock" && bedrock_has_inference_profile_arn {
93 if seen_bedrock {
94 return Err(LlmError::DuplicateProvider {
95 provider: "bedrock".to_string(),
96 field: "inferenceProfileArn".to_string(),
97 });
98 }
99 seen_bedrock = true;
100 }
101
102 let connection = connections.config_for(provider_name);
103 if connection.request_model.is_some() && !seen_request_model_providers.insert(provider_name) {
104 return Err(LlmError::DuplicateProvider {
105 provider: provider_name.to_string(),
106 field: "requestModel".to_string(),
107 });
108 }
109
110 let factory = self
111 .factories
112 .get(provider_name)
113 .ok_or_else(|| LlmError::UnknownProvider { provider: provider_name.to_string() })?;
114 providers.push(factory(model, connection).await?);
115
116 if first_identity.is_none() {
117 first_identity = Some(pair.parse::<LlmModel>().map_err(LlmError::InvalidModelSpec)?);
118 }
119 }
120
121 let identity = first_identity.ok_or(LlmError::EmptyModelSpec)?;
122
123 let provider: Box<dyn StreamingModelProvider> = if providers.len() == 1 {
124 providers.into_iter().next().ok_or(LlmError::EmptyModelSpec)?
125 } else {
126 Box::new(AlloyedModelProvider::new(providers))
127 };
128
129 Ok((provider, identity))
130 }
131}
132
133type CreateProviderFn = Box<
134 dyn Fn(&str, ProviderConnectionConfig) -> BoxFuture<'static, Result<Box<dyn StreamingModelProvider>>> + Send + Sync,
135>;
136
137struct ProviderPool<T, U> {
138 connect: U,
139 providers: Mutex<HashMap<ProviderConnectionConfig, Arc<OnceCell<T>>>>,
140}
141
142impl<T, U, V> ProviderPool<T, U>
143where
144 T: Clone,
145 U: Fn(ProviderConnectionConfig) -> V,
146 V: Future<Output = Result<T>>,
147{
148 async fn connect(&self, config: ProviderConnectionConfig) -> Result<T> {
149 let cell = Arc::clone(
150 self.providers.lock().unwrap_or_else(PoisonError::into_inner).entry(config.clone()).or_default(),
151 );
152 cell.get_or_try_init(|| (self.connect)(config)).await.cloned()
153 }
154}
155
156impl ModelProviderParser {
157 fn register<T, U, V>(mut self, name: &str, connect: U, with_model: fn(T, &str) -> T) -> Self
158 where
159 T: StreamingModelProvider + Clone + 'static,
160 U: Fn(ProviderConnectionConfig) -> V + Send + Sync + 'static,
161 V: Future<Output = Result<T>> + Send + 'static,
162 {
163 let pool = Arc::new(ProviderPool { connect, providers: Mutex::default() });
164 let create: CreateProviderFn = Box::new(move |model, config| {
165 let (pool, model) = (Arc::clone(&pool), model.to_string());
166 Box::pin(async move {
167 let provider = pool.connect(config).await?;
168 Ok(Box::new(with_model(provider, &model)) as Box<dyn StreamingModelProvider>)
169 })
170 });
171 self.factories.insert(name.to_string(), create);
172 self
173 }
174
175 fn register_sync<T: StreamingModelProvider + Clone + 'static>(
176 self,
177 name: &str,
178 connect: impl Fn(ProviderConnectionConfig) -> Result<T> + Send + Sync + 'static,
179 with_model: fn(T, &str) -> T,
180 ) -> Self {
181 self.register(name, move |config| ready(connect(config)), with_model)
182 }
183}
184
185#[cfg(test)]
186mod tests {
187 use std::collections::BTreeMap;
188
189 use super::*;
190
191 #[tokio::test]
192 async fn test_parse_llamacpp() {
193 let result = parse("llamacpp").await;
194 assert!(result.is_ok());
195 let (_, model) = result.unwrap();
196 assert_eq!(model, LlmModel::LlamaCpp(String::new()));
197 }
198
199 #[tokio::test]
200 async fn test_parse_anthropic() {
201 let result = parse("anthropic:claude-sonnet-4-6").await;
202 match result {
203 Ok((_, model)) => {
204 assert_eq!(model, LlmModel::Anthropic(crate::catalog::AnthropicModel::ClaudeSonnet46));
205 }
206 Err(e) => {
207 let err = e.to_string();
208 assert!(
209 err.contains("API")
210 || err.contains("ANTHROPIC")
211 || err.contains("credentials")
212 || err.contains("JSON"),
213 "Should fail on API key or credentials, not parsing. Got: {err}"
214 );
215 }
216 }
217 }
218
219 #[tokio::test]
220 async fn test_parse_ollama() {
221 let result = parse("ollama:llama3.2").await;
222 assert!(result.is_ok());
223 let (_, model) = result.unwrap();
224 assert_eq!(model, LlmModel::Ollama("llama3.2".to_string()));
225 }
226
227 #[tokio::test]
228 async fn test_parse_openai() {
229 let result = parse("openai:gpt-4.1").await;
230 if let Err(e) = result {
231 let err = e.to_string();
232 assert!(err.contains("API") || err.contains("OPENAI"), "Should fail on API key, not parsing. Got: {err}");
233 }
234 }
235
236 #[tokio::test]
237 async fn test_parse_openrouter() {
238 let result = parse("openrouter:google/gemini-2.5-flash").await;
239 if let Err(e) = result {
240 let err = e.to_string();
241 assert!(err.contains("API") || err.contains("OPENROUTER"), "Should fail on API key, not parsing");
242 }
243 }
244
245 #[tokio::test]
246 async fn test_parse_gemini() {
247 let result = parse("gemini:gemini-2.5-flash").await;
248 if let Err(e) = result {
249 let err = e.to_string();
250 assert!(err.contains("API") || err.contains("GEMINI"), "Should fail on API key, not parsing");
251 }
252 }
253
254 #[tokio::test]
255 async fn test_parse_provider_without_model() {
256 let result = parse("anthropic").await;
257 assert!(result.is_err());
258 }
259
260 #[cfg(feature = "bedrock")]
261 #[tokio::test]
262 async fn test_parse_rejects_bedrock_inference_profile_arn() {
263 let spec = "bedrock:arn:aws:bedrock:us-west-2:000000000000:inference-profile/us.anthropic.claude-opus-4-7";
264
265 let error = match parse(spec).await {
266 Ok(_) => panic!("Bedrock ARN should be rejected"),
267 Err(error) => error.to_string(),
268 };
269
270 assert!(error.contains("providers.bedrock.inferenceProfileArn"), "{error}");
271 }
272
273 #[cfg(feature = "bedrock")]
274 #[tokio::test]
275 async fn test_parse_rejects_bedrock_application_inference_profile_arn() {
276 let spec = "bedrock:arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/000000000000";
277
278 let error = match parse(spec).await {
279 Ok(_) => panic!("Bedrock ARN should be rejected"),
280 Err(error) => error.to_string(),
281 };
282
283 assert!(error.contains("providers.bedrock.inferenceProfileArn"), "{error}");
284 }
285
286 #[tokio::test]
287 async fn test_parse_rejects_repeated_request_model_provider() {
288 let parser = ModelProviderParser::default();
289 let connections = ProviderConnectionOverrides::new(BTreeMap::from([(
290 "azure-foundry".to_string(),
291 crate::ProviderConnectionOverride {
292 base_url: Some("http://127.0.0.1:8787".to_string()),
293 auth_mode: Some(crate::ProviderAuthMode::None),
294 request_model: Some("production-coding".to_string()),
295 inference_profile_arn: None,
296 idle_timeout_secs: None,
297 },
298 )]));
299
300 let Err(error) = parser.parse("azure-foundry:gpt-5.5,azure-foundry:gpt-5.5", &connections).await else {
301 panic!("repeated request-model provider should be rejected");
302 };
303 assert!(error.to_string().contains("providers.azure-foundry.requestModel"));
304 }
305
306 #[tokio::test]
307 async fn test_parse_unknown_provider() {
308 let result = parse("unknown:model").await;
309 assert!(result.is_err());
310 if let Err(e) = result {
311 assert!(e.to_string().contains("Unknown provider"));
312 }
313 }
314
315 #[tokio::test]
316 async fn test_parse_single_provider() {
317 let result = parse("llamacpp").await;
318 assert!(result.is_ok());
319 }
320
321 #[tokio::test]
322 async fn test_parse_multiple_providers() {
323 let result = parse("llamacpp,ollama:llama3.2").await;
324 assert!(result.is_ok());
325 let (_, model) = result.unwrap();
326 assert_eq!(model, LlmModel::LlamaCpp(String::new()));
327 }
328
329 #[tokio::test]
330 async fn test_parse_with_spaces() {
331 let result = parse("llamacpp , ollama:llama3.2").await;
332 assert!(result.is_ok());
333 }
334
335 #[test]
336 fn test_parser_is_send_sync() {
337 fn assert_send_sync<T: Send + Sync>() {}
338 assert_send_sync::<ModelProviderParser>();
339 }
340
341 async fn parse(spec: &str) -> Result<(Box<dyn StreamingModelProvider>, LlmModel)> {
342 ModelProviderParser::default().parse(spec, &ProviderConnectionOverrides::default()).await
343 }
344}