use std::collections::BTreeMap;
use serde_json::{Map, Value, json};
use super::{Counts, Merged, server_ms, server_tokens};
pub(super) fn counts(fields: &Value) -> Counts {
Counts {
prompt_tokens: server_tokens(fields, "/usage/prompt_tokens"),
completion_tokens: server_tokens(fields, "/usage/completion_tokens")
.or_else(|| server_tokens(fields, "/timings/predicted_n")),
reused_tokens: server_tokens(fields, "/timings/cache_n")
.or_else(|| server_tokens(fields, "/usage/prompt_tokens_details/cached_tokens")),
evaluated_tokens: server_tokens(fields, "/timings/prompt_n"),
prefill_ms: server_ms(fields, "/timings/prompt_ms", 1.0),
decode_ms: server_ms(fields, "/timings/predicted_ms", 1.0),
load_ms: None,
}
}
#[derive(Default)]
struct Choice {
role: Option<String>,
content: Option<String>,
reasoning_key: Option<&'static str>,
reasoning: String,
tool_calls: BTreeMap<u64, ToolCall>,
finish_reason: Option<Value>,
}
#[derive(Default)]
struct ToolCall {
id: Option<String>,
name: String,
arguments: String,
}
pub(super) fn merge_sse(raw: &[u8]) -> Merged {
let text = String::from_utf8_lossy(raw);
let mut chunks = 0u32;
let mut bad_chunks = 0u32;
let mut top = Map::new();
let mut choices: BTreeMap<u64, Choice> = BTreeMap::new();
for line in text.lines() {
let Some(data) = line.strip_prefix("data:") else {
continue;
};
let data = data.trim();
if data == "[DONE]" {
continue;
}
chunks += 1;
let Ok(Value::Object(event)) = serde_json::from_str::<Value>(data) else {
bad_chunks += 1;
continue;
};
for (key, value) in event {
match (key.as_str(), value) {
("choices", Value::Array(items)) => {
for (position, item) in items.iter().enumerate() {
merge_choice(&mut choices, position, item);
}
}
("choices", _) => {}
(_, Value::Null) => {}
(_, value) => {
top.insert(key, value);
}
}
}
}
if chunks == bad_chunks {
return Merged {
body: Value::Null,
chunks,
bad_chunks,
};
}
if top.contains_key("object") {
top.insert("object".into(), Value::from("chat.completion"));
}
let choices = choices
.into_iter()
.map(|(index, choice)| finish_choice(index, choice))
.collect();
top.insert("choices".into(), Value::Array(choices));
Merged {
body: Value::Object(top),
chunks,
bad_chunks,
}
}
fn merge_choice(choices: &mut BTreeMap<u64, Choice>, position: usize, item: &Value) {
let index = item
.get("index")
.and_then(Value::as_u64)
.unwrap_or(position as u64);
let choice = choices.entry(index).or_default();
if let Some(reason) = item.get("finish_reason").filter(|r| !r.is_null()) {
choice.finish_reason = Some(reason.clone());
}
let Some(delta) = item.get("delta").or_else(|| item.get("message")) else {
return;
};
if let Some(role) = delta.get("role").and_then(Value::as_str) {
choice.role = Some(role.to_string());
}
if let Some(content) = delta.get("content").and_then(Value::as_str) {
choice
.content
.get_or_insert_with(String::new)
.push_str(content);
}
for key in ["reasoning_content", "reasoning"] {
if let Some(reasoning) = delta.get(key).and_then(Value::as_str) {
choice.reasoning_key.get_or_insert(key);
choice.reasoning.push_str(reasoning);
break;
}
}
if let Some(calls) = delta.get("tool_calls").and_then(Value::as_array) {
for (position, call) in calls.iter().enumerate() {
let slot = call
.get("index")
.and_then(Value::as_u64)
.unwrap_or(position as u64);
let acc = choice.tool_calls.entry(slot).or_default();
if let Some(id) = call.get("id").and_then(Value::as_str) {
acc.id = Some(id.to_string());
}
if let Some(name) = call.pointer("/function/name").and_then(Value::as_str)
&& acc.name.is_empty()
{
acc.name = name.to_string();
}
if let Some(arguments) = call.pointer("/function/arguments").and_then(Value::as_str) {
acc.arguments.push_str(arguments);
}
}
}
}
fn finish_choice(index: u64, choice: Choice) -> Value {
let mut message = Map::new();
message.insert(
"role".into(),
Value::from(choice.role.unwrap_or_else(|| "assistant".into())),
);
message.insert("content".into(), Value::from(choice.content));
if let Some(key) = choice.reasoning_key {
message.insert(key.into(), Value::from(choice.reasoning));
}
if !choice.tool_calls.is_empty() {
let calls = choice
.tool_calls
.into_values()
.map(|call| {
json!({
"id": call.id,
"type": "function",
"function": { "name": call.name, "arguments": call.arguments }
})
})
.collect();
message.insert("tool_calls".into(), Value::Array(calls));
}
json!({
"index": index,
"message": message,
"finish_reason": choice.finish_reason,
})
}
#[cfg(test)]
mod tests {
use super::*;
fn sse(events: &[Value]) -> Vec<u8> {
let mut out = String::new();
for event in events {
out.push_str(&format!("data: {event}\n\n"));
}
out.push_str("data: [DONE]\n\n");
out.into_bytes()
}
#[test]
fn content_and_reasoning_are_concatenated() {
let raw = sse(&[
json!({"id": "c", "object": "chat.completion.chunk", "model": "m",
"choices": [{"index": 0, "delta": {"role": "assistant", "content": null, "reasoning_content": "thi"}}]}),
json!({"id": "c", "object": "chat.completion.chunk",
"choices": [{"index": 0, "delta": {"reasoning_content": "nk"}}]}),
json!({"choices": [{"index": 0, "delta": {"content": "Hel"}}]}),
json!({"choices": [{"index": 0, "delta": {"content": "lo"}, "finish_reason": "stop"}]}),
]);
let merged = merge_sse(&raw);
assert_eq!((merged.chunks, merged.bad_chunks), (4, 0));
assert_eq!(merged.body["object"], "chat.completion");
assert_eq!(merged.body["model"], "m");
let message = &merged.body["choices"][0]["message"];
assert_eq!(message["content"], "Hello");
assert_eq!(message["reasoning_content"], "think");
assert_eq!(message["role"], "assistant");
assert_eq!(merged.body["choices"][0]["finish_reason"], "stop");
}
#[test]
fn reasoning_alias_keeps_its_name() {
let raw = sse(&[
json!({"choices": [{"delta": {"reasoning": "a"}}]}),
json!({"choices": [{"delta": {"reasoning": "b"}}]}),
]);
let message = &merge_sse(&raw).body["choices"][0]["message"];
assert_eq!(message["reasoning"], "ab");
assert!(message.get("reasoning_content").is_none());
}
#[test]
fn tool_calls_merge_by_index() {
let raw = sse(&[
json!({"choices": [{"delta": {"tool_calls": [{"index": 0, "id": "a", "type": "function", "function": {"name": "read_file", "arguments": "{\"pa"}}]}}]}),
json!({"choices": [{"delta": {"tool_calls": [{"index": 0, "function": {"arguments": "th\":1}"}}]}}]}),
json!({"choices": [{"delta": {"tool_calls": [{"index": 1, "id": "b", "function": {"name": "ls", "arguments": "{}"}}]}}]}),
]);
let calls = merge_sse(&raw).body["choices"][0]["message"]["tool_calls"].clone();
assert_eq!(
calls,
json!([
{"id": "a", "type": "function", "function": {"name": "read_file", "arguments": "{\"path\":1}"}},
{"id": "b", "type": "function", "function": {"name": "ls", "arguments": "{}"}}
])
);
}
#[test]
fn last_usage_and_timings_win_and_nulls_are_ignored() {
let raw = sse(&[
json!({"choices": [{"delta": {"content": "x"}}], "usage": {"prompt_tokens": 1}}),
json!({"choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 2}, "timings": {"cache_n": 5}}),
json!({"choices": [], "usage": null}),
]);
let body = merge_sse(&raw).body;
assert_eq!(
body["usage"],
json!({"prompt_tokens": 10, "completion_tokens": 2})
);
assert_eq!(body["timings"], json!({"cache_n": 5}));
}
#[test]
fn bad_chunks_are_counted_and_skipped() {
let raw = b"data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\ndata: {broken\n\n: comment\nevent: x\ndata: [DONE]\n\n";
let merged = merge_sse(raw);
assert_eq!((merged.chunks, merged.bad_chunks), (2, 1));
assert_eq!(merged.body["choices"][0]["message"]["content"], "a");
}
#[test]
fn crlf_lines_and_no_space_after_colon() {
let raw =
b"data:{\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\r\n\r\ndata: [DONE]\r\n\r\n";
assert_eq!(merge_sse(raw).body["choices"][0]["message"]["content"], "a");
}
#[test]
fn nothing_parsed_gives_null() {
let merged = merge_sse(b"data: nope\n\n");
assert_eq!(merged.body, Value::Null);
assert_eq!((merged.chunks, merged.bad_chunks), (1, 1));
assert_eq!(merge_sse(b"").body, Value::Null);
}
#[test]
fn merged_stream_has_the_non_streamed_shape() {
let raw = sse(&[
json!({"choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]}),
]);
let body = merge_sse(&raw).body;
assert_eq!(
body["choices"],
json!([{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}])
);
}
}