litellm_rust/
stream.rs

1use crate::error::{LiteLLMError, Result};
2use crate::http::MAX_SSE_BUFFER_SIZE;
3use crate::types::{Reasoning, Usage};
4use bytes::Bytes;
5use futures_util::stream::{Stream, StreamExt, TryStreamExt};
6use serde_json::Value;
7use std::pin::Pin;
8use tokio::io::{AsyncBufReadExt, BufReader};
9use tokio_util::io::StreamReader;
10
11#[derive(Debug, Clone)]
12pub struct ChatStreamChunk {
13    /// User-visible answer text for this chunk.
14    pub content: String,
15    /// Provider reasoning, kept separate from answer content.
16    pub reasoning: Option<Reasoning>,
17    pub raw: Option<Value>,
18    pub usage: Option<Usage>,
19}
20
21pub type ChatStream = Pin<Box<dyn Stream<Item = Result<ChatStreamChunk>> + Send>>;
22
23#[derive(Debug, Clone)]
24struct SseEvent {
25    event: Option<String>,
26    data: String,
27}
28
29type SseEventStream = Pin<Box<dyn Stream<Item = Result<SseEvent>> + Send>>;
30
31fn sse_event_stream<S>(stream: S) -> SseEventStream
32where
33    S: Stream<Item = std::result::Result<Bytes, reqwest::Error>> + Send + Unpin + 'static,
34{
35    let s = async_stream::try_stream! {
36        let stream = stream.map_err(std::io::Error::other);
37        let reader = StreamReader::new(stream);
38        let mut lines = BufReader::new(reader).lines();
39
40        let mut event_name: Option<String> = None;
41        let mut data_buf = String::new();
42
43        while let Some(line) = lines.next_line().await.map_err(LiteLLMError::from)? {
44            if line.is_empty() {
45                if !data_buf.is_empty() {
46                    let data = std::mem::take(&mut data_buf);
47                    let event = event_name.take();
48                    yield SseEvent { event, data };
49                } else {
50                    event_name = None;
51                }
52                continue;
53            }
54
55            if line.starts_with(':') {
56                continue;
57            }
58
59            let (field, value) = if let Some((field, value)) = line.split_once(':') {
60                (field, value.strip_prefix(' ').unwrap_or(value))
61            } else {
62                (line.as_str(), "")
63            };
64
65            match field {
66                "event" => {
67                    event_name = Some(value.to_string());
68                }
69                "data" => {
70                    if !data_buf.is_empty() {
71                        data_buf.push('\n');
72                    }
73                    data_buf.push_str(value);
74                    if data_buf.len() > MAX_SSE_BUFFER_SIZE {
75                        Err(LiteLLMError::http(format!(
76                            "SSE data buffer exceeded maximum size of {} bytes",
77                            MAX_SSE_BUFFER_SIZE
78                        )))?;
79                    }
80                }
81                _ => {}
82            }
83        }
84
85        if !data_buf.is_empty() {
86            let data = std::mem::take(&mut data_buf);
87            let event = event_name.take();
88            yield SseEvent { event, data };
89        }
90    };
91    Box::pin(s)
92}
93
94/// Parse an OpenAI-compatible SSE stream into chat chunks.
95///
96/// This function includes protection against unbounded memory growth by limiting
97/// the internal buffer size to `MAX_SSE_BUFFER_SIZE`.
98pub fn parse_sse_stream<S>(stream: S) -> ChatStream
99where
100    S: Stream<Item = std::result::Result<Bytes, reqwest::Error>> + Send + Unpin + 'static,
101{
102    let s = async_stream::try_stream! {
103        let mut events = sse_event_stream(stream);
104        while let Some(event) = events.next().await {
105            let event = event?;
106            let data = event.data.trim();
107            if data == "[DONE]" {
108                return;
109            }
110            let value: Value = serde_json::from_str(data)
111                .map_err(|e| LiteLLMError::Parse(e.to_string()))?;
112            let usage = parse_usage(&value);
113            let content = value
114                .pointer("/choices/0/delta/content")
115                .and_then(|v| v.as_str())
116                .map(|s| s.to_string());
117            let reasoning_text = value
118                .pointer("/choices/0/delta/reasoning_content")
119                .and_then(|v| v.as_str())
120                .or_else(|| {
121                    value
122                        .pointer("/choices/0/delta/reasoning")
123                        .and_then(|v| v.as_str())
124                })
125                .map(|s| s.to_string());
126            let reasoning_details = value
127                .pointer("/choices/0/delta/reasoning_details")
128                .and_then(|v| v.as_array())
129                .cloned();
130            let reasoning = reasoning_details
131                .and_then(Reasoning::from_details)
132                .or_else(|| reasoning_text.and_then(Reasoning::from_text));
133            let content = content.unwrap_or_default();
134            yield ChatStreamChunk {
135                content,
136                reasoning,
137                raw: Some(value),
138                usage,
139            };
140        }
141    };
142    Box::pin(s)
143}
144
145/// Parse Responses API SSE events into the existing text/reasoning stream.
146/// Function-call events remain available in `raw`; callers that need a final
147/// normalized `tool_calls` array should use non-streaming completion.
148pub fn parse_responses_sse_stream<S>(stream: S) -> ChatStream
149where
150    S: Stream<Item = std::result::Result<Bytes, reqwest::Error>> + Send + Unpin + 'static,
151{
152    let s = async_stream::try_stream! {
153        let mut events = sse_event_stream(stream);
154        let mut accumulated_content = String::new();
155        let mut refusal = String::new();
156        let mut usage = Usage::default();
157        while let Some(event) = events.next().await {
158            let event = event?;
159            let data = event.data.trim();
160            if data == "[DONE]" {
161                if !refusal.is_empty() {
162                    Err(LiteLLMError::Refusal { text: refusal.clone(), usage: usage.clone() })?;
163                }
164                return;
165            }
166            let value: Value = serde_json::from_str(data)
167                .map_err(|e| LiteLLMError::Parse(e.to_string()))?;
168            let kind = value.get("type").and_then(Value::as_str).or(event.event.as_deref());
169            if kind == Some("error") || value.get("type").and_then(Value::as_str) == Some("error") {
170                let message = value
171                    .get("message")
172                    .or_else(|| value.pointer("/error/message"))
173                    .and_then(Value::as_str)
174                    .unwrap_or("Responses API stream error");
175                Err(LiteLLMError::http(message))?;
176            }
177            if matches!(kind, Some("response.failed" | "response.incomplete")) {
178                if kind == Some("response.incomplete") {
179                    if let Some(incomplete_usage) = value.get("response").and_then(parse_responses_usage) {
180                        usage = incomplete_usage;
181                    }
182                }
183                let message = value
184                    .pointer("/response/error/message")
185                    .and_then(Value::as_str)
186                    .or_else(|| value.pointer("/response/incomplete_details/reason").and_then(Value::as_str))
187                    .unwrap_or("Responses API did not complete the response");
188                if kind == Some("response.incomplete") {
189                    Err(LiteLLMError::Truncated { text: accumulated_content.clone(), usage: usage.clone() })?;
190                }
191                Err(LiteLLMError::http(message))?;
192            }
193            if kind == Some("response.refusal.delta") {
194                refusal.push_str(value.get("delta").and_then(Value::as_str).unwrap_or_default());
195                continue;
196            }
197            if kind == Some("response.refusal.done") {
198                if let Some(text) = value.get("refusal").and_then(Value::as_str) {
199                    refusal = text.to_owned();
200                }
201                if !refusal.is_empty() {
202                    Err(LiteLLMError::Refusal { text: refusal.clone(), usage: usage.clone() })?;
203                }
204                continue;
205            }
206            let content = if kind == Some("response.output_text.delta") {
207                value.get("delta").and_then(Value::as_str).unwrap_or_default().to_string()
208            } else { String::new() };
209            accumulated_content.push_str(&content);
210            let reasoning = if kind == Some("response.reasoning_summary_text.delta") {
211                value.get("delta").and_then(Value::as_str).map(str::to_owned).and_then(Reasoning::from_text)
212            } else { None };
213            let event_usage = if kind == Some("response.completed") {
214                value.get("response").and_then(parse_responses_usage)
215            } else { None };
216            if let Some(completed_usage) = event_usage.clone() {
217                usage = completed_usage;
218            }
219            if kind == Some("response.completed") {
220                let response = value.get("response").unwrap_or(&Value::Null);
221                let completed_refusal = response
222                    .get("output").and_then(Value::as_array).into_iter().flatten()
223                    .flat_map(|item| item.get("content").and_then(Value::as_array).into_iter().flatten())
224                    .filter(|part| part.get("type").and_then(Value::as_str) == Some("refusal"))
225                    .filter_map(|part| part.get("refusal").and_then(Value::as_str))
226                    .collect::<String>();
227                if !completed_refusal.is_empty() {
228                    Err(LiteLLMError::Refusal { text: completed_refusal, usage: usage.clone() })?;
229                }
230                if !refusal.is_empty() {
231                    Err(LiteLLMError::Refusal { text: refusal.clone(), usage: usage.clone() })?;
232                }
233            }
234            if !content.is_empty() || reasoning.is_some() || event_usage.is_some()
235                || kind.map(|k| k.starts_with("response.function_call_arguments.") || k.starts_with("response.output_item.")).unwrap_or(false) {
236                yield ChatStreamChunk { content, reasoning, raw: Some(value), usage: event_usage };
237            }
238        }
239        if !refusal.is_empty() {
240            Err(LiteLLMError::Refusal { text: refusal, usage })?;
241        }
242    };
243    Box::pin(s)
244}
245
246fn parse_responses_usage(value: &Value) -> Option<Usage> {
247    let usage = value.get("usage")?.as_object()?;
248    Some(Usage {
249        prompt_tokens: usage.get("input_tokens").and_then(Value::as_u64),
250        completion_tokens: usage.get("output_tokens").and_then(Value::as_u64),
251        thoughts_tokens: usage
252            .get("output_tokens_details")
253            .and_then(|v| v.get("reasoning_tokens"))
254            .and_then(Value::as_u64),
255        total_tokens: usage.get("total_tokens").and_then(Value::as_u64),
256        cost_usd: None,
257    })
258}
259
260/// Parse an Anthropic SSE stream into chat chunks.
261///
262/// This function includes protection against unbounded memory growth by limiting
263/// the internal buffer size to `MAX_SSE_BUFFER_SIZE`.
264pub fn parse_anthropic_sse_stream<S>(stream: S) -> ChatStream
265where
266    S: Stream<Item = std::result::Result<Bytes, reqwest::Error>> + Send + Unpin + 'static,
267{
268    let s = async_stream::try_stream! {
269        let mut events = sse_event_stream(stream);
270        while let Some(event) = events.next().await {
271            let event = event?;
272            let data = event.data.trim();
273            if data == "[DONE]" {
274                return;
275            }
276            let value: Value = serde_json::from_str(data)
277                .map_err(|e| LiteLLMError::Parse(e.to_string()))?;
278            let usage = parse_usage(&value);
279            if event.event.as_deref() == Some("content_block_delta") {
280                let text = value
281                    .pointer("/delta/text")
282                    .and_then(|v| v.as_str())
283                    .map(|s| s.to_string());
284                let reasoning = value
285                    .pointer("/delta/thinking")
286                    .and_then(|v| v.as_str())
287                    .map(|s| s.to_string())
288                    .and_then(Reasoning::from_text);
289                let content = text.unwrap_or_default();
290                if !content.is_empty() || reasoning.is_some() {
291                    yield ChatStreamChunk {
292                        content,
293                        reasoning,
294                        raw: Some(value),
295                        usage,
296                    };
297                }
298            }
299        }
300    };
301    Box::pin(s)
302}
303
304fn parse_usage(value: &Value) -> Option<Usage> {
305    let usage = value.get("usage")?.as_object()?;
306    let prompt_tokens = usage.get("prompt_tokens").and_then(|v| v.as_u64());
307    let completion_tokens = usage.get("completion_tokens").and_then(|v| v.as_u64());
308    let total_tokens = usage.get("total_tokens").and_then(|v| v.as_u64());
309    let cost_usd = usage
310        .get("cost")
311        .and_then(|v| v.as_f64())
312        .or_else(|| usage.get("cost").and_then(|v| v.as_str())?.parse().ok())
313        .or_else(|| usage.get("cost_usd").and_then(|v| v.as_f64()))
314        .or_else(|| usage.get("total_cost").and_then(|v| v.as_f64()));
315    Some(Usage {
316        prompt_tokens,
317        completion_tokens,
318        thoughts_tokens: None,
319        total_tokens,
320        cost_usd,
321    })
322}
323
324#[cfg(test)]
325mod tests {
326    use super::*;
327    use bytes::Bytes;
328    use futures_util::stream;
329
330    #[tokio::test]
331    async fn parse_sse_basic() {
332        let data = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\
333                    data: {\"choices\":[{\"delta\":{\"content\":\" World\"}}]}\n\n\
334                    data: [DONE]\n\n";
335        let bytes_stream = stream::iter(vec![Ok(Bytes::from(data))]);
336        let mut chat_stream = parse_sse_stream(bytes_stream);
337
338        let chunk1 = chat_stream.next().await.unwrap().unwrap();
339        assert_eq!(chunk1.content, "Hello");
340
341        let chunk2 = chat_stream.next().await.unwrap().unwrap();
342        assert_eq!(chunk2.content, " World");
343
344        assert!(chat_stream.next().await.is_none());
345    }
346
347    #[tokio::test]
348    async fn parse_responses_sse_text_function_and_usage() {
349        let data = "event: response.output_text.delta\n\
350                    data: {\"type\":\"response.output_text.delta\",\"delta\":\"Hello\"}\n\n\
351                    event: response.function_call_arguments.delta\n\
352                    data: {\"type\":\"response.function_call_arguments.delta\",\"delta\":\"{}\"}\n\n\
353                    event: response.completed\n\
354                    data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"total_tokens\":8,\"output_tokens_details\":{\"reasoning_tokens\":2}}}}\n\n";
355        let mut stream = parse_responses_sse_stream(stream::iter(vec![Ok(Bytes::from(data))]));
356        assert_eq!(stream.next().await.unwrap().unwrap().content, "Hello");
357        assert_eq!(
358            stream.next().await.unwrap().unwrap().raw.unwrap()["delta"],
359            "{}"
360        );
361        let usage = stream.next().await.unwrap().unwrap().usage.unwrap();
362        assert_eq!(usage.prompt_tokens, Some(3));
363        assert_eq!(usage.thoughts_tokens, Some(2));
364    }
365
366    #[tokio::test]
367    async fn parse_responses_sse_preserves_function_items_and_fails_status_events() {
368        let data = "data: {\"type\":\"response.output_item.added\",\"item\":{\"type\":\"function_call\",\"call_id\":\"call_1\",\"name\":\"weather\"}}\n\n\
369                    data: {\"type\":\"response.function_call_arguments.done\",\"call_id\":\"call_1\",\"arguments\":\"{}\"}\n\n";
370        let mut stream = parse_responses_sse_stream(stream::iter(vec![Ok(Bytes::from(data))]));
371        assert_eq!(
372            stream.next().await.unwrap().unwrap().raw.unwrap()["item"]["call_id"],
373            "call_1"
374        );
375        assert_eq!(
376            stream.next().await.unwrap().unwrap().raw.unwrap()["arguments"],
377            "{}"
378        );
379
380        let failed = "data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"message\":\"blocked\"}}}\n\n";
381        let mut stream = parse_responses_sse_stream(stream::iter(vec![Ok(Bytes::from(failed))]));
382        assert!(stream
383            .next()
384            .await
385            .unwrap()
386            .unwrap_err()
387            .to_string()
388            .contains("blocked"));
389
390        let event_error =
391            "event: error\ndata: {\"type\":\"error\",\"message\":\"transport failed\"}\n\n";
392        let mut stream =
393            parse_responses_sse_stream(stream::iter(vec![Ok(Bytes::from(event_error))]));
394        assert!(stream
395            .next()
396            .await
397            .unwrap()
398            .unwrap_err()
399            .to_string()
400            .contains("transport failed"));
401
402        let refusal = "data: {\"type\":\"response.refusal.delta\",\"delta\":\"no\"}\n\n\
403                       data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"refusal\",\"refusal\":\"no\"}]}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n";
404        let mut stream = parse_responses_sse_stream(stream::iter(vec![Ok(Bytes::from(refusal))]));
405        match stream.next().await.unwrap().unwrap_err() {
406            LiteLLMError::Refusal { text, usage } => {
407                assert_eq!(text, "no");
408                assert_eq!(usage.total_tokens, Some(2));
409            }
410            other => panic!("unexpected error: {other:?}"),
411        }
412
413        let refusal_done = concat!(
414            "data: {\"type\":\"response.refusal.delta\",\"delta\":\"no\"}\n\n",
415            "data: {\"type\":\"response.refusal.done\",\"refusal\":\"no\"}\n\n",
416        );
417        let mut stream =
418            parse_responses_sse_stream(stream::iter(vec![Ok(Bytes::from(refusal_done))]));
419        match stream.next().await.unwrap().unwrap_err() {
420            LiteLLMError::Refusal { text, .. } => assert_eq!(text, "no"),
421            other => panic!("unexpected error: {other:?}"),
422        }
423
424        let refusal_completed = concat!(
425            "data: {\"type\":\"response.refusal.delta\",\"delta\":\"no\"}\n\n",
426            "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n",
427        );
428        let mut stream =
429            parse_responses_sse_stream(stream::iter(vec![Ok(Bytes::from(refusal_completed))]));
430        match stream.next().await.unwrap().unwrap_err() {
431            LiteLLMError::Refusal { text, usage } => {
432                assert_eq!(text, "no");
433                assert_eq!(usage.total_tokens, Some(2));
434            }
435            other => panic!("unexpected error: {other:?}"),
436        }
437
438        let refusal_eof = "data: {\"type\":\"response.refusal.delta\",\"delta\":\"no\"}\n\n";
439        let mut stream =
440            parse_responses_sse_stream(stream::iter(vec![Ok(Bytes::from(refusal_eof))]));
441        match stream.next().await.unwrap().unwrap_err() {
442            LiteLLMError::Refusal { text, .. } => assert_eq!(text, "no"),
443            other => panic!("unexpected error: {other:?}"),
444        }
445
446        let incomplete = concat!(
447            "data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n",
448            "data: {\"type\":\"response.incomplete\",\"response\":{\"status\":\"incomplete\",\"incomplete_details\":{\"reason\":\"max_output_tokens\"},\"usage\":{\"input_tokens\":4,\"output_tokens\":6,\"total_tokens\":10}}}\n\n",
449        );
450        let mut stream =
451            parse_responses_sse_stream(stream::iter(vec![Ok(Bytes::from(incomplete))]));
452        assert_eq!(stream.next().await.unwrap().unwrap().content, "partial");
453        match stream.next().await.unwrap().unwrap_err() {
454            LiteLLMError::Truncated { text, usage } => {
455                assert_eq!(text, "partial");
456                assert_eq!(usage.prompt_tokens, Some(4));
457                assert_eq!(usage.total_tokens, Some(10));
458            }
459            other => panic!("unexpected error: {other:?}"),
460        }
461    }
462
463    #[tokio::test]
464    async fn parse_anthropic_sse_basic() {
465        let data = "event: content_block_delta\n\
466                    data: {\"delta\":{\"text\":\"Hello\"}}\n\n\
467                    event: content_block_delta\n\
468                    data: {\"delta\":{\"text\":\" World\"}}\n\n";
469        let bytes_stream = stream::iter(vec![Ok(Bytes::from(data))]);
470        let mut chat_stream = parse_anthropic_sse_stream(bytes_stream);
471
472        let chunk1 = chat_stream.next().await.unwrap().unwrap();
473        assert_eq!(chunk1.content, "Hello");
474
475        let chunk2 = chat_stream.next().await.unwrap().unwrap();
476        assert_eq!(chunk2.content, " World");
477    }
478
479    #[tokio::test]
480    async fn parse_sse_handles_split_chunks() {
481        // Simulate data coming in multiple network chunks
482        let chunk1 = "data: {\"choices\":[{\"delta\":{\"con";
483        let chunk2 = "tent\":\"Split\"}}]}\n\ndata: [DONE]\n\n";
484        let bytes_stream = stream::iter(vec![Ok(Bytes::from(chunk1)), Ok(Bytes::from(chunk2))]);
485        let mut chat_stream = parse_sse_stream(bytes_stream);
486
487        let chunk = chat_stream.next().await.unwrap().unwrap();
488        assert_eq!(chunk.content, "Split");
489
490        assert!(chat_stream.next().await.is_none());
491    }
492}
litellm_rust/stream.rs

litellm_rust/
stream.rs
