mistralrs_server_core/
speech_generation.rs

1//! ## Speech generation functionality and route handler.
2
3use std::{error::Error, sync::Arc};
4
5use anyhow::Result;
6use axum::{
7    body::Bytes,
8    extract::{Json, State},
9    http::{self, HeaderMap, HeaderValue, StatusCode},
10    response::IntoResponse,
11};
12use mistralrs_core::{
13    speech_utils::{self, Sample},
14    Constraint, MistralRs, NormalRequest, Request, RequestMessage, Response, SamplingParams,
15};
16use tokio::sync::mpsc::{Receiver, Sender};
17
18use crate::{
19    handler_core::{create_response_channel, send_request, ErrorToResponse, JsonError},
20    openai::{AudioResponseFormat, SpeechGenerationRequest},
21    types::SharedMistralRsState,
22    util::{sanitize_error_message, validate_model_name},
23};
24
25/// Represents different types of speech generation responses.
26pub enum SpeechGenerationResponder {
27    InternalError(Box<dyn Error>),
28    ValidationError(Box<dyn Error>),
29    RawResponse(axum::response::Response),
30}
31
32impl IntoResponse for SpeechGenerationResponder {
33    /// Converts the speech generation responder into an HTTP response.
34    fn into_response(self) -> axum::response::Response {
35        match self {
36            SpeechGenerationResponder::InternalError(e) => {
37                JsonError::new(sanitize_error_message(e.as_ref()))
38                    .to_response(http::StatusCode::INTERNAL_SERVER_ERROR)
39            }
40            SpeechGenerationResponder::ValidationError(e) => {
41                JsonError::new(sanitize_error_message(e.as_ref()))
42                    .to_response(http::StatusCode::UNPROCESSABLE_ENTITY)
43            }
44            SpeechGenerationResponder::RawResponse(resp) => resp,
45        }
46    }
47}
48
49/// Parses and validates a speech generation request.
50///
51/// This function transforms a speech generation request into the
52/// request format used by mistral.rs.
53pub fn parse_request(
54    oairequest: SpeechGenerationRequest,
55    state: Arc<MistralRs>,
56    tx: Sender<Response>,
57) -> Result<(Request, AudioResponseFormat)> {
58    let repr = serde_json::to_string(&oairequest).expect("Serialization of request failed.");
59    MistralRs::maybe_log_request(state.clone(), repr);
60
61    // Validate that the requested model matches the loaded model
62    validate_model_name(&oairequest.model, state.clone())?;
63
64    let request = Request::Normal(Box::new(NormalRequest {
65        id: state.next_request_id(),
66        messages: RequestMessage::SpeechGeneration {
67            prompt: oairequest.input,
68        },
69        sampling_params: SamplingParams::deterministic(),
70        response: tx,
71        return_logprobs: false,
72        is_streaming: false,
73        suffix: None,
74        constraint: Constraint::None,
75        tool_choice: None,
76        tools: None,
77        logits_processors: None,
78        return_raw_logits: false,
79        web_search_options: None,
80        model_id: if oairequest.model == "default" {
81            None
82        } else {
83            Some(oairequest.model.clone())
84        },
85    }));
86
87    Ok((request, oairequest.response_format))
88}
89
90/// Speech generation endpoint handler.
91#[utoipa::path(
92    post,
93    tag = "Mistral.rs",
94    path = "/v1/audio/speech",
95    request_body = SpeechGenerationRequest,
96    responses((status = 200, description = "Speech generation"))
97)]
98pub async fn speech_generation(
99    State(state): State<Arc<MistralRs>>,
100    Json(oairequest): Json<SpeechGenerationRequest>,
101) -> SpeechGenerationResponder {
102    let (tx, mut rx) = create_response_channel(None);
103
104    let (request, response_format) = match parse_request(oairequest, state.clone(), tx) {
105        Ok(x) => x,
106        Err(e) => return handle_error(state, e.into()),
107    };
108
109    // Validate response format here
110    if !matches!(
111        response_format,
112        AudioResponseFormat::Wav | AudioResponseFormat::Pcm
113    ) {
114        return SpeechGenerationResponder::ValidationError(Box::new(JsonError::new(
115            "Only support wav/pcm response format.".to_string(),
116        )));
117    }
118
119    if let Err(e) = send_request(&state, request).await {
120        return handle_error(state, e.into());
121    }
122
123    process_non_streaming_response(&mut rx, state, response_format).await
124}
125
126/// Helper function to handle speech generation errors and logging them.
127pub fn handle_error(
128    state: SharedMistralRsState,
129    e: Box<dyn std::error::Error + Send + Sync + 'static>,
130) -> SpeechGenerationResponder {
131    let sanitized_msg = sanitize_error_message(&*e);
132    let e = anyhow::Error::msg(sanitized_msg);
133    MistralRs::maybe_log_error(state, &*e);
134    SpeechGenerationResponder::InternalError(e.into())
135}
136
137/// Process non-streaming speech generation responses.
138pub async fn process_non_streaming_response(
139    rx: &mut Receiver<Response>,
140    state: SharedMistralRsState,
141    response_format: AudioResponseFormat,
142) -> SpeechGenerationResponder {
143    let response = match rx.recv().await {
144        Some(response) => response,
145        None => {
146            let e = anyhow::Error::msg("No response received from the model.");
147            return handle_error(state, e.into());
148        }
149    };
150
151    match_responses(state, response, response_format)
152}
153
154/// Matches and processes different types of model responses into appropriate speech generation responses.
155pub fn match_responses(
156    state: SharedMistralRsState,
157    response: Response,
158    response_format: AudioResponseFormat,
159) -> SpeechGenerationResponder {
160    match response {
161        Response::InternalError(e) => {
162            MistralRs::maybe_log_error(state, &*e);
163            SpeechGenerationResponder::InternalError(e)
164        }
165        Response::ValidationError(e) => SpeechGenerationResponder::ValidationError(e),
166        Response::ImageGeneration(_) => unreachable!(),
167        Response::CompletionModelError(m, _) => {
168            let e = anyhow::Error::msg(m.to_string());
169            MistralRs::maybe_log_error(state, &*e);
170            SpeechGenerationResponder::InternalError(e.into())
171        }
172        Response::CompletionDone(_) => unreachable!(),
173        Response::CompletionChunk(_) => unreachable!(),
174        Response::Chunk(_) => unreachable!(),
175        Response::Done(_) => unreachable!(),
176        Response::ModelError(_, _) => unreachable!(),
177        Response::Speech {
178            pcm,
179            rate,
180            channels,
181        } => {
182            let pcm_endianness = "s16le";
183
184            let content_type = response_format.audio_content_type(rate, channels, pcm_endianness);
185            let mut headers = HeaderMap::new();
186            headers.insert(
187                http::header::CONTENT_TYPE,
188                HeaderValue::from_str(&content_type).unwrap(),
189            );
190
191            let encoded = match response_format {
192                AudioResponseFormat::Pcm => {
193                    let samples: &[f32] = &pcm;
194                    let mut buf = Vec::with_capacity(samples.len() * std::mem::size_of::<i64>());
195                    for &sample in samples {
196                        buf.extend_from_slice(&sample.to_i16().to_le_bytes());
197                    }
198                    buf
199                }
200                AudioResponseFormat::Wav => {
201                    // Write WAV data into an in-memory buffer
202                    let mut buf = Vec::new();
203                    speech_utils::write_pcm_as_wav(&mut buf, &pcm, rate as u32, channels as u16)
204                        .unwrap();
205                    buf
206                }
207                _ => unreachable!("Should be validated above."),
208            };
209
210            let bytes = Bytes::from(encoded);
211
212            SpeechGenerationResponder::RawResponse((StatusCode::OK, headers, bytes).into_response())
213        }
214        Response::Raw { .. } => unreachable!(),
215    }
216}
mistralrs_server_core/speech_generation.rs

mistralrs_server_core/
speech_generation.rs
