Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: return streaming errors as an event formatted for openai's client #2668

Merged
24 changes: 24 additions & 0 deletions router/src/infer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,12 @@ use crate::{
};
use async_stream::stream;
use async_trait::async_trait;
use axum::response::sse::Event;
use chat_template::ChatTemplate;
use futures::future::try_join_all;
use futures::Stream;
use minijinja::ErrorKind;
use serde::Serialize;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use thiserror::Error;
Expand Down Expand Up @@ -373,4 +375,26 @@ impl InferError {
InferError::StreamSerializationError(_) => "stream_serialization_error",
}
}

pub(crate) fn into_openai_event(self) -> Event {
Event::default()
.json_data(OpenaiErrorEvent {
error: APIError {
message: self.to_string(),
http_status_code: 422,
},
})
.unwrap()
}
}

#[derive(Serialize)]
pub struct APIError {
message: String,
http_status_code: usize,
}

#[derive(Serialize)]
pub struct OpenaiErrorEvent {
error: APIError,
}
7 changes: 5 additions & 2 deletions router/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -866,7 +866,7 @@ pub(crate) async fn completions(

yield Ok(event);
}
Err(err) => yield Ok(Event::from(err)),
Err(err) => yield Ok(err.into_openai_event()),
}
}
};
Expand Down Expand Up @@ -1274,7 +1274,8 @@ pub(crate) async fn chat_completions(
};
let mut response_as_tool = using_tools;
while let Some(result) = response_stream.next().await {
if let Ok(stream_token) = result {
match result{
Ok(stream_token) => {
let token_text = &stream_token.token.text.clone();
match state {
StreamState::Buffering => {
Expand Down Expand Up @@ -1368,6 +1369,8 @@ pub(crate) async fn chat_completions(
}
}
}
Err(err) => yield Ok(err.into_openai_event())
}
}
yield Ok::<Event, Infallible>(Event::default().data("[DONE]"));
};
Expand Down
Loading