litellm_rust/
stream.rs

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

litellm_rust/
stream.rs
