mod ollama;
mod openai;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::measured::{Measured, Source};
use crate::record::{self, Output, Params, Record, RequestDigest, StreamKind, ToolCall};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum Api {
#[serde(rename = "openai_chat")]
OpenAiChat,
#[serde(rename = "ollama_chat")]
OllamaChat,
#[serde(rename = "other")]
Other,
}
#[derive(Debug, Clone, Copy, PartialEq, Default, Serialize, Deserialize)]
pub struct Counts {
pub prompt_tokens: Option<Measured<u32>>,
pub completion_tokens: Option<Measured<u32>>,
pub reused_tokens: Option<Measured<u32>>,
pub evaluated_tokens: Option<Measured<u32>>,
pub prefill_ms: Option<Measured<f64>>,
pub decode_ms: Option<Measured<f64>>,
pub load_ms: Option<Measured<f64>>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct Exchange {
pub id: u64,
pub t_start_ms: i64,
pub t_end_ms: i64,
pub api: Api,
pub model: String,
pub user_agent: Option<String>,
pub messages: Vec<Msg>,
pub tools: Vec<ToolDef>,
pub tools_hash: Option<String>,
pub params: Params,
pub counts: Counts,
pub output: Output,
pub ok: bool,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct Msg {
pub role: String,
pub hash: String,
pub chars: u32,
pub tokens: Option<Measured<u32>>,
pub head: String,
pub tail: String,
pub text: Option<String>,
pub tool_calls: Vec<ToolCall>,
pub tool_call_id: Option<String>,
pub name: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct ToolDef {
pub name: String,
pub hash: String,
pub chars: u32,
pub tokens: Option<Measured<u32>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum SkipReason {
NotChat,
Error,
Non2xx,
Unparsed,
}
#[derive(Debug, Clone, PartialEq, Default)]
pub(crate) struct Normalized {
pub exchanges: Vec<Exchange>,
pub skipped: Vec<Skipped>,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct Skipped {
pub path: String,
pub reason: SkipReason,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct Merged {
pub body: Value,
pub chunks: u32,
pub bad_chunks: u32,
}
pub(crate) fn detect_api(path: &str) -> Api {
match path.split('?').next().unwrap_or_default() {
"/v1/chat/completions" => Api::OpenAiChat,
"/api/chat" => Api::OllamaChat,
_ => Api::Other,
}
}
pub(crate) fn normalize_all(records: &[Record]) -> Normalized {
let mut normalized = Normalized::default();
for record in records {
match normalize(record) {
Ok(exchange) => normalized.exchanges.push(exchange),
Err(reason) => normalized.skipped.push(Skipped {
path: record.request.path.clone(),
reason,
}),
}
}
normalized
}
fn normalize(record: &Record) -> Result<Exchange, SkipReason> {
let api = detect_api(&record.request.path);
if api == Api::Other {
return Err(SkipReason::NotChat);
}
if record.error.is_some() {
return Err(SkipReason::Error);
}
let Some(response) = record.response.as_ref() else {
return Err(SkipReason::Error);
};
if !(200..300).contains(&response.status) {
return Err(SkipReason::Non2xx);
}
let (request, texts) = request_parts(api, record)?;
let response_digest = match (&response.digest, &response.body) {
(Some(digest), _) => Some(digest.clone()),
(None, Some(body)) => Some(record::response_digest(body)),
(None, None) => None,
};
let ok = response_digest.is_some();
let (fields, output) = response_digest.map_or((Value::Null, Output::default()), |digest| {
(Value::Object(digest.usage), digest.output)
});
let server = match api {
Api::OllamaChat => ollama::counts(&fields),
Api::OpenAiChat | Api::Other => openai::counts(&fields),
};
let ttfb_ms = record
.t_first_byte_ms
.map(|first_byte| (first_byte - record.t_start_ms).max(0));
let counts = derive(server, response.stream.is_some(), ttfb_ms);
let tokens = record.tokens.as_ref();
let tokenizer = |value: &u32| Measured {
value: *value,
source: Source::Tokenizer,
};
let messages = request
.messages
.into_iter()
.enumerate()
.map(|(index, message)| Msg {
tokens: tokens
.and_then(|t| t.messages.get(&message.hash))
.map(tokenizer),
text: texts.as_ref().and_then(|t| t.get(index).cloned()),
role: message.role,
hash: message.hash,
chars: message.chars,
head: message.head,
tail: message.tail,
tool_calls: message.tool_calls,
tool_call_id: message.tool_call_id,
name: message.name,
})
.collect();
let tools = request
.tools
.into_iter()
.map(|tool| ToolDef {
tokens: tokens.and_then(|t| t.tools.get(&tool.hash)).map(tokenizer),
name: tool.name,
hash: tool.hash,
chars: tool.chars,
})
.collect();
Ok(Exchange {
id: record.id,
t_start_ms: record.t_start_ms,
t_end_ms: record.t_end_ms,
api,
model: request.model.unwrap_or_default(),
user_agent: record.request.headers.get("user-agent").cloned(),
messages,
tools,
tools_hash: request.tools_hash,
params: request.params,
counts,
output,
ok,
})
}
fn request_parts(
api: Api,
record: &Record,
) -> Result<(RequestDigest, Option<Vec<String>>), SkipReason> {
if let Some(digest) = &record.request.digest {
return Ok((digest.clone(), None));
}
let body = record.request.body.as_ref().ok_or(SkipReason::Unparsed)?;
let digest = record::request_digest(api, body).ok_or(SkipReason::Unparsed)?;
let texts = body
.get("messages")
.and_then(Value::as_array)
.map(|messages| messages.iter().map(record::message_text).collect());
Ok((digest, texts))
}
fn derive(mut counts: Counts, streamed: bool, ttfb_ms: Option<i64>) -> Counts {
if counts.prompt_tokens.is_none()
&& let (Some(reused), Some(evaluated)) = (counts.reused_tokens, counts.evaluated_tokens)
{
counts.prompt_tokens = Some(Measured {
value: reused.value.saturating_add(evaluated.value),
source: reused.source.max(evaluated.source),
});
}
if let Some(prompt) = counts.prompt_tokens {
match (counts.reused_tokens, counts.evaluated_tokens) {
(Some(reused), None) => {
counts.evaluated_tokens = Some(Measured {
value: prompt.value.saturating_sub(reused.value),
source: reused.source,
});
}
(None, Some(evaluated)) => {
counts.reused_tokens = Some(Measured {
value: prompt.value.saturating_sub(evaluated.value),
source: evaluated.source,
});
}
_ => {}
}
}
if counts.prefill_ms.is_none()
&& streamed
&& let Some(ttfb_ms) = ttfb_ms
{
counts.prefill_ms = Some(Measured {
value: ttfb_ms as f64,
source: Source::Proxy,
});
}
counts
}
fn server_tokens(fields: &Value, pointer: &str) -> Option<Measured<u32>> {
let value = fields.pointer(pointer)?.as_u64()?;
Some(Measured {
value: u32::try_from(value).ok()?,
source: Source::Server,
})
}
fn server_ms(fields: &Value, pointer: &str, divisor: f64) -> Option<Measured<f64>> {
let value = fields.pointer(pointer)?.as_f64()?;
Some(Measured {
value: value / divisor,
source: Source::Server,
})
}
#[cfg_attr(not(feature = "net"), allow(dead_code))]
pub(crate) fn merge_stream(kind: StreamKind, raw: &[u8]) -> Merged {
match kind {
StreamKind::Sse => openai::merge_sse(raw),
StreamKind::Ndjson => ollama::merge_ndjson(raw),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::RECORD_VERSION;
use crate::record::{ExchangeError, RequestRecord, ResponseRecord, TokenCounts};
use serde_json::json;
use std::collections::BTreeMap;
const T0: i64 = 1_789_984_894_228;
fn chat_body() -> Value {
json!({
"model": "gemma",
"max_tokens": 32,
"messages": [
{"role": "system", "content": "Be brief."},
{"role": "user", "content": "Say hi."},
{"role": "assistant", "content": null,
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "ls", "arguments": "{}"}}]},
{"role": "tool", "tool_call_id": "c1", "content": "a.txt"}
],
"tools": [{"type": "function", "function": {"name": "ls", "parameters": {"type": "object"}}}]
})
}
fn record(path: &str, request: Value, response: Value, stream: Option<StreamKind>) -> Record {
Record {
v: RECORD_VERSION,
id: 7,
conn: 1,
t_start_ms: T0,
t_first_byte_ms: Some(T0 + 25),
t_end_ms: T0 + 54,
request: RequestRecord {
method: "POST".into(),
path: path.into(),
headers: BTreeMap::new(),
body_bytes: 1,
body: Some(request),
body_text: None,
digest: None,
},
response: Some(ResponseRecord {
status: 200,
headers: BTreeMap::new(),
stream,
chunks: stream.map(|_| 6),
bad_chunks: stream.map(|_| 0),
body_bytes: 1,
truncated: false,
body: Some(response),
body_text: None,
digest: None,
}),
error: None,
tokens: None,
}
}
fn exchange(record: &Record) -> Exchange {
normalize(record).unwrap()
}
fn server<T>(value: T) -> Option<Measured<T>> {
Some(Measured {
value,
source: Source::Server,
})
}
fn llama_cpp_response() -> Value {
json!({
"choices": [{"index": 0, "finish_reason": "stop",
"message": {"role": "assistant", "content": "Hello there."}}],
"object": "chat.completion",
"timings": {"cache_n": 20, "prompt_n": 5, "prompt_ms": 21.987, "predicted_n": 4, "predicted_ms": 28.827},
"usage": {"completion_tokens": 4, "prompt_tokens": 25, "total_tokens": 29,
"prompt_tokens_details": {"cached_tokens": 20}}
})
}
#[test]
fn llama_cpp_fields() {
let counts = exchange(&record(
"/v1/chat/completions",
chat_body(),
llama_cpp_response(),
Some(StreamKind::Sse),
))
.counts;
assert_eq!(
counts,
Counts {
prompt_tokens: server(25),
completion_tokens: server(4),
reused_tokens: server(20),
evaluated_tokens: server(5),
prefill_ms: server(21.987),
decode_ms: server(28.827),
load_ms: None,
}
);
}
#[test]
fn llama_cpp_timings_win_over_cached_tokens() {
let mut response = llama_cpp_response();
response["usage"]["prompt_tokens_details"]["cached_tokens"] = json!(0);
let counts = exchange(&record("/v1/chat/completions", chat_body(), response, None)).counts;
assert_eq!(counts.reused_tokens, server(20));
}
fn vllm_response() -> Value {
json!({
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}}],
"usage": {"prompt_tokens": 100, "completion_tokens": 3, "prompt_tokens_details": {"cached_tokens": 64}}
})
}
#[test]
fn vllm_style_fields_derive_evaluated_and_use_ttfb_when_streamed() {
let counts = exchange(&record(
"/v1/chat/completions",
chat_body(),
vllm_response(),
Some(StreamKind::Sse),
))
.counts;
assert_eq!(
counts,
Counts {
prompt_tokens: server(100),
completion_tokens: server(3),
reused_tokens: server(64),
evaluated_tokens: server(36),
prefill_ms: Some(Measured {
value: 25.0,
source: Source::Proxy
}),
decode_ms: None,
load_ms: None,
}
);
}
#[test]
fn non_streamed_time_to_first_byte_is_never_prefill() {
let counts = exchange(&record(
"/v1/chat/completions",
chat_body(),
vllm_response(),
None,
))
.counts;
assert_eq!(counts.prefill_ms, None);
}
#[test]
fn server_prefill_beats_proxy_ttfb() {
let counts = exchange(&record(
"/v1/chat/completions",
chat_body(),
llama_cpp_response(),
Some(StreamKind::Sse),
))
.counts;
assert_eq!(counts.prefill_ms.map(|m| m.source), Some(Source::Server));
}
#[test]
fn no_first_byte_means_no_proxy_prefill() {
let mut r = record(
"/v1/chat/completions",
chat_body(),
vllm_response(),
Some(StreamKind::Sse),
);
r.t_first_byte_ms = None;
assert_eq!(exchange(&r).counts.prefill_ms, None);
}
#[test]
fn streamed_llama_cpp_without_usage_derives_the_totals_from_timings() {
let response = json!({
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}}],
"timings": {"cache_n": 605, "prompt_n": 16, "prompt_ms": 35.2, "predicted_n": 160, "predicted_ms": 1331.9}
});
let counts = exchange(&record(
"/v1/chat/completions",
chat_body(),
response,
Some(StreamKind::Sse),
))
.counts;
assert_eq!(counts.prompt_tokens, server(621));
assert_eq!(counts.completion_tokens, server(160));
assert_eq!(
(counts.reused_tokens, counts.evaluated_tokens),
(server(605), server(16))
);
}
#[test]
fn usage_wins_over_timings_for_totals() {
let counts = exchange(&record(
"/v1/chat/completions",
chat_body(),
llama_cpp_response(),
None,
))
.counts;
assert_eq!(
(counts.prompt_tokens, counts.completion_tokens),
(server(25), server(4))
);
}
#[test]
fn reused_is_derived_from_evaluated() {
let response =
json!({"choices": [], "usage": {"prompt_tokens": 30}, "timings": {"prompt_n": 12}});
let counts = exchange(&record("/v1/chat/completions", chat_body(), response, None)).counts;
assert_eq!(counts.reused_tokens, server(18));
assert_eq!(counts.evaluated_tokens, server(12));
}
#[test]
fn derivation_needs_the_prompt_total() {
let response = json!({"choices": [], "timings": {"prompt_n": 12}});
let counts = exchange(&record("/v1/chat/completions", chat_body(), response, None)).counts;
assert_eq!(counts.reused_tokens, None);
assert_eq!(counts.prompt_tokens, None);
}
#[test]
fn derivation_never_goes_negative() {
let response = json!({"choices": [], "usage": {"prompt_tokens": 10, "prompt_tokens_details": {"cached_tokens": 50}}});
let counts = exchange(&record("/v1/chat/completions", chat_body(), response, None)).counts;
assert_eq!(counts.evaluated_tokens, server(0));
}
#[test]
fn ollama_native_fields() {
let response = json!({
"model": "m", "message": {"role": "assistant", "content": "hi"}, "done": true,
"prompt_eval_count": 12, "prompt_eval_duration": 3_500_000, "eval_count": 3,
"eval_duration": 9_000_000, "load_duration": 250_000, "total_duration": 20_000_000
});
let counts = exchange(&record(
"/api/chat",
chat_body(),
response,
Some(StreamKind::Ndjson),
))
.counts;
assert_eq!(
counts,
Counts {
prompt_tokens: server(12),
completion_tokens: server(3),
reused_tokens: None,
evaluated_tokens: None,
prefill_ms: server(3.5),
decode_ms: server(9.0),
load_ms: server(0.25),
}
);
}
#[test]
fn ollama_openai_compatible_fields() {
let response = json!({
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}}],
"system_fingerprint": "fp_ollama",
"usage": {"prompt_tokens": 40, "completion_tokens": 2, "total_tokens": 42}
});
let counts = exchange(&record(
"/v1/chat/completions",
chat_body(),
response.clone(),
Some(StreamKind::Sse),
))
.counts;
assert_eq!(
counts,
Counts {
prompt_tokens: server(40),
completion_tokens: server(2),
reused_tokens: None,
evaluated_tokens: None,
prefill_ms: Some(Measured {
value: 25.0,
source: Source::Proxy
}),
decode_ms: None,
load_ms: None,
}
);
let counts = exchange(&record("/v1/chat/completions", chat_body(), response, None)).counts;
assert_eq!(counts.prefill_ms, None);
}
#[test]
fn nothing_reported_is_all_null() {
let response = json!({"choices": [{"message": {"content": "x"}}]});
let counts = exchange(&record("/v1/chat/completions", chat_body(), response, None)).counts;
assert_eq!(counts, Counts::default());
}
#[test]
fn exchange_carries_the_request_and_the_output() {
let response = json!({
"choices": [{"message": {"role": "assistant", "content": "ok", "reasoning_content": "hm",
"tool_calls": [{"id": "c2", "function": {"name": "read", "arguments": "{}"}}]}}]
});
let body = chat_body();
let x = exchange(&record(
"/v1/chat/completions",
body.clone(),
response,
None,
));
assert_eq!(
(x.id, x.t_start_ms, x.t_end_ms, x.api),
(7, T0, T0 + 54, Api::OpenAiChat)
);
assert_eq!(x.model, "gemma");
assert!(x.ok);
assert_eq!(x.params.max_tokens, Some(32));
let expected = record::request_digest(Api::OpenAiChat, &body).unwrap();
assert_eq!(x.messages.len(), 4);
for (msg, digest) in x.messages.iter().zip(&expected.messages) {
assert_eq!(
(&msg.role, &msg.hash, msg.chars),
(&digest.role, &digest.hash, digest.chars)
);
assert_eq!(msg.tokens, None);
}
assert_eq!(x.messages[0].text.as_deref(), Some("Be brief."));
assert_eq!(x.messages[2].tool_calls[0].name, "ls");
assert_eq!(x.messages[3].tool_call_id.as_deref(), Some("c1"));
assert_eq!(x.tools.len(), 1);
assert_eq!(x.tools[0].name, "ls");
assert_eq!(x.tools_hash, expected.tools_hash);
assert_eq!(x.output.content_chars, 2);
assert_eq!(x.output.reasoning_chars, 2);
assert_eq!(x.output.tool_calls[0].id.as_deref(), Some("c2"));
}
#[test]
fn full_and_digest_recordings_give_the_same_exchange() {
for (path, response, stream) in [
(
"/v1/chat/completions",
llama_cpp_response(),
Some(StreamKind::Sse),
),
("/v1/chat/completions", vllm_response(), None),
(
"/api/chat",
json!({"message": {"content": "hi"}, "prompt_eval_count": 9, "eval_duration": 1000}),
Some(StreamKind::Ndjson),
),
] {
let full = record(path, chat_body(), response, stream);
let mut digest = full.clone();
record::apply_digest(&mut digest);
assert!(digest.request.body.is_none() && digest.request.digest.is_some());
let mut from_full = exchange(&full);
let from_digest = exchange(&digest);
assert!(from_full.messages.iter().all(|m| m.text.is_some()));
assert!(from_digest.messages.iter().all(|m| m.text.is_none()));
for m in &mut from_full.messages {
m.text = None;
}
assert_eq!(from_full, from_digest);
}
}
#[test]
fn tokenizer_counts_are_attached_by_hash() {
let mut r = record("/v1/chat/completions", chat_body(), vllm_response(), None);
let digest = record::request_digest(Api::OpenAiChat, &chat_body()).unwrap();
r.tokens = Some(TokenCounts {
messages: BTreeMap::from([(digest.messages[1].hash.clone(), 6)]),
tools: BTreeMap::from([(digest.tools[0].hash.clone(), 40)]),
});
let x = exchange(&r);
let tokenizer = |value| {
Some(Measured {
value,
source: Source::Tokenizer,
})
};
assert_eq!(x.messages[0].tokens, None);
assert_eq!(x.messages[1].tokens, tokenizer(6));
assert_eq!(x.tools[0].tokens, tokenizer(40));
}
#[test]
fn unreadable_response_is_still_analysed() {
let mut r = record("/v1/chat/completions", chat_body(), json!({}), None);
if let Some(response) = r.response.as_mut() {
response.body = None;
response.body_text = Some("{\"choices\": [tru".into());
response.truncated = true;
}
let x = exchange(&r);
assert!(!x.ok);
assert_eq!(x.counts, Counts::default());
assert_eq!(x.output, Output::default());
assert_eq!(x.messages.len(), 4);
}
#[test]
fn missing_model_is_empty() {
let mut body = chat_body();
body.as_object_mut().unwrap().remove("model");
assert_eq!(
exchange(&record("/v1/chat/completions", body, vllm_response(), None)).model,
""
);
}
#[test]
fn skip_reasons() {
let ok = record("/v1/chat/completions", chat_body(), vllm_response(), None);
let skip = |r: &Record| normalize(r).err();
let mut r = ok.clone();
r.request.path = "/v1/embeddings".into();
assert_eq!(skip(&r), Some(SkipReason::NotChat));
let mut r = ok.clone();
r.error = Some(ExchangeError::ClientDisconnected);
assert_eq!(skip(&r), Some(SkipReason::Error));
let mut r = ok.clone();
r.response = None;
assert_eq!(skip(&r), Some(SkipReason::Error));
let mut r = ok.clone();
if let Some(response) = r.response.as_mut() {
response.status = 500;
}
assert_eq!(skip(&r), Some(SkipReason::Non2xx));
let mut r = ok.clone();
r.request.body = None;
r.request.body_text = Some("not json".into());
assert_eq!(skip(&r), Some(SkipReason::Unparsed));
let mut r = ok.clone();
r.request.body = Some(json!({"model": "m", "prompt": "x"}));
assert_eq!(skip(&r), Some(SkipReason::Unparsed));
let mut r = ok.clone();
r.request.body = None;
assert_eq!(skip(&r), Some(SkipReason::Unparsed));
assert_eq!(skip(&ok), None);
}
#[test]
fn normalize_all_keeps_order_and_lists_skips() {
let a = record("/v1/chat/completions", chat_body(), vllm_response(), None);
let mut b = a.clone();
b.id = 8;
b.request.path = "/v1/models".into();
let mut c = a.clone();
c.id = 9;
let normalized = normalize_all(&[a, b, c]);
assert_eq!(
normalized
.exchanges
.iter()
.map(|x| x.id)
.collect::<Vec<_>>(),
vec![7, 9]
);
assert_eq!(
normalized.skipped,
vec![Skipped {
path: "/v1/models".into(),
reason: SkipReason::NotChat
}]
);
}
#[test]
fn api_detection_uses_the_path_only() {
assert_eq!(detect_api("/v1/chat/completions"), Api::OpenAiChat);
assert_eq!(detect_api("/v1/chat/completions?x=1"), Api::OpenAiChat);
assert_eq!(detect_api("/api/chat"), Api::OllamaChat);
assert_eq!(detect_api("/v1/completions"), Api::Other);
assert_eq!(detect_api("/v1/models"), Api::Other);
assert_eq!(detect_api("/api/generate"), Api::Other);
}
#[test]
fn api_names() {
assert_eq!(
serde_json::to_string(&Api::OpenAiChat).unwrap(),
"\"openai_chat\""
);
assert_eq!(
serde_json::to_string(&Api::OllamaChat).unwrap(),
"\"ollama_chat\""
);
}
}