1use std::{
4 pin::Pin,
5 sync::Arc,
6 task::{Context, Poll},
7 time::Duration,
8};
9
10use anyhow::Result;
11use axum::{
12 extract::{Json, State},
13 http::{self},
14 response::{
15 sse::{Event, KeepAlive},
16 IntoResponse, Sse,
17 },
18};
19use mistralrs_core::{
20 CompletionChunkResponse, CompletionResponse, Constraint, MistralRs, NormalRequest, Request,
21 RequestMessage, Response, SamplingParams,
22};
23use tokio::sync::mpsc::{Receiver, Sender};
24use tracing::warn;
25
26use crate::{
27 completion_core::{
28 convert_stop_tokens, get_dry_sampling_params, handle_completion_error,
29 BaseCompletionResponder,
30 },
31 handler_core::{
32 base_process_non_streaming_response, create_response_channel, send_request,
33 BaseJsonModelError, ErrorToResponse, JsonError, ModelErrorMessage,
34 },
35 openai::{CompletionRequest, Grammar},
36 streaming::{base_create_streamer, get_keep_alive_interval, BaseStreamer, DoneState},
37 types::{ExtractedMistralRsState, OnChunkCallback, OnDoneCallback, SharedMistralRsState},
38 util::{sanitize_error_message, validate_model_name},
39};
40
41pub type CompletionOnChunkCallback = OnChunkCallback<CompletionChunkResponse>;
59
60pub type CompletionOnDoneCallback = OnDoneCallback<CompletionChunkResponse>;
76
77pub type CompletionStreamer =
82 BaseStreamer<CompletionChunkResponse, CompletionOnChunkCallback, CompletionOnDoneCallback>;
83
84impl futures::Stream for CompletionStreamer {
85 type Item = Result<Event, axum::Error>;
86
87 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
95 match self.done_state {
96 DoneState::SendingDone => {
97 self.done_state = DoneState::Done;
100 return Poll::Ready(Some(Ok(Event::default().data("[DONE]"))));
101 }
102 DoneState::Done => {
103 if let Some(on_done) = &self.on_done {
104 on_done(&self.chunks);
105 }
106 return Poll::Ready(None);
107 }
108 DoneState::Running => (),
109 }
110
111 match self.rx.poll_recv(cx) {
112 Poll::Ready(Some(resp)) => match resp {
113 Response::CompletionModelError(msg, _) => {
114 MistralRs::maybe_log_error(
115 self.state.clone(),
116 &ModelErrorMessage(msg.to_string()),
117 );
118 self.done_state = DoneState::SendingDone;
120 Poll::Ready(Some(Ok(Event::default().data(msg))))
121 }
122 Response::ValidationError(e) => Poll::Ready(Some(Ok(
123 Event::default().data(sanitize_error_message(e.as_ref()))
124 ))),
125 Response::InternalError(e) => {
126 MistralRs::maybe_log_error(self.state.clone(), &*e);
127 Poll::Ready(Some(Ok(
128 Event::default().data(sanitize_error_message(e.as_ref()))
129 )))
130 }
131 Response::CompletionChunk(mut response) => {
132 if response.choices.iter().all(|x| x.finish_reason.is_some()) {
133 self.done_state = DoneState::SendingDone;
134 }
135 MistralRs::maybe_log_response(self.state.clone(), &response);
137
138 if let Some(on_chunk) = &self.on_chunk {
139 response = on_chunk(response);
140 }
141
142 if self.store_chunks {
143 self.chunks.push(response.clone());
144 }
145
146 Poll::Ready(Some(Event::default().json_data(response)))
147 }
148 Response::Done(_) => unreachable!(),
149 Response::CompletionDone(_) => unreachable!(),
150 Response::Chunk(_) => unreachable!(),
151 Response::ImageGeneration(_) => unreachable!(),
152 Response::ModelError(_, _) => unreachable!(),
153 Response::Speech { .. } => unreachable!(),
154 Response::Raw { .. } => unreachable!(),
155 },
156 Poll::Pending | Poll::Ready(None) => Poll::Pending,
157 }
158 }
159}
160
161pub type CompletionResponder = BaseCompletionResponder<CompletionResponse, CompletionStreamer>;
163
164type JsonModelError = BaseJsonModelError<CompletionResponse>;
166impl ErrorToResponse for JsonModelError {}
167
168impl IntoResponse for CompletionResponder {
169 fn into_response(self) -> axum::response::Response {
171 match self {
172 CompletionResponder::Sse(s) => s.into_response(),
173 CompletionResponder::Json(s) => Json(s).into_response(),
174 CompletionResponder::InternalError(e) => {
175 JsonError::new(sanitize_error_message(e.as_ref()))
176 .to_response(http::StatusCode::INTERNAL_SERVER_ERROR)
177 }
178 CompletionResponder::ValidationError(e) => {
179 JsonError::new(sanitize_error_message(e.as_ref()))
180 .to_response(http::StatusCode::UNPROCESSABLE_ENTITY)
181 }
182 CompletionResponder::ModelError(msg, response) => JsonModelError::new(msg, response)
183 .to_response(http::StatusCode::INTERNAL_SERVER_ERROR),
184 }
185 }
186}
187
188pub fn parse_request(
193 oairequest: CompletionRequest,
194 state: Arc<MistralRs>,
195 tx: Sender<Response>,
196) -> Result<(Request, bool)> {
197 let repr = serde_json::to_string(&oairequest).expect("Serialization of request failed.");
198 MistralRs::maybe_log_request(state.clone(), repr);
199
200 validate_model_name(&oairequest.model, state.clone())?;
202
203 let stop_toks = convert_stop_tokens(oairequest.stop_seqs);
204
205 if oairequest.logprobs.is_some() {
206 warn!("Completion requests do not support logprobs.");
207 }
208
209 let is_streaming = oairequest.stream.unwrap_or(false);
210
211 let dry_params = get_dry_sampling_params(
212 oairequest.dry_multiplier,
213 oairequest.dry_sequence_breakers,
214 oairequest.dry_base,
215 oairequest.dry_allowed_length,
216 )?;
217
218 Ok((
219 Request::Normal(Box::new(NormalRequest {
220 id: state.next_request_id(),
221 messages: RequestMessage::Completion {
222 text: oairequest.prompt,
223 echo_prompt: oairequest.echo_prompt,
224 best_of: oairequest.best_of,
225 },
226 sampling_params: SamplingParams {
227 temperature: oairequest.temperature,
228 top_k: oairequest.top_k,
229 top_p: oairequest.top_p,
230 min_p: oairequest.min_p,
231 top_n_logprobs: 1,
232 frequency_penalty: oairequest.frequency_penalty,
233 presence_penalty: oairequest.presence_penalty,
234 repetition_penalty: oairequest.repetition_penalty,
235 max_len: oairequest.max_tokens,
236 stop_toks,
237 logits_bias: oairequest.logit_bias,
238 n_choices: oairequest.n_choices,
239 dry_params,
240 },
241 response: tx,
242 return_logprobs: false,
243 is_streaming,
244 suffix: oairequest.suffix,
245 constraint: match oairequest.grammar {
246 Some(Grammar::Regex(regex)) => Constraint::Regex(regex),
247 Some(Grammar::Lark(lark)) => Constraint::Lark(lark),
248 Some(Grammar::JsonSchema(schema)) => Constraint::JsonSchema(schema),
249 Some(Grammar::Llguidance(llguidance)) => Constraint::Llguidance(llguidance),
250 None => Constraint::None,
251 },
252 tool_choice: oairequest.tool_choice,
253 tools: oairequest.tools,
254 logits_processors: None,
255 return_raw_logits: false,
256 web_search_options: None,
257 model_id: if oairequest.model == "default" {
258 None
259 } else {
260 Some(oairequest.model.clone())
261 },
262 })),
263 is_streaming,
264 ))
265}
266
267#[utoipa::path(
269 post,
270 tag = "Mistral.rs",
271 path = "/v1/completions",
272 request_body = CompletionRequest,
273 responses((status = 200, description = "Completions"))
274)]
275pub async fn completions(
276 State(state): ExtractedMistralRsState,
277 Json(oairequest): Json<CompletionRequest>,
278) -> CompletionResponder {
279 let (tx, mut rx) = create_response_channel(None);
280
281 let (request, is_streaming) = match parse_request(oairequest, state.clone(), tx) {
282 Ok(x) => x,
283 Err(e) => return handle_error(state, e.into()),
284 };
285
286 if let Err(e) = send_request(&state, request).await {
287 return handle_error(state, e.into());
288 }
289
290 if is_streaming {
291 CompletionResponder::Sse(create_streamer(rx, state, None, None))
292 } else {
293 process_non_streaming_response(&mut rx, state).await
294 }
295}
296
297pub fn handle_error(
299 state: SharedMistralRsState,
300 e: Box<dyn std::error::Error + Send + Sync + 'static>,
301) -> CompletionResponder {
302 handle_completion_error(state, e)
303}
304
305pub fn create_streamer(
307 rx: Receiver<Response>,
308 state: SharedMistralRsState,
309 on_chunk: Option<CompletionOnChunkCallback>,
310 on_done: Option<CompletionOnDoneCallback>,
311) -> Sse<CompletionStreamer> {
312 let streamer = base_create_streamer(rx, state, on_chunk, on_done);
313 let keep_alive_interval = get_keep_alive_interval();
314
315 Sse::new(streamer)
316 .keep_alive(KeepAlive::new().interval(Duration::from_millis(keep_alive_interval)))
317}
318
319pub async fn process_non_streaming_response(
321 rx: &mut Receiver<Response>,
322 state: SharedMistralRsState,
323) -> CompletionResponder {
324 base_process_non_streaming_response(rx, state, match_responses, handle_error).await
325}
326
327pub fn match_responses(state: SharedMistralRsState, response: Response) -> CompletionResponder {
329 match response {
330 Response::InternalError(e) => {
331 MistralRs::maybe_log_error(state, &*e);
332 CompletionResponder::InternalError(e)
333 }
334 Response::CompletionModelError(msg, response) => {
335 MistralRs::maybe_log_error(state.clone(), &ModelErrorMessage(msg.to_string()));
336 MistralRs::maybe_log_response(state, &response);
337 CompletionResponder::ModelError(msg, response)
338 }
339 Response::ValidationError(e) => CompletionResponder::ValidationError(e),
340 Response::CompletionDone(response) => {
341 MistralRs::maybe_log_response(state, &response);
342 CompletionResponder::Json(response)
343 }
344 Response::CompletionChunk(_) => unreachable!(),
345 Response::Chunk(_) => unreachable!(),
346 Response::Done(_) => unreachable!(),
347 Response::ModelError(_, _) => unreachable!(),
348 Response::ImageGeneration(_) => unreachable!(),
349 Response::Speech { .. } => unreachable!(),
350 Response::Raw { .. } => unreachable!(),
351 }
352}