rig/providers/openrouter/
embedding.rs

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

rig/providers/openrouter/
embedding.rs
