rig_core/providers/mistral/
embedding.rs

1use serde::Deserialize;
2use serde_json::json;
3
4use crate::{
5    embeddings::{self, EmbeddingError},
6    http_client::{self, HttpClientExt},
7};
8
9use super::client::{ApiResponse, Client, Usage};
10
11// ================================================================
12// Mistral Embedding API
13// ================================================================
14pub const MISTRAL_EMBED: &str = "mistral-embed";
15
16pub const MAX_DOCUMENTS: usize = 1024;
17
18#[derive(Clone)]
19pub struct EmbeddingModel<T = reqwest::Client> {
20    client: Client<T>,
21    pub model: String,
22    ndims: usize,
23}
24
25impl<T> EmbeddingModel<T> {
26    pub fn new(client: Client<T>, model: impl Into<String>, ndims: usize) -> Self {
27        Self {
28            client,
29            model: model.into(),
30            ndims,
31        }
32    }
33
34    pub fn with_model(client: Client<T>, model: &str, ndims: usize) -> Self {
35        Self {
36            client,
37            model: model.to_string(),
38            ndims,
39        }
40    }
41}
42
43impl<T> embeddings::EmbeddingModel for EmbeddingModel<T>
44where
45    T: HttpClientExt + Clone + 'static,
46{
47    type Client = Client<T>;
48
49    const MAX_DOCUMENTS: usize = MAX_DOCUMENTS;
50
51    fn make(client: &Self::Client, model: impl Into<String>, dims: Option<usize>) -> Self {
52        Self::new(client.clone(), model, dims.unwrap_or_default())
53    }
54
55    fn ndims(&self) -> usize {
56        self.ndims
57    }
58
59    async fn embed_texts(
60        &self,
61        documents: impl IntoIterator<Item = String>,
62    ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
63        let documents = documents.into_iter().collect::<Vec<_>>();
64
65        let body = serde_json::to_vec(&json!({
66            "model": self.model,
67            "input": documents
68        }))?;
69
70        let req = self
71            .client
72            .post("v1/embeddings")?
73            .header("Content-Type", "application/json")
74            .body(body)
75            .map_err(|e| EmbeddingError::HttpError(e.into()))?;
76
77        let response = self.client.send(req).await?;
78
79        let status = response.status();
80        if status.is_success() {
81            let response_body: Vec<u8> = response.into_body().await?;
82            let parsed: ApiResponse<EmbeddingResponse> = serde_json::from_slice(&response_body)?;
83
84            match parsed {
85                ApiResponse::Ok(response) => {
86                    tracing::debug!(target: "rig",
87                        "Mistral embedding token usage: {}",
88                        response.usage
89                    );
90
91                    if response.data.len() != documents.len() {
92                        return Err(EmbeddingError::ResponseError(
93                            "Response data length does not match input length".into(),
94                        ));
95                    }
96
97                    Ok(response
98                        .data
99                        .into_iter()
100                        .zip(documents.into_iter())
101                        .map(|(embedding, document)| embeddings::Embedding {
102                            document,
103                            vec: embedding
104                                .embedding
105                                .into_iter()
106                                .filter_map(|n| n.as_f64())
107                                .collect(),
108                        })
109                        .collect())
110                }
111                ApiResponse::Err(err) => {
112                    tracing::warn!(message = %err.message, "provider returned an error response");
113                    Err(EmbeddingError::from_http_response(
114                        status,
115                        String::from_utf8_lossy(&response_body),
116                    ))
117                }
118            }
119        } else {
120            let text = http_client::text(response).await?;
121            Err(EmbeddingError::from_http_response(status, text))
122        }
123    }
124}
125
126#[derive(Debug, Deserialize)]
127pub struct EmbeddingResponse {
128    pub id: String,
129    pub object: String,
130    pub model: String,
131    pub usage: Usage,
132    pub data: Vec<EmbeddingData>,
133}
134
135#[derive(Debug, Deserialize)]
136pub struct EmbeddingData {
137    pub object: String,
138    pub embedding: Vec<serde_json::Number>,
139    pub index: usize,
140}
rig_core/providers/mistral/embedding.rs

rig_core/providers/mistral/
embedding.rs
