rig_core/providers/doubleword/
embedding.rs

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

rig_core/providers/doubleword/
embedding.rs
