lc_embeddings/
qwen.rs

1// lc-embeddings/src/qwen.rs
2//! Qwen (Alibaba Cloud) embeddings implementation.
3//!
4//! Qwen (DashScope compatible mode) speaks the OpenAI-compatible `/embeddings` protocol and
5//! reuses the [`crate::openai_compat`] shared base class (P1-5); this file only configures
6//! the spec (URL / model / dimension / batch size).
7
8use crate::openai_compat::{CompatConfigAccess, CompatSpec, OpenAICompatEmbeddings};
9use crate::EmbeddingError;
10
11/// Default base URL for the Qwen (DashScope) API.
12pub const QWEN_BASE_URL: &str = "https://dashscope.aliyuncs.com/compatible-mode/v1";
13
14/// Default embedding model for Qwen.
15pub const QWEN_EMBED_MODEL: &str = "text-embedding-v1";
16
17/// Configuration for Qwen embeddings API.
18#[derive(Debug, Clone)]
19pub struct QwenEmbeddingsConfig {
20    /// Qwen API key.
21    pub api_key: String,
22    /// Base URL for the Qwen (DashScope) embeddings API.
23    pub base_url: String,
24    /// Embedding model name.
25    pub model: String,
26}
27
28impl Default for QwenEmbeddingsConfig {
29    fn default() -> Self {
30        Self {
31            api_key: std::env::var("QWEN_API_KEY").unwrap_or_default(),
32            base_url: QWEN_BASE_URL.to_string(),
33            model: QWEN_EMBED_MODEL.to_string(),
34        }
35    }
36}
37
38impl QwenEmbeddingsConfig {
39    /// Creates a new QwenEmbeddingsConfig with the given API key.
40    pub fn new(api_key: impl Into<String>) -> Self {
41        Self {
42            api_key: api_key.into(),
43            ..Default::default()
44        }
45    }
46
47    /// Creates a QwenEmbeddingsConfig from environment variables, returning a Result.
48    ///
49    /// Environment variables:
50    /// - `QWEN_API_KEY`: API key (required)
51    /// - `QWEN_BASE_URL`: API endpoint (optional)
52    /// - `QWEN_EMBED_MODEL`: Model name (optional)
53    pub fn from_env_result() -> Result<Self, EmbeddingError> {
54        let api_key = std::env::var("QWEN_API_KEY").map_err(|_| {
55            EmbeddingError::Config("QWEN_API_KEY environment variable not set".to_string())
56        })?;
57        let base_url = std::env::var("QWEN_BASE_URL").unwrap_or_else(|_| QWEN_BASE_URL.to_string());
58        let model =
59            std::env::var("QWEN_EMBED_MODEL").unwrap_or_else(|_| QWEN_EMBED_MODEL.to_string());
60        Ok(Self {
61            api_key,
62            base_url,
63            model,
64        })
65    }
66
67    /// Sets the embedding model.
68    pub fn with_model(mut self, model: impl Into<String>) -> Self {
69        self.model = model.into();
70        self
71    }
72}
73
74impl CompatConfigAccess for QwenEmbeddingsConfig {
75    fn api_key(&self) -> &str {
76        &self.api_key
77    }
78    fn base_url(&self) -> &str {
79        &self.base_url
80    }
81    fn model(&self) -> &str {
82        &self.model
83    }
84}
85
86impl CompatSpec for QwenEmbeddingsConfig {
87    fn api_key_env() -> &'static str {
88        "QWEN_API_KEY"
89    }
90    fn batch_size() -> usize {
91        64
92    }
93    fn dimension_for(model: &str) -> Result<usize, EmbeddingError> {
94        if model == QWEN_EMBED_MODEL {
95            Ok(1536)
96        } else {
97            Err(EmbeddingError::Config(format!(
98                "unknown embedding dimension for Qwen model '{model}' (supported: '{QWEN_EMBED_MODEL}')"
99            )))
100        }
101    }
102    fn from_env_result() -> Result<Self, EmbeddingError> {
103        Self::from_env_result()
104    }
105}
106
107/// Qwen embeddings client for generating vector embeddings.
108///
109/// Reuses the OpenAI-compatible-protocol shared base class (P1-5): fails fast at construction
110/// validating a non-empty API key and known model dimension (P1-2/P1-3), batch alignment errors
111/// explicitly (P0-1), and error bodies are not swallowed (P1-4).
112pub type QwenEmbeddings = OpenAICompatEmbeddings<QwenEmbeddingsConfig>;
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use crate::test_support::spawn_embeddings_stub;
118    use crate::Embeddings;
119    use std::env;
120    use std::sync::Arc;
121
122    /// P0-1: provider returns fewer entries → explicit `EmptyVectorInBatch`, not a silent empty vector.
123    #[tokio::test]
124    async fn test_embed_documents_truncated_errors() {
125        let base_url = spawn_embeddings_stub(Arc::new(|n| n.saturating_sub(1))).await;
126        let config = QwenEmbeddingsConfig {
127            api_key: "test-key".into(),
128            base_url,
129            model: QWEN_EMBED_MODEL.into(),
130        };
131        let embeddings = QwenEmbeddings::new(config).unwrap();
132
133        let result = embeddings.embed_documents(&["a", "b"]).await;
134        assert!(
135            matches!(result, Err(EmbeddingError::EmptyVectorInBatch)),
136            "truncated response should report EmptyVectorInBatch, got: {:?}",
137            result
138        );
139    }
140
141    /// P1-3: an empty API key → `Config` error at construction (fail fast), not a delayed 401.
142    #[test]
143    fn test_new_rejects_empty_api_key() {
144        let config = QwenEmbeddingsConfig {
145            api_key: String::new(),
146            base_url: QWEN_BASE_URL.into(),
147            model: QWEN_EMBED_MODEL.into(),
148        };
149        let err = QwenEmbeddings::new(config).unwrap_err();
150        assert!(matches!(err, EmbeddingError::Config(_)));
151    }
152
153    /// P1-2: unknown model → construction-time error, never lying with a default 1536.
154    #[test]
155    fn test_new_rejects_unknown_model() {
156        let config = QwenEmbeddingsConfig {
157            api_key: "test-key".into(),
158            base_url: QWEN_BASE_URL.into(),
159            model: "some-unknown-model".into(),
160        };
161        let err = QwenEmbeddings::new(config).unwrap_err();
162        assert!(matches!(err, EmbeddingError::Config(_)));
163    }
164
165    fn save_and_set(key: &str, value: &str) -> Option<String> {
166        let old = env::var(key).ok();
167        env::set_var(key, value);
168        old
169    }
170
171    fn restore(key: &str, old: Option<String>) {
172        match old {
173            Some(v) => env::set_var(key, v),
174            None => env::remove_var(key),
175        }
176    }
177
178    #[test]
179    fn test_from_env_result_ok_when_key_set() {
180        let _lock = crate::ENV_TEST_LOCK
181            .lock()
182            .unwrap_or_else(|e| e.into_inner());
183        let old = save_and_set("QWEN_API_KEY", "test-key-123");
184        let result = QwenEmbeddingsConfig::from_env_result();
185        assert!(result.is_ok());
186        assert_eq!(result.unwrap().api_key, "test-key-123");
187        restore("QWEN_API_KEY", old);
188    }
189
190    #[test]
191    fn test_from_env_result_err_when_key_missing() {
192        let _lock = crate::ENV_TEST_LOCK
193            .lock()
194            .unwrap_or_else(|e| e.into_inner());
195        let old = env::var("QWEN_API_KEY").ok();
196        env::remove_var("QWEN_API_KEY");
197        let result = QwenEmbeddingsConfig::from_env_result();
198        assert!(result.is_err());
199        assert!(result.unwrap_err().to_string().contains("QWEN_API_KEY"));
200        restore("QWEN_API_KEY", old);
201    }
202
203    #[test]
204    fn test_from_env_result_uses_optional_vars() {
205        let _lock = crate::ENV_TEST_LOCK
206            .lock()
207            .unwrap_or_else(|e| e.into_inner());
208        let old_key = save_and_set("QWEN_API_KEY", "key");
209        let old_url = save_and_set("QWEN_BASE_URL", "https://custom.api.com");
210        let old_model = save_and_set("QWEN_EMBED_MODEL", "custom-model");
211        let config = QwenEmbeddingsConfig::from_env_result().unwrap();
212        assert_eq!(config.base_url, "https://custom.api.com");
213        assert_eq!(config.model, "custom-model");
214        restore("QWEN_API_KEY", old_key);
215        restore("QWEN_BASE_URL", old_url);
216        restore("QWEN_EMBED_MODEL", old_model);
217    }
218
219    #[test]
220    fn test_from_env_result_uses_defaults_for_optional_vars() {
221        let _lock = crate::ENV_TEST_LOCK
222            .lock()
223            .unwrap_or_else(|e| e.into_inner());
224        let old_key = save_and_set("QWEN_API_KEY", "key");
225        let old_url = env::var("QWEN_BASE_URL").ok();
226        env::remove_var("QWEN_BASE_URL");
227        let old_model = env::var("QWEN_EMBED_MODEL").ok();
228        env::remove_var("QWEN_EMBED_MODEL");
229        let config = QwenEmbeddingsConfig::from_env_result().unwrap();
230        assert_eq!(config.base_url, QWEN_BASE_URL.to_string());
231        assert_eq!(config.model, QWEN_EMBED_MODEL);
232        restore("QWEN_API_KEY", old_key);
233        restore("QWEN_BASE_URL", old_url);
234        restore("QWEN_EMBED_MODEL", old_model);
235    }
236
237    #[test]
238    fn test_embeddings_from_env_result_ok() {
239        let _lock = crate::ENV_TEST_LOCK
240            .lock()
241            .unwrap_or_else(|e| e.into_inner());
242        let old = save_and_set("QWEN_API_KEY", "test-key");
243        assert!(QwenEmbeddings::from_env_result().is_ok());
244        restore("QWEN_API_KEY", old);
245    }
246
247    #[test]
248    fn test_embeddings_from_env_result_err_when_key_missing() {
249        let _lock = crate::ENV_TEST_LOCK
250            .lock()
251            .unwrap_or_else(|e| e.into_inner());
252        let old = env::var("QWEN_API_KEY").ok();
253        env::remove_var("QWEN_API_KEY");
254        assert!(QwenEmbeddings::from_env_result().is_err());
255        restore("QWEN_API_KEY", old);
256    }
257}
lc_embeddings/qwen.rs

lc_embeddings/
qwen.rs
