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<f64>,
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.embedding,
135                        })
136                        .collect())
137                }
138                ApiResponse::Err(err) => Err(EmbeddingError::ProviderError(err.message)),
139            }
140        } else {
141            let text = http_client::text(response).await?;
142            Err(EmbeddingError::ProviderError(text))
143        }
144    }
145}
146
147impl<T> EmbeddingModel<T> {
148    pub fn new(client: Client<T>, model: impl Into<String>, ndims: usize) -> Self {
149        Self {
150            client,
151            model: model.into(),
152            encoding_format: None,
153            ndims,
154            user: None,
155        }
156    }
157
158    pub fn with_model(client: Client<T>, model: &str, ndims: usize) -> Self {
159        Self {
160            client,
161            model: model.into(),
162            encoding_format: None,
163            ndims,
164            user: None,
165        }
166    }
167
168    pub fn with_encoding_format(
169        client: Client<T>,
170        model: &str,
171        ndims: usize,
172        encoding_format: EncodingFormat,
173    ) -> Self {
174        Self {
175            client,
176            model: model.into(),
177            encoding_format: Some(encoding_format),
178            ndims,
179            user: None,
180        }
181    }
182
183    pub fn encoding_format(mut self, encoding_format: EncodingFormat) -> Self {
184        self.encoding_format = Some(encoding_format);
185        self
186    }
187
188    pub fn user(mut self, user: impl Into<String>) -> Self {
189        self.user = Some(user.into());
190        self
191    }
192}
rig/providers/openrouter/embedding.rs

rig/providers/openrouter/
embedding.rs
