agent_framework_mistral/
embeddings.rs

1//! Mistral embeddings client.
2//!
3//! Rust equivalent of upstream's `MistralEmbeddingClient`
4//! (`agent_framework_mistral/_embedding_client.py`):
5//! `POST {base_url}/embeddings` (`https://api.mistral.ai/v1` by default),
6//! OpenAI-shaped request/response, default model `mistral-embed`. Response
7//! parsing is shared with `agent-framework-openai`.
8
9use std::sync::Arc;
10
11use agent_framework_core::client::EmbeddingClient;
12use agent_framework_core::error::{Error, Result};
13use agent_framework_core::types::{EmbeddingGenerationOptions, GeneratedEmbeddings};
14use serde_json::{json, Map, Value};
15
16use crate::DEFAULT_BASE_URL;
17
18/// Upstream's default Mistral embedding model.
19pub const DEFAULT_EMBEDDING_MODEL: &str = "mistral-embed";
20
21/// A Mistral embeddings client (`POST {base_url}/embeddings`).
22///
23/// ```no_run
24/// # use agent_framework_mistral::MistralEmbeddingClient;
25/// # use agent_framework_core::client::EmbeddingClient;
26/// # async fn demo() -> agent_framework_core::error::Result<()> {
27/// let client = MistralEmbeddingClient::from_env(None)?; // mistral-embed
28/// let batch = client.get_embeddings(vec!["hello".into()], None).await?;
29/// println!("{} dims", batch[0].dimensions());
30/// # Ok(())
31/// # }
32/// ```
33#[derive(Clone)]
34pub struct MistralEmbeddingClient {
35    inner: Arc<Inner>,
36}
37
38#[derive(Clone)]
39struct Inner {
40    http: reqwest::Client,
41    api_key: String,
42    base_url: String,
43    model: String,
44}
45
46impl std::fmt::Debug for MistralEmbeddingClient {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        f.debug_struct("MistralEmbeddingClient")
49            .field("base_url", &self.inner.base_url)
50            .field("model", &self.inner.model)
51            .finish_non_exhaustive()
52    }
53}
54
55impl MistralEmbeddingClient {
56    /// Create a client for the given API key and embedding model
57    /// (`None` → [`DEFAULT_EMBEDDING_MODEL`]).
58    pub fn new(api_key: impl Into<String>, model: Option<String>) -> Self {
59        Self {
60            inner: Arc::new(Inner {
61                http: reqwest::Client::new(),
62                api_key: api_key.into(),
63                base_url: DEFAULT_BASE_URL.to_string(),
64                model: model.unwrap_or_else(|| DEFAULT_EMBEDDING_MODEL.to_string()),
65            }),
66        }
67    }
68
69    /// Build a client from the `MISTRAL_API_KEY` (and optional
70    /// `MISTRAL_BASE_URL`) environment variables.
71    pub fn from_env(model: Option<String>) -> Result<Self> {
72        let key = std::env::var("MISTRAL_API_KEY")
73            .map_err(|_| Error::Configuration("MISTRAL_API_KEY is not set".into()))?;
74        let mut client = Self::new(key, model);
75        if let Ok(base) = std::env::var("MISTRAL_BASE_URL") {
76            client = client.with_base_url(base);
77        }
78        Ok(client)
79    }
80
81    /// Override the base URL.
82    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
83        Arc::make_mut(&mut self.inner).base_url = base_url.into();
84        self
85    }
86}
87
88#[async_trait::async_trait]
89impl EmbeddingClient for MistralEmbeddingClient {
90    async fn get_embeddings(
91        &self,
92        values: Vec<String>,
93        options: Option<EmbeddingGenerationOptions>,
94    ) -> Result<GeneratedEmbeddings> {
95        let mut body = Map::new();
96        let model = options
97            .as_ref()
98            .and_then(|o| o.model.clone())
99            .unwrap_or_else(|| self.inner.model.clone());
100        body.insert("model".into(), json!(model));
101        body.insert("input".into(), json!(values));
102        // Mistral spells requested dimensionality `output_dimension`
103        // (upstream maps options["dimensions"] onto it).
104        if let Some(dimensions) = options.as_ref().and_then(|o| o.dimensions) {
105            body.insert("output_dimension".into(), json!(dimensions));
106        }
107        if let Some(dtype) = options
108            .as_ref()
109            .and_then(|o| o.additional_properties.get("output_dtype"))
110        {
111            body.insert("output_dtype".into(), dtype.clone());
112        }
113        let url = format!("{}/embeddings", self.inner.base_url.trim_end_matches('/'));
114        let resp = self
115            .inner
116            .http
117            .post(&url)
118            .bearer_auth(&self.inner.api_key)
119            .json(&Value::Object(body))
120            .send()
121            .await
122            .map_err(|e| Error::service(format!("request failed: {e}")))?;
123        if !resp.status().is_success() {
124            let status = resp.status();
125            let text = resp.text().await.unwrap_or_default();
126            return Err(Error::service_status(
127                status.as_u16(),
128                format!("Mistral API error {status}: {text}"),
129                None,
130            ));
131        }
132        let value: Value = resp
133            .json()
134            .await
135            .map_err(|e| Error::service(format!("invalid response json: {e}")))?;
136        agent_framework_openai::embeddings::parse_embeddings_response(&value)
137    }
138
139    fn model(&self) -> Option<&str> {
140        Some(&self.inner.model)
141    }
142}
agent_framework_mistral/embeddings.rs

agent_framework_mistral/
embeddings.rs
