use std::collections::BTreeMap;
use std::io::{self, BufRead, Write};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::RECORD_VERSION;
use crate::hash;
use crate::normalize::{self, Api};
const EDGE_CHARS: usize = 256;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Record {
pub v: u32,
pub id: u64,
pub conn: u64,
pub t_start_ms: i64,
pub t_first_byte_ms: Option<i64>,
pub t_end_ms: i64,
pub request: RequestRecord,
pub response: Option<ResponseRecord>,
pub error: Option<ExchangeError>,
pub tokens: Option<TokenCounts>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RequestRecord {
pub method: String,
pub path: String,
pub headers: BTreeMap<String, String>,
pub body_bytes: u64,
pub body: Option<Value>,
pub body_text: Option<String>,
pub digest: Option<RequestDigest>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ResponseRecord {
pub status: u16,
pub headers: BTreeMap<String, String>,
pub stream: Option<StreamKind>,
pub chunks: Option<u32>,
pub bad_chunks: Option<u32>,
pub body_bytes: u64,
pub truncated: bool,
pub body: Option<Value>,
pub body_text: Option<String>,
pub digest: Option<ResponseDigest>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StreamKind {
Sse,
Ndjson,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ExchangeError {
UpstreamUnreachable,
UpstreamReset,
ClientDisconnected,
Timeout,
BodyTooLarge,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RequestDigest {
pub model: Option<String>,
pub params: Params,
pub messages: Vec<MessageDigest>,
pub tools: Vec<ToolDigest>,
pub tools_hash: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct Params {
pub max_tokens: Option<u32>,
pub response_format: bool,
pub grammar: bool,
pub stream: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MessageDigest {
pub role: String,
pub hash: String,
pub chars: u32,
pub head: String,
pub tail: String,
pub tool_calls: Vec<ToolCall>,
pub tool_call_id: Option<String>,
pub name: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolCall {
pub id: Option<String>,
pub name: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolDigest {
pub name: String,
pub hash: String,
pub chars: u32,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ResponseDigest {
pub usage: Map<String, Value>,
pub output: Output,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct Output {
pub content_chars: u32,
pub reasoning_chars: u32,
pub tool_calls: Vec<ToolCall>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct TokenCounts {
pub messages: BTreeMap<String, u32>,
pub tools: BTreeMap<String, u32>,
}
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct ReadResult {
pub records: Vec<Record>,
pub bad_lines: Vec<BadLine>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BadLine {
pub line: usize,
pub error: String,
}
pub fn read_jsonl(reader: impl BufRead) -> ReadResult {
let mut result = ReadResult::default();
for (index, line) in reader.lines().enumerate() {
let number = index + 1;
let line = match line {
Ok(line) => line,
Err(e) => {
result.bad_lines.push(BadLine {
line: number,
error: e.to_string(),
});
break;
}
};
if line.trim().is_empty() {
continue;
}
match parse_line(&line) {
Ok(record) => result.records.push(record),
Err(error) => result.bad_lines.push(BadLine {
line: number,
error,
}),
}
}
result.records.sort_by_key(|r| r.id);
result
}
fn parse_line(line: &str) -> Result<Record, String> {
let value: Value = serde_json::from_str(line).map_err(|e| e.to_string())?;
match value.get("v").and_then(Value::as_u64) {
Some(v) if v == u64::from(RECORD_VERSION) => {}
Some(v) => return Err(format!("unsupported record version {v}")),
None => return Err("missing record version".to_string()),
}
serde_json::from_value(value).map_err(|e| e.to_string())
}
pub fn write_jsonl(mut writer: impl Write, record: &Record) -> io::Result<()> {
let mut line = serde_json::to_vec(record).map_err(io::Error::other)?;
line.push(b'\n');
writer.write_all(&line)
}
#[cfg_attr(not(feature = "net"), allow(dead_code))]
pub(crate) fn apply_digest(record: &mut Record) {
let api = normalize::detect_api(&record.request.path);
let request = &mut record.request;
request.digest = match (api, request.body.as_ref()) {
(Api::Other, _) | (_, None) => None,
(api, Some(body)) => request_digest(api, body),
};
request.body = None;
request.body_text = None;
if let Some(response) = record.response.as_mut() {
let success = (200..300).contains(&response.status);
response.digest = match response.body.as_ref() {
Some(body) if api != Api::Other && success => Some(response_digest(body)),
_ => None,
};
response.body = None;
response.body_text = None;
}
}
pub(crate) fn request_digest(api: Api, body: &Value) -> Option<RequestDigest> {
let messages = body.get("messages")?.as_array()?;
let tools: &[Value] = body
.get("tools")
.and_then(Value::as_array)
.map(Vec::as_slice)
.unwrap_or_default();
let tools_hash = if tools.is_empty() {
None
} else {
Some(hash::hex_value(&Value::Array(tools.to_vec())))
};
Some(RequestDigest {
model: str_field(body, "model"),
params: params(api, body),
messages: messages.iter().map(message_digest).collect(),
tools: tools.iter().map(tool_digest).collect(),
tools_hash,
})
}
pub(crate) fn response_digest(body: &Value) -> ResponseDigest {
let usage = body
.as_object()
.map(|object| {
object
.iter()
.filter(|(key, _)| key.as_str() != "choices" && key.as_str() != "message")
.map(|(key, value)| (key.clone(), value.clone()))
.collect()
})
.unwrap_or_default();
let output = body
.pointer("/choices/0/message")
.or_else(|| body.get("message"))
.map(|message| Output {
content_chars: char_count(&content_text(message.get("content"))),
reasoning_chars: ["reasoning_content", "reasoning", "thinking"]
.iter()
.find_map(|key| message.get(*key).and_then(Value::as_str))
.map(char_count)
.unwrap_or_default(),
tool_calls: tool_calls(message),
})
.unwrap_or_default();
ResponseDigest { usage, output }
}
pub(crate) fn message_digest(message: &Value) -> MessageDigest {
let role = str_field(message, "role").unwrap_or_default();
let name = str_field(message, "name").or_else(|| str_field(message, "tool_name"));
let tool_call_id = str_field(message, "tool_call_id");
let identity = Value::Array(vec![
Value::from(role.clone()),
message.get("content").cloned().unwrap_or(Value::Null),
raw_tool_calls(message).cloned().unwrap_or(Value::Null),
Value::from(name.clone()),
Value::from(tool_call_id.clone()),
]);
let text = message_text(message);
let (head, tail) = edges(&text);
MessageDigest {
role,
hash: hash::hex_value(&identity),
chars: char_count(&text),
head: hash::hex(head.as_bytes()),
tail: hash::hex(tail.as_bytes()),
tool_calls: tool_calls(message),
tool_call_id,
name,
}
}
pub(crate) fn message_text(message: &Value) -> String {
let mut text = content_text(message.get("content"));
if let Some(calls) = raw_tool_calls(message) {
text.push_str(&hash::canonical_json(calls));
}
text
}
fn tool_digest(tool: &Value) -> ToolDigest {
let canonical = hash::canonical_json(tool);
let name = tool
.pointer("/function/name")
.or_else(|| tool.get("name"))
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
ToolDigest {
name,
hash: hash::hex(canonical.as_bytes()),
chars: char_count(&canonical),
}
}
fn params(api: Api, body: &Value) -> Params {
let max_tokens = ["max_tokens", "max_completion_tokens", "n_predict"]
.iter()
.find_map(|key| body.get(*key).and_then(Value::as_u64))
.or_else(|| body.pointer("/options/num_predict").and_then(Value::as_u64))
.and_then(|n| u32::try_from(n).ok());
let response_format = is_set(body.get("format"))
|| (is_set(body.get("response_format"))
&& body
.pointer("/response_format/type")
.and_then(Value::as_str)
!= Some("text"));
let grammar = is_set(body.get("grammar")) || is_set(body.get("json_schema"));
let stream = body
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(api == Api::OllamaChat);
Params {
max_tokens,
response_format,
grammar,
stream,
}
}
fn is_set(value: Option<&Value>) -> bool {
match value {
None | Some(Value::Null) => false,
Some(Value::String(s)) => !s.is_empty(),
Some(_) => true,
}
}
fn raw_tool_calls(message: &Value) -> Option<&Value> {
message
.get("tool_calls")
.filter(|calls| calls.as_array().is_some_and(|a| !a.is_empty()))
}
fn tool_calls(message: &Value) -> Vec<ToolCall> {
raw_tool_calls(message)
.and_then(Value::as_array)
.map(|calls| {
calls
.iter()
.map(|call| ToolCall {
id: str_field(call, "id"),
name: call
.pointer("/function/name")
.or_else(|| call.get("name"))
.and_then(Value::as_str)
.unwrap_or_default()
.to_string(),
})
.collect()
})
.unwrap_or_default()
}
fn content_text(content: Option<&Value>) -> String {
match content {
Some(Value::String(s)) => s.clone(),
Some(Value::Array(parts)) => parts
.iter()
.filter_map(|part| part.get("text").and_then(Value::as_str))
.collect(),
_ => String::new(),
}
}
fn edges(text: &str) -> (String, String) {
let total = text.chars().count();
let head = text.chars().take(EDGE_CHARS).collect();
let tail = text
.chars()
.skip(total.saturating_sub(EDGE_CHARS))
.collect();
(head, tail)
}
fn char_count(text: &str) -> u32 {
u32::try_from(text.chars().count()).unwrap_or(u32::MAX)
}
fn str_field(value: &Value, key: &str) -> Option<String> {
value.get(key).and_then(Value::as_str).map(str::to_string)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn same_message_same_hash_regardless_of_key_order() {
let a = json!({"role": "user", "content": "hello"});
let b: Value = serde_json::from_str(r#"{"content":"hello","role":"user"}"#).unwrap();
assert_eq!(message_digest(&a).hash, message_digest(&b).hash);
}
#[test]
fn role_content_name_and_call_id_are_part_of_identity() {
let base = message_digest(&json!({"role": "tool", "content": "x", "tool_call_id": "1"}));
for other in [
json!({"role": "user", "content": "x", "tool_call_id": "1"}),
json!({"role": "tool", "content": "y", "tool_call_id": "1"}),
json!({"role": "tool", "content": "x", "tool_call_id": "2"}),
json!({"role": "tool", "content": "x", "tool_call_id": "1", "name": "n"}),
] {
assert_ne!(base.hash, message_digest(&other).hash);
}
}
#[test]
fn empty_tool_calls_equal_absent_tool_calls() {
let a = message_digest(&json!({"role": "assistant", "content": "x"}));
let b = message_digest(&json!({"role": "assistant", "content": "x", "tool_calls": []}));
assert_eq!(a.hash, b.hash);
assert_eq!(a.chars, b.chars);
}
#[test]
fn text_parts_are_joined_and_tool_calls_counted() {
let m = json!({
"role": "assistant",
"content": [{"type": "text", "text": "ab"}, {"type": "image_url", "image_url": {"url": "x"}}, {"type": "text", "text": "cd"}],
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "read_file", "arguments": "{}"}}]
});
let d = message_digest(&m);
let calls = hash::canonical_json(m.get("tool_calls").unwrap());
assert_eq!(d.chars as usize, 4 + calls.chars().count());
assert_eq!(
d.tool_calls,
vec![ToolCall {
id: Some("c1".into()),
name: "read_file".into()
}]
);
}
#[test]
fn head_and_tail_cover_256_chars() {
let text: String = "é".repeat(300) + "end";
let d = message_digest(&json!({"role": "user", "content": text}));
assert_eq!(d.chars, 303);
assert_eq!(d.head, hash::hex("é".repeat(256).as_bytes()));
assert_eq!(d.tail, hash::hex(("é".repeat(253) + "end").as_bytes()));
let short = message_digest(&json!({"role": "user", "content": "hi"}));
assert_eq!(short.head, short.tail);
}
#[test]
fn ollama_tool_name_is_the_name() {
let d = message_digest(&json!({"role": "tool", "content": "ok", "tool_name": "ls"}));
assert_eq!(d.name.as_deref(), Some("ls"));
}
#[test]
fn params_openai() {
let body = json!({"model": "m", "messages": [], "max_tokens": 8, "response_format": {"type": "json_object"}, "stream": true});
let d = request_digest(Api::OpenAiChat, &body).unwrap();
assert_eq!(
d.params,
Params {
max_tokens: Some(8),
response_format: true,
grammar: false,
stream: true
}
);
assert_eq!(d.model.as_deref(), Some("m"));
}
#[test]
fn params_defaults_and_text_format() {
let body = json!({"messages": [], "response_format": {"type": "text"}, "n_predict": -1});
let p = request_digest(Api::OpenAiChat, &body).unwrap().params;
assert_eq!(p, Params::default());
let p = request_digest(Api::OllamaChat, &json!({"messages": []}))
.unwrap()
.params;
assert!(p.stream);
}
#[test]
fn params_ollama_and_grammar() {
let body = json!({"messages": [], "format": "json", "options": {"num_predict": 16}, "stream": false});
let p = request_digest(Api::OllamaChat, &body).unwrap().params;
assert_eq!(p.max_tokens, Some(16));
assert!(p.response_format);
assert!(!p.stream);
let p = request_digest(
Api::OpenAiChat,
&json!({"messages": [], "grammar": "root ::= x"}),
)
.unwrap()
.params;
assert!(p.grammar);
}
#[test]
fn tools_hash_and_tool_digests() {
let tool = json!({"type": "function", "function": {"name": "read_file", "parameters": {"type": "object"}}});
let body = json!({"messages": [], "tools": [tool.clone()]});
let d = request_digest(Api::OpenAiChat, &body).unwrap();
assert_eq!(d.tools.len(), 1);
assert_eq!(d.tools[0].name, "read_file");
assert_eq!(d.tools[0].hash, hash::hex_value(&tool));
assert_eq!(d.tools_hash, Some(hash::hex_value(&json!([tool]))));
let none = request_digest(Api::OpenAiChat, &json!({"messages": [], "tools": []})).unwrap();
assert_eq!(none.tools_hash, None);
}
#[test]
fn not_a_chat_body() {
assert!(request_digest(Api::OpenAiChat, &json!({"input": "x"})).is_none());
}
#[test]
fn response_digest_openai() {
let body = json!({
"id": "x",
"choices": [{"message": {"role": "assistant", "content": "hello", "reasoning_content": "think",
"tool_calls": [{"id": "c", "function": {"name": "ls", "arguments": "{}"}}]}}],
"usage": {"prompt_tokens": 10},
"timings": {"cache_n": 4}
});
let d = response_digest(&body);
assert_eq!(d.output.content_chars, 5);
assert_eq!(d.output.reasoning_chars, 5);
assert_eq!(d.output.tool_calls[0].name, "ls");
assert!(d.usage.contains_key("usage"));
assert!(d.usage.contains_key("timings"));
assert!(!d.usage.contains_key("choices"));
}
#[test]
fn response_digest_ollama() {
let body = json!({
"model": "m",
"message": {"role": "assistant", "content": "hi", "thinking": "hmm",
"tool_calls": [{"function": {"name": "ls", "arguments": {}}}]},
"done": true,
"prompt_eval_count": 12,
"eval_duration": 5000000
});
let d = response_digest(&body);
assert_eq!(d.output.content_chars, 2);
assert_eq!(d.output.reasoning_chars, 3);
assert_eq!(
d.output.tool_calls,
vec![ToolCall {
id: None,
name: "ls".into()
}]
);
assert_eq!(d.usage.get("prompt_eval_count"), Some(&json!(12)));
assert!(!d.usage.contains_key("message"));
}
fn full_record(path: &str, status: u16) -> Record {
Record {
v: RECORD_VERSION,
id: 1,
conn: 1,
t_start_ms: 0,
t_first_byte_ms: Some(1),
t_end_ms: 2,
request: RequestRecord {
method: "POST".into(),
path: path.into(),
headers: BTreeMap::new(),
body_bytes: 10,
body: Some(
json!({"model": "m", "messages": [{"role": "user", "content": "secret"}]}),
),
body_text: None,
digest: None,
},
response: Some(ResponseRecord {
status,
headers: BTreeMap::new(),
stream: None,
chunks: None,
bad_chunks: None,
body_bytes: 10,
truncated: false,
body: Some(
json!({"choices": [{"message": {"content": "secret"}}], "usage": {"prompt_tokens": 3}}),
),
body_text: None,
digest: None,
}),
error: None,
tokens: None,
}
}
#[test]
fn apply_digest_keeps_no_text() {
let mut r = full_record("/v1/chat/completions", 200);
apply_digest(&mut r);
let line = serde_json::to_string(&r).unwrap();
assert!(!line.contains("secret"));
assert_eq!(r.request.digest.as_ref().unwrap().messages.len(), 1);
assert!(r.request.body.is_none());
let response = r.response.unwrap();
assert!(response.body.is_none());
assert_eq!(response.digest.unwrap().output.content_chars, 6);
}
#[test]
fn apply_digest_drops_other_paths_and_errors() {
let mut r = full_record("/v1/embeddings", 200);
apply_digest(&mut r);
assert!(r.request.digest.is_none() && r.request.body.is_none());
assert!(r.response.as_ref().unwrap().digest.is_none());
let mut r = full_record("/v1/chat/completions", 500);
apply_digest(&mut r);
assert!(r.request.digest.is_some());
let response = r.response.unwrap();
assert!(response.digest.is_none() && response.body.is_none());
}
}