next_web_ai/ai/deep_seek/api/
deep_seek_api.rs

1use futures_core::stream::BoxStream;
2use futures_util::StreamExt;
3use next_web_core::error::BoxError;
4
5use crate::ai::deep_seek::chat_model::ChatModel;
6
7const ERROR: [u8; 8] = [123, 34, 101, 114, 114, 111, 114, 34];
8const DATA: [u8; 6] = [100, 97, 116, 97, 58, 32];
9const DONE: [u8; 12] = [100, 97, 116, 97, 58, 32, 91, 68, 79, 78, 69, 93];
10
11#[derive(Clone)]
12pub struct DeepSeekApi {
13    pub(crate) base_url: Box<str>,
14    pub(crate) api_key: Box<str>,
15
16    pub(crate) chat_model: ChatModel,
17
18    pub(crate) client: reqwest::Client,
19}
20
21impl DeepSeekApi {
22    pub fn new(api_key: impl Into<Box<str>>, chat_model: ChatModel) -> Self {
23        let client = reqwest::Client::builder().build().unwrap();
24        let api_key = api_key.into();
25        Self {
26            base_url: "https://api.deepseek.com/chat/completions".into(),
27            api_key,
28            chat_model,
29            client,
30        }
31    }
32
33    pub async fn send(&self, req: &ChatCompletionRequest) -> Result<ChatApiRespnose, BoxError> {
34        let resp = self
35            .client
36            .post(self.base_url.as_ref())
37            .header("Content-Type", "application/json")
38            .bearer_auth(self.api_key.as_ref())
39            .body(serde_json::to_string(req)?)
40            .send()
41            .await?;
42        if !req.stream {
43            return Ok(ChatApiRespnose::Data(resp.json().await?));
44        }
45
46        // SSE stream
47        let stream = resp.bytes_stream().then(|data| async move {
48            data.map(|data| {
49                // Check for error message
50                if data.starts_with(&ERROR) {
51                    return Err(String::from_utf8(data.to_vec())
52                        .unwrap_or("Unknown error".into())
53                        .into());
54                }
55
56                if data.starts_with(&DONE) {
57                    println!("\n\nEnd of stream\n\n")
58                }
59
60                data.split(|&s| s == b'\n')
61                    .filter(|line| line.starts_with(&DATA))
62                    .map(|line| {
63                        serde_json::from_slice::<ChatCompletion>(&line[6..]).map_err(Into::into)
64                    })
65                    .collect::<Result<Vec<ChatCompletion>, BoxError>>()
66                    .map_err(Into::into)
67            })
68            .unwrap_or_else(|err| Err(err.into()))
69        });
70
71        Ok(ChatApiRespnose::DataStream(Box::pin(stream)))
72    }
73}
74
75pub enum ChatApiRespnose {
76    Data(ChatCompletion),
77    DataStream(BoxStream<'static, Result<Vec<ChatCompletion>, BoxError>>),
78}
79
80impl Default for DeepSeekApi {
81    fn default() -> Self {
82        let api_key = std::env::var("DEEPSEEK_API_KEY").unwrap();
83        Self::new(api_key, ChatModel::Chat)
84    }
85}
86
87#[derive(serde::Serialize)]
88pub struct ChatCompletionRequest {
89    pub(crate) messages: Vec<ChatCompletionMessage>,
90    pub(crate) model: Box<str>,
91    pub(crate) stream: bool,
92    pub(crate) temperature: Option<f32>,
93}
94
95impl ChatCompletionRequest {
96    pub fn new(
97        messages: Vec<ChatCompletionMessage>,
98        model: impl Into<Box<str>>,
99        stream: bool,
100    ) -> Self {
101        Self {
102            messages,
103            model: model.into(),
104            stream,
105            temperature: None,
106        }
107    }
108}
109
110#[derive(Debug, serde::Serialize, serde::Deserialize)]
111pub struct ChatCompletionMessage {
112    pub(crate) role: Box<str>,
113    pub(crate) content: Box<str>,
114}
115
116impl ChatCompletionMessage {
117    pub fn new(role: impl Into<Box<str>>, content: impl Into<Box<str>>) -> Self {
118        Self {
119            role: role.into(),
120            content: content.into(),
121        }
122    }
123}
124///
125/// {
126///  "id": "9710a6c0-1b51-427b-b95d-b6734ec46270",
127///  "object": "chat.completion",
128///  "created": 1754571706,
129///  "model": "deepseek-chat",
130///  "choices": [
131///    {
132///      "index": 0,
133///      "message": {
134///        "role": "assistant",
135///        "content": "Hello! 😊 How can I assist you today?"
136///      },
137///      "logprobs": null,
138///      "finish_reason": "stop"
139///    }
140///  ],
141///  "usage": {
142///    "prompt_tokens": 11,
143///    "completion_tokens": 11,
144///    "total_tokens": 22,
145///    "prompt_tokens_details": {
146///      "cached_tokens": 0
147///    },
148///    "prompt_cache_hit_tokens": 0,
149///    "prompt_cache_miss_tokens": 11
150///  },
151///  "system_fingerprint": "fp_8802369eaa_prod0623_fp8_kvcache"
152/// }
153///
154///
155///
156#[derive(Debug, serde::Deserialize)]
157pub struct ChatCompletion {
158    pub id: Box<str>,
159    pub object: Box<str>,
160    pub created: u64,
161    pub model: Box<str>,
162    pub choices: Vec<Choice>,
163    pub usage: Option<Usage>,
164    pub system_fingerprint: Box<str>,
165}
166
167#[derive(Debug, serde::Deserialize)]
168pub struct Usage {
169    pub prompt_tokens: u32,
170    pub completion_tokens: u32,
171    pub total_tokens: u64,
172    pub prompt_tokens_details: PromptTokensDetails,
173    pub prompt_cache_hit_tokens: u32,
174    pub prompt_cache_miss_tokens: u32,
175}
176
177#[derive(Debug, serde::Deserialize)]
178pub struct PromptTokensDetails {
179    pub cached_tokens: u32,
180}
181
182#[derive(Debug, serde::Deserialize)]
183pub struct Choice {
184    pub index: u32,
185    pub delta: Option<DeltaContent>,
186    pub message: Option<ChatCompletionMessage>,
187    pub logprobs: Option<Box<str>>,
188    pub finish_reason: Option<Box<str>>,
189}
190
191#[derive(Debug, serde::Deserialize)]
192pub struct DeltaContent {
193    pub content: Box<str>,
194}
195
196#[derive(Clone)]
197pub struct DefaultUsage {
198    pub prompt_tokens: u32,
199    pub completion_tokens: u32,
200    pub total_tokens: u64,
201}
202
203impl crate::chat::meta_data::usage::Usage for DefaultUsage {
204    fn get_prompt_tokens(&self) -> u32 {
205        self.prompt_tokens
206    }
207
208    fn get_completion_tokens(&self) -> u32 {
209        self.completion_tokens
210    }
211}