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

af_llm/
stream.rs
