1use async_trait::async_trait;
27use serde::Deserialize;
28use serde_json::json;
29
30use crate::retry::{post_json_with_retry, DEFAULT_RETRY};
31use crate::vision::{ImageInput, VisionEmbeddings};
32use crate::{l2_normalize, EmbeddingError};
33
34pub const QWEN_VISION_BASE_URL: &str =
36 "https://dashscope.aliyuncs.com/api/v1/services/embeddings/multimodal-embedding";
37
38pub const QWEN_VISION_EMBED_MODEL: &str = "multimodal-embedding-v1";
40
41pub const QWEN_VISION_DIMENSION: usize = 1024;
43
44const MAX_CONTENTS_PER_REQUEST: usize = 10;
46
47#[derive(Debug, Clone)]
49pub struct QwenVisionEmbeddingsConfig {
50 pub api_key: String,
52 pub base_url: String,
54 pub model: String,
56}
57
58impl Default for QwenVisionEmbeddingsConfig {
59 fn default() -> Self {
60 Self {
61 api_key: std::env::var("QWEN_API_KEY").unwrap_or_default(),
62 base_url: QWEN_VISION_BASE_URL.to_string(),
63 model: QWEN_VISION_EMBED_MODEL.to_string(),
64 }
65 }
66}
67
68impl QwenVisionEmbeddingsConfig {
69 pub fn new(api_key: impl Into<String>) -> Self {
71 Self {
72 api_key: api_key.into(),
73 ..Default::default()
74 }
75 }
76
77 pub fn from_env_result() -> Result<Self, EmbeddingError> {
82 let api_key = std::env::var("QWEN_API_KEY").map_err(|_| {
83 EmbeddingError::Config("QWEN_API_KEY environment variable not set".to_string())
84 })?;
85 let base_url = std::env::var("QWEN_VISION_BASE_URL")
86 .unwrap_or_else(|_| QWEN_VISION_BASE_URL.to_string());
87 let model = std::env::var("QWEN_VISION_EMBED_MODEL")
88 .unwrap_or_else(|_| QWEN_VISION_EMBED_MODEL.to_string());
89 Ok(Self {
90 api_key,
91 base_url,
92 model,
93 })
94 }
95
96 pub fn with_model(mut self, model: impl Into<String>) -> Self {
98 self.model = model.into();
99 self
100 }
101}
102
103pub struct QwenVisionEmbeddings {
105 config: QwenVisionEmbeddingsConfig,
106 client: reqwest::Client,
107 dimension: usize,
108}
109
110impl std::fmt::Debug for QwenVisionEmbeddings {
111 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112 f.debug_struct("QwenVisionEmbeddings")
113 .field("model", &self.config.model)
114 .finish()
115 }
116}
117
118#[derive(Debug, Deserialize)]
119struct DashScopeMmResponse {
120 output: Option<DashScopeMmOutput>,
121 #[serde(default)]
122 code: Option<String>,
123 #[serde(default)]
124 message: Option<String>,
125}
126
127#[derive(Debug, Deserialize)]
128struct DashScopeMmOutput {
129 embeddings: Vec<DashScopeMmEmbedding>,
130}
131
132#[derive(Debug, Deserialize)]
133struct DashScopeMmEmbedding {
134 embedding: Vec<f32>,
135 #[serde(default)]
136 index: usize,
137}
138
139impl QwenVisionEmbeddings {
140 pub fn new(config: QwenVisionEmbeddingsConfig) -> Result<Self, EmbeddingError> {
143 if config.api_key.trim().is_empty() {
144 return Err(EmbeddingError::Config("QWEN_API_KEY is empty".to_string()));
145 }
146 let dimension = Self::dimension_for(&config.model)?;
147 Ok(Self {
148 config,
149 client: reqwest::Client::new(),
150 dimension,
151 })
152 }
153
154 pub fn from_env_result() -> Result<Self, EmbeddingError> {
156 Self::new(QwenVisionEmbeddingsConfig::from_env_result()?)
157 }
158
159 fn dimension_for(model: &str) -> Result<usize, EmbeddingError> {
160 match model {
161 QWEN_VISION_EMBED_MODEL => Ok(QWEN_VISION_DIMENSION),
162 other => Err(EmbeddingError::Config(format!(
163 "unknown embedding dimension for Qwen multimodal model '{other}' \
164 (supported: '{QWEN_VISION_EMBED_MODEL}')"
165 ))),
166 }
167 }
168
169 fn endpoint(&self) -> String {
170 format!(
171 "{}/{}",
172 self.config.base_url.trim_end_matches('/'),
173 self.config.model
174 )
175 }
176
177 async fn post_contents(
180 &self,
181 contents: &[serde_json::Value],
182 ) -> Result<Vec<Vec<f32>>, EmbeddingError> {
183 let body = json!({
184 "model": self.config.model,
185 "input": {"contents": contents},
186 });
187
188 let response = post_json_with_retry(
189 &self.client,
190 &self.endpoint(),
191 &self.config.api_key,
192 &body,
193 &DEFAULT_RETRY,
194 )
195 .await
196 .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
197
198 let status = response.status();
199 if !status.is_success() {
200 let error_text = response.text().await.map_err(|e| {
201 EmbeddingError::HttpError(format!("failed to read error response body: {e}"))
202 })?;
203 return Err(EmbeddingError::ApiError(format!(
204 "HTTP {}: {}",
205 status, error_text
206 )));
207 }
208
209 let parsed: DashScopeMmResponse = response
210 .json()
211 .await
212 .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
213
214 let output = match parsed.output {
215 Some(output) if !output.embeddings.is_empty() => output,
216 _ => {
217 return Err(EmbeddingError::ApiError(
218 parsed
219 .code
220 .zip(parsed.message)
221 .map(|(code, message)| format!("{code}: {message}"))
222 .unwrap_or_else(|| "No embeddings in DashScope response".to_string()),
223 ));
224 }
225 };
226
227 let mut indexed = output.embeddings;
229 indexed.sort_by_key(|e| e.index);
230 Ok(indexed.into_iter().map(|e| e.embedding).collect())
231 }
232}
233
234#[async_trait]
235impl VisionEmbeddings for QwenVisionEmbeddings {
236 async fn embed_image(&self, image: &ImageInput) -> Result<Vec<f32>, EmbeddingError> {
237 let mut batch = self.embed_images(std::slice::from_ref(image)).await?;
238 batch
239 .pop()
240 .ok_or_else(|| EmbeddingError::ApiError("No embedding in response".to_string()))
241 }
242
243 async fn embed_images(&self, images: &[ImageInput]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
244 if images.is_empty() {
245 return Ok(Vec::new());
246 }
247
248 let entries: Result<Vec<serde_json::Value>, EmbeddingError> = images
250 .iter()
251 .map(|image| {
252 let reference = image.reference()?;
253 Ok(json!([{"image": reference}]))
254 })
255 .collect();
256 let entries = entries?;
257
258 let mut all = Vec::with_capacity(images.len());
259 for chunk in entries.chunks(MAX_CONTENTS_PER_REQUEST) {
260 all.extend(self.post_contents(chunk).await?);
261 }
262
263 if all.len() != images.len() {
264 return Err(EmbeddingError::BatchMismatch {
265 expected: images.len(),
266 actual: all.len(),
267 });
268 }
269 for vector in all.iter_mut() {
270 l2_normalize(vector);
271 }
272 Ok(all)
273 }
274
275 async fn embed_text(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
276 if text.trim().is_empty() {
277 return Err(EmbeddingError::EmptyInput);
278 }
279 let contents = [json!([{"text": text}])];
280 let mut batch = self.post_contents(&contents).await?;
281 if batch.len() != 1 {
282 return Err(EmbeddingError::BatchMismatch {
283 expected: 1,
284 actual: batch.len(),
285 });
286 }
287 let mut embedding = batch.remove(0);
288 l2_normalize(&mut embedding);
289 Ok(embedding)
290 }
291
292 fn dimension(&self) -> usize {
293 self.dimension
294 }
295
296 fn model_name(&self) -> &str {
297 &self.config.model
298 }
299}
300
301#[cfg(test)]
302mod tests {
303 use super::*;
304 use crate::test_support::{
305 spawn_json_handler_stub, spawn_json_recording_stub, spawn_status_stub,
306 };
307 use std::sync::atomic::Ordering;
308 use std::sync::Arc;
309
310 fn config_with(base_url: String) -> QwenVisionEmbeddingsConfig {
311 QwenVisionEmbeddingsConfig {
312 api_key: "test-key".into(),
313 base_url,
314 model: QWEN_VISION_EMBED_MODEL.into(),
315 }
316 }
317
318 fn two_vector_response() -> serde_json::Value {
319 json!({
320 "output": {
321 "embeddings": [
322 {"embedding": [0.8, -0.6], "index": 1, "type": 0},
324 {"embedding": [0.6, 0.8], "index": 0, "type": 1}
325 ]
326 },
327 "request_id": "abc"
328 })
329 }
330
331 #[test]
332 fn construction_validates_key_and_model() {
333 let mut bad_key = config_with("http://127.0.0.1:1".into());
334 bad_key.api_key = String::new();
335 assert!(matches!(
336 QwenVisionEmbeddings::new(bad_key),
337 Err(EmbeddingError::Config(_))
338 ));
339
340 let mut bad_model = config_with("http://127.0.0.1:1".into());
341 bad_model.model = "some-other-model".into();
342 assert!(matches!(
343 QwenVisionEmbeddings::new(bad_model),
344 Err(EmbeddingError::Config(_))
345 ));
346
347 let ok = QwenVisionEmbeddings::new(config_with("http://127.0.0.1:1".into())).unwrap();
348 assert_eq!(ok.dimension(), QWEN_VISION_DIMENSION);
349 assert_eq!(ok.model_name(), QWEN_VISION_EMBED_MODEL);
350 }
351
352 #[tokio::test]
354 async fn image_request_body_matches_native_snapshot() {
355 let (base_url, bodies) = spawn_json_recording_stub(two_vector_response()).await;
356 let embeddings = QwenVisionEmbeddings::new(config_with(base_url)).unwrap();
357
358 let images = vec![
359 ImageInput::from_url("https://example.com/cat.png"),
360 ImageInput::from_base64("aW1n", "image/png"),
361 ];
362 let vectors = embeddings.embed_images(&images).await.unwrap();
363 assert_eq!(vectors.len(), 2);
364 assert_eq!(vectors[0], vec![0.6, 0.8]);
366 assert_eq!(vectors[1], vec![0.8, -0.6]);
367
368 let recorded = bodies.lock().unwrap();
369 assert_eq!(recorded.len(), 1);
370 assert_eq!(recorded[0]["model"], QWEN_VISION_EMBED_MODEL);
371 assert_eq!(
372 recorded[0]["input"]["contents"],
373 json!([
374 [{"image": "https://example.com/cat.png"}],
375 [{"image": "data:image/png;base64,aW1n"}],
376 ])
377 );
378 }
379
380 #[tokio::test]
381 async fn text_request_uses_text_content() {
382 let response = json!({"output": {"embeddings": [{"embedding": [0.6, 0.8], "index": 0}]}});
383 let (base_url, bodies) = spawn_json_recording_stub(response).await;
384 let embeddings = QwenVisionEmbeddings::new(config_with(base_url)).unwrap();
385
386 let v = embeddings.embed_text("一只猫").await.unwrap();
387 assert_eq!(v.len(), 2);
388 let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
389 assert!((norm - 1.0).abs() < 1e-5);
390
391 let recorded = bodies.lock().unwrap();
392 assert_eq!(
393 recorded[0]["input"]["contents"],
394 json!([[{"text": "一只猫"}]])
395 );
396 }
397
398 #[tokio::test]
399 async fn batches_larger_than_ten_contents_are_chunked_in_order() {
400 let handler = Arc::new(|request: serde_json::Value| {
404 let n = request["input"]["contents"]
405 .as_array()
406 .map(|a| a.len())
407 .unwrap_or(0);
408 json!({
409 "output": {
410 "embeddings": (0..n)
411 .map(|i| json!({"embedding": [i as f32, 0.0], "index": i}))
412 .collect::<Vec<_>>()
413 }
414 })
415 });
416 let (base_url, bodies) = spawn_json_handler_stub(handler).await;
417 let embeddings = QwenVisionEmbeddings::new(config_with(base_url)).unwrap();
418
419 let images: Vec<ImageInput> = (0..21)
420 .map(|i| ImageInput::from_url(format!("https://example.com/img{i}.png")))
421 .collect();
422 let vectors = embeddings.embed_images(&images).await.unwrap();
423 assert_eq!(vectors.len(), 21, "10 + 10 + 1 across three requests");
424 assert_eq!(vectors[0], vectors[10]);
427 assert_eq!(vectors[10], vectors[20]);
428 assert_eq!(vectors[9], vectors[19]);
429 assert_ne!(vectors[9], vectors[10], "chunk boundary must not reorder");
430
431 let recorded = bodies.lock().unwrap();
432 assert_eq!(recorded.len(), 3);
433 assert_eq!(
434 recorded[0]["input"]["contents"].as_array().unwrap().len(),
435 10
436 );
437 assert_eq!(
438 recorded[1]["input"]["contents"].as_array().unwrap().len(),
439 10
440 );
441 assert_eq!(
442 recorded[2]["input"]["contents"].as_array().unwrap().len(),
443 1
444 );
445 }
446
447 #[tokio::test]
448 async fn retries_on_429_and_enforces_empty_contract() {
449 let success_body = r#"{"output":{"embeddings":[{"embedding":[0.6,0.8],"index":0}]}}"#;
450 let (base_url, requests) = spawn_status_stub(429, 2, 200, success_body).await;
451 let embeddings = QwenVisionEmbeddings::new(config_with(base_url)).unwrap();
452
453 let v = embeddings
454 .embed_text("hello")
455 .await
456 .expect("should retry after two 429s");
457 assert_eq!(v.len(), 2);
458 assert_eq!(requests.load(Ordering::SeqCst), 3, "1 initial + 2 retries");
459
460 assert!(matches!(
461 embeddings.embed_text(" ").await,
462 Err(EmbeddingError::EmptyInput)
463 ));
464 assert_eq!(
465 embeddings.embed_images(&[]).await.unwrap(),
466 Vec::<Vec<f32>>::new()
467 );
468 }
469
470 #[tokio::test]
471 async fn provider_error_payload_is_surfaced() {
472 let response = json!({"code": "InvalidApiKey", "message": "key invalid"});
473 let (base_url, _bodies) = spawn_json_recording_stub(response).await;
474 let embeddings = QwenVisionEmbeddings::new(config_with(base_url)).unwrap();
475
476 let err = embeddings.embed_text("hello").await.unwrap_err();
477 match err {
478 EmbeddingError::ApiError(message) => assert!(message.contains("InvalidApiKey")),
479 other => panic!("expected ApiError, got {other:?}"),
480 }
481 }
482}