rig_core/providers/openrouter/
embedding.rs

1use super::{Client, Usage, client::ApiResponse};
2use crate::embeddings::EmbeddingError;
3use crate::http_client::HttpClientExt;
4use crate::wasm_compat::WasmCompatSend;
5use crate::{embeddings, http_client};
6use serde::{Deserialize, Serialize};
7use serde_json::json;
8
9#[derive(Debug, Deserialize)]
10pub struct EmbeddingResponse {
11    pub object: String,
12    pub data: Vec<EmbeddingData>,
13    pub model: String,
14    pub usage: Option<Usage>,
15    pub id: Option<String>,
16}
17
18#[derive(Debug, Deserialize, Clone, Serialize)]
19#[serde(rename_all = "snake_case")]
20pub enum EncodingFormat {
21    Float,
22    Base64,
23}
24
25#[derive(Debug, Deserialize)]
26pub struct EmbeddingData {
27    pub object: String,
28    pub embedding: Vec<serde_json::Number>,
29    pub index: usize,
30}
31
32#[derive(Clone)]
33pub struct EmbeddingModel<T = reqwest::Client> {
34    client: Client<T>,
35    pub model: String,
36    pub encoding_format: Option<EncodingFormat>,
37    pub user: Option<String>,
38    ndims: usize,
39}
40
41impl<T> embeddings::EmbeddingModel for EmbeddingModel<T>
42where
43    T: HttpClientExt + Clone + std::fmt::Debug + Default + WasmCompatSend + 'static,
44{
45    const MAX_DOCUMENTS: usize = 1024;
46
47    type Client = Client<T>;
48
49    fn make(client: &Self::Client, model: impl Into<String>, ndims: Option<usize>) -> Self {
50        let model = model.into();
51        let dims = ndims.unwrap_or_default();
52
53        Self::new(client.clone(), model, dims)
54    }
55
56    fn ndims(&self) -> usize {
57        self.ndims
58    }
59
60    async fn embed_texts(
61        &self,
62        documents: impl IntoIterator<Item = String>,
63    ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
64        let documents = documents.into_iter().collect::<Vec<_>>();
65
66        let mut body = json!({
67            "model": self.model,
68            "input": documents,
69        });
70
71        let body_object = body.as_object_mut().ok_or_else(|| {
72            EmbeddingError::ResponseError("embedding request body must be a JSON object".into())
73        })?;
74
75        if self.ndims > 0 {
76            body_object.insert("dimensions".to_owned(), json!(self.ndims));
77        }
78
79        if let Some(encoding_format) = &self.encoding_format {
80            body_object.insert("encoding_format".to_owned(), json!(encoding_format));
81        }
82
83        if let Some(user) = &self.user {
84            body_object.insert("user".to_owned(), json!(user));
85        }
86
87        let body = serde_json::to_vec(&body)?;
88
89        let req = self
90            .client
91            .post("/embeddings")?
92            .body(body)
93            .map_err(|e| EmbeddingError::HttpError(e.into()))?;
94
95        let response = self.client.send(req).await?;
96
97        let status = response.status();
98        if status.is_success() {
99            let response_body: Vec<u8> = response.into_body().await?;
100            let parsed: ApiResponse<EmbeddingResponse> = serde_json::from_slice(&response_body)?;
101
102            match parsed {
103                ApiResponse::Ok(response) => {
104                    tracing::info!(target: "rig",
105                        "OpenRouter embedding token usage: {:?}",
106                        response.usage
107                    );
108
109                    if response.data.len() != documents.len() {
110                        return Err(EmbeddingError::ResponseError(
111                            "Response data length does not match input length".into(),
112                        ));
113                    }
114
115                    Ok(response
116                        .data
117                        .into_iter()
118                        .zip(documents.into_iter())
119                        .map(|(embedding, document)| embeddings::Embedding {
120                            document,
121                            vec: embedding
122                                .embedding
123                                .into_iter()
124                                .filter_map(|n| n.as_f64())
125                                .collect(),
126                        })
127                        .collect())
128                }
129                ApiResponse::Err(err) => {
130                    tracing::warn!(message = %err.message, "provider returned an error response");
131                    Err(EmbeddingError::from_http_response(
132                        status,
133                        String::from_utf8_lossy(&response_body).into_owned(),
134                    ))
135                }
136            }
137        } else {
138            let text = http_client::text(response).await?;
139            Err(EmbeddingError::from_http_response(status, text))
140        }
141    }
142}
143
144impl<T> EmbeddingModel<T> {
145    pub fn new(client: Client<T>, model: impl Into<String>, ndims: usize) -> Self {
146        Self {
147            client,
148            model: model.into(),
149            encoding_format: None,
150            ndims,
151            user: None,
152        }
153    }
154
155    pub fn with_model(client: Client<T>, model: &str, ndims: usize) -> Self {
156        Self {
157            client,
158            model: model.into(),
159            encoding_format: None,
160            ndims,
161            user: None,
162        }
163    }
164
165    pub fn with_encoding_format(
166        client: Client<T>,
167        model: &str,
168        ndims: usize,
169        encoding_format: EncodingFormat,
170    ) -> Self {
171        Self {
172            client,
173            model: model.into(),
174            encoding_format: Some(encoding_format),
175            ndims,
176            user: None,
177        }
178    }
179
180    pub fn encoding_format(mut self, encoding_format: EncodingFormat) -> Self {
181        self.encoding_format = Some(encoding_format);
182        self
183    }
184
185    pub fn user(mut self, user: impl Into<String>) -> Self {
186        self.user = Some(user.into());
187        self
188    }
189}
rig_core/providers/openrouter/embedding.rs

rig_core/providers/openrouter/
embedding.rs
