agent_framework_mistral/
lib.rs

1//! # agent-framework-mistral
2//!
3//! A [Mistral AI](https://mistral.ai) [`ChatClient`]
4//! for `agent-framework-rs`.
5//!
6//! Talks to the Mistral Chat Completions API (`POST /v1/chat/completions`),
7//! whose wire format is OpenAI-Chat-compatible: message shape, function-tool
8//! shape, tool-call/response shape, and usage shape all match
9//! [`agent-framework-openai`](agent_framework_openai)'s Chat Completions
10//! client, so this crate reuses that conversion code directly (the same
11//! approach `agent-framework-azure` takes for Azure OpenAI) rather than
12//! duplicating it. Only the parts that genuinely differ — the base URL,
13//! bearer-token auth, the supported request-option set (see
14//! [`convert::apply_options`]), and error classification — are implemented
15//! here; see [`convert`] for the details.
16//!
17//! Upstream (the Python/.NET `agent-framework`) ships no dedicated Mistral
18//! chat connector at all — Mistral only appears there as an
19//! embeddings/text-embedding provider. This crate instead provides a full
20//! Mistral **chat** client, since Mistral's hosted models are chat models
21//! first and foremost and the framework's [`ChatClient`] trait is the
22//! natural fit. It can gain an embeddings client of its own once
23//! `agent-framework-core` grows a shared embeddings trait to implement
24//! against; until then, chat is the complete scope of this crate.
25//!
26//! ```no_run
27//! use agent_framework_mistral::MistralChatClient;
28//! use agent_framework_core::prelude::*;
29//!
30//! # async fn demo() -> Result<()> {
31//! let client = MistralChatClient::new("...", "mistral-large-latest");
32//! let agent = Agent::builder(client)
33//!     .instructions("You are concise.")
34//!     .build();
35//! let reply = agent.run_once("Say hi").await?;
36//! println!("{}", reply.text());
37//! # Ok(())
38//! # }
39//! ```
40
41pub mod convert;
42pub mod embeddings;
43pub use embeddings::{MistralEmbeddingClient, DEFAULT_EMBEDDING_MODEL};
44
45use std::sync::Arc;
46
47use agent_framework_core::client::{ChatClient, ChatStream};
48use agent_framework_core::error::{Error, Result};
49use agent_framework_core::types::{ChatOptions, ChatResponse, Message};
50use futures::StreamExt;
51use serde_json::Value;
52
53pub(crate) const DEFAULT_BASE_URL: &str = "https://api.mistral.ai/v1";
54
55/// Parse a `Retry-After` header into a delay in seconds.
56///
57/// Mistral returns the integer/decimal-seconds form on `429`, mirroring
58/// OpenAI/Anthropic; a date-form or unparseable value is treated as absent.
59fn parse_retry_after(headers: &reqwest::header::HeaderMap) -> Option<f64> {
60    headers
61        .get(reqwest::header::RETRY_AFTER)
62        .and_then(|v| v.to_str().ok())
63        .and_then(|s| s.trim().parse::<f64>().ok())
64        .filter(|s| s.is_finite() && *s >= 0.0)
65}
66
67/// A Mistral AI chat client (`POST {base_url}/chat/completions`).
68#[derive(Clone)]
69pub struct MistralChatClient {
70    inner: Arc<Inner>,
71}
72
73#[derive(Clone)]
74struct Inner {
75    http: reqwest::Client,
76    api_key: String,
77    base_url: String,
78    model: String,
79}
80
81impl std::fmt::Debug for MistralChatClient {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        f.debug_struct("MistralChatClient")
84            .field("base_url", &self.inner.base_url)
85            .field("model", &self.inner.model)
86            .finish_non_exhaustive()
87    }
88}
89
90impl MistralChatClient {
91    /// Create a client for the given API key and default model.
92    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
93        Self {
94            inner: Arc::new(Inner {
95                http: reqwest::Client::new(),
96                api_key: api_key.into(),
97                base_url: DEFAULT_BASE_URL.to_string(),
98                model: model.into(),
99            }),
100        }
101    }
102
103    /// Build a client from the `MISTRAL_API_KEY` (and optional
104    /// `MISTRAL_BASE_URL`) environment variables.
105    ///
106    /// Unlike Ollama (which runs unauthenticated by default and so has no
107    /// API-key environment variable at all), Mistral's hosted API always
108    /// requires a bearer token, so `MISTRAL_API_KEY` is required here.
109    pub fn from_env(model: impl Into<String>) -> Result<Self> {
110        let key = std::env::var("MISTRAL_API_KEY")
111            .map_err(|_| Error::Configuration("MISTRAL_API_KEY is not set".into()))?;
112        let mut client = Self::new(key, model);
113        if let Ok(base) = std::env::var("MISTRAL_BASE_URL") {
114            client = client.with_base_url(base);
115        }
116        Ok(client)
117    }
118
119    /// Override the base URL (for proxies or private deployments).
120    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
121        Arc::make_mut(&mut self.inner).base_url = base_url.into();
122        self
123    }
124
125    /// The default model id.
126    pub fn model(&self) -> &str {
127        &self.inner.model
128    }
129
130    fn url(&self) -> String {
131        format!(
132            "{}/chat/completions",
133            self.inner.base_url.trim_end_matches('/')
134        )
135    }
136
137    fn build_body(&self, messages: &[Message], options: &ChatOptions, stream: bool) -> Value {
138        let model = options
139            .model
140            .clone()
141            .unwrap_or_else(|| self.inner.model.clone());
142        convert::build_request(messages, options, &model, stream)
143    }
144
145    async fn post(&self, body: &Value) -> Result<reqwest::Response> {
146        let resp = self
147            .inner
148            .http
149            .post(self.url())
150            .bearer_auth(&self.inner.api_key)
151            .json(body)
152            .send()
153            .await
154            .map_err(|e| Error::service(format!("request failed: {e}")))?;
155        if !resp.status().is_success() {
156            let status = resp.status();
157            let retry_after = parse_retry_after(resp.headers());
158            let text = resp.text().await.unwrap_or_default();
159            return Err(convert::classify_mistral_error(
160                status.as_u16(),
161                format!("Mistral API error {status}: {text}"),
162                retry_after,
163            ));
164        }
165        Ok(resp)
166    }
167}
168
169#[async_trait::async_trait]
170impl ChatClient for MistralChatClient {
171    async fn get_response(
172        &self,
173        messages: Vec<Message>,
174        options: ChatOptions,
175    ) -> Result<ChatResponse> {
176        let body = self.build_body(&messages, &options, false);
177        let resp = self.post(&body).await?;
178        let value: Value = resp
179            .json()
180            .await
181            .map_err(|e| Error::service(format!("invalid response json: {e}")))?;
182        Ok(convert::parse_response(&value))
183    }
184
185    async fn get_streaming_response(
186        &self,
187        messages: Vec<Message>,
188        options: ChatOptions,
189    ) -> Result<ChatStream> {
190        let body = self.build_body(&messages, &options, true);
191        let resp = self.post(&body).await?;
192        // Mistral's streaming Chat Completions wire shape (SSE
193        // `chat.completion.chunk` objects, `[DONE]` terminator, a trailing
194        // usage-only chunk when `stream_options.include_usage` is set, and
195        // mid-stream `{"error": {...}}` objects) is identical to OpenAI's, so
196        // parsing is reused verbatim rather than duplicated.
197        Ok(agent_framework_openai::parse_sse_stream(resp).boxed())
198    }
199
200    fn model(&self) -> Option<&str> {
201        Some(&self.inner.model)
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    fn client() -> MistralChatClient {
210        MistralChatClient::new("test-key", "mistral-large-latest")
211    }
212
213    // region: URL building
214
215    #[test]
216    fn url_uses_default_base_url() {
217        let c = client();
218        assert_eq!(c.url(), "https://api.mistral.ai/v1/chat/completions");
219    }
220
221    #[test]
222    fn with_base_url_overrides_default_and_trims_trailing_slash() {
223        let c = client().with_base_url("https://proxy.example.com/v1/");
224        assert_eq!(c.url(), "https://proxy.example.com/v1/chat/completions");
225    }
226
227    // endregion
228
229    // region: request body building
230
231    #[test]
232    fn build_body_defaults_to_client_model() {
233        let c = client();
234        let body = c.build_body(&[Message::user("hi")], &ChatOptions::new(), false);
235        assert_eq!(body["model"], serde_json::json!("mistral-large-latest"));
236    }
237
238    #[test]
239    fn build_body_per_request_model_overrides_client_default() {
240        let c = client();
241        let options = ChatOptions::new().with_model("mistral-small-latest");
242        let body = c.build_body(&[Message::user("hi")], &options, false);
243        assert_eq!(body["model"], serde_json::json!("mistral-small-latest"));
244    }
245
246    #[test]
247    fn build_body_stream_flag() {
248        let c = client();
249        let body = c.build_body(&[Message::user("hi")], &ChatOptions::new(), true);
250        assert_eq!(body["stream"], serde_json::json!(true));
251    }
252
253    // endregion
254
255    // Streaming: `get_streaming_response` delegates directly to
256    // `agent_framework_openai::parse_sse_stream`, with no Mistral-specific
257    // logic of its own -- the wiring is verified at compile time (this crate
258    // wouldn't type-check against `ChatStream` otherwise), and SSE chunk
259    // parsing itself (text deltas, tool-call argument accumulation, `[DONE]`
260    // handling, trailing usage chunk, mid-stream error surfacing) is already
261    // covered by `agent-framework-openai`'s own test suite, which this crate
262    // depends on unchanged. `reqwest::Response` can't be constructed from raw
263    // bytes outside an actual HTTP exchange, so reproducing those fixtures
264    // here would require standing up a mock server rather than a plain unit
265    // test (`agent-framework-azure`, which reuses the very same function,
266    // documents this identically).
267
268    // region: env-var constructor
269
270    /// Guards `MISTRAL_API_KEY` / `MISTRAL_BASE_URL` mutation: tests within a
271    /// crate run on multiple threads, and env vars are process-global, so
272    /// this serializes access across the tests below.
273    static ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
274
275    #[test]
276    fn from_env_reads_api_key_and_base_url() {
277        let _guard = ENV_MUTEX.lock().unwrap();
278        // SAFETY: serialized by ENV_MUTEX against the other env-var test in
279        // this module; no other test in this crate touches these variables.
280        unsafe {
281            std::env::set_var("MISTRAL_API_KEY", "test-key-123");
282            std::env::set_var("MISTRAL_BASE_URL", "https://example.test/v1");
283        }
284        let client = MistralChatClient::from_env("mistral-large-latest").unwrap();
285        assert_eq!(client.inner.api_key, "test-key-123");
286        assert_eq!(client.inner.base_url, "https://example.test/v1");
287        unsafe {
288            std::env::remove_var("MISTRAL_API_KEY");
289            std::env::remove_var("MISTRAL_BASE_URL");
290        }
291    }
292
293    #[test]
294    fn from_env_errors_when_api_key_missing() {
295        let _guard = ENV_MUTEX.lock().unwrap();
296        // SAFETY: serialized by ENV_MUTEX; see above.
297        unsafe {
298            std::env::remove_var("MISTRAL_API_KEY");
299            std::env::remove_var("MISTRAL_BASE_URL");
300        }
301        let result = MistralChatClient::from_env("mistral-large-latest");
302        assert!(result.is_err());
303    }
304
305    #[test]
306    fn from_env_defaults_base_url_when_unset() {
307        let _guard = ENV_MUTEX.lock().unwrap();
308        // SAFETY: serialized by ENV_MUTEX; see above.
309        unsafe {
310            std::env::set_var("MISTRAL_API_KEY", "test-key-123");
311            std::env::remove_var("MISTRAL_BASE_URL");
312        }
313        let client = MistralChatClient::from_env("mistral-large-latest").unwrap();
314        assert_eq!(client.inner.base_url, DEFAULT_BASE_URL);
315        unsafe {
316            std::env::remove_var("MISTRAL_API_KEY");
317        }
318    }
319
320    // endregion
321
322    #[test]
323    fn model_returns_default_model() {
324        let c = client();
325        assert_eq!(c.model(), "mistral-large-latest");
326        assert_eq!(ChatClient::model(&c), Some("mistral-large-latest"));
327    }
328}
agent_framework_mistral/lib.rs

agent_framework_mistral/
lib.rs
