lc_embeddings/
deepseek.rs

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

lc_embeddings/
deepseek.rs
