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                    images: Vec::new(),
153                    role: Role::Assistant,
154                    content: (!self.content.is_empty()).then_some(self.content),
155                    tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
156                    tool_call_id: None,
157                    name: None,
158                },
159                finish_reason: self.finish_reason,
160                output_blocks: self.output_blocks,
161            }],
162            usage: self.usage,
163        })
164    }
165}
166
167fn stable_hash(left: &str, right: &str) -> u32 {
168    let mut hash = 0x811c_9dc5u32;
169    for byte in left.bytes().chain(right.bytes()) {
170        hash ^= u32::from(byte);
171        hash = hash.wrapping_mul(0x0100_0193);
172    }
173    hash
174}
175
176fn merge_tool_name(current: &mut String, incoming: &str) {
177    if current.is_empty() {
178        current.push_str(incoming);
179    } else if incoming.starts_with(current.as_str()) {
180        *current = incoming.to_string();
181    } else if !current.starts_with(incoming) {
182        current.push_str(incoming);
183    }
184}
185
186#[derive(Debug, Deserialize)]
187struct StreamChunk {
188    #[serde(default)]
189    id: String,
190    #[serde(default)]
191    choices: Vec<StreamChoice>,
192    #[serde(default)]
193    usage: Option<Usage>,
194}
195
196#[derive(Debug, Deserialize)]
197struct StreamChoice {
198    #[serde(default)]
199    index: u32,
200    #[serde(default)]
201    delta: Option<DeltaBody>,
202    #[serde(default)]
203    finish_reason: Option<FinishReason>,
204}
205
206#[derive(Debug, Deserialize)]
207struct DeltaBody {
208    #[serde(default)]
209    content: Option<String>,
210    #[serde(default)]
211    tool_calls: Option<Vec<ToolCallDelta>>,
212    #[serde(default, alias = "content_blocks")]
213    output_blocks: Vec<AssistantBlock>,
214}
215
216#[derive(Debug, Deserialize)]
217struct ToolCallDelta {
218    #[serde(default)]
219    index: u32,
220    #[serde(default)]
221    id: Option<String>,
222    #[serde(default, rename = "type")]
223    kind: Option<String>,
224    #[serde(default)]
225    function: Option<ToolFunctionDelta>,
226}
227
228#[derive(Debug, Deserialize)]
229struct ToolFunctionDelta {
230    #[serde(default)]
231    name: Option<String>,
232    #[serde(default)]
233    arguments: Option<String>,
234}
235
236/// Incremental SSE decoder. It keeps split UTF-8 code points and partial lines
237/// as bytes until a complete line arrives, then joins all `data:` lines in one
238/// event as required by the SSE format.
239#[derive(Debug, Default)]
240pub(crate) struct SseDecoder {
241    buffer: Vec<u8>,
242    data_lines: Vec<String>,
243}
244
245impl SseDecoder {
246    pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<String>> {
247        self.buffer.extend_from_slice(chunk);
248        let mut events = Vec::new();
249        while let Some(newline) = self.buffer.iter().position(|byte| *byte == b'\n') {
250            let mut line = self.buffer.drain(..=newline).collect::<Vec<_>>();
251            line.pop();
252            if line.last() == Some(&b'\r') {
253                line.pop();
254            }
255            let line = String::from_utf8(line)?;
256            if line.is_empty() {
257                if !self.data_lines.is_empty() {
258                    events.push(self.data_lines.join("\n"));
259                    self.data_lines.clear();
260                }
261            } else if let Some(data) = line.strip_prefix("data:") {
262                self.data_lines
263                    .push(data.strip_prefix(' ').unwrap_or(data).to_string());
264            }
265        }
266        Ok(events)
267    }
268
269    pub fn finish(self) -> Result<()> {
270        if self.buffer.is_empty() && self.data_lines.is_empty() {
271            return Ok(());
272        }
273        String::from_utf8(self.buffer)?;
274        Err(LlmError::StreamProtocol(
275            "stream ended with an incomplete SSE frame".into(),
276        ))
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use super::*;
283
284    #[test]
285    fn assembles_content_and_tool_calls() {
286        let mut assembler = StreamAssembler::default();
287        assert!(assembler
288            .apply_json(r#"{"id":"chatcmpl-1","choices":[{"delta":{"content":"Hi"}}]}"#,)
289            .unwrap()
290            .is_some());
291        assembler
292            .apply_json(r#"{"choices":[{"delta":{"content":" there"}}]}"#)
293            .unwrap();
294        assembler.apply_json(
295            r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"lookup_weather","arguments":""}}]}}]}"#,
296        ).unwrap();
297        assembler.apply_json(
298            r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"city\":\"Paris\"}"}}]}}]}"#,
299        ).unwrap();
300        assembler
301            .apply_json(r#"{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}"#)
302            .unwrap();
303
304        let response = assembler.finish().unwrap();
305        assert_eq!(response.first_content(), Some("Hi there"));
306        let calls = response.first_tool_calls().unwrap();
307        assert_eq!(calls[0].function.name, "lookup_weather");
308        assert_eq!(calls[0].function.arguments, r#"{"city":"Paris"}"#);
309    }
310
311    #[test]
312    fn drains_lf_and_crlf_frames() {
313        let mut decoder = SseDecoder::default();
314        assert_eq!(
315            decoder
316                .push(b"data: {\"a\":1}\r\n\r\ndata: [DONE]\n\npartial")
317                .unwrap(),
318            [r#"{"a":1}"#, "[DONE]"]
319        );
320        assert!(matches!(decoder.finish(), Err(LlmError::StreamProtocol(_))));
321    }
322
323    #[test]
324    fn joins_multiline_data_and_preserves_usage() {
325        let mut decoder = SseDecoder::default();
326        let frames = decoder
327            .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")
328            .unwrap();
329        let mut assembler = StreamAssembler::default();
330        assembler.apply_json(&frames[0]).unwrap();
331        let response = assembler.finish().unwrap();
332        assert_eq!(response.first_content(), Some("hi"));
333        assert_eq!(response.usage.unwrap().total_tokens, 3);
334    }
335
336    #[test]
337    fn streams_typed_output_blocks_and_rejects_unknown_kinds() {
338        let mut assembler = StreamAssembler::default();
339        assembler.apply_json(r#"{"choices":[{"delta":{"output_blocks":[{"type":"citation","resource_id":"doc-1","label":"Doc","uri":"docs://doc-1"}]},"finish_reason":"stop"}]}"#).unwrap();
340        let response = assembler.finish().unwrap();
341        assert!(matches!(
342            &response.choices[0].output_blocks[0],
343            AssistantBlock::Citation { resource_id, .. } if resource_id == "doc-1"
344        ));
345        assert!(StreamAssembler::default()
346            .apply_json(
347                r#"{"choices":[{"delta":{"output_blocks":[{"type":"html","html":"bad"}]}}]}"#
348            )
349            .is_err());
350    }
351
352    #[test]
353    fn malformed_json_and_utf8_fail_closed() {
354        assert!(StreamAssembler::default().apply_json("{").is_err());
355        let mut decoder = SseDecoder::default();
356        assert!(matches!(
357            decoder.push(b"data: \xff\n\n"),
358            Err(LlmError::InvalidUtf8(_))
359        ));
360    }
361
362    #[test]
363    fn empty_and_unfinished_assemblies_fail_closed() {
364        assert!(matches!(
365            StreamAssembler::default().finish(),
366            Err(LlmError::StreamProtocol(_))
367        ));
368        let mut assembler = StreamAssembler::default();
369        assembler
370            .apply_json(r#"{"choices":[{"delta":{"content":"partial"}}]}"#)
371            .unwrap();
372        assert!(matches!(
373            assembler.finish(),
374            Err(LlmError::StreamProtocol(_))
375        ));
376    }
377
378    #[test]
379    fn repeated_tool_names_do_not_duplicate() {
380        let mut name = String::new();
381        merge_tool_name(&mut name, "fetch_report");
382        merge_tool_name(&mut name, "fetch_report");
383        assert_eq!(name, "fetch_report");
384    }
385
386    #[test]
387    fn incremental_tool_name_fragments_append() {
388        let mut name = String::new();
389        merge_tool_name(&mut name, "fetch_");
390        merge_tool_name(&mut name, "report");
391        assert_eq!(name, "fetch_report");
392    }
393}
af_llm/stream.rs

af_llm/
stream.rs
