use serde::Deserialize;
use crate::error::{LlmError, Result};
use crate::types::{AssistantBlock, FinishReason, FunctionCall, ToolCall, Usage};
use af_context::ToolCallId;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StreamDelta {
pub content: String,
pub has_tool_calls: bool,
}
#[derive(Debug, Default)]
pub(crate) struct StreamAssembler {
id: String,
content: String,
tool_slots: Vec<Option<ToolCallBuilder>>,
usage: Option<Usage>,
finish_reason: Option<FinishReason>,
output_blocks: Vec<AssistantBlock>,
saw_choice: bool,
}
#[derive(Debug, Default, Clone)]
struct ToolCallBuilder {
id: String,
kind: String,
name: String,
arguments: String,
}
impl StreamAssembler {
pub fn apply_json(&mut self, data: &str) -> Result<Option<StreamDelta>> {
let chunk: StreamChunk = serde_json::from_str(data)?;
if !chunk.id.is_empty() {
self.id = chunk.id;
}
if let Some(usage) = chunk.usage {
self.usage = Some(usage);
}
let mut content_changed = false;
for choice in chunk.choices {
if choice.index != 0 {
continue;
}
self.saw_choice = true;
if let Some(reason) = choice.finish_reason {
self.finish_reason = Some(reason);
}
let Some(delta) = choice.delta else {
continue;
};
self.output_blocks.extend(delta.output_blocks);
if let Some(piece) = delta.content.filter(|piece| !piece.is_empty()) {
self.content.push_str(&piece);
content_changed = true;
}
for tool_delta in delta.tool_calls.unwrap_or_default() {
self.merge_tool_delta(tool_delta);
}
}
Ok(content_changed.then(|| StreamDelta {
content: self.content.clone(),
has_tool_calls: self.has_tool_calls(),
}))
}
fn has_tool_calls(&self) -> bool {
self.tool_slots
.iter()
.flatten()
.any(|tool| !tool.id.is_empty() || !tool.name.is_empty() || !tool.arguments.is_empty())
}
fn merge_tool_delta(&mut self, delta: ToolCallDelta) {
let index = delta.index as usize;
if self.tool_slots.len() <= index {
self.tool_slots.resize_with(index + 1, || None);
}
let tool = self.tool_slots[index].get_or_insert_with(ToolCallBuilder::default);
if let Some(id) = delta.id.filter(|value| !value.is_empty()) {
tool.id = id;
}
if let Some(kind) = delta.kind.filter(|value| !value.is_empty()) {
tool.kind = kind;
}
if let Some(function) = delta.function {
if let Some(name) = function.name.filter(|value| !value.is_empty()) {
merge_tool_name(&mut tool.name, &name);
}
if let Some(arguments) = function.arguments {
tool.arguments.push_str(&arguments);
}
}
}
pub fn finish(self) -> Result<crate::types::CompletionResponse> {
use crate::types::{ChatMessage, Choice, CompletionResponse, Role};
if !self.saw_choice {
return Err(LlmError::StreamProtocol(
"stream completed without choice 0".into(),
));
}
if self.finish_reason.is_none() {
return Err(LlmError::StreamProtocol(
"stream completed without finish_reason".into(),
));
}
let tool_calls = self
.tool_slots
.into_iter()
.flatten()
.filter(|tool| !tool.name.is_empty() || !tool.arguments.is_empty())
.map(|tool| ToolCall {
id: ToolCallId::parse(tool.id).unwrap_or_else(|_| {
ToolCallId::parse(format!(
"call_{:08x}",
stable_hash(&tool.name, &tool.arguments)
))
.expect("synthesized tool call id is non-empty")
}),
kind: if tool.kind.is_empty() {
"function".into()
} else {
tool.kind
},
function: FunctionCall {
name: tool.name,
arguments: tool.arguments,
},
})
.collect::<Vec<_>>();
Ok(CompletionResponse {
id: self.id,
choices: vec![Choice {
index: 0,
message: ChatMessage {
role: Role::Assistant,
content: (!self.content.is_empty()).then_some(self.content),
tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
tool_call_id: None,
name: None,
},
finish_reason: self.finish_reason,
output_blocks: self.output_blocks,
}],
usage: self.usage,
})
}
}
fn stable_hash(left: &str, right: &str) -> u32 {
let mut hash = 0x811c_9dc5u32;
for byte in left.bytes().chain(right.bytes()) {
hash ^= u32::from(byte);
hash = hash.wrapping_mul(0x0100_0193);
}
hash
}
fn merge_tool_name(current: &mut String, incoming: &str) {
if current.is_empty() {
current.push_str(incoming);
} else if incoming.starts_with(current.as_str()) {
*current = incoming.to_string();
} else if !current.starts_with(incoming) {
current.push_str(incoming);
}
}
#[derive(Debug, Deserialize)]
struct StreamChunk {
#[serde(default)]
id: String,
#[serde(default)]
choices: Vec<StreamChoice>,
#[serde(default)]
usage: Option<Usage>,
}
#[derive(Debug, Deserialize)]
struct StreamChoice {
#[serde(default)]
index: u32,
#[serde(default)]
delta: Option<DeltaBody>,
#[serde(default)]
finish_reason: Option<FinishReason>,
}
#[derive(Debug, Deserialize)]
struct DeltaBody {
#[serde(default)]
content: Option<String>,
#[serde(default)]
tool_calls: Option<Vec<ToolCallDelta>>,
#[serde(default, alias = "content_blocks")]
output_blocks: Vec<AssistantBlock>,
}
#[derive(Debug, Deserialize)]
struct ToolCallDelta {
#[serde(default)]
index: u32,
#[serde(default)]
id: Option<String>,
#[serde(default, rename = "type")]
kind: Option<String>,
#[serde(default)]
function: Option<ToolFunctionDelta>,
}
#[derive(Debug, Deserialize)]
struct ToolFunctionDelta {
#[serde(default)]
name: Option<String>,
#[serde(default)]
arguments: Option<String>,
}
#[derive(Debug, Default)]
pub(crate) struct SseDecoder {
buffer: Vec<u8>,
data_lines: Vec<String>,
}
impl SseDecoder {
pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<String>> {
self.buffer.extend_from_slice(chunk);
let mut events = Vec::new();
while let Some(newline) = self.buffer.iter().position(|byte| *byte == b'\n') {
let mut line = self.buffer.drain(..=newline).collect::<Vec<_>>();
line.pop();
if line.last() == Some(&b'\r') {
line.pop();
}
let line = String::from_utf8(line)?;
if line.is_empty() {
if !self.data_lines.is_empty() {
events.push(self.data_lines.join("\n"));
self.data_lines.clear();
}
} else if let Some(data) = line.strip_prefix("data:") {
self.data_lines
.push(data.strip_prefix(' ').unwrap_or(data).to_string());
}
}
Ok(events)
}
pub fn finish(self) -> Result<()> {
if self.buffer.is_empty() && self.data_lines.is_empty() {
return Ok(());
}
String::from_utf8(self.buffer)?;
Err(LlmError::StreamProtocol(
"stream ended with an incomplete SSE frame".into(),
))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn assembles_content_and_tool_calls() {
let mut assembler = StreamAssembler::default();
assert!(assembler
.apply_json(r#"{"id":"chatcmpl-1","choices":[{"delta":{"content":"Hi"}}]}"#,)
.unwrap()
.is_some());
assembler
.apply_json(r#"{"choices":[{"delta":{"content":" there"}}]}"#)
.unwrap();
assembler.apply_json(
r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"lookup_weather","arguments":""}}]}}]}"#,
).unwrap();
assembler.apply_json(
r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"city\":\"Paris\"}"}}]}}]}"#,
).unwrap();
assembler
.apply_json(r#"{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}"#)
.unwrap();
let response = assembler.finish().unwrap();
assert_eq!(response.first_content(), Some("Hi there"));
let calls = response.first_tool_calls().unwrap();
assert_eq!(calls[0].function.name, "lookup_weather");
assert_eq!(calls[0].function.arguments, r#"{"city":"Paris"}"#);
}
#[test]
fn drains_lf_and_crlf_frames() {
let mut decoder = SseDecoder::default();
assert_eq!(
decoder
.push(b"data: {\"a\":1}\r\n\r\ndata: [DONE]\n\npartial")
.unwrap(),
[r#"{"a":1}"#, "[DONE]"]
);
assert!(matches!(decoder.finish(), Err(LlmError::StreamProtocol(_))));
}
#[test]
fn joins_multiline_data_and_preserves_usage() {
let mut decoder = SseDecoder::default();
let frames = decoder
.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")
.unwrap();
let mut assembler = StreamAssembler::default();
assembler.apply_json(&frames[0]).unwrap();
let response = assembler.finish().unwrap();
assert_eq!(response.first_content(), Some("hi"));
assert_eq!(response.usage.unwrap().total_tokens, 3);
}
#[test]
fn streams_typed_output_blocks_and_rejects_unknown_kinds() {
let mut assembler = StreamAssembler::default();
assembler.apply_json(r#"{"choices":[{"delta":{"output_blocks":[{"type":"citation","resource_id":"doc-1","label":"Doc","uri":"docs://doc-1"}]},"finish_reason":"stop"}]}"#).unwrap();
let response = assembler.finish().unwrap();
assert!(matches!(
&response.choices[0].output_blocks[0],
AssistantBlock::Citation { resource_id, .. } if resource_id == "doc-1"
));
assert!(StreamAssembler::default()
.apply_json(
r#"{"choices":[{"delta":{"output_blocks":[{"type":"html","html":"bad"}]}}]}"#
)
.is_err());
}
#[test]
fn malformed_json_and_utf8_fail_closed() {
assert!(StreamAssembler::default().apply_json("{").is_err());
let mut decoder = SseDecoder::default();
assert!(matches!(
decoder.push(b"data: \xff\n\n"),
Err(LlmError::InvalidUtf8(_))
));
}
#[test]
fn empty_and_unfinished_assemblies_fail_closed() {
assert!(matches!(
StreamAssembler::default().finish(),
Err(LlmError::StreamProtocol(_))
));
let mut assembler = StreamAssembler::default();
assembler
.apply_json(r#"{"choices":[{"delta":{"content":"partial"}}]}"#)
.unwrap();
assert!(matches!(
assembler.finish(),
Err(LlmError::StreamProtocol(_))
));
}
#[test]
fn repeated_tool_names_do_not_duplicate() {
let mut name = String::new();
merge_tool_name(&mut name, "fetch_report");
merge_tool_name(&mut name, "fetch_report");
assert_eq!(name, "fetch_report");
}
#[test]
fn incremental_tool_name_fragments_append() {
let mut name = String::new();
merge_tool_name(&mut name, "fetch_");
merge_tool_name(&mut name, "report");
assert_eq!(name, "fetch_report");
}
}