next_web_ai/ai/deep_seek/api/
deep_seek_api.rs1use 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 let stream = resp.bytes_stream().then(|data| async move {
48 data.map(|data| {
49 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#[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}