mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
refactor(rust): one stream codec shape per API format (#45996)
Chat Completions, Messages and Responses each get a streaming module with the same three names, event_stream, decode and encode, built on a shared SSE helper in base_llm/sse.rs that owns the [DONE] marker. Callers use the module path instead of vendor-prefixed names. Encoding matches what the Python proxy writes to clients: Messages frames carry an event line, the OpenAI formats are bare data frames. Responses gets one event type for SSE and WebSocket. ResponsesEventType names every event on OpenAI's streaming-events reference, with unknown names riding in Other, and is_terminal and is_failure follow Python's iterator. The framer adds a bulk decode_all and an SseBlockCodec that keeps each event's wire bytes. An event the stream ends inside is dropped, as the SSE spec and the openai SDK's decoder do. ChatCompletionChunk no longer requires id, created or object, since compatible hosts omit them.
This commit is contained in:
yujonglee
GitHub
parent
7bd108139f
commit
271b4f0306
22 changed files with 1219 additions and 266 deletions
|
|
@ -1,6 +1,6 @@
|
|||
use std::str;
|
||||
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use bytes::{Buf, BufMut, Bytes, BytesMut};
|
||||
use tokio_util::codec::{Decoder, Encoder};
|
||||
|
||||
use crate::SseError;
|
||||
|
|
@ -29,38 +29,107 @@ impl Decoder for SseCodec {
|
|||
return Ok(None);
|
||||
}
|
||||
while let Some(end) = block_end(src) {
|
||||
let block = src.split_to(end);
|
||||
let pending = lines(&block)
|
||||
.map(|(line, _)| line)
|
||||
.take_while(|line| !line.is_empty())
|
||||
.try_fold(Pending::default(), Pending::apply)?;
|
||||
if let Some(event) = pending.dispatch() {
|
||||
if let Some(event) = event_in(&src.split_to(end))? {
|
||||
return Ok(Some(event));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
|
||||
Ok(None)
|
||||
/// An event the stream ended inside of is dropped, as the SSE spec requires.
|
||||
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
|
||||
self.decode(pending)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decode_all(wire: &[u8]) -> Result<Vec<SseEvent>, SseError> {
|
||||
let mut codec = SseCodec::default();
|
||||
let mut buffer = BytesMut::from(wire);
|
||||
let mut events = Vec::new();
|
||||
while let Some(event) = codec.decode_eof(&mut buffer)? {
|
||||
events.push(event);
|
||||
}
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
impl SseCodec {
|
||||
fn skip_bom(&mut self, src: &mut BytesMut) -> bool {
|
||||
if self.past_bom {
|
||||
return true;
|
||||
}
|
||||
if src.starts_with(BOM) {
|
||||
src.advance(BOM.len());
|
||||
} else if BOM.starts_with(src) {
|
||||
return false;
|
||||
}
|
||||
self.past_bom = true;
|
||||
true
|
||||
skip_bom(&mut self.past_bom, src)
|
||||
}
|
||||
}
|
||||
|
||||
fn event_in(block: &[u8]) -> Result<Option<SseEvent>, SseError> {
|
||||
Ok(lines(block)
|
||||
.map(|(line, _)| line)
|
||||
.take_while(|line| !line.is_empty())
|
||||
.try_fold(Pending::default(), Pending::apply)?
|
||||
.dispatch())
|
||||
}
|
||||
|
||||
/// One event together with the exact wire bytes that carried it, for relays that forward a
|
||||
/// provider's bytes untouched but still need to see each event go by.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct SseBlock {
|
||||
pub bytes: Bytes,
|
||||
pub event: SseEvent,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct SseBlockCodec {
|
||||
past_bom: bool,
|
||||
}
|
||||
|
||||
impl Decoder for SseBlockCodec {
|
||||
type Item = SseBlock;
|
||||
type Error = SseError;
|
||||
|
||||
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
|
||||
self.decode_block(src, false)
|
||||
}
|
||||
|
||||
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
|
||||
self.decode_block(pending, true)
|
||||
}
|
||||
}
|
||||
|
||||
impl SseBlockCodec {
|
||||
fn decode_block(
|
||||
&mut self,
|
||||
src: &mut BytesMut,
|
||||
eof: bool,
|
||||
) -> Result<Option<SseBlock>, SseError> {
|
||||
if !skip_bom(&mut self.past_bom, src) {
|
||||
return Ok(None);
|
||||
}
|
||||
while let Some(end) = block_end(src) {
|
||||
if !eof && end == src.len() && src.ends_with(b"\r") {
|
||||
return Ok(None);
|
||||
}
|
||||
let block = src.split_to(end);
|
||||
if let Some(event) = event_in(&block)? {
|
||||
return Ok(Some(SseBlock {
|
||||
bytes: block.freeze(),
|
||||
event,
|
||||
}));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
fn skip_bom(past_bom: &mut bool, src: &mut BytesMut) -> bool {
|
||||
if *past_bom {
|
||||
return true;
|
||||
}
|
||||
if src.starts_with(BOM) {
|
||||
src.advance(BOM.len());
|
||||
} else if BOM.starts_with(src) {
|
||||
return false;
|
||||
}
|
||||
*past_bom = true;
|
||||
true
|
||||
}
|
||||
|
||||
fn block_end(bytes: &[u8]) -> Option<usize> {
|
||||
lines(bytes)
|
||||
.find(|(line, _)| line.is_empty())
|
||||
|
|
@ -168,3 +237,179 @@ impl Encoder<SseEvent> for SseCodec {
|
|||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::frames;
|
||||
|
||||
fn pieces(wire: &[u8], size: usize) -> Vec<Bytes> {
|
||||
wire.chunks(size).map(Bytes::copy_from_slice).collect()
|
||||
}
|
||||
|
||||
async fn blocks(wire: &[u8], size: usize) -> Result<Vec<SseBlock>, SseError> {
|
||||
frames(
|
||||
stream::iter(pieces(wire, size).into_iter().map(Ok::<_, std::io::Error>)),
|
||||
SseBlockCodec::default(),
|
||||
)
|
||||
.try_collect()
|
||||
.await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lf(
|
||||
b": keepalive\n\nevent: delta\ndata: a\ndata: b\n\n",
|
||||
b"event: delta\ndata: a\ndata: b\n\n"
|
||||
)]
|
||||
#[case::crlf(
|
||||
b": keepalive\r\n\r\nevent: delta\r\ndata: a\r\ndata: b\r\n\r\n",
|
||||
b"event: delta\r\ndata: a\r\ndata: b\r\n\r\n"
|
||||
)]
|
||||
#[case::cr(
|
||||
b": keepalive\r\revent: delta\rdata: a\rdata: b\r\r",
|
||||
b"event: delta\rdata: a\rdata: b\r\r"
|
||||
)]
|
||||
#[case::bom(b"\xEF\xBB\xBFdata: \xe2\x82\xac\n\n", b"data: \xe2\x82\xac\n\n")]
|
||||
#[case::empty_data(b"event: ignored\n\ndata:\n\n", b"data:\n\n")]
|
||||
#[case::extensions(
|
||||
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n",
|
||||
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn blocks_carry_the_wire_bytes_and_the_event_they_dispatch(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: &[u8],
|
||||
#[values(1, 3, 4096)] size: usize,
|
||||
) {
|
||||
let blocks = blocks(wire, size).await.unwrap();
|
||||
let expected_events = decode_all(expected).unwrap();
|
||||
assert_eq!(expected_events.len(), 1);
|
||||
assert_eq!(
|
||||
blocks,
|
||||
vec![SseBlock {
|
||||
bytes: Bytes::copy_from_slice(expected),
|
||||
event: expected_events[0].clone(),
|
||||
}]
|
||||
);
|
||||
assert_eq!(expected_events, decode_all(wire).unwrap());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty(b"", vec![])]
|
||||
#[case::unterminated(b"data: partial\n", vec![])]
|
||||
#[case::partial_tail(b"data: complete\n\ndata: partial\n", vec![&b"data: complete\n\n"[..]])]
|
||||
#[case::multiple(b"data: first\n\ndata: second\n\n", vec![&b"data: first\n\n"[..], &b"data: second\n\n"[..]])]
|
||||
#[case::trailing_comment(b"data: first\n\n: bye\n\n", vec![&b"data: first\n\n"[..]])]
|
||||
#[tokio::test]
|
||||
async fn blocks_dispatch_complete_events_and_drop_a_cut_tail(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<&[u8]>,
|
||||
) {
|
||||
let mut blocks = Box::pin(frames(
|
||||
stream::iter(pieces(wire, 1).into_iter().map(Ok::<_, std::io::Error>)),
|
||||
SseBlockCodec::default(),
|
||||
));
|
||||
for expected in expected {
|
||||
assert_eq!(blocks.next().await.unwrap().unwrap().bytes, expected);
|
||||
}
|
||||
assert!(blocks.next().await.is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn bulk_decode_preserves_every_field() {
|
||||
let events = decode_all(b"id: 1\nretry: 20\nevent: delta\ndata: a\ndata: b\n\n").unwrap();
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![SseEvent {
|
||||
event: Some("delta".into()),
|
||||
data: "a\nb".into(),
|
||||
id: Some("1".into()),
|
||||
retry: Some(20),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::only_a_tail(b"data: partial\n", vec![])]
|
||||
#[case::tail_after_an_event(b"data: complete\n\ndata: partial\n", vec!["complete"])]
|
||||
#[case::no_terminator_at_all(b"data: partial", vec![])]
|
||||
fn bulk_decode_drops_the_event_a_stream_ends_inside(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<&str>,
|
||||
) {
|
||||
let data: Vec<String> = decode_all(wire)
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.map(|e| e.data)
|
||||
.collect();
|
||||
assert_eq!(data, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::trailing_blank_lines(b"data: x\n\n\n\n")]
|
||||
#[case::trailing_comment(b"data: x\n\n: bye\n\n")]
|
||||
#[case::lone_cr_at_eof(b"data: x\r\r")]
|
||||
fn bulk_decode_accepts_a_stream_that_ends_between_events(#[case] wire: &[u8]) {
|
||||
assert_eq!(decode_all(wire).unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn bulk_decode_rejects_invalid_utf8() {
|
||||
assert!(matches!(
|
||||
decode_all(b"data: \xff\n\n"),
|
||||
Err(SseError::InvalidUtf8(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn block_framing_preserves_transport_errors() {
|
||||
let input = stream::iter([
|
||||
Ok(Bytes::from_static(b"data: first\n\n")),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset",
|
||||
)),
|
||||
]);
|
||||
let mut blocks = Box::pin(frames(input, SseBlockCodec::default()));
|
||||
assert_eq!(
|
||||
blocks.next().await.unwrap().unwrap().bytes,
|
||||
Bytes::from_static(b"data: first\n\n")
|
||||
);
|
||||
let Some(Err(SseError::Body(error))) = blocks.next().await else {
|
||||
panic!("expected body error");
|
||||
};
|
||||
let cause = error
|
||||
.get_ref()
|
||||
.unwrap()
|
||||
.downcast_ref::<std::io::Error>()
|
||||
.unwrap();
|
||||
assert_eq!(cause.kind(), std::io::ErrorKind::ConnectionReset);
|
||||
assert!(blocks.next().await.is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn invalid_utf8_stops_block_framing_after_prior_events() {
|
||||
let mut blocks = Box::pin(frames(
|
||||
stream::iter(
|
||||
pieces(b"data: first\n\ndata: \xff\n\n", 1)
|
||||
.into_iter()
|
||||
.map(Ok::<_, std::io::Error>),
|
||||
),
|
||||
SseBlockCodec::default(),
|
||||
));
|
||||
assert_eq!(
|
||||
blocks.next().await.unwrap().unwrap().bytes,
|
||||
Bytes::from_static(b"data: first\n\n")
|
||||
);
|
||||
assert!(matches!(
|
||||
blocks.next().await,
|
||||
Some(Err(SseError::InvalidUtf8(_)))
|
||||
));
|
||||
assert!(blocks.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -97,12 +97,11 @@ async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected:
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unterminated_single(b"data: partial\n", vec![])]
|
||||
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
|
||||
#[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])]
|
||||
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
|
||||
#[case::trailing_blank_lines(b"data: x\n\n\n", vec![event(None, "x")])]
|
||||
#[case::trailing_comment(b"data: x\n\n: bye\n\n", vec![event(None, "x")])]
|
||||
#[tokio::test]
|
||||
async fn eof_dispatches_only_terminated_events(
|
||||
async fn eof_between_events_ends_the_stream_cleanly(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<SseEvent>,
|
||||
) {
|
||||
|
|
@ -112,6 +111,19 @@ async fn eof_dispatches_only_terminated_events(
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unterminated_single(b"data: partial\n", vec![])]
|
||||
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
|
||||
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
|
||||
#[case::partial_bom(b"\xEF\xBB", vec![])]
|
||||
#[tokio::test]
|
||||
async fn eof_inside_an_event_drops_it_after_the_terminated_events(
|
||||
#[case] wire: &[u8],
|
||||
#[case] expected: Vec<SseEvent>,
|
||||
) {
|
||||
assert_eq!(collect(every(wire, 1)).await.unwrap(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])]
|
||||
#[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])]
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ use litellm_http::transport::Error as TransportError;
|
|||
use litellm_llms::base_llm::{
|
||||
auth::{Authenticated, resolve_auth},
|
||||
messages::{
|
||||
streaming::{ByteStream, StreamDecoder, encode_anthropic_sse},
|
||||
streaming::{self, ByteStream, StreamDecoder},
|
||||
transformation::BaseMessagesConfig,
|
||||
},
|
||||
};
|
||||
|
|
@ -245,7 +245,7 @@ fn decoded_chunks(
|
|||
let Some(event) = events.try_next().await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let chunk = encode_anthropic_sse(&event)?;
|
||||
let chunk = streaming::encode(&event)?;
|
||||
log_chunk(provider, "client_response", &chunk);
|
||||
Ok(Some((chunk, events)))
|
||||
})
|
||||
|
|
@ -271,7 +271,6 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_llms::base_llm::messages::streaming::anthropic_sse_event_stream;
|
||||
use rstest::rstest;
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
|
||||
|
||||
|
|
@ -300,7 +299,7 @@ mod tests {
|
|||
.await
|
||||
.unwrap();
|
||||
let MessagesCallResponse::Stream { mut chunks, .. } =
|
||||
streaming_response(response, Some(anthropic_sse_event_stream), "test")
|
||||
streaming_response(response, Some(streaming::event_stream), "test")
|
||||
else {
|
||||
panic!("a streaming response returns chunks");
|
||||
};
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
|
|||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_http::websocket::{UpstreamWebSocket, connect_upstream};
|
||||
use litellm_llms_types::formats::responses::streaming_websocket::ResponsesWsEventType;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio_tungstenite::tungstenite::{
|
||||
Message,
|
||||
|
|
@ -12,17 +11,6 @@ use tokio_tungstenite::tungstenite::{
|
|||
|
||||
use super::Error;
|
||||
|
||||
pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool {
|
||||
matches!(
|
||||
event_type,
|
||||
ResponsesWsEventType::ResponseCreated
|
||||
| ResponsesWsEventType::ResponseCompleted
|
||||
| ResponsesWsEventType::ResponseFailed
|
||||
| ResponsesWsEventType::ResponseIncomplete
|
||||
| ResponsesWsEventType::Error
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponsesWebSocketConnection {
|
||||
socket: Arc<Mutex<Option<UpstreamWebSocket>>>,
|
||||
|
|
|
|||
|
|
@ -33,16 +33,82 @@ pub struct ChatCompletionStreamingChoice {
|
|||
pub logprobs: Option<Value>,
|
||||
}
|
||||
|
||||
/// Compatible hosts omit `id`, `created` and `object` on chunks, and Python's
|
||||
/// `ModelResponseStream` fills them in, so none of the envelope fields is required here.
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
#[derive(Default)]
|
||||
pub struct ChatCompletionChunk {
|
||||
pub id: String,
|
||||
pub created: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub created: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
pub object: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub object: Option<String>,
|
||||
#[serde(default)]
|
||||
pub choices: Vec<ChatCompletionStreamingChoice>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub usage: Option<ChatCompletionsUsage>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_specific_fields: Option<Map<String, Value>>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::openai(json!({
|
||||
"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m",
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi", "future": 1}, "finish_reason": null, "logprobs": null}],
|
||||
"usage": null
|
||||
}))]
|
||||
#[case::usage_only(json!({
|
||||
"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3,
|
||||
"prompt_tokens_details": {"cached_tokens": 0, "cache_creation_tokens": 0, "text_tokens": 1}
|
||||
}
|
||||
}))]
|
||||
#[case::bare_host(json!({"choices": [{"index": 0, "delta": {"content": "hi"}}]}))]
|
||||
#[case::host_extensions(json!({"id": "x", "choices": [], "citations": ["a"], "future": {"k": null}}))]
|
||||
fn chunks_round_trip_without_losing_fields(#[case] wire: Value) {
|
||||
let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap();
|
||||
assert_eq!(serde_json::to_value(chunk).unwrap(), without_nulls(wire));
|
||||
}
|
||||
|
||||
/// Typed optional fields serialize by omission; extension fields keep their nulls.
|
||||
fn without_nulls(value: Value) -> Value {
|
||||
let Value::Object(fields) = value else {
|
||||
return value;
|
||||
};
|
||||
Value::Object(
|
||||
fields
|
||||
.into_iter()
|
||||
.filter(|(_, value)| !value.is_null())
|
||||
.map(|(key, value)| match (key.as_str(), value) {
|
||||
("choices", Value::Array(choices)) => (
|
||||
key,
|
||||
Value::Array(choices.into_iter().map(without_nulls).collect()),
|
||||
),
|
||||
(_, value) => (key, value),
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::choices_not_a_list(json!({"choices": {}}))]
|
||||
#[case::choice_without_delta(json!({"choices": [{"index": 0}]}))]
|
||||
#[case::created_not_a_number(json!({"created": "now", "choices": []}))]
|
||||
#[case::non_object(json!([]))]
|
||||
fn malformed_chunks_are_rejected(#[case] wire: Value) {
|
||||
assert!(serde_json::from_value::<ChatCompletionChunk>(wire).is_err());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -109,7 +109,9 @@ pub struct MessagesStreamError {
|
|||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
#[derive(strum::AsRefStr)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum MessagesStreamEvent {
|
||||
MessageStart {
|
||||
message: MessagesStreamMessage,
|
||||
|
|
@ -175,6 +177,23 @@ mod tests {
|
|||
assert_eq!(serde_json::to_value(event).unwrap(), wire);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::start(json!({"type": "message_start", "message": {"id": "m", "type": "message", "role": "assistant", "model": "x", "content": [], "stop_reason": null, "stop_sequence": null, "usage": {}}}))]
|
||||
#[case::block_start(json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}))]
|
||||
#[case::block_delta(json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "a"}}))]
|
||||
#[case::block_stop(json!({"type": "content_block_stop", "index": 0}))]
|
||||
#[case::delta(json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}))]
|
||||
#[case::stop(json!({"type": "message_stop"}))]
|
||||
#[case::ping(json!({"type": "ping"}))]
|
||||
#[case::error(json!({"type": "error", "error": {"type": "overloaded_error", "message": "busy"}}))]
|
||||
fn the_event_name_is_the_serialized_type_tag(#[case] wire: Value) {
|
||||
let event: MessagesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
|
||||
assert_eq!(
|
||||
Some(event.as_ref()),
|
||||
wire.get("type").and_then(Value::as_str)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_index(json!({"type": "content_block_stop", "index": "zero"}))]
|
||||
#[case::missing_delta(json!({"type": "content_block_delta", "index": 0}))]
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# rules
|
||||
|
||||
- Responses originated at OpenAI and is now offered by other hosts and gateways. Types here describe the format, not OpenAI's deployment of it
|
||||
- Served over HTTP and over the WebSocket mode. `ResponsesWsEvent` is a wire event and lives here
|
||||
- Served over HTTP and over the WebSocket mode. `ResponsesStreamEvent` and `ResponsesEventType` in `streaming.rs` are the wire event for both
|
||||
- `ResponsesWsTransformResult` wraps a provider transformation's output, so it lives in `llms/src/base_llm/responses/`
|
||||
- Hosts that accept only part of the format, or add output item types, are handled in `llms/src/<provider>/responses/`. Keep unknown output items passing through instead of rejecting them
|
||||
- Responses-to-Chat emulation and model-specific parameter rewriting are transformations, not format rules
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
mod output;
|
||||
mod response;
|
||||
pub mod streaming;
|
||||
pub mod streaming_websocket;
|
||||
|
||||
pub use output::*;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,305 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
/// Every event name on https://developers.openai.com/api/reference/resources/responses/streaming-events
|
||||
/// as of 2026-10-11, plus `response.create` from the WebSocket mode. Names OpenAI adds later ride in
|
||||
/// `Other` untouched.
|
||||
#[derive(
|
||||
Clone,
|
||||
Debug,
|
||||
PartialEq,
|
||||
Eq,
|
||||
strum::AsRefStr,
|
||||
strum::EnumIter,
|
||||
strum::EnumString,
|
||||
strum::Display,
|
||||
serde_with::DeserializeFromStr,
|
||||
serde_with::SerializeDisplay,
|
||||
)]
|
||||
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
|
||||
#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))]
|
||||
pub enum ResponsesEventType {
|
||||
#[strum(serialize = "response.create")]
|
||||
ResponseCreate,
|
||||
#[strum(serialize = "response.created")]
|
||||
ResponseCreated,
|
||||
#[strum(serialize = "response.queued")]
|
||||
ResponseQueued,
|
||||
#[strum(serialize = "response.in_progress")]
|
||||
ResponseInProgress,
|
||||
#[strum(serialize = "response.completed")]
|
||||
ResponseCompleted,
|
||||
#[strum(serialize = "response.failed")]
|
||||
ResponseFailed,
|
||||
#[strum(serialize = "response.incomplete")]
|
||||
ResponseIncomplete,
|
||||
#[strum(serialize = "error")]
|
||||
Error,
|
||||
#[strum(serialize = "response.audio.delta")]
|
||||
AudioDelta,
|
||||
#[strum(serialize = "response.audio.done")]
|
||||
AudioDone,
|
||||
#[strum(serialize = "response.audio.transcript.delta")]
|
||||
AudioTranscriptDelta,
|
||||
#[strum(serialize = "response.audio.transcript.done")]
|
||||
AudioTranscriptDone,
|
||||
#[strum(serialize = "response.code_interpreter_call.completed")]
|
||||
CodeInterpreterCallCompleted,
|
||||
#[strum(serialize = "response.code_interpreter_call.in_progress")]
|
||||
CodeInterpreterCallInProgress,
|
||||
#[strum(serialize = "response.code_interpreter_call.interpreting")]
|
||||
CodeInterpreterCallInterpreting,
|
||||
#[strum(serialize = "response.code_interpreter_call_code.delta")]
|
||||
CodeInterpreterCallCodeDelta,
|
||||
#[strum(serialize = "response.code_interpreter_call_code.done")]
|
||||
CodeInterpreterCallCodeDone,
|
||||
#[strum(serialize = "response.compaction.compacting")]
|
||||
CompactionCompacting,
|
||||
#[strum(serialize = "response.content_part.added")]
|
||||
ContentPartAdded,
|
||||
#[strum(serialize = "response.content_part.done")]
|
||||
ContentPartDone,
|
||||
#[strum(serialize = "response.custom_tool_call_input.delta")]
|
||||
CustomToolCallInputDelta,
|
||||
#[strum(serialize = "response.custom_tool_call_input.done")]
|
||||
CustomToolCallInputDone,
|
||||
#[strum(serialize = "response.file_search_call.completed")]
|
||||
FileSearchCallCompleted,
|
||||
#[strum(serialize = "response.file_search_call.in_progress")]
|
||||
FileSearchCallInProgress,
|
||||
#[strum(serialize = "response.file_search_call.searching")]
|
||||
FileSearchCallSearching,
|
||||
#[strum(serialize = "response.function_call_arguments.delta")]
|
||||
FunctionCallArgumentsDelta,
|
||||
#[strum(serialize = "response.function_call_arguments.done")]
|
||||
FunctionCallArgumentsDone,
|
||||
#[strum(serialize = "response.image_generation_call.completed")]
|
||||
ImageGenerationCallCompleted,
|
||||
#[strum(serialize = "response.image_generation_call.generating")]
|
||||
ImageGenerationCallGenerating,
|
||||
#[strum(serialize = "response.image_generation_call.in_progress")]
|
||||
ImageGenerationCallInProgress,
|
||||
#[strum(serialize = "response.image_generation_call.partial_image")]
|
||||
ImageGenerationCallPartialImage,
|
||||
#[strum(serialize = "response.mcp_call.completed")]
|
||||
McpCallCompleted,
|
||||
#[strum(serialize = "response.mcp_call.failed")]
|
||||
McpCallFailed,
|
||||
#[strum(serialize = "response.mcp_call.in_progress")]
|
||||
McpCallInProgress,
|
||||
#[strum(serialize = "response.mcp_call_arguments.delta")]
|
||||
McpCallArgumentsDelta,
|
||||
#[strum(serialize = "response.mcp_call_arguments.done")]
|
||||
McpCallArgumentsDone,
|
||||
#[strum(serialize = "response.mcp_list_tools.completed")]
|
||||
McpListToolsCompleted,
|
||||
#[strum(serialize = "response.mcp_list_tools.failed")]
|
||||
McpListToolsFailed,
|
||||
#[strum(serialize = "response.mcp_list_tools.in_progress")]
|
||||
McpListToolsInProgress,
|
||||
#[strum(serialize = "response.output_item.added")]
|
||||
OutputItemAdded,
|
||||
#[strum(serialize = "response.output_item.done")]
|
||||
OutputItemDone,
|
||||
#[strum(serialize = "response.output_text.annotation.added")]
|
||||
OutputTextAnnotationAdded,
|
||||
#[strum(serialize = "response.output_text.delta")]
|
||||
OutputTextDelta,
|
||||
#[strum(serialize = "response.output_text.done")]
|
||||
OutputTextDone,
|
||||
#[strum(serialize = "response.reasoning_summary_part.added")]
|
||||
ReasoningSummaryPartAdded,
|
||||
#[strum(serialize = "response.reasoning_summary_part.done")]
|
||||
ReasoningSummaryPartDone,
|
||||
#[strum(serialize = "response.reasoning_summary_text.delta")]
|
||||
ReasoningSummaryTextDelta,
|
||||
#[strum(serialize = "response.reasoning_summary_text.done")]
|
||||
ReasoningSummaryTextDone,
|
||||
#[strum(serialize = "response.reasoning_text.delta")]
|
||||
ReasoningTextDelta,
|
||||
#[strum(serialize = "response.reasoning_text.done")]
|
||||
ReasoningTextDone,
|
||||
#[strum(serialize = "response.refusal.delta")]
|
||||
RefusalDelta,
|
||||
#[strum(serialize = "response.refusal.done")]
|
||||
RefusalDone,
|
||||
#[strum(serialize = "response.shell_call_command.added")]
|
||||
ShellCallCommandAdded,
|
||||
#[strum(serialize = "response.shell_call_command.delta")]
|
||||
ShellCallCommandDelta,
|
||||
#[strum(serialize = "response.shell_call_command.done")]
|
||||
ShellCallCommandDone,
|
||||
#[strum(serialize = "response.shell_call_output_content.delta")]
|
||||
ShellCallOutputContentDelta,
|
||||
#[strum(serialize = "response.shell_call_output_content.done")]
|
||||
ShellCallOutputContentDone,
|
||||
#[strum(serialize = "response.web_search_call.completed")]
|
||||
WebSearchCallCompleted,
|
||||
#[strum(serialize = "response.web_search_call.in_progress")]
|
||||
WebSearchCallInProgress,
|
||||
#[strum(serialize = "response.web_search_call.searching")]
|
||||
WebSearchCallSearching,
|
||||
#[strum(default, transparent)]
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl ResponsesEventType {
|
||||
/// The last event of a response. `response.completed`, `response.incomplete` and
|
||||
/// `response.failed` each close the response and carry its final body; `error` aborts it.
|
||||
pub fn is_terminal(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::ResponseCompleted | Self::ResponseFailed | Self::ResponseIncomplete | Self::Error
|
||||
)
|
||||
}
|
||||
|
||||
/// A terminal event after which the response is not a success.
|
||||
pub fn is_failure(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::ResponseFailed | Self::ResponseIncomplete | Self::Error
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
pub struct ResponsesStreamEvent {
|
||||
#[serde(rename = "type")]
|
||||
pub event_type: ResponsesEventType,
|
||||
#[serde(flatten)]
|
||||
pub data: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl ResponsesStreamEvent {
|
||||
pub fn model(&self) -> Option<&str> {
|
||||
let model = self.data.get("model").and_then(Value::as_str);
|
||||
if model.is_some() {
|
||||
return model;
|
||||
}
|
||||
self.data
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|response| response.get("model"))
|
||||
.and_then(Value::as_str)
|
||||
}
|
||||
|
||||
pub fn is_response_create(&self) -> bool {
|
||||
self.event_type == ResponsesEventType::ResponseCreate
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use strum::IntoEnumIterator;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
fn every_named_event_type_round_trips() {
|
||||
let named: Vec<_> = ResponsesEventType::iter()
|
||||
.filter(|event_type| !matches!(event_type, ResponsesEventType::Other(_)))
|
||||
.collect();
|
||||
assert!(named.len() > 50);
|
||||
for event_type in named {
|
||||
let wire = event_type.to_string();
|
||||
assert_eq!(wire, event_type.as_ref());
|
||||
assert_eq!(serde_json::to_value(&event_type).unwrap(), json!(wire));
|
||||
assert_eq!(wire.parse::<ResponsesEventType>().unwrap(), event_type);
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ResponsesEventType>(json!(wire)).unwrap(),
|
||||
event_type
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unknown("future.event", ResponsesEventType::Other("future.event".into()))]
|
||||
#[case::empty("", ResponsesEventType::Other(String::new()))]
|
||||
#[case::case_sensitive(
|
||||
"Response.Completed",
|
||||
ResponsesEventType::Other("Response.Completed".into())
|
||||
)]
|
||||
#[case::escaped("future\"\\\n", ResponsesEventType::Other("future\"\\\n".into()))]
|
||||
fn unknown_event_types_round_trip_through_other(
|
||||
#[case] wire: &str,
|
||||
#[case] expected: ResponsesEventType,
|
||||
) {
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ResponsesEventType>(json!(wire)).unwrap(),
|
||||
expected
|
||||
);
|
||||
assert_eq!(serde_json::to_value(&expected).unwrap(), json!(wire));
|
||||
assert_eq!(expected.as_ref(), wire);
|
||||
assert_eq!(expected.to_string(), wire);
|
||||
}
|
||||
|
||||
/// Python's `ResponsesAPIStreamingIterator` stores the response on `response.completed`,
|
||||
/// `response.incomplete` and `response.failed`, and raises on `response.failed` and `error`.
|
||||
#[rstest]
|
||||
#[case::completed(ResponsesEventType::ResponseCompleted, true, false)]
|
||||
#[case::incomplete(ResponsesEventType::ResponseIncomplete, true, true)]
|
||||
#[case::failed(ResponsesEventType::ResponseFailed, true, true)]
|
||||
#[case::error(ResponsesEventType::Error, true, true)]
|
||||
#[case::create(ResponsesEventType::ResponseCreate, false, false)]
|
||||
#[case::created(ResponsesEventType::ResponseCreated, false, false)]
|
||||
#[case::in_progress(ResponsesEventType::ResponseInProgress, false, false)]
|
||||
#[case::delta(ResponsesEventType::OutputTextDelta, false, false)]
|
||||
#[case::item_done(ResponsesEventType::OutputItemDone, false, false)]
|
||||
#[case::unknown(ResponsesEventType::Other("response.future".into()), false, false)]
|
||||
#[case::tool_failed(ResponsesEventType::McpCallFailed, false, false)]
|
||||
fn terminal_and_failure_events_match_the_python_iterator(
|
||||
#[case] event_type: ResponsesEventType,
|
||||
#[case] terminal: bool,
|
||||
#[case] failure: bool,
|
||||
) {
|
||||
assert_eq!(event_type.is_terminal(), terminal);
|
||||
assert_eq!(event_type.is_failure(), failure);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::number("17")]
|
||||
#[case::boolean("true")]
|
||||
#[case::null("null")]
|
||||
#[case::array("[]")]
|
||||
#[case::object("{}")]
|
||||
fn event_types_reject_non_strings(#[case] wire: &str) {
|
||||
assert!(serde_json::from_str::<ResponsesEventType>(wire).is_err());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::text(json!({"type":"response.output_text.delta","delta":"hello","sequence_number":3,"future":null}))]
|
||||
#[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1","output_index":0}))]
|
||||
#[case::unknown(json!({"type":"future.event","payload":{"items":[1,null]},"future":true}))]
|
||||
#[case::terminal(json!({"type":"response.completed","response":{"id":"r","output":[],"future":null}}))]
|
||||
fn events_preserve_every_payload_field(#[case] wire: Value) {
|
||||
let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
|
||||
assert_eq!(serde_json::to_value(event).unwrap(), wire);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::missing_type(json!({"delta":"hello"}))]
|
||||
#[case::numeric_type(json!({"type":7}))]
|
||||
#[case::null_type(json!({"type":null}))]
|
||||
#[case::non_object(json!([]))]
|
||||
fn malformed_events_are_rejected(#[case] wire: Value) {
|
||||
assert!(serde_json::from_value::<ResponsesStreamEvent>(wire).is_err());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::flat(json!({"type":"response.create","model":"model-a"}), Some("model-a"))]
|
||||
#[case::nested(json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))]
|
||||
#[case::none(json!({"type":"response.created","response":{}}), None)]
|
||||
fn model_reads_flat_and_nested_shapes(#[case] wire: Value, #[case] expected: Option<&str>) {
|
||||
let event: ResponsesStreamEvent = serde_json::from_value(wire).unwrap();
|
||||
assert_eq!(event.model(), expected);
|
||||
}
|
||||
|
||||
#[cfg(feature = "schema")]
|
||||
#[rstest]
|
||||
fn event_type_schema_accepts_unknown_strings() {
|
||||
let schema = schemars::schema_for!(ResponsesEventType).to_value();
|
||||
assert_eq!(schema.get("type"), Some(&json!("string")));
|
||||
assert!(schema.get("enum").is_none());
|
||||
}
|
||||
}
|
||||
|
|
@ -2,68 +2,6 @@ use super::ResponsesOutputItem;
|
|||
use crate::recognized::Recognized;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
#[derive(
|
||||
Clone,
|
||||
Debug,
|
||||
PartialEq,
|
||||
Eq,
|
||||
strum::AsRefStr,
|
||||
strum::EnumString,
|
||||
strum::Display,
|
||||
serde_with::DeserializeFromStr,
|
||||
serde_with::SerializeDisplay,
|
||||
)]
|
||||
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
|
||||
#[cfg_attr(feature = "schema", schemars(with = "String"))]
|
||||
pub enum ResponsesWsEventType {
|
||||
#[strum(serialize = "response.create")]
|
||||
ResponseCreate,
|
||||
#[strum(serialize = "response.created")]
|
||||
ResponseCreated,
|
||||
#[strum(serialize = "response.completed")]
|
||||
ResponseCompleted,
|
||||
#[strum(serialize = "response.failed")]
|
||||
ResponseFailed,
|
||||
#[strum(serialize = "response.incomplete")]
|
||||
ResponseIncomplete,
|
||||
#[strum(serialize = "error")]
|
||||
Error,
|
||||
#[strum(default, transparent)]
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl ResponsesWsEventType {
|
||||
pub fn as_str(&self) -> &str {
|
||||
self.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
pub struct ResponsesWsEvent {
|
||||
#[serde(rename = "type")]
|
||||
pub event_type: ResponsesWsEventType,
|
||||
#[serde(flatten)]
|
||||
pub data: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl ResponsesWsEvent {
|
||||
pub fn model(&self) -> Option<&str> {
|
||||
let model = self.data.get("model").and_then(Value::as_str);
|
||||
if model.is_some() {
|
||||
return model;
|
||||
}
|
||||
self.data
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|response| response.get("model"))
|
||||
.and_then(Value::as_str)
|
||||
}
|
||||
|
||||
pub fn is_response_create(&self) -> bool {
|
||||
self.event_type == ResponsesWsEventType::ResponseCreate
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
#[derive(Eq)]
|
||||
pub struct ResponsesErrorFrame {
|
||||
|
|
@ -112,8 +50,7 @@ mod tests {
|
|||
use crate::formats::responses::ResponsesOutputItem;
|
||||
|
||||
use crate::{
|
||||
formats::responses::streaming_websocket::{ResponsesEventResponse, ResponsesWsEventType},
|
||||
recognized::Recognized,
|
||||
formats::responses::streaming_websocket::ResponsesEventResponse, recognized::Recognized,
|
||||
test_support::*,
|
||||
};
|
||||
|
||||
|
|
@ -134,66 +71,6 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::flat(serde_json::json!({"type":"response.create","model":"model-a"}), Some("model-a"))]
|
||||
#[case::nested(serde_json::json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))]
|
||||
fn model_reads_flat_and_nested_create_shapes(
|
||||
#[case] payload: serde_json::Value,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let event: ResponsesWsEvent = serde_json::from_value(payload).expect("valid event");
|
||||
assert_eq!(event.model(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::create("response.create", ResponsesWsEventType::ResponseCreate)]
|
||||
#[case::created("response.created", ResponsesWsEventType::ResponseCreated)]
|
||||
#[case::completed("response.completed", ResponsesWsEventType::ResponseCompleted)]
|
||||
#[case::failed("response.failed", ResponsesWsEventType::ResponseFailed)]
|
||||
#[case::incomplete("response.incomplete", ResponsesWsEventType::ResponseIncomplete)]
|
||||
#[case::error("error", ResponsesWsEventType::Error)]
|
||||
#[case::unknown(
|
||||
"response.output_text.delta",
|
||||
ResponsesWsEventType::Other("response.output_text.delta".to_string())
|
||||
)]
|
||||
#[case::empty("", ResponsesWsEventType::Other(String::new()))]
|
||||
#[case::case_sensitive(
|
||||
"Response.Completed",
|
||||
ResponsesWsEventType::Other("Response.Completed".into())
|
||||
)]
|
||||
#[case::escaped("future\"\\\n", ResponsesWsEventType::Other("future\"\\\n".into()))]
|
||||
fn websocket_event_type_round_trips(
|
||||
#[case] wire: &str,
|
||||
#[case] expected: ResponsesWsEventType,
|
||||
) {
|
||||
let serialized = serde_json::to_string(&expected).unwrap();
|
||||
assert_eq!(serialized, serde_json::to_string(wire).unwrap());
|
||||
assert_eq!(
|
||||
serde_json::from_str::<ResponsesWsEventType>(&serialized).unwrap(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::number("17")]
|
||||
#[case::boolean("true")]
|
||||
#[case::null("null")]
|
||||
#[case::array("[]")]
|
||||
#[case::object("{}")]
|
||||
fn websocket_event_type_rejects_non_strings(#[case] wire: &str) {
|
||||
assert!(serde_json::from_str::<ResponsesWsEventType>(wire).is_err());
|
||||
}
|
||||
|
||||
#[cfg(feature = "schema")]
|
||||
#[rstest]
|
||||
fn websocket_event_type_schema_is_open_string() {
|
||||
let schema = schemars::schema_for!(ResponsesWsEventType);
|
||||
assert_eq!(
|
||||
schema.to_value().get("type"),
|
||||
Some(&serde_json::json!("string"))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn nested_event_response_exposes_typed_output_and_preserves_extensions() {
|
||||
let wire = json!({
|
||||
|
|
|
|||
|
|
@ -32,6 +32,13 @@
|
|||
- Provider-wide docs (auth, errors, regions): `src/<provider>/`
|
||||
- One host's endpoint docs: `src/<provider>/<format>/`
|
||||
|
||||
## Streaming
|
||||
|
||||
- Each format's `src/base_llm/<format>/streaming.rs` exposes the same three names: `event_stream` (bytes to typed events), `decode` (a whole stream to typed events) and `encode` (one typed event to the SSE frame LiteLLM's proxy sends clients). Callers use the module path, never a vendor prefix
|
||||
- SSE framing, `[DONE]` and the generic helpers live in `src/base_llm/sse.rs`. Messages frames are `event:` plus `data:`; Chat Completions and Responses frames are bare `data:`, as Python's proxy writes them
|
||||
- Decoding is strict: a malformed event, or `[DONE]` on a Messages stream, is `InvalidResponse` and ends the stream. Python never fails a stream this way; the divergence is deliberate and `src/base_llm/sse.rs` lists what Python does per format. A stream cut inside an event drops that event, as the SSE spec requires
|
||||
- What ends or fails a stream belongs on the format's event type (`is_terminal`, `is_failure`). The cache's success check in `inference/src/caching.rs` still matches type strings; move it onto the typed events rather than adding new string matches
|
||||
|
||||
## Python pairs
|
||||
|
||||
- Python paths identify counterparts but do not dictate Rust module names or class hierarchy. Preserve behavior and concepts, not structure
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ use crate::{
|
|||
Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param,
|
||||
},
|
||||
},
|
||||
messages::streaming::anthropic_sse_event_stream,
|
||||
messages::streaming::event_stream as messages_event_stream,
|
||||
},
|
||||
};
|
||||
|
||||
|
|
@ -196,7 +196,7 @@ impl BaseConfig for AnthropicConfig {
|
|||
|
||||
fn model_response_iterator(&self, shape: StreamShape) -> Option<ChatStream> {
|
||||
Some(ChatStream::new(
|
||||
anthropic_sse_event_stream,
|
||||
messages_event_stream,
|
||||
ModelResponseIterator::new(shape),
|
||||
))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,15 +1,32 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, stream::BoxStream};
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionChunk;
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
base_llm::base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream},
|
||||
base_llm::{
|
||||
base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream},
|
||||
sse::{self, is_done},
|
||||
},
|
||||
};
|
||||
|
||||
pub type ChatChunkStream = BoxStream<'static, Result<ChatCompletionChunk, Error>>;
|
||||
|
||||
/// OpenAI-compatible SSE: bare `data:` frames closed by `data: [DONE]`.
|
||||
pub fn event_stream(bytes: ByteStream) -> ChatChunkStream {
|
||||
sse::event_stream(bytes, is_done)
|
||||
}
|
||||
|
||||
pub fn decode(wire: &[u8]) -> Result<Vec<ChatCompletionChunk>, Error> {
|
||||
sse::decode(wire, is_done)
|
||||
}
|
||||
|
||||
pub fn encode(chunk: &ChatCompletionChunk) -> Result<Bytes, Error> {
|
||||
sse::encode_data(chunk)
|
||||
}
|
||||
|
||||
/// What Python's `map_openai_params` decides about the stream and `completion`
|
||||
/// hands to `ModelResponseIterator`: it is settled while the request is built,
|
||||
/// never re-derived from the body.
|
||||
|
|
@ -52,3 +69,109 @@ impl ChatStream {
|
|||
(self.run)(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
|
||||
stream::iter(pieces.into_iter().map(Ok)).boxed()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::role(json!({"id":"c","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}))]
|
||||
#[case::text(json!({"id":"c","choices":[{"index":0,"delta":{"content":"hi\nthere"}}]}))]
|
||||
#[case::tool(json!({"id":"c","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"t","type":"function","function":{"name":"f","arguments":"{\"x\":"}}]}}]}))]
|
||||
#[case::finish(json!({"id":"c","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}))]
|
||||
#[case::usage(json!({"id":"c","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"prompt_tokens_details":{"cached_tokens":0,"cache_creation_tokens":0,"text_tokens":1}}}))]
|
||||
#[case::host_extension(json!({"id":"c","choices":[],"citations":[{"url":"u"}]}))]
|
||||
fn encoded_chunks_are_bare_data_frames_that_decode_back(#[case] wire: Value) {
|
||||
let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap();
|
||||
let encoded = encode(&chunk).unwrap();
|
||||
assert!(encoded.starts_with(b"data: "));
|
||||
assert!(encoded.ends_with(b"\n\n"));
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(),
|
||||
wire
|
||||
);
|
||||
assert_eq!(decode(&encoded).unwrap(), vec![chunk]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lf("\n")]
|
||||
#[case::crlf("\r\n")]
|
||||
#[case::cr("\r")]
|
||||
#[tokio::test]
|
||||
async fn streams_split_anywhere_skip_comments_and_done(
|
||||
#[case] newline: &str,
|
||||
#[values(1, 5, 4096)] size: usize,
|
||||
) {
|
||||
let wire = format!(
|
||||
": keepalive{newline}{newline}data: {{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{\"content\":\"a\"}}}}]}}{newline}{newline}data:{{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{}},\"finish_reason\":\"stop\"}}]}}{newline}{newline}data: [DONE]{newline}{newline}"
|
||||
);
|
||||
let streamed = event_stream(in_pieces(wire.as_bytes(), size))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(streamed, decode(wire.as_bytes()).unwrap());
|
||||
assert_eq!(streamed.len(), 2);
|
||||
assert_eq!(streamed[0].choices[0].delta.content.as_deref(), Some("a"));
|
||||
assert_eq!(
|
||||
streamed[1].choices[0].finish_reason.as_deref(),
|
||||
Some("stop")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::spaced(b"data: [DONE]\n\n")]
|
||||
#[case::unspaced(b"data:[DONE]\n\n")]
|
||||
#[case::trailing_space(b"data: [DONE] \n\n")]
|
||||
fn every_done_spelling_closes_the_stream_without_a_chunk(#[case] wire: &[u8]) {
|
||||
assert_eq!(decode(wire).unwrap(), vec![]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_json(b"data: {\n\n")]
|
||||
#[case::choices_not_a_list(b"data: {\"choices\":{}}\n\n")]
|
||||
#[case::invalid_utf8(b"data: \xff\n\n")]
|
||||
#[tokio::test]
|
||||
async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
|
||||
assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
|
||||
let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
|
||||
assert!(matches!(
|
||||
streamed.last(),
|
||||
Some(Err(Error::InvalidResponse(_))) | None
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_body_error_ends_the_stream_after_delivered_chunks() {
|
||||
let mut chunks = event_stream(
|
||||
stream::iter([
|
||||
Ok(Bytes::from_static(
|
||||
b"data: {\"id\":\"c\",\"choices\":[]}\n\n",
|
||||
)),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset",
|
||||
)),
|
||||
])
|
||||
.boxed(),
|
||||
);
|
||||
assert_eq!(
|
||||
chunks.next().await.unwrap().unwrap().id.as_deref(),
|
||||
Some("c")
|
||||
);
|
||||
assert!(matches!(
|
||||
chunks.next().await,
|
||||
Some(Err(Error::InvalidResponse(_)))
|
||||
));
|
||||
assert!(chunks.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,38 +1,32 @@
|
|||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, stream::BoxStream};
|
||||
use litellm_framer::{frames, sse::SseCodec};
|
||||
use litellm_llms_types::formats::messages::streaming::MessagesStreamEvent;
|
||||
|
||||
use crate::Error;
|
||||
pub use crate::base_llm::base_model_iterator::ByteStream;
|
||||
use crate::{Error, base_llm::sse};
|
||||
|
||||
pub type EventStream = BoxStream<'static, Result<MessagesStreamEvent, Error>>;
|
||||
pub type EventStream = sse::EventStream<MessagesStreamEvent>;
|
||||
pub type StreamDecoder = fn(ByteStream) -> EventStream;
|
||||
|
||||
pub fn anthropic_sse_event_stream(bytes: ByteStream) -> EventStream {
|
||||
Box::pin(frames(bytes, SseCodec::default()).map(|event| {
|
||||
let event = event.map_err(|error| {
|
||||
Error::InvalidResponse(crate::ErrorDetail::failed("stream framing", error))
|
||||
})?;
|
||||
serde_json::from_str(&event.data).map_err(|error| {
|
||||
Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error))
|
||||
})
|
||||
}))
|
||||
fn no_control_frames(_: &litellm_framer::sse::SseEvent) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
pub fn encode_anthropic_sse(event: &MessagesStreamEvent) -> Result<Bytes, Error> {
|
||||
let data = serde_json::to_value(event).map_err(|error| {
|
||||
Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error))
|
||||
})?;
|
||||
let name = data
|
||||
.get("type")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.ok_or_else(|| {
|
||||
Error::InvalidResponse(
|
||||
"Anthropic stream event is invalid: stream event has no type".into(),
|
||||
)
|
||||
})?;
|
||||
Ok(Bytes::from(format!("event: {name}\ndata: {data}\n\n")))
|
||||
pub fn event_stream(bytes: ByteStream) -> EventStream {
|
||||
sse::event_stream(bytes, no_control_frames)
|
||||
}
|
||||
|
||||
pub fn decode(wire: &[u8]) -> Result<Vec<MessagesStreamEvent>, Error> {
|
||||
sse::decode(wire, no_control_frames)
|
||||
}
|
||||
|
||||
/// Anthropic names every frame after the event's `type`, and LiteLLM's proxy does the same
|
||||
/// when it re-emits translated events.
|
||||
pub fn encode(event: &MessagesStreamEvent) -> Result<Bytes, Error> {
|
||||
let data = serde_json::to_string(event).map_err(sse::invalid)?;
|
||||
Ok(Bytes::from(format!(
|
||||
"event: {}\ndata: {data}\n\n",
|
||||
event.as_ref()
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -41,6 +35,7 @@ mod tests {
|
|||
use litellm_llms_types::formats::messages::streaming::{
|
||||
MessagesContentBlockDelta, MessagesStreamUsage,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
|
@ -48,51 +43,11 @@ mod tests {
|
|||
const TEXT_DELTA: &str =
|
||||
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
|
||||
|
||||
fn in_pieces(wire: &[u8]) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(3).map(Bytes::copy_from_slice).collect();
|
||||
fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
|
||||
stream::iter(pieces.into_iter().map(Ok)).boxed()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sse_frames_split_anywhere_decode_into_typed_events() {
|
||||
let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n");
|
||||
let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes()))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![MessagesStreamEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: MessagesContentBlockDelta::TextDelta {
|
||||
text: "hello".into(),
|
||||
},
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decodes_citations_delta_events() {
|
||||
let wire = concat!(
|
||||
"event: content_block_delta\n",
|
||||
r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#,
|
||||
"\n\n",
|
||||
);
|
||||
let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes()))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
events.as_slice(),
|
||||
[MessagesStreamEvent::ContentBlockDelta {
|
||||
delta: MessagesContentBlockDelta::Citations { .. },
|
||||
..
|
||||
}]
|
||||
));
|
||||
}
|
||||
|
||||
fn events() -> Vec<MessagesStreamEvent> {
|
||||
vec![
|
||||
MessagesStreamEvent::Ping,
|
||||
|
|
@ -110,27 +65,80 @@ mod tests {
|
|||
]
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn encoded_events_decode_back_to_themselves() {
|
||||
let wire = events()
|
||||
.iter()
|
||||
.map(encode_anthropic_sse)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.unwrap();
|
||||
fn wire(newline: &str) -> String {
|
||||
format!(
|
||||
"event: content_block_delta{newline}data: {TEXT_DELTA}{newline}{newline}: keepalive{newline}{newline}event: message_stop{newline}data: {{\"type\":\"message_stop\"}}{newline}{newline}"
|
||||
)
|
||||
}
|
||||
|
||||
let decoded = anthropic_sse_event_stream(stream::iter(wire.into_iter().map(Ok)).boxed())
|
||||
#[rstest]
|
||||
#[case::lf("\n")]
|
||||
#[case::crlf("\r\n")]
|
||||
#[case::cr("\r")]
|
||||
#[tokio::test]
|
||||
async fn frames_split_anywhere_decode_into_typed_events(
|
||||
#[case] newline: &str,
|
||||
#[values(1, 3, 4096)] size: usize,
|
||||
) {
|
||||
let wire = wire(newline);
|
||||
let expected = vec![
|
||||
MessagesStreamEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: MessagesContentBlockDelta::TextDelta {
|
||||
text: "hello".into(),
|
||||
},
|
||||
},
|
||||
MessagesStreamEvent::MessageStop { usage: None },
|
||||
];
|
||||
let streamed = event_stream(in_pieces(wire.as_bytes(), size))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(streamed, expected);
|
||||
assert_eq!(decode(wire.as_bytes()).unwrap(), expected);
|
||||
}
|
||||
|
||||
assert_eq!(decoded, events());
|
||||
#[tokio::test]
|
||||
async fn decodes_citations_delta_events() {
|
||||
let wire = concat!(
|
||||
"event: content_block_delta\n",
|
||||
r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#,
|
||||
"\n\n",
|
||||
);
|
||||
let events = event_stream(in_pieces(wire.as_bytes(), 3))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
events.as_slice(),
|
||||
[MessagesStreamEvent::ContentBlockDelta {
|
||||
delta: MessagesContentBlockDelta::Citations { .. },
|
||||
..
|
||||
}]
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn encoded_events_decode_back_to_themselves() {
|
||||
let wire: Vec<u8> = events()
|
||||
.iter()
|
||||
.map(encode)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect();
|
||||
let streamed = event_stream(in_pieces(&wire, 3))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(streamed, events());
|
||||
assert_eq!(decode(&wire).unwrap(), events());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_event_is_named_by_its_type() {
|
||||
let encoded =
|
||||
encode_anthropic_sse(&MessagesStreamEvent::MessageStop { usage: None }).unwrap();
|
||||
|
||||
let encoded = encode(&MessagesStreamEvent::MessageStop { usage: None }).unwrap();
|
||||
assert_eq!(
|
||||
encoded,
|
||||
Bytes::from(format!(
|
||||
|
|
@ -139,4 +147,68 @@ mod tests {
|
|||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::json(b"data: {\n\n")]
|
||||
#[case::unknown_event(b"data: {\"type\":\"future_event\"}\n\n")]
|
||||
#[case::malformed_tool(b"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":3}}\n\n")]
|
||||
#[case::utf8(b"data: \xff\n\n")]
|
||||
#[case::done_is_not_an_anthropic_frame(b"data: [DONE]\n\n")]
|
||||
#[tokio::test]
|
||||
async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
|
||||
assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
|
||||
let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
|
||||
assert!(matches!(
|
||||
streamed.last(),
|
||||
Some(Err(Error::InvalidResponse(_))) | None
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn decoding_keeps_tool_deltas_and_block_extensions() {
|
||||
let wire = concat!(
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"t\",\"name\":\"lookup\",\"input\":{},\"future\":null}}\n\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"x\\\":\"}}\n\n",
|
||||
);
|
||||
let decoded = decode(wire.as_bytes()).unwrap();
|
||||
let [
|
||||
MessagesStreamEvent::ContentBlockStart { content_block, .. },
|
||||
MessagesStreamEvent::ContentBlockDelta {
|
||||
delta: MessagesContentBlockDelta::InputJsonDelta { partial_json },
|
||||
..
|
||||
},
|
||||
] = decoded.as_slice()
|
||||
else {
|
||||
panic!("expected tool events");
|
||||
};
|
||||
assert_eq!(
|
||||
content_block.extra.get("future"),
|
||||
Some(&serde_json::Value::Null)
|
||||
);
|
||||
assert_eq!(partial_json, "{\"x\":");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_body_error_ends_the_stream_after_delivered_events() {
|
||||
let mut events = event_stream(
|
||||
stream::iter([
|
||||
Ok(Bytes::from_static(b"data: {\"type\":\"ping\"}\n\n")),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset",
|
||||
)),
|
||||
])
|
||||
.boxed(),
|
||||
);
|
||||
assert_eq!(
|
||||
events.next().await.unwrap().unwrap(),
|
||||
MessagesStreamEvent::Ping
|
||||
);
|
||||
assert!(matches!(
|
||||
events.next().await,
|
||||
Some(Err(Error::InvalidResponse(_)))
|
||||
));
|
||||
assert!(events.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,3 +5,4 @@ pub mod chat;
|
|||
pub mod messages;
|
||||
pub mod ocr;
|
||||
pub mod responses;
|
||||
pub mod sse;
|
||||
|
|
|
|||
|
|
@ -1 +1,2 @@
|
|||
pub mod streaming;
|
||||
pub mod transformation;
|
||||
|
|
|
|||
135
litellm-rust/crates/llms/src/base_llm/responses/streaming.rs
Normal file
135
litellm-rust/crates/llms/src/base_llm/responses/streaming.rs
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
use bytes::Bytes;
|
||||
use litellm_llms_types::formats::responses::streaming::ResponsesStreamEvent;
|
||||
|
||||
pub use crate::base_llm::base_model_iterator::ByteStream;
|
||||
use crate::{
|
||||
Error,
|
||||
base_llm::sse::{self, is_done},
|
||||
};
|
||||
|
||||
pub type EventStream = sse::EventStream<ResponsesStreamEvent>;
|
||||
|
||||
pub fn event_stream(bytes: ByteStream) -> EventStream {
|
||||
sse::event_stream(bytes, is_done)
|
||||
}
|
||||
|
||||
pub fn decode(wire: &[u8]) -> Result<Vec<ResponsesStreamEvent>, Error> {
|
||||
sse::decode(wire, is_done)
|
||||
}
|
||||
|
||||
/// A bare `data:` frame. Python's proxy serializes every Responses event this way and never
|
||||
/// writes an `event:` line, so clients see one shape whichever runtime served them.
|
||||
pub fn encode(event: &ResponsesStreamEvent) -> Result<Bytes, Error> {
|
||||
sse::encode_data(event)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use litellm_llms_types::formats::responses::streaming::ResponsesEventType;
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
|
||||
stream::iter(pieces.into_iter().map(Ok)).boxed()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::created(json!({"type":"response.created","response":{"id":"r","output":[]}}))]
|
||||
#[case::text(json!({"type":"response.output_text.delta","delta":"hello\nworld","sequence_number":3}))]
|
||||
#[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1"}))]
|
||||
#[case::unknown(json!({"type":"future.event","payload":[1,null],"future":{"nested":true}}))]
|
||||
#[case::completed(json!({"type":"response.completed","response":{"id":"r","usage":{"input_tokens":3}}}))]
|
||||
#[case::failed(json!({"type":"response.failed","response":{"error":{"code":"failure"}}}))]
|
||||
#[case::incomplete(json!({"type":"response.incomplete","response":{"incomplete_details":{"reason":"limit"}}}))]
|
||||
#[case::error(json!({"type":"error","error":{"message":"boom"}}))]
|
||||
fn encoded_events_are_bare_data_frames_that_decode_back(#[case] wire: Value) {
|
||||
let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
|
||||
let encoded = encode(&event).unwrap();
|
||||
assert!(encoded.starts_with(b"data: "));
|
||||
assert!(encoded.ends_with(b"\n\n"));
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(),
|
||||
wire
|
||||
);
|
||||
assert_eq!(decode(&encoded).unwrap(), vec![event]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lf("\n")]
|
||||
#[case::crlf("\r\n")]
|
||||
#[case::cr("\r")]
|
||||
#[tokio::test]
|
||||
async fn streams_split_anywhere_skip_comments_and_done(
|
||||
#[case] newline: &str,
|
||||
#[values(1, 5, 4096)] size: usize,
|
||||
) {
|
||||
let wire = format!(
|
||||
"\u{feff}event: response.created{newline}data: {{\"type\":\"response.created\"}}{newline}{newline}: keepalive{newline}{newline}event: future.event{newline}data: {{\"type\":\"future.event\",\"payload\":null}}{newline}{newline}data: [DONE]{newline}{newline}"
|
||||
);
|
||||
let streamed = event_stream(in_pieces(wire.as_bytes(), size))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(streamed, decode(wire.as_bytes()).unwrap());
|
||||
assert_eq!(streamed.len(), 2);
|
||||
assert_eq!(streamed[0].event_type, ResponsesEventType::ResponseCreated);
|
||||
assert_eq!(
|
||||
streamed[1].event_type,
|
||||
ResponsesEventType::Other("future.event".into())
|
||||
);
|
||||
assert_eq!(streamed[1].data.get("payload"), Some(&Value::Null));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::spaced(b"data: [DONE]\n\n")]
|
||||
#[case::unspaced(b"data:[DONE]\n\n")]
|
||||
#[case::trailing_space(b"data: [DONE] \n\n")]
|
||||
fn every_done_spelling_closes_the_stream_without_an_event(#[case] wire: &[u8]) {
|
||||
assert_eq!(decode(wire).unwrap(), vec![]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_json(b"data: {\n\n")]
|
||||
#[case::missing_type(b"data: {}\n\n")]
|
||||
#[case::numeric_type(b"data: {\"type\":3}\n\n")]
|
||||
#[case::invalid_utf8(b"data: \xff\n\n")]
|
||||
#[tokio::test]
|
||||
async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
|
||||
assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
|
||||
let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
|
||||
assert!(matches!(
|
||||
streamed.last(),
|
||||
Some(Err(Error::InvalidResponse(_))) | None
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_body_error_ends_the_stream_after_delivered_events() {
|
||||
let mut events = event_stream(
|
||||
stream::iter([
|
||||
Ok(Bytes::from_static(
|
||||
b"data: {\"type\":\"response.created\"}\n\n",
|
||||
)),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset",
|
||||
)),
|
||||
])
|
||||
.boxed(),
|
||||
);
|
||||
assert_eq!(
|
||||
events.next().await.unwrap().unwrap().event_type,
|
||||
ResponsesEventType::ResponseCreated
|
||||
);
|
||||
assert!(matches!(
|
||||
events.next().await,
|
||||
Some(Err(Error::InvalidResponse(_)))
|
||||
));
|
||||
assert!(events.next().await.is_none());
|
||||
}
|
||||
}
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
use litellm_llms_types::formats::responses::{
|
||||
ResponsesApiResponse, streaming_websocket::ResponsesWsEvent,
|
||||
ResponsesApiResponse, streaming::ResponsesStreamEvent,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
|
@ -38,11 +38,11 @@ pub trait BaseResponsesApiConfig: Sync {
|
|||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ResponsesWsTransformResult {
|
||||
pub events: Vec<ResponsesWsEvent>,
|
||||
pub events: Vec<ResponsesStreamEvent>,
|
||||
}
|
||||
|
||||
impl ResponsesWsTransformResult {
|
||||
pub fn passthrough(event: ResponsesWsEvent) -> Self {
|
||||
pub fn passthrough(event: ResponsesStreamEvent) -> Self {
|
||||
Self {
|
||||
events: vec![event],
|
||||
}
|
||||
|
|
@ -58,18 +58,18 @@ pub trait ResponsesWebSocketProviderConfig: Sync {
|
|||
|
||||
fn transform_ws_request(
|
||||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
event: &ResponsesStreamEvent,
|
||||
model: &str,
|
||||
) -> Result<ResponsesWsTransformResult, Error>;
|
||||
|
||||
fn transform_ws_response(
|
||||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
event: &ResponsesStreamEvent,
|
||||
model: &str,
|
||||
) -> Result<ResponsesWsTransformResult, Error>;
|
||||
}
|
||||
|
||||
pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent {
|
||||
pub fn enforce_model(event: &ResponsesStreamEvent, model: &str) -> ResponsesStreamEvent {
|
||||
if !event.is_response_create() {
|
||||
return event.clone();
|
||||
}
|
||||
|
|
@ -103,7 +103,7 @@ pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent
|
|||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn event(value: serde_json::Value) -> ResponsesWsEvent {
|
||||
fn event(value: serde_json::Value) -> ResponsesStreamEvent {
|
||||
serde_json::from_value(value).expect("valid event")
|
||||
}
|
||||
|
||||
|
|
|
|||
102
litellm-rust/crates/llms/src/base_llm/sse.rs
Normal file
102
litellm-rust/crates/llms/src/base_llm/sse.rs
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
//! SSE decoding and encoding shared by every format's `streaming` module.
|
||||
//!
|
||||
//! DIVERGENCE: decoding is strict. An event whose `data` is not a valid event of the format is
|
||||
//! `InvalidResponse` and ends the stream, and so is `[DONE]` on a Messages stream. Python
|
||||
//! never fails a stream on a malformed event:
|
||||
//! - Chat Completions (`BaseModelResponseIterator`) and Responses
|
||||
//! (`BaseResponsesAPIStreamingIterator`) skip the event and keep going
|
||||
//! - Messages (`BaseAnthropicMessagesStreamingIterator`) forwards the provider's bytes to the
|
||||
//! client without parsing them
|
||||
//! - Anthropic to Chat (`anthropic/chat/handler.py` `ModelResponseIterator`) buffers the invalid
|
||||
//! lines on the guess that one payload was split across them
|
||||
//!
|
||||
//! Skipping loses content or a tool-call fragment with nothing to show for it, so Rust reports it.
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, TryStreamExt, future::ready, stream::BoxStream};
|
||||
use litellm_framer::{
|
||||
frames,
|
||||
sse::{SseCodec, SseEvent, decode_all},
|
||||
};
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
|
||||
use crate::{Error, ErrorDetail, base_llm::base_model_iterator::ByteStream};
|
||||
|
||||
pub type EventStream<E> = BoxStream<'static, Result<E, Error>>;
|
||||
|
||||
/// Which frames are control frames rather than events of the format.
|
||||
pub type IsControl = fn(&SseEvent) -> bool;
|
||||
|
||||
/// OpenAI's end-of-stream marker on Chat Completions and Responses streams. It closes the
|
||||
/// stream and is not an event; Anthropic Messages streams never carry it.
|
||||
pub const DONE_FRAME: &[u8] = b"data: [DONE]\n\n";
|
||||
|
||||
pub fn is_done(event: &SseEvent) -> bool {
|
||||
event.data.trim() == "[DONE]"
|
||||
}
|
||||
|
||||
pub fn event_stream<E: DeserializeOwned + Send + 'static>(
|
||||
bytes: ByteStream,
|
||||
is_control: IsControl,
|
||||
) -> EventStream<E> {
|
||||
frames(bytes, SseCodec::default())
|
||||
.map_err(framing)
|
||||
.try_filter(move |event| ready(!is_control(event)))
|
||||
.and_then(|event| ready(parse(&event)))
|
||||
.boxed()
|
||||
}
|
||||
|
||||
/// Decodes a whole stream at once. An event the stream ends inside is dropped.
|
||||
pub fn decode<E: DeserializeOwned>(wire: &[u8], is_control: IsControl) -> Result<Vec<E>, Error> {
|
||||
decode_all(wire)
|
||||
.map_err(framing)?
|
||||
.iter()
|
||||
.filter(|event| !is_control(event))
|
||||
.map(parse)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// A bare `data:` frame, which is what LiteLLM's proxy sends to clients on OpenAI formats.
|
||||
pub fn encode_data<E: Serialize>(event: &E) -> Result<Bytes, Error> {
|
||||
let data = serde_json::to_string(event).map_err(invalid)?;
|
||||
Ok(Bytes::from(format!("data: {data}\n\n")))
|
||||
}
|
||||
|
||||
fn parse<E: DeserializeOwned>(event: &SseEvent) -> Result<E, Error> {
|
||||
serde_json::from_str(&event.data).map_err(invalid)
|
||||
}
|
||||
|
||||
pub fn invalid(error: serde_json::Error) -> Error {
|
||||
Error::InvalidResponse(ErrorDetail::invalid("stream event", error))
|
||||
}
|
||||
|
||||
fn framing(error: litellm_framer::SseError) -> Error {
|
||||
Error::InvalidResponse(ErrorDetail::failed("stream framing", error))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::exact("[DONE]", true)]
|
||||
#[case::padded(" [DONE] ", true)]
|
||||
#[case::json("{\"type\":\"x\"}", false)]
|
||||
#[case::mentions("[DONE] was sent", false)]
|
||||
#[case::empty("", false)]
|
||||
fn done_is_an_exact_marker(#[case] data: &str, #[case] done: bool) {
|
||||
let event = SseEvent {
|
||||
data: data.into(),
|
||||
..SseEvent::default()
|
||||
};
|
||||
assert_eq!(is_done(&event), done);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn the_done_frame_we_send_decodes_as_a_control_frame() {
|
||||
let events: Vec<serde_json::Value> = decode(DONE_FRAME, is_done).unwrap();
|
||||
assert!(events.is_empty());
|
||||
}
|
||||
}
|
||||
|
|
@ -102,7 +102,7 @@ mod tests {
|
|||
use litellm_llms_types::formats::messages::streaming::MessagesContentBlockDelta;
|
||||
|
||||
use super::*;
|
||||
use crate::base_llm::messages::streaming::anthropic_sse_event_stream;
|
||||
use crate::base_llm::messages::streaming;
|
||||
|
||||
fn in_pieces(wire: &[u8]) -> ByteStream {
|
||||
let pieces: Vec<Bytes> = wire.chunks(3).map(Bytes::copy_from_slice).collect();
|
||||
|
|
@ -147,7 +147,7 @@ mod tests {
|
|||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let from_sse = anthropic_sse_event_stream(in_pieces(sse.as_bytes()))
|
||||
let from_sse = streaming::event_stream(in_pieces(sse.as_bytes()))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
|
|
|||
|
|
@ -332,7 +332,7 @@ mod tests {
|
|||
use litellm_auth_aws::constants::DEFAULT_BEDROCK_REGION;
|
||||
|
||||
use super::*;
|
||||
use crate::base_llm::messages::streaming::encode_anthropic_sse;
|
||||
use crate::base_llm::messages::streaming;
|
||||
|
||||
fn event(value: Value) -> MessagesStreamEvent {
|
||||
serde_json::from_value(value).unwrap()
|
||||
|
|
@ -493,7 +493,7 @@ mod tests {
|
|||
.boxed();
|
||||
|
||||
let sse = bedrock_anthropic_messages_event_stream(bytes)
|
||||
.map_ok(|event| encode_anthropic_sse(&event).unwrap())
|
||||
.map_ok(|event| streaming::encode(&event).unwrap())
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap()
|
||||
|
|
@ -508,7 +508,7 @@ mod tests {
|
|||
)),
|
||||
]
|
||||
.iter()
|
||||
.flat_map(|event| encode_anthropic_sse(event).unwrap())
|
||||
.flat_map(|event| streaming::encode(event).unwrap())
|
||||
.collect();
|
||||
assert_eq!(sse, expected);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
use litellm_llms_types::formats::responses::{
|
||||
ResponsesApiResponse, streaming_websocket::ResponsesWsEvent,
|
||||
ResponsesApiResponse, streaming::ResponsesStreamEvent,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -34,7 +34,7 @@ impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig {
|
|||
|
||||
fn transform_ws_request(
|
||||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
event: &ResponsesStreamEvent,
|
||||
model: &str,
|
||||
) -> Result<ResponsesWsTransformResult, Error> {
|
||||
Ok(ResponsesWsTransformResult::passthrough(enforce_model(
|
||||
|
|
@ -44,7 +44,7 @@ impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig {
|
|||
|
||||
fn transform_ws_response(
|
||||
&self,
|
||||
event: &ResponsesWsEvent,
|
||||
event: &ResponsesStreamEvent,
|
||||
_model: &str,
|
||||
) -> Result<ResponsesWsTransformResult, Error> {
|
||||
Ok(ResponsesWsTransformResult::passthrough(event.clone()))
|
||||
|
|
@ -230,7 +230,7 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn openai_config_is_native_and_enforces_model() {
|
||||
let event: ResponsesWsEvent =
|
||||
let event: ResponsesStreamEvent =
|
||||
serde_json::from_value(serde_json::json!({"type":"response.create"}))
|
||||
.expect("valid event");
|
||||
let result = OPENAI_RESPONSES_WS_CONFIG
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue