uni_xervo/provider/
mistral.rs

1use crate::api::{ModelAliasSpec, ModelTask};
2use crate::error::{Result, RuntimeError};
3use crate::provider::remote_common::{RemoteProviderBase, check_http_status, resolve_api_key};
4use crate::traits::{
5    EmbedResult, EmbeddingModel, GenerationOptions, GenerationResult, GeneratorModel,
6    LoadedModelHandle, Message, MessageRole, ModelProvider, ProviderCapabilities, ProviderHealth,
7    TokenUsage,
8};
9use async_trait::async_trait;
10use reqwest::Client;
11use serde_json::json;
12use std::sync::Arc;
13
14/// Remote provider that calls the [Mistral AI API](https://docs.mistral.ai/api/)
15/// for embedding and text generation (chat completions).
16///
17/// Requires the `MISTRAL_API_KEY` environment variable (or a custom env var
18/// name via the `api_key_env` option).
19pub struct RemoteMistralProvider {
20    base: RemoteProviderBase,
21}
22
23impl Default for RemoteMistralProvider {
24    fn default() -> Self {
25        Self {
26            base: RemoteProviderBase::new(),
27        }
28    }
29}
30
31impl RemoteMistralProvider {
32    pub fn new() -> Self {
33        Self::default()
34    }
35
36    #[cfg(test)]
37    fn insert_test_breaker(&self, key: crate::api::ModelRuntimeKey, age: std::time::Duration) {
38        self.base.insert_test_breaker(key, age);
39    }
40
41    #[cfg(test)]
42    fn breaker_count(&self) -> usize {
43        self.base.breaker_count()
44    }
45
46    #[cfg(test)]
47    fn force_cleanup_now_for_test(&self) {
48        self.base.force_cleanup_now_for_test();
49    }
50}
51
52#[async_trait]
53impl ModelProvider for RemoteMistralProvider {
54    fn provider_id(&self) -> &'static str {
55        "remote/mistral"
56    }
57
58    fn capabilities(&self) -> ProviderCapabilities {
59        ProviderCapabilities {
60            supported_tasks: vec![ModelTask::Embed, ModelTask::Generate],
61        }
62    }
63
64    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
65        let cb = self.base.circuit_breaker_for(spec);
66        let api_key = resolve_api_key(&spec.options, "api_key_env", "MISTRAL_API_KEY")?;
67        let embedding_dimensions = spec
68            .options
69            .get("embedding_dimensions")
70            .and_then(|v| v.as_u64())
71            .map(|v| v as u32);
72
73        match spec.task {
74            ModelTask::Embed => {
75                let default_dims = match spec.model_id.as_str() {
76                    "codestral-embed" => 1536,
77                    _ => 1024,
78                };
79                let model = MistralEmbeddingModel {
80                    client: self.base.client.clone(),
81                    cb: cb.clone(),
82                    model_id: spec.model_id.clone(),
83                    api_key,
84                    dimensions: embedding_dimensions.unwrap_or(default_dims),
85                };
86                let handle: Arc<dyn EmbeddingModel> = Arc::new(model);
87                Ok(Arc::new(handle) as LoadedModelHandle)
88            }
89            ModelTask::Generate => {
90                let model = MistralGeneratorModel {
91                    client: self.base.client.clone(),
92                    cb,
93                    model_id: spec.model_id.clone(),
94                    api_key,
95                };
96                let handle: Arc<dyn GeneratorModel> = Arc::new(model);
97                Ok(Arc::new(handle) as LoadedModelHandle)
98            }
99            ModelTask::Raw => Err(RuntimeError::CapabilityMismatch(
100                "Mistral provider does not support task Raw".to_string(),
101            )),
102            _ => Err(RuntimeError::CapabilityMismatch(format!(
103                "Mistral provider does not support task {:?}",
104                spec.task
105            ))),
106        }
107    }
108
109    async fn health(&self) -> ProviderHealth {
110        ProviderHealth::Healthy
111    }
112}
113
114struct MistralEmbeddingModel {
115    client: Client,
116    cb: crate::reliability::CircuitBreakerWrapper,
117    model_id: String,
118    api_key: String,
119    dimensions: u32,
120}
121
122#[async_trait]
123impl EmbeddingModel for MistralEmbeddingModel {
124    async fn embed(&self, texts: &[&str]) -> Result<EmbedResult> {
125        let texts: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
126
127        self.cb
128            .call(move || async move {
129                let response = self
130                    .client
131                    .post("https://api.mistral.ai/v1/embeddings")
132                    .header("Authorization", format!("Bearer {}", self.api_key))
133                    .json(&json!({
134                        "model": self.model_id,
135                        "input": texts
136                    }))
137                    .send()
138                    .await
139                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
140
141                let body: serde_json::Value = check_http_status("Mistral", response)?
142                    .json()
143                    .await
144                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
145
146                let mut vectors = Vec::new();
147                if let Some(data) = body.get("data").and_then(|d| d.as_array()) {
148                    for item in data {
149                        if let Some(embedding) = item.get("embedding").and_then(|e| e.as_array()) {
150                            let vec: Vec<f32> = embedding
151                                .iter()
152                                .filter_map(|v| v.as_f64().map(|f| f as f32))
153                                .collect();
154                            vectors.push(vec);
155                        }
156                    }
157                }
158
159                // Mistral embeddings response (OpenAI-compatible) carries a
160                // `usage` object with `prompt_tokens` and `total_tokens`.
161                let usage = body.get("usage").map(|u| {
162                    let prompt = u["prompt_tokens"].as_u64().unwrap_or(0) as usize;
163                    let total = u["total_tokens"].as_u64().unwrap_or(prompt as u64) as usize;
164                    TokenUsage {
165                        prompt_tokens: prompt,
166                        completion_tokens: 0,
167                        total_tokens: total,
168                    }
169                });
170
171                Ok(EmbedResult { vectors, usage })
172            })
173            .await
174    }
175
176    fn dimensions(&self) -> u32 {
177        self.dimensions
178    }
179}
180
181impl crate::traits::ModelInfo for MistralEmbeddingModel {
182    fn model_id(&self) -> &str {
183        &self.model_id
184    }
185}
186
187struct MistralGeneratorModel {
188    client: Client,
189    cb: crate::reliability::CircuitBreakerWrapper,
190    model_id: String,
191    api_key: String,
192}
193
194impl crate::traits::ModelInfo for MistralGeneratorModel {
195    fn model_id(&self) -> &str {
196        &self.model_id
197    }
198}
199
200#[async_trait]
201impl GeneratorModel for MistralGeneratorModel {
202    async fn generate(
203        &self,
204        messages: &[Message],
205        options: GenerationOptions,
206    ) -> Result<GenerationResult> {
207        let messages: Vec<serde_json::Value> = messages
208            .iter()
209            .map(|msg| {
210                let role = match msg.role {
211                    MessageRole::System => "system",
212                    MessageRole::User => "user",
213                    MessageRole::Assistant => "assistant",
214                };
215                json!({ "role": role, "content": msg.text() })
216            })
217            .collect();
218
219        self.cb
220            .call(move || async move {
221                let mut body = json!({
222                    "model": self.model_id,
223                    "messages": messages,
224                });
225
226                if let Some(max_tokens) = options.max_tokens {
227                    body["max_tokens"] = json!(max_tokens);
228                }
229                if let Some(temperature) = options.temperature {
230                    body["temperature"] = json!(temperature);
231                }
232                if let Some(top_p) = options.top_p {
233                    body["top_p"] = json!(top_p);
234                }
235
236                let response = self
237                    .client
238                    .post("https://api.mistral.ai/v1/chat/completions")
239                    .header("Authorization", format!("Bearer {}", self.api_key))
240                    .json(&body)
241                    .send()
242                    .await
243                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
244
245                let body: serde_json::Value = check_http_status("Mistral", response)?
246                    .json()
247                    .await
248                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
249
250                let text = body["choices"][0]["message"]["content"]
251                    .as_str()
252                    .unwrap_or("")
253                    .to_string();
254
255                let usage = body.get("usage").map(|u| TokenUsage {
256                    prompt_tokens: u["prompt_tokens"].as_u64().unwrap_or(0) as usize,
257                    completion_tokens: u["completion_tokens"].as_u64().unwrap_or(0) as usize,
258                    total_tokens: u["total_tokens"].as_u64().unwrap_or(0) as usize,
259                });
260
261                Ok(GenerationResult {
262                    text,
263                    usage,
264                    images: vec![],
265                    audio: None,
266                })
267            })
268            .await
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275    use crate::api::ModelRuntimeKey;
276    use crate::provider::remote_common::RemoteProviderBase;
277    use crate::traits::ModelProvider;
278    use std::time::Duration;
279
280    static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
281
282    fn spec(alias: &str, task: ModelTask, model_id: &str) -> ModelAliasSpec {
283        ModelAliasSpec {
284            alias: alias.to_string(),
285            task,
286            provider_id: "remote/mistral".to_string(),
287            model_id: model_id.to_string(),
288            revision: None,
289            warmup: crate::api::WarmupPolicy::Lazy,
290            required: false,
291            timeout: None,
292            load_timeout: None,
293            retry: None,
294            options: serde_json::Value::Null,
295        }
296    }
297
298    #[tokio::test]
299    async fn breaker_reused_for_same_runtime_key() {
300        let _lock = ENV_LOCK.lock().await;
301        unsafe { std::env::set_var("MISTRAL_API_KEY", "test-key") };
302
303        let provider = RemoteMistralProvider::new();
304        let s1 = spec("embed/a", ModelTask::Embed, "mistral-embed");
305        let s2 = spec("embed/b", ModelTask::Embed, "mistral-embed");
306
307        let _ = provider.load(&s1).await.unwrap();
308        let _ = provider.load(&s2).await.unwrap();
309
310        assert_eq!(provider.breaker_count(), 1);
311
312        unsafe { std::env::remove_var("MISTRAL_API_KEY") };
313    }
314
315    #[tokio::test]
316    async fn breaker_isolated_by_task_and_model() {
317        let _lock = ENV_LOCK.lock().await;
318        unsafe { std::env::set_var("MISTRAL_API_KEY", "test-key") };
319
320        let provider = RemoteMistralProvider::new();
321        let embed = spec("embed/a", ModelTask::Embed, "mistral-embed");
322        let gen_spec = spec("chat/a", ModelTask::Generate, "mistral-small-latest");
323
324        let _ = provider.load(&embed).await.unwrap();
325        let _ = provider.load(&gen_spec).await.unwrap();
326
327        assert_eq!(provider.breaker_count(), 2);
328
329        unsafe { std::env::remove_var("MISTRAL_API_KEY") };
330    }
331
332    #[tokio::test]
333    async fn breaker_cleanup_evicts_stale_entries() {
334        let _lock = ENV_LOCK.lock().await;
335        unsafe { std::env::set_var("MISTRAL_API_KEY", "test-key") };
336
337        let provider = RemoteMistralProvider::new();
338        let stale = spec("embed/stale", ModelTask::Embed, "mistral-embed");
339        let fresh = spec("chat/fresh", ModelTask::Generate, "mistral-small-latest");
340        provider.insert_test_breaker(
341            ModelRuntimeKey::new(&stale),
342            RemoteProviderBase::BREAKER_TTL + Duration::from_secs(5),
343        );
344        provider.insert_test_breaker(ModelRuntimeKey::new(&fresh), Duration::from_secs(1));
345        assert_eq!(provider.breaker_count(), 2);
346
347        provider.force_cleanup_now_for_test();
348        let _ = provider.load(&fresh).await.unwrap();
349
350        assert_eq!(provider.breaker_count(), 1);
351
352        unsafe { std::env::remove_var("MISTRAL_API_KEY") };
353    }
354
355    #[tokio::test]
356    async fn rerank_capability_mismatch() {
357        let _lock = ENV_LOCK.lock().await;
358        unsafe { std::env::set_var("MISTRAL_API_KEY", "test-key") };
359
360        let provider = RemoteMistralProvider::new();
361        let s = spec("rerank/a", ModelTask::Rerank, "mistral-embed");
362        let result = provider.load(&s).await;
363        assert!(result.is_err());
364        assert!(
365            result
366                .unwrap_err()
367                .to_string()
368                .contains("does not support task")
369        );
370
371        unsafe { std::env::remove_var("MISTRAL_API_KEY") };
372    }
373
374    #[tokio::test]
375    async fn default_embedding_dimensions() {
376        let _lock = ENV_LOCK.lock().await;
377        unsafe { std::env::set_var("MISTRAL_API_KEY", "test-key") };
378
379        let provider = RemoteMistralProvider::new();
380        let s = spec("embed/dim", ModelTask::Embed, "mistral-embed");
381
382        let handle = provider.load(&s).await.unwrap();
383        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
384        assert_eq!(model.dimensions(), 1024);
385
386        unsafe { std::env::remove_var("MISTRAL_API_KEY") };
387    }
388
389    #[tokio::test]
390    async fn custom_embedding_dimensions() {
391        let _lock = ENV_LOCK.lock().await;
392        unsafe { std::env::set_var("MISTRAL_API_KEY", "test-key") };
393
394        let provider = RemoteMistralProvider::new();
395        let mut s = spec("embed/dim-custom", ModelTask::Embed, "mistral-embed");
396        s.options = serde_json::json!({"embedding_dimensions": 256});
397
398        let handle = provider.load(&s).await.unwrap();
399        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
400        assert_eq!(model.dimensions(), 256);
401
402        unsafe { std::env::remove_var("MISTRAL_API_KEY") };
403    }
404}
uni_xervo/provider/mistral.rs

uni_xervo/provider/
mistral.rs
