1use std::collections::HashMap;
6
7use crate::DeepSeekRequest;
8use crate::chat::request::{Stop, StreamOptions, is_none_or_empty_stop};
9use crate::chat::response::ChatGeneric;
10use crate::error::DeepSeekError;
11use crate::{DeepSeekClient, api_post, api_request_stream, consume_sse, spawn_blocking_stream};
12use derive_builder::Builder;
13use reqwest::Method;
14use serde::{Deserialize, Serialize};
15use tokio::sync::mpsc;
16
17pub type Completion = ChatGeneric<CompletionChoice>;
19
20#[derive(Clone, Debug, PartialEq, Serialize, Builder)]
22#[builder(
23 pattern = "owned",
24 setter(into, strip_option),
25 build_fn(validate = "Self::validate"),
26 name = "FIMCompletionRequestBuilder"
27)]
28pub struct FIMCompletionRequest {
29 #[serde(skip_serializing)]
30 pub client: DeepSeekClient,
31
32 pub model: String,
36
37 pub prompt: String,
39
40 #[builder(default)]
42 #[serde(skip_serializing_if = "Option::is_none")]
43 pub echo: Option<bool>,
44
45 #[builder(default)]
52 #[serde(skip_serializing_if = "Option::is_none")]
53 pub logprobs: Option<u32>,
54
55 #[builder(default)]
57 #[serde(skip_serializing_if = "Option::is_none")]
58 pub max_tokens: Option<u32>,
59
60 #[builder(default)]
63 #[serde(skip_serializing_if = "is_none_or_empty_stop")]
64 pub stop: Option<Stop>,
65
66 #[builder(default)]
69 #[serde(skip_serializing_if = "Option::is_none")]
70 pub stream: Option<bool>,
71
72 #[builder(default)]
74 #[serde(skip_serializing_if = "Option::is_none")]
75 pub stream_options: Option<StreamOptions>,
76
77 #[builder(default)]
79 #[serde(skip_serializing_if = "Option::is_none")]
80 pub suffix: Option<String>,
81
82 #[builder(default)]
90 #[serde(skip_serializing_if = "Option::is_none")]
91 pub temperature: Option<f64>,
92
93 #[builder(default)]
102 #[serde(skip_serializing_if = "Option::is_none")]
103 pub top_p: Option<f64>,
104}
105
106impl FIMCompletionRequestBuilder {
107 fn validate(&self) -> Result<(), String> {
108 if let Some(temperature) = self.temperature.flatten()
109 && !(0.0..=2.0).contains(&temperature)
110 {
111 return Err("temperature must be between 0 and 2".to_string());
112 }
113 if let Some(logprobs) = self.logprobs.flatten()
114 && logprobs > 20
115 {
116 return Err("logprobs must be <= 20".to_string());
117 }
118
119 if let Some(top_p) = self.top_p.flatten()
120 && !(0.0..=1.0).contains(&top_p)
121 {
122 return Err("top_p must be between 0 and 1".to_string());
123 }
124
125 if let Some(stream) = self.stream.flatten()
126 && !stream
127 && self.stream_options.is_some()
128 {
129 return Err("stream_options cannot be set when stream is false".to_string());
130 }
131
132 if let Some(stop) = self.stop.as_ref().and_then(|s| s.as_ref())
133 && let Stop::Many(values) = stop
134 && values.len() > 16
135 {
136 return Err("a maximum of 16 stop sequences are allowed".to_string());
137 }
138
139 Ok(())
140 }
141}
142
143#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
145pub struct CompletionChoice {
146 pub finish_reason: FinishReason,
154 pub index: u64,
155 pub text: String,
156 #[serde(skip_serializing_if = "Option::is_none")]
157 pub logprobs: Option<Logprobs>,
158}
159
160#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
162#[serde(rename_all = "snake_case")]
163pub enum FinishReason {
164 Stop,
165 Length,
166 ContentFilter,
167 InsufficientSystemResources,
168}
169
170#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
172pub struct Logprobs {
173 pub text_offset: Vec<u64>,
174 pub token_logprobs: Vec<f64>,
175 pub tokens: Vec<String>,
176 pub top_logprobs: Option<Vec<HashMap<String, f64>>>,
177}
178#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
180pub struct CompletionChoiceStream {
181 pub finish_reason: Option<FinishReason>,
182 pub index: u64,
183 pub text: String,
184 #[serde(skip_serializing_if = "Option::is_none")]
185 pub logprobs: Option<Logprobs>,
186}
187pub type CompletionStream = ChatGeneric<CompletionChoiceStream>;
189pub type CompletionStreamItem = Result<CompletionStream, DeepSeekError>;
191pub struct CompletionStreamBlocking {
193 rx: std::sync::mpsc::Receiver<CompletionStreamItem>,
194}
195
196impl Iterator for CompletionStreamBlocking {
197 type Item = CompletionStreamItem;
198
199 fn next(&mut self) -> Option<Self::Item> {
200 self.rx.recv().ok()
201 }
202}
203impl DeepSeekRequest for FIMCompletionRequest {
204 type Response = Completion;
205 type StreamItem = CompletionStreamItem;
206 type BlockingStream = CompletionStreamBlocking;
207
208 async fn send(self) -> Result<Self::Response, DeepSeekError> {
209 let client = self.client.clone();
210 api_post("/completions", &self, client).await
211 }
212
213 async fn stream(self) -> Result<mpsc::Receiver<Self::StreamItem>, DeepSeekError> {
214 let mut request = self;
215 request.stream = Some(true);
216
217 let client = request.client.clone();
218 let event_source = api_request_stream(
219 Method::POST,
220 "/completions",
221 |builder| builder.json(&request),
222 client,
223 )
224 .await?;
225
226 Ok(consume_sse(event_source, |data| {
227 serde_json::from_str::<CompletionStream>(&data)
228 .map(Some)
229 .map_err(|err| DeepSeekError::decode(err.to_string(), data))
230 }))
231 }
232
233 fn stream_blocking(self) -> Result<CompletionStreamBlocking, DeepSeekError> {
234 let rx = spawn_blocking_stream(self.stream())?;
235 Ok(CompletionStreamBlocking { rx })
236 }
237}
238
239#[cfg(test)]
240mod tests {
241 use super::*;
242 use crate::DEFAULT_BETA_BASE_URL;
243
244 fn get_client() -> DeepSeekClient {
245 DeepSeekClient::new(
246 std::env::var("DEEPSEEK_API_KEY").expect("DEEPSEEK_API_KEY is not set"),
247 DEFAULT_BETA_BASE_URL.clone(),
248 )
249 }
250
251 fn get_fim_builder() -> FIMCompletionRequestBuilder {
252 FIMCompletionRequestBuilder::default()
253 .client(get_client())
254 .model("deepseek-v4-flash")
255 .max_tokens(64_u32)
256 }
257
258 #[tokio::test]
259 async fn test_fim_completion() {
260 let fim_request = get_fim_builder()
261 .prompt("def fib(a):")
262 .suffix(" return fib(a-1) + fib(a-2)")
263 .build()
264 .unwrap();
265 let response = fim_request.send().await.unwrap();
266 println!("{:#?}", response);
267 assert_eq!(response.object, "text_completion");
268 assert_eq!(response.model, "deepseek-v4-flash");
269 assert_eq!(response.choices.len(), 1);
270 }
271
272 #[tokio::test]
273 async fn test_fim_completion_stream() {
274 let fim_request = get_fim_builder()
275 .prompt("def fib(a):")
276 .suffix(" return fib(a-1) + fib(a-2)")
277 .stream(true)
278 .build()
279 .unwrap();
280 let mut stream = fim_request.stream().await.unwrap();
281 while let Some(item) = stream.recv().await {
282 match item {
283 Ok(chunk) => println!("Received chunk: {:#?}", chunk),
284 Err(err) => eprintln!("Stream error: {}", err),
285 }
286 }
287 }
288
289 #[tokio::test]
290 async fn test_fim_completion_stream_blocking() {
291 let fim_request = get_fim_builder()
292 .prompt("def fib(a):")
293 .suffix(" return fib(a-1) + fib(a-2)")
294 .stream(true)
295 .build()
296 .unwrap();
297 let stream = fim_request.stream_blocking().unwrap();
298 for item in stream {
299 match item {
300 Ok(chunk) => println!("Received chunk: {:#?}", chunk),
301 Err(err) => eprintln!("Stream error: {}", err),
302 }
303 }
304 }
305}