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 an Anthropic SSE stream into chat chunks.
146///
147/// This function includes protection against unbounded memory growth by limiting
148/// the internal buffer size to `MAX_SSE_BUFFER_SIZE`.
149pub fn parse_anthropic_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        while let Some(event) = events.next().await {
156            let event = event?;
157            let data = event.data.trim();
158            if data == "[DONE]" {
159                return;
160            }
161            let value: Value = serde_json::from_str(data)
162                .map_err(|e| LiteLLMError::Parse(e.to_string()))?;
163            let usage = parse_usage(&value);
164            if event.event.as_deref() == Some("content_block_delta") {
165                let text = value
166                    .pointer("/delta/text")
167                    .and_then(|v| v.as_str())
168                    .map(|s| s.to_string());
169                let reasoning = value
170                    .pointer("/delta/thinking")
171                    .and_then(|v| v.as_str())
172                    .map(|s| s.to_string())
173                    .and_then(Reasoning::from_text);
174                let content = text.unwrap_or_default();
175                if !content.is_empty() || reasoning.is_some() {
176                    yield ChatStreamChunk {
177                        content,
178                        reasoning,
179                        raw: Some(value),
180                        usage,
181                    };
182                }
183            }
184        }
185    };
186    Box::pin(s)
187}
188
189fn parse_usage(value: &Value) -> Option<Usage> {
190    let usage = value.get("usage")?.as_object()?;
191    let prompt_tokens = usage.get("prompt_tokens").and_then(|v| v.as_u64());
192    let completion_tokens = usage.get("completion_tokens").and_then(|v| v.as_u64());
193    let total_tokens = usage.get("total_tokens").and_then(|v| v.as_u64());
194    let cost_usd = usage
195        .get("cost")
196        .and_then(|v| v.as_f64())
197        .or_else(|| usage.get("cost").and_then(|v| v.as_str())?.parse().ok())
198        .or_else(|| usage.get("cost_usd").and_then(|v| v.as_f64()))
199        .or_else(|| usage.get("total_cost").and_then(|v| v.as_f64()));
200    Some(Usage {
201        prompt_tokens,
202        completion_tokens,
203        thoughts_tokens: None,
204        total_tokens,
205        cost_usd,
206    })
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212    use bytes::Bytes;
213    use futures_util::stream;
214
215    #[tokio::test]
216    async fn parse_sse_basic() {
217        let data = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\
218                    data: {\"choices\":[{\"delta\":{\"content\":\" World\"}}]}\n\n\
219                    data: [DONE]\n\n";
220        let bytes_stream = stream::iter(vec![Ok(Bytes::from(data))]);
221        let mut chat_stream = parse_sse_stream(bytes_stream);
222
223        let chunk1 = chat_stream.next().await.unwrap().unwrap();
224        assert_eq!(chunk1.content, "Hello");
225
226        let chunk2 = chat_stream.next().await.unwrap().unwrap();
227        assert_eq!(chunk2.content, " World");
228
229        assert!(chat_stream.next().await.is_none());
230    }
231
232    #[tokio::test]
233    async fn parse_anthropic_sse_basic() {
234        let data = "event: content_block_delta\n\
235                    data: {\"delta\":{\"text\":\"Hello\"}}\n\n\
236                    event: content_block_delta\n\
237                    data: {\"delta\":{\"text\":\" World\"}}\n\n";
238        let bytes_stream = stream::iter(vec![Ok(Bytes::from(data))]);
239        let mut chat_stream = parse_anthropic_sse_stream(bytes_stream);
240
241        let chunk1 = chat_stream.next().await.unwrap().unwrap();
242        assert_eq!(chunk1.content, "Hello");
243
244        let chunk2 = chat_stream.next().await.unwrap().unwrap();
245        assert_eq!(chunk2.content, " World");
246    }
247
248    #[tokio::test]
249    async fn parse_sse_handles_split_chunks() {
250        // Simulate data coming in multiple network chunks
251        let chunk1 = "data: {\"choices\":[{\"delta\":{\"con";
252        let chunk2 = "tent\":\"Split\"}}]}\n\ndata: [DONE]\n\n";
253        let bytes_stream = stream::iter(vec![Ok(Bytes::from(chunk1)), Ok(Bytes::from(chunk2))]);
254        let mut chat_stream = parse_sse_stream(bytes_stream);
255
256        let chunk = chat_stream.next().await.unwrap().unwrap();
257        assert_eq!(chunk.content, "Split");
258
259        assert!(chat_stream.next().await.is_none());
260    }
261}
litellm_rust/stream.rs

litellm_rust/
stream.rs
