use serde_json::{Map, Value};
use super::{Counts, Merged, server_ms, server_tokens};
const NANOS_PER_MS: f64 = 1e6;
pub(super) fn counts(fields: &Value) -> Counts {
Counts {
prompt_tokens: server_tokens(fields, "/prompt_eval_count"),
completion_tokens: server_tokens(fields, "/eval_count"),
reused_tokens: None,
evaluated_tokens: None,
prefill_ms: server_ms(fields, "/prompt_eval_duration", NANOS_PER_MS),
decode_ms: server_ms(fields, "/eval_duration", NANOS_PER_MS),
load_ms: server_ms(fields, "/load_duration", NANOS_PER_MS),
}
}
pub(super) fn merge_ndjson(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 role: Option<String> = None;
let mut content = String::new();
let mut thinking: Option<String> = None;
let mut tool_calls: Vec<Value> = Vec::new();
for line in text.lines() {
let line = line.trim();
if line.is_empty() {
continue;
}
chunks += 1;
let Ok(Value::Object(object)) = serde_json::from_str::<Value>(line) else {
bad_chunks += 1;
continue;
};
for (key, value) in object {
if key != "message" {
top.insert(key, value);
continue;
}
if let Some(r) = value.get("role").and_then(Value::as_str) {
role = Some(r.to_string());
}
if let Some(c) = value.get("content").and_then(Value::as_str) {
content.push_str(c);
}
if let Some(t) = value.get("thinking").and_then(Value::as_str) {
thinking.get_or_insert_with(String::new).push_str(t);
}
if let Some(calls) = value.get("tool_calls").and_then(Value::as_array) {
tool_calls.extend(calls.iter().cloned());
}
}
}
if chunks == bad_chunks {
return Merged {
body: Value::Null,
chunks,
bad_chunks,
};
}
let mut message = Map::new();
message.insert(
"role".into(),
Value::from(role.unwrap_or_else(|| "assistant".into())),
);
message.insert("content".into(), Value::from(content));
if let Some(thinking) = thinking {
message.insert("thinking".into(), Value::from(thinking));
}
if !tool_calls.is_empty() {
message.insert("tool_calls".into(), Value::Array(tool_calls));
}
top.insert("message".into(), Value::Object(message));
Merged {
body: Value::Object(top),
chunks,
bad_chunks,
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn ndjson(lines: &[Value]) -> Vec<u8> {
lines
.iter()
.map(|l| format!("{l}\n"))
.collect::<String>()
.into_bytes()
}
#[test]
fn content_thinking_and_final_counters() {
let raw = ndjson(&[
json!({"model": "m", "message": {"role": "assistant", "content": "", "thinking": "hm"}, "done": false}),
json!({"model": "m", "message": {"role": "assistant", "content": "Hel"}, "done": false}),
json!({"model": "m", "message": {"role": "assistant", "content": "lo"}, "done": false}),
json!({"model": "m", "message": {"role": "assistant", "content": ""}, "done": true, "done_reason": "stop",
"prompt_eval_count": 12, "prompt_eval_duration": 3000000, "eval_count": 3, "eval_duration": 9000000,
"load_duration": 100, "total_duration": 5}),
]);
let merged = merge_ndjson(&raw);
assert_eq!((merged.chunks, merged.bad_chunks), (4, 0));
let body = merged.body;
assert_eq!(body["message"]["content"], "Hello");
assert_eq!(body["message"]["thinking"], "hm");
assert_eq!(body["done"], true);
assert_eq!(body["prompt_eval_count"], 12);
assert_eq!(body["eval_duration"], 9000000);
assert_eq!(body["load_duration"], 100);
}
#[test]
fn tool_calls_are_collected() {
let raw = ndjson(&[
json!({"message": {"role": "assistant", "content": "", "tool_calls": [{"function": {"name": "ls", "arguments": {}}}]}, "done": false}),
json!({"message": {"role": "assistant", "content": ""}, "done": true}),
]);
let body = merge_ndjson(&raw).body;
assert_eq!(
body["message"]["tool_calls"],
json!([{"function": {"name": "ls", "arguments": {}}}])
);
assert!(body["message"].get("thinking").is_none());
}
#[test]
fn bad_lines_are_counted() {
let raw =
b"{\"message\":{\"content\":\"a\"},\"done\":false}\nnot json\n\n{\"done\":true}\n";
let merged = merge_ndjson(raw);
assert_eq!((merged.chunks, merged.bad_chunks), (3, 1));
assert_eq!(merged.body["message"]["content"], "a");
}
#[test]
fn nothing_parsed_gives_null() {
assert_eq!(merge_ndjson(b"oops\n").body, Value::Null);
assert_eq!(merge_ndjson(b"").body, Value::Null);
}
}