af_llm/
stream.rs

1//! OpenAI-compatible SSE parsing for streaming chat completions.
2
3use serde::Deserialize;
4
5use crate::error::{LlmError, Result};
6use crate::types::{AssistantBlock, FinishReason, FunctionCall, ToolCall, Usage};
7
8/// One cumulative content update while a completion is streaming.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct StreamDelta {
11    pub content: String,
12    pub has_tool_calls: bool,
13}
14
15#[derive(Debug, Default)]
16pub(crate) struct StreamAssembler {
17    id: String,
18    content: String,
19    tool_slots: Vec<Option<ToolCallBuilder>>,
20    usage: Option<Usage>,
21    finish_reason: Option<FinishReason>,
22    output_blocks: Vec<AssistantBlock>,
23    saw_choice: bool,
24}
25
26#[derive(Debug, Default, Clone)]
27struct ToolCallBuilder {
28    id: String,
29    kind: String,
30    name: String,
31    arguments: String,
32}
33
34impl StreamAssembler {
35    pub fn apply_json(&mut self, data: &str) -> Result<Option<StreamDelta>> {
36        let chunk: StreamChunk = serde_json::from_str(data)?;
37        if !chunk.id.is_empty() {
38            self.id = chunk.id;
39        }
40        if let Some(usage) = chunk.usage {
41            self.usage = Some(usage);
42        }
43
44        let mut content_changed = false;
45        for choice in chunk.choices {
46            if choice.index != 0 {
47                continue;
48            }
49            self.saw_choice = true;
50            if let Some(reason) = choice.finish_reason {
51                self.finish_reason = Some(reason);
52            }
53            let Some(delta) = choice.delta else {
54                continue;
55            };
56            self.output_blocks.extend(delta.output_blocks);
57            if let Some(piece) = delta.content.filter(|piece| !piece.is_empty()) {
58                self.content.push_str(&piece);
59                content_changed = true;
60            }
61            for tool_delta in delta.tool_calls.unwrap_or_default() {
62                self.merge_tool_delta(tool_delta);
63            }
64        }
65
66        Ok(content_changed.then(|| StreamDelta {
67            content: self.content.clone(),
68            has_tool_calls: self.has_tool_calls(),
69        }))
70    }
71
72    fn has_tool_calls(&self) -> bool {
73        self.tool_slots
74            .iter()
75            .flatten()
76            .any(|tool| !tool.id.is_empty() || !tool.name.is_empty() || !tool.arguments.is_empty())
77    }
78
79    fn merge_tool_delta(&mut self, delta: ToolCallDelta) {
80        let index = delta.index as usize;
81        if self.tool_slots.len() <= index {
82            self.tool_slots.resize_with(index + 1, || None);
83        }
84        let tool = self.tool_slots[index].get_or_insert_with(ToolCallBuilder::default);
85        if let Some(id) = delta.id.filter(|value| !value.is_empty()) {
86            tool.id = id;
87        }
88        if let Some(kind) = delta.kind.filter(|value| !value.is_empty()) {
89            tool.kind = kind;
90        }
91        if let Some(function) = delta.function {
92            if let Some(name) = function.name.filter(|value| !value.is_empty()) {
93                merge_tool_name(&mut tool.name, &name);
94            }
95            if let Some(arguments) = function.arguments {
96                tool.arguments.push_str(&arguments);
97            }
98        }
99    }
100
101    pub fn finish(self) -> Result<crate::types::CompletionResponse> {
102        use crate::types::{ChatMessage, Choice, CompletionResponse, Role};
103
104        if !self.saw_choice {
105            return Err(LlmError::StreamProtocol(
106                "stream completed without choice 0".into(),
107            ));
108        }
109        if self.finish_reason.is_none() {
110            return Err(LlmError::StreamProtocol(
111                "stream completed without finish_reason".into(),
112            ));
113        }
114
115        let tool_calls = self
116            .tool_slots
117            .into_iter()
118            .flatten()
119            .filter(|tool| !tool.name.is_empty() || !tool.arguments.is_empty())
120            .map(|tool| ToolCall {
121                id: if tool.id.is_empty() {
122                    format!("call_{:08x}", stable_hash(&tool.name, &tool.arguments))
123                } else {
124                    tool.id
125                },
126                kind: if tool.kind.is_empty() {
127                    "function".into()
128                } else {
129                    tool.kind
130                },
131                function: FunctionCall {
132                    name: tool.name,
133                    arguments: tool.arguments,
134                },
135            })
136            .collect::<Vec<_>>();
137
138        Ok(CompletionResponse {
139            id: self.id,
140            choices: vec![Choice {
141                index: 0,
142                message: ChatMessage {
143                    role: Role::Assistant,
144                    content: (!self.content.is_empty()).then_some(self.content),
145                    tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
146                    tool_call_id: None,
147                    name: None,
148                },
149                finish_reason: self.finish_reason,
150                output_blocks: self.output_blocks,
151            }],
152            usage: self.usage,
153        })
154    }
155}
156
157fn stable_hash(left: &str, right: &str) -> u32 {
158    let mut hash = 0x811c_9dc5u32;
159    for byte in left.bytes().chain(right.bytes()) {
160        hash ^= u32::from(byte);
161        hash = hash.wrapping_mul(0x0100_0193);
162    }
163    hash
164}
165
166fn merge_tool_name(current: &mut String, incoming: &str) {
167    if current.is_empty() {
168        current.push_str(incoming);
169    } else if incoming.starts_with(current.as_str()) {
170        *current = incoming.to_string();
171    } else if !current.starts_with(incoming) {
172        current.push_str(incoming);
173    }
174}
175
176#[derive(Debug, Deserialize)]
177struct StreamChunk {
178    #[serde(default)]
179    id: String,
180    #[serde(default)]
181    choices: Vec<StreamChoice>,
182    #[serde(default)]
183    usage: Option<Usage>,
184}
185
186#[derive(Debug, Deserialize)]
187struct StreamChoice {
188    #[serde(default)]
189    index: u32,
190    #[serde(default)]
191    delta: Option<DeltaBody>,
192    #[serde(default)]
193    finish_reason: Option<FinishReason>,
194}
195
196#[derive(Debug, Deserialize)]
197struct DeltaBody {
198    #[serde(default)]
199    content: Option<String>,
200    #[serde(default)]
201    tool_calls: Option<Vec<ToolCallDelta>>,
202    #[serde(default, alias = "content_blocks")]
203    output_blocks: Vec<AssistantBlock>,
204}
205
206#[derive(Debug, Deserialize)]
207struct ToolCallDelta {
208    #[serde(default)]
209    index: u32,
210    #[serde(default)]
211    id: Option<String>,
212    #[serde(default, rename = "type")]
213    kind: Option<String>,
214    #[serde(default)]
215    function: Option<ToolFunctionDelta>,
216}
217
218#[derive(Debug, Deserialize)]
219struct ToolFunctionDelta {
220    #[serde(default)]
221    name: Option<String>,
222    #[serde(default)]
223    arguments: Option<String>,
224}
225
226/// Incremental SSE decoder. It keeps split UTF-8 code points and partial lines
227/// as bytes until a complete line arrives, then joins all `data:` lines in one
228/// event as required by the SSE format.
229#[derive(Debug, Default)]
230pub(crate) struct SseDecoder {
231    buffer: Vec<u8>,
232    data_lines: Vec<String>,
233}
234
235impl SseDecoder {
236    pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<String>> {
237        self.buffer.extend_from_slice(chunk);
238        let mut events = Vec::new();
239        while let Some(newline) = self.buffer.iter().position(|byte| *byte == b'\n') {
240            let mut line = self.buffer.drain(..=newline).collect::<Vec<_>>();
241            line.pop();
242            if line.last() == Some(&b'\r') {
243                line.pop();
244            }
245            let line = String::from_utf8(line)?;
246            if line.is_empty() {
247                if !self.data_lines.is_empty() {
248                    events.push(self.data_lines.join("\n"));
249                    self.data_lines.clear();
250                }
251            } else if let Some(data) = line.strip_prefix("data:") {
252                self.data_lines
253                    .push(data.strip_prefix(' ').unwrap_or(data).to_string());
254            }
255        }
256        Ok(events)
257    }
258
259    pub fn finish(self) -> Result<()> {
260        if self.buffer.is_empty() && self.data_lines.is_empty() {
261            return Ok(());
262        }
263        String::from_utf8(self.buffer)?;
264        Err(LlmError::StreamProtocol(
265            "stream ended with an incomplete SSE frame".into(),
266        ))
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273
274    #[test]
275    fn assembles_content_and_tool_calls() {
276        let mut assembler = StreamAssembler::default();
277        assert!(assembler
278            .apply_json(r#"{"id":"chatcmpl-1","choices":[{"delta":{"content":"Hi"}}]}"#,)
279            .unwrap()
280            .is_some());
281        assembler
282            .apply_json(r#"{"choices":[{"delta":{"content":" there"}}]}"#)
283            .unwrap();
284        assembler.apply_json(
285            r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"lookup_weather","arguments":""}}]}}]}"#,
286        ).unwrap();
287        assembler.apply_json(
288            r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"city\":\"Paris\"}"}}]}}]}"#,
289        ).unwrap();
290        assembler
291            .apply_json(r#"{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}"#)
292            .unwrap();
293
294        let response = assembler.finish().unwrap();
295        assert_eq!(response.first_content(), Some("Hi there"));
296        let calls = response.first_tool_calls().unwrap();
297        assert_eq!(calls[0].function.name, "lookup_weather");
298        assert_eq!(calls[0].function.arguments, r#"{"city":"Paris"}"#);
299    }
300
301    #[test]
302    fn drains_lf_and_crlf_frames() {
303        let mut decoder = SseDecoder::default();
304        assert_eq!(
305            decoder
306                .push(b"data: {\"a\":1}\r\n\r\ndata: [DONE]\n\npartial")
307                .unwrap(),
308            [r#"{"a":1}"#, "[DONE]"]
309        );
310        assert!(matches!(decoder.finish(), Err(LlmError::StreamProtocol(_))));
311    }
312
313    #[test]
314    fn joins_multiline_data_and_preserves_usage() {
315        let mut decoder = SseDecoder::default();
316        let frames = decoder
317            .push(b"data: {\"id\":\"x\",\"choices\":[\r\ndata: {\"index\":0,\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}],\r\ndata: \"usage\":{\"prompt_tokens\":2,\"completion_tokens\":1,\"total_tokens\":3}}\r\n\r\n")
318            .unwrap();
319        let mut assembler = StreamAssembler::default();
320        assembler.apply_json(&frames[0]).unwrap();
321        let response = assembler.finish().unwrap();
322        assert_eq!(response.first_content(), Some("hi"));
323        assert_eq!(response.usage.unwrap().total_tokens, 3);
324    }
325
326    #[test]
327    fn streams_typed_output_blocks_and_rejects_unknown_kinds() {
328        let mut assembler = StreamAssembler::default();
329        assembler.apply_json(r#"{"choices":[{"delta":{"output_blocks":[{"type":"citation","resource_id":"doc-1","label":"Doc","uri":"docs://doc-1"}]},"finish_reason":"stop"}]}"#).unwrap();
330        let response = assembler.finish().unwrap();
331        assert!(matches!(
332            &response.choices[0].output_blocks[0],
333            AssistantBlock::Citation { resource_id, .. } if resource_id == "doc-1"
334        ));
335        assert!(StreamAssembler::default()
336            .apply_json(
337                r#"{"choices":[{"delta":{"output_blocks":[{"type":"html","html":"bad"}]}}]}"#
338            )
339            .is_err());
340    }
341
342    #[test]
343    fn malformed_json_and_utf8_fail_closed() {
344        assert!(StreamAssembler::default().apply_json("{").is_err());
345        let mut decoder = SseDecoder::default();
346        assert!(matches!(
347            decoder.push(b"data: \xff\n\n"),
348            Err(LlmError::InvalidUtf8(_))
349        ));
350    }
351
352    #[test]
353    fn empty_and_unfinished_assemblies_fail_closed() {
354        assert!(matches!(
355            StreamAssembler::default().finish(),
356            Err(LlmError::StreamProtocol(_))
357        ));
358        let mut assembler = StreamAssembler::default();
359        assembler
360            .apply_json(r#"{"choices":[{"delta":{"content":"partial"}}]}"#)
361            .unwrap();
362        assert!(matches!(
363            assembler.finish(),
364            Err(LlmError::StreamProtocol(_))
365        ));
366    }
367
368    #[test]
369    fn repeated_tool_names_do_not_duplicate() {
370        let mut name = String::new();
371        merge_tool_name(&mut name, "fetch_report");
372        merge_tool_name(&mut name, "fetch_report");
373        assert_eq!(name, "fetch_report");
374    }
375
376    #[test]
377    fn incremental_tool_name_fragments_append() {
378        let mut name = String::new();
379        merge_tool_name(&mut name, "fetch_");
380        merge_tool_name(&mut name, "report");
381        assert_eq!(name, "fetch_report");
382    }
383}
af_llm/stream.rs

af_llm/
stream.rs
