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::{
8    anthropic::AnthropicProvider,
9    gemini::GeminiProvider,
10    local::{llama_cpp::LlamaCppProvider, ollama::OllamaProvider},
11    openai::OpenAiProvider,
12    openai_compatible::generic::{self, GenericOpenAiProvider},
13    openrouter::OpenRouterProvider,
14};
15use crate::{
16    LlmError, ProviderConnectionConfig, ProviderConnectionOverrides, ProviderFactory, StreamingModelProvider,
17    alloyed::AlloyedModelProvider,
18};
19#[cfg(feature = "codex")]
20use aether_auth::OAuthCredentialStorage;
21use futures::future::BoxFuture;
22use std::collections::HashMap;
23#[cfg(feature = "codex")]
24use std::sync::Arc;
25
26#[doc = include_str!("docs/parser.md")]
27pub struct ModelProviderParser {
28    factories: HashMap<String, CreateProviderFn>,
29    provider_connections: ProviderConnectionOverrides,
30}
31
32impl ModelProviderParser {
33    /// Create a new parser with custom provider factories
34    pub fn new(factories: HashMap<String, CreateProviderFn>) -> Self {
35        Self { factories, provider_connections: ProviderConnectionOverrides::default() }
36    }
37}
38
39impl Default for ModelProviderParser {
40    /// Create a parser with all built-in providers registered
41    fn default() -> Self {
42        let parser = Self::new(HashMap::new())
43            .with_provider::<AnthropicProvider>("anthropic")
44            .with_provider::<GeminiProvider>("gemini")
45            .with_provider::<OpenRouterProvider>("openrouter")
46            .with_provider::<OllamaProvider>("ollama")
47            .with_provider::<LlamaCppProvider>("llamacpp")
48            .with_provider::<OpenAiProvider>("openai")
49            .with_openai_provider("deepseek", &generic::DEEPSEEK)
50            .with_openai_provider("moonshot", &generic::MOONSHOT)
51            .with_openai_provider("zai", &generic::ZAI);
52
53        #[cfg(feature = "bedrock")]
54        let parser = parser.with_provider::<BedrockProvider>("bedrock");
55
56        parser
57    }
58}
59
60impl ModelProviderParser {
61    pub fn with_provider_connections(mut self, connections: ProviderConnectionOverrides) -> Self {
62        self.provider_connections = connections;
63        self
64    }
65
66    pub fn with_provider<P: ProviderFactory + StreamingModelProvider + 'static>(
67        mut self,
68        name: impl Into<String>,
69    ) -> Self {
70        self.factories.insert(
71            name.into(),
72            Box::new(|model: &str, connection: ProviderConnectionConfig| {
73                let model = model.to_string();
74                Box::pin(
75                    async move { Ok(Box::new(P::from_env_with_connection(connection).await?.with_model(&model)) as _) },
76                )
77            }),
78        );
79        self
80    }
81
82    #[cfg(feature = "codex")]
83    pub fn with_codex_provider(mut self, store: Arc<dyn OAuthCredentialStorage>) -> Self {
84        self.factories.insert(
85            "codex".to_string(),
86            Box::new(move |model: &str, connection: ProviderConnectionConfig| {
87                let store = Arc::clone(&store);
88                let model = model.to_string();
89                Box::pin(async move {
90                    Ok(Box::new(CodexProvider::new(store).with_connection(connection).with_model(&model)) as _)
91                })
92            }),
93        );
94        self
95    }
96
97    pub fn with_openai_provider(mut self, name: impl Into<String>, config: &'static generic::ProviderConfig) -> Self {
98        self.factories.insert(
99            name.into(),
100            Box::new(move |model: &str, connection: ProviderConnectionConfig| {
101                let model = model.to_string();
102                Box::pin(async move {
103                    Ok(
104                        Box::new(
105                            GenericOpenAiProvider::from_env_with_connection(config, connection)?.with_model(&model),
106                        ) as _,
107                    )
108                })
109            }),
110        );
111        self
112    }
113
114    /// Create a provider from a typed `LlmModel`
115    pub async fn create_provider(&self, model: &LlmModel) -> Result<Box<dyn StreamingModelProvider>> {
116        let key = model.provider();
117        let factory = self.factories.get(key).ok_or_else(|| LlmError::Other(format!("Unknown provider: {key}")))?;
118        factory(&model.model_id(), self.provider_connections.config_for(key)).await
119    }
120
121    /// Parse a model specification string and create a provider instance.
122    ///
123    /// Returns both the provider and an `LlmModel` describing the identity
124    /// of the first (or only) provider in the spec.
125    ///
126    /// # Format
127    ///
128    /// - `"provider:model"` - Single provider (e.g., "anthropic:claude-3.5-sonnet")
129    /// - `"provider1:model1,provider2:model2"` - Multiple providers create an `AlloyedModelProvider`
130    ///
131    pub async fn parse(&self, models_str: &str) -> Result<(Box<dyn StreamingModelProvider>, LlmModel)> {
132        let provider_model_pairs: Vec<&str> = models_str.split(',').map(str::trim).collect();
133        if provider_model_pairs.is_empty() {
134            return Err(LlmError::Other("No models provided".to_string()));
135        }
136
137        let bedrock_has_inference_profile_arn =
138            self.provider_connections.config_for("bedrock").inference_profile_arn.is_some();
139        let mut seen_bedrock = false;
140        let mut providers = Vec::new();
141        let mut first_identity: Option<LlmModel> = None;
142
143        for pair in provider_model_pairs {
144            let (provider_name, model) = pair.split_once(':').unwrap_or((pair, ""));
145
146            if provider_name == "bedrock" && bedrock_has_inference_profile_arn {
147                if seen_bedrock {
148                    return Err(LlmError::Other(
149                        "providers.bedrock.inferenceProfileArn cannot be used with multiple bedrock models in one alloy spec"
150                            .to_string(),
151                    ));
152                }
153                seen_bedrock = true;
154            }
155
156            let factory = self
157                .factories
158                .get(provider_name)
159                .ok_or_else(|| LlmError::Other(format!("Unknown provider: {provider_name}")))?;
160
161            let connection = self.provider_connections.config_for(provider_name);
162            providers.push(factory(model, connection).await?);
163
164            if first_identity.is_none() {
165                first_identity = Some(pair.parse::<LlmModel>().map_err(LlmError::Other)?);
166            }
167        }
168
169        let identity = first_identity.ok_or_else(|| LlmError::Other("No providers parsed".to_string()))?;
170
171        let provider: Box<dyn StreamingModelProvider> = if providers.len() == 1 {
172            providers.into_iter().next().ok_or_else(|| LlmError::Other("No providers available".to_string()))?
173        } else {
174            Box::new(AlloyedModelProvider::new(providers))
175        };
176
177        Ok((provider, identity))
178    }
179}
180
181/// Factory function type for creating model providers
182///
183/// Takes a model name and returns a boxed future that resolves to a `StreamingModelProvider`
184pub type CreateProviderFn = Box<
185    dyn Fn(&str, ProviderConnectionConfig) -> BoxFuture<'static, Result<Box<dyn StreamingModelProvider>>> + Send + Sync,
186>;
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    #[tokio::test]
193    async fn test_parse_llamacpp() {
194        let parser = ModelProviderParser::default();
195        let result = parser.parse("llamacpp").await;
196        assert!(result.is_ok());
197        let (_, model) = result.unwrap();
198        assert_eq!(model, LlmModel::LlamaCpp(String::new()));
199    }
200
201    #[tokio::test]
202    async fn test_parse_anthropic() {
203        let parser = ModelProviderParser::default();
204        let result = parser.parse("anthropic:claude-sonnet-4-6").await;
205        match result {
206            Ok((_, model)) => {
207                assert_eq!(model, LlmModel::Anthropic(crate::catalog::AnthropicModel::ClaudeSonnet46));
208            }
209            Err(e) => {
210                let err = e.to_string();
211                assert!(
212                    err.contains("API")
213                        || err.contains("ANTHROPIC")
214                        || err.contains("credentials")
215                        || err.contains("JSON"),
216                    "Should fail on API key or credentials, not parsing. Got: {err}"
217                );
218            }
219        }
220    }
221
222    #[tokio::test]
223    async fn test_parse_ollama() {
224        let parser = ModelProviderParser::default();
225        let result = parser.parse("ollama:llama3.2").await;
226        assert!(result.is_ok());
227        let (_, model) = result.unwrap();
228        assert_eq!(model, LlmModel::Ollama("llama3.2".to_string()));
229    }
230
231    #[tokio::test]
232    async fn test_parse_openai() {
233        let parser = ModelProviderParser::default();
234        let result = parser.parse("openai:gpt-4.1").await;
235        if let Err(e) = result {
236            let err = e.to_string();
237            assert!(err.contains("API") || err.contains("OPENAI"), "Should fail on API key, not parsing. Got: {err}");
238        }
239    }
240
241    #[tokio::test]
242    async fn test_parse_openrouter() {
243        let parser = ModelProviderParser::default();
244        let result = parser.parse("openrouter:google/gemini-2.5-flash").await;
245        if let Err(e) = result {
246            let err = e.to_string();
247            assert!(err.contains("API") || err.contains("OPENROUTER"), "Should fail on API key, not parsing");
248        }
249    }
250
251    #[tokio::test]
252    async fn test_parse_gemini() {
253        let parser = ModelProviderParser::default();
254        let result = parser.parse("gemini:gemini-2.5-flash").await;
255        if let Err(e) = result {
256            let err = e.to_string();
257            assert!(err.contains("API") || err.contains("GEMINI"), "Should fail on API key, not parsing");
258        }
259    }
260
261    #[tokio::test]
262    async fn test_parse_provider_without_model() {
263        let parser = ModelProviderParser::default();
264        let result = parser.parse("anthropic").await;
265        assert!(result.is_err());
266    }
267
268    #[cfg(feature = "bedrock")]
269    #[tokio::test]
270    async fn test_parse_rejects_bedrock_inference_profile_arn() {
271        let parser = ModelProviderParser::default();
272        let spec = "bedrock:arn:aws:bedrock:us-west-2:000000000000:inference-profile/us.anthropic.claude-opus-4-7";
273
274        let error = match parser.parse(spec).await {
275            Ok(_) => panic!("Bedrock ARN should be rejected"),
276            Err(error) => error.to_string(),
277        };
278
279        assert!(error.contains("providers.bedrock.inferenceProfileArn"), "{error}");
280    }
281
282    #[cfg(feature = "bedrock")]
283    #[tokio::test]
284    async fn test_parse_rejects_bedrock_application_inference_profile_arn() {
285        let parser = ModelProviderParser::default();
286        let spec = "bedrock:arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/000000000000";
287
288        let error = match parser.parse(spec).await {
289            Ok(_) => panic!("Bedrock ARN should be rejected"),
290            Err(error) => error.to_string(),
291        };
292
293        assert!(error.contains("providers.bedrock.inferenceProfileArn"), "{error}");
294    }
295
296    #[tokio::test]
297    async fn test_parse_unknown_provider() {
298        let parser = ModelProviderParser::default();
299        let result = parser.parse("unknown:model").await;
300        assert!(result.is_err());
301        if let Err(e) = result {
302            assert!(e.to_string().contains("Unknown provider"));
303        }
304    }
305
306    #[tokio::test]
307    async fn test_with_custom_provider() {
308        let parser = ModelProviderParser::default().with_provider::<OllamaProvider>("custom");
309
310        let model = LlmModel::Ollama("test-model".to_string());
311        let result = parser.create_provider(&model).await;
312        assert!(result.is_ok());
313    }
314
315    #[tokio::test]
316    async fn test_parse_single_provider() {
317        let parser = ModelProviderParser::default();
318        let result = parser.parse("llamacpp").await;
319        assert!(result.is_ok());
320    }
321
322    #[tokio::test]
323    async fn test_parse_multiple_providers() {
324        let parser = ModelProviderParser::default();
325        let result = parser.parse("llamacpp,ollama:llama3.2").await;
326        assert!(result.is_ok());
327        let (_, model) = result.unwrap();
328        assert_eq!(model, LlmModel::LlamaCpp(String::new()));
329    }
330
331    #[tokio::test]
332    async fn test_parse_with_spaces() {
333        let parser = ModelProviderParser::default();
334        let result = parser.parse("llamacpp , ollama:llama3.2").await;
335        assert!(result.is_ok());
336    }
337
338    #[test]
339    fn test_parser_is_send_sync() {
340        fn assert_send_sync<T: Send + Sync>() {}
341        assert_send_sync::<ModelProviderParser>();
342    }
343}
llm/parser.rs

llm/
parser.rs
