langchain_rust/llm/deepseek/
client.rs

1use crate::{
2    language_models::{llm::LLM, options::CallOptions, GenerateResult, LLMError, TokenUsage},
3    llm::DeepseekError,
4    schemas::{Message, StreamData},
5};
6use async_trait::async_trait;
7use futures::{Stream, StreamExt};
8use reqwest::Client;
9use std::pin::Pin;
10
11use super::models::{ApiResponse, DeepseekMessage, Payload, ResponseFormat};
12use crate::llm::sse::{sse_json_events, usage_from_chunk};
13use crate::observability::traced_generate;
14
15pub enum DeepseekModel {
16    DeepseekChat,
17    DeepseekReasoner,
18}
19
20impl ToString for DeepseekModel {
21    fn to_string(&self) -> String {
22        match self {
23            DeepseekModel::DeepseekChat => "deepseek-chat".to_string(),
24            DeepseekModel::DeepseekReasoner => "deepseek-reasoner".to_string(),
25        }
26    }
27}
28
29#[derive(Clone)]
30pub struct Deepseek {
31    model: String,
32    options: CallOptions,
33    api_key: String,
34    base_url: String,
35    json_mode: bool,
36    include_reasoning: bool,
37}
38
39impl Default for Deepseek {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl Deepseek {
46    pub fn new() -> Self {
47        Self {
48            model: DeepseekModel::DeepseekChat.to_string(),
49            options: CallOptions::default(),
50            api_key: std::env::var("DEEPSEEK_API_KEY").unwrap_or_default(),
51            base_url: "https://api.deepseek.com".to_string(),
52            json_mode: false,
53            include_reasoning: false,
54        }
55    }
56
57    pub fn with_model<S: Into<String>>(mut self, model: S) -> Self {
58        self.model = model.into();
59        self
60    }
61
62    pub fn with_options(mut self, options: CallOptions) -> Self {
63        self.options = options;
64        self
65    }
66
67    pub fn with_api_key<S: Into<String>>(mut self, api_key: S) -> Self {
68        self.api_key = api_key.into();
69        self
70    }
71
72    pub fn with_base_url<S: Into<String>>(mut self, base_url: S) -> Self {
73        self.base_url = base_url.into();
74        self
75    }
76
77    pub fn with_json_mode(mut self, json_mode: bool) -> Self {
78        self.json_mode = json_mode;
79        self
80    }
81
82    pub fn with_include_reasoning(mut self, include_reasoning: bool) -> Self {
83        self.include_reasoning = include_reasoning;
84        self
85    }
86
87    async fn generate(&self, messages: &[Message]) -> Result<GenerateResult, LLMError> {
88        let client = Client::new();
89        let is_stream = self.options.streaming_func.is_some();
90
91        let payload = self.build_payload(messages, is_stream);
92        let res = client
93            .post(&format!("{}/v1/chat/completions", self.base_url))
94            .header("Authorization", format!("Bearer {}", self.api_key))
95            .header("Content-Type", "application/json")
96            .json(&payload)
97            .send()
98            .await?;
99
100        let status = res.status().as_u16();
101
102        let res = match status {
103            400 => Err(LLMError::DeepseekError(DeepseekError::InvalidFormatError(
104                "Invalid request format".to_string(),
105            ))),
106            401 => Err(LLMError::DeepseekError(DeepseekError::AuthenticationError(
107                "Invalid API Key".to_string(),
108            ))),
109            402 => Err(LLMError::DeepseekError(
110                DeepseekError::InsufficientBalanceError("Insufficient balance".to_string()),
111            )),
112            422 => Err(LLMError::DeepseekError(
113                DeepseekError::InvalidParametersError("Invalid parameters".to_string()),
114            )),
115            429 => Err(LLMError::DeepseekError(DeepseekError::RateLimitError(
116                "Rate limit reached".to_string(),
117            ))),
118            500 => Err(LLMError::DeepseekError(DeepseekError::ServerError(
119                "Server error".to_string(),
120            ))),
121            503 => Err(LLMError::DeepseekError(
122                DeepseekError::ServerOverloadedError("Server overloaded".to_string()),
123            )),
124            _ => Ok(res.json::<ApiResponse>().await?),
125        }?;
126
127        let choice = res.choices.first();
128
129        let mut generation = choice
130            .map(|c| c.message.content.clone())
131            .unwrap_or_default();
132
133        // If include_reasoning is enabled and the model is deepseek-reasoner,
134        // append the reasoning content to the generation if available
135        if self.include_reasoning && self.model == DeepseekModel::DeepseekReasoner.to_string() {
136            if let Some(reasoning) = choice.and_then(|c| c.message.reasoning_content.clone()) {
137                generation = format!("Reasoning:\n{}\n\nAnswer:\n{}", reasoning, generation);
138            }
139        }
140
141        let tokens = Some(TokenUsage {
142            prompt_tokens: res.usage.prompt_tokens,
143            completion_tokens: res.usage.completion_tokens,
144            total_tokens: res.usage.total_tokens,
145        });
146
147        Ok(GenerateResult {
148            tokens,
149            generation,
150            ..Default::default()
151        })
152    }
153
154    fn build_payload(&self, messages: &[Message], stream: bool) -> Payload {
155        let mut response_format = None;
156        if self.json_mode {
157            response_format = Some(ResponseFormat {
158                format_type: "json_object".to_string(),
159            });
160        }
161
162        let mut payload = Payload {
163            model: self.model.clone(),
164            messages: messages
165                .iter()
166                .map(DeepseekMessage::from_message)
167                .collect::<Vec<_>>(),
168            max_tokens: self.options.max_tokens,
169            stream: None,
170            temperature: self.options.temperature,
171            top_p: self.options.top_p,
172            frequency_penalty: None,
173            presence_penalty: None,
174            stop: self.options.stop_words.clone(),
175            response_format,
176        };
177
178        if stream {
179            payload.stream = Some(true);
180        }
181
182        // Apply frequency_penalty if it's in the options range
183        if let Some(fp) = self.options.frequency_penalty {
184            if fp >= -2.0 && fp <= 2.0 {
185                payload.frequency_penalty = Some(fp);
186            }
187        }
188
189        // Apply presence_penalty if it's in the options range
190        if let Some(pp) = self.options.presence_penalty {
191            if pp >= -2.0 && pp <= 2.0 {
192                payload.presence_penalty = Some(pp);
193            }
194        }
195
196        payload
197    }
198}
199
200impl Deepseek {
201    /// `generate` without the tracing span: honours `streaming_func`.
202    async fn generate_dispatch(&self, messages: &[Message]) -> Result<GenerateResult, LLMError> {
203        match &self.options.streaming_func {
204            Some(func) => {
205                let mut complete_response = String::new();
206                let mut stream = self.stream(messages).await?;
207                while let Some(data) = stream.next().await {
208                    match data {
209                        Ok(value) => {
210                            let mut func = func.lock().await;
211                            complete_response.push_str(&value.content);
212                            let _ = func(value.content).await;
213                        }
214                        Err(e) => return Err(e),
215                    }
216                }
217                let mut generate_result = GenerateResult::default();
218                generate_result.generation = complete_response;
219                Ok(generate_result)
220            }
221            None => self.generate(messages).await,
222        }
223    }
224}
225
226#[async_trait]
227impl LLM for Deepseek {
228    async fn generate(&self, messages: &[Message]) -> Result<GenerateResult, LLMError> {
229        traced_generate("deepseek", &self.model, self.generate_dispatch(messages)).await
230    }
231
232    async fn stream(
233        &self,
234        messages: &[Message],
235    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamData, LLMError>> + Send>>, LLMError> {
236        let client = Client::new();
237        let payload = self.build_payload(messages, true);
238        let request = client
239            .post(&format!("{}/v1/chat/completions", self.base_url))
240            .header("Authorization", format!("Bearer {}", self.api_key))
241            .header("Content-Type", "application/json")
242            .json(&payload)
243            .build()?;
244
245        let stream = client.execute(request).await?;
246        let stream = stream.bytes_stream();
247
248        let include_reasoning = self.include_reasoning;
249        let is_reasoner = self.model == DeepseekModel::DeepseekReasoner.to_string();
250
251        let processed_stream = sse_json_events(stream).filter_map(move |event| async move {
252            let chunk = match event {
253                Ok(chunk) => chunk,
254                Err(e) => return Some(Err(e)),
255            };
256            let delta = chunk.get("choices")?.as_array()?.first()?.get("delta")?;
257            let usage = usage_from_chunk(&chunk);
258            if include_reasoning && is_reasoner {
259                if let Some(reasoning) = delta
260                    .get("reasoning_content")
261                    .and_then(|c| c.as_str())
262                    .filter(|r| !r.is_empty())
263                {
264                    let content = format!("Reasoning: {}", reasoning);
265                    return Some(Ok(StreamData::new(chunk.clone(), usage, content)));
266                }
267            }
268            let content = delta
269                .get("content")
270                .and_then(|c| c.as_str())
271                .filter(|c| !c.is_empty())?
272                .to_string();
273            Some(Ok(StreamData::new(chunk, usage, content)))
274        });
275
276        Ok(Box::pin(processed_stream))
277    }
278
279    fn add_options(&mut self, options: CallOptions) {
280        self.options = options;
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287    use crate::schemas::{Message, MessageType};
288
289    #[tokio::test]
290    #[ignore]
291    async fn test_deepseek_generate() {
292        let messages = vec![Message {
293            content: "Hello".to_string(),
294            message_type: MessageType::HumanMessage,
295            id: Some("test_id".to_string()),
296            images: None,
297            tool_calls: None,
298        }];
299
300        let client = Deepseek::new();
301        let res = client.generate(&messages).await;
302        assert!(res.is_ok());
303    }
304
305    #[tokio::test]
306    #[ignore]
307    async fn test_deepseek_stream() {
308        let messages = vec![Message {
309            content: "Hello".to_string(),
310            message_type: MessageType::HumanMessage,
311            id: Some("test_id".to_string()),
312            images: None,
313            tool_calls: None,
314        }];
315
316        let client = Deepseek::new();
317        let res = client.stream(&messages).await;
318        assert!(res.is_ok());
319    }
320
321    #[tokio::test]
322    #[ignore]
323    async fn test_deepseek_reasoner() {
324        let messages = vec![Message {
325            content: "9.11 and 9.8, which is greater?".to_string(),
326            message_type: MessageType::HumanMessage,
327            id: Some("test_id".to_string()),
328            images: None,
329            tool_calls: None,
330        }];
331
332        // Create a client with the DeepseekReasoner model and enable reasoning content
333        let client = Deepseek::new()
334            .with_model(DeepseekModel::DeepseekReasoner.to_string())
335            .with_include_reasoning(true);
336
337        let res = client.generate(&messages).await;
338        assert!(res.is_ok());
339
340        // The response will contain both the reasoning and answer content
341        if let Ok(result) = res {
342            println!("Generation result: {}", result.generation);
343        }
344    }
345}
langchain_rust/llm/deepseek/client.rs

langchain_rust/llm/deepseek/
client.rs
