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

langchain_rust/llm/deepseek/
client.rs
