Skip to main content

llm/
parser.rs

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    /// Create a parser with all built-in providers registered
25    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    /// Parse a model specification string and create a provider instance.
64    ///
65    /// Returns both the provider and an `LlmModel` describing the identity
66    /// of the first (or only) provider in the spec.
67    ///
68    /// # Format
69    ///
70    /// - `"provider:model"` - Single provider (e.g., "anthropic:claude-3.5-sonnet")
71    /// - `"provider1:model1,provider2:model2"` - Multiple providers create an `AlloyedModelProvider`
72    ///
73    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}