use serde_json::Value;
use crate::memory::token_utils::count_tokens;
use crate::runtime::normalize_token_usage_with_hints;
use crate::types::{LLMResponse, Metadata, TokenUsage, ToolCall, UsageSource};
#[derive(Debug, Clone)]
pub(super) struct UsageEstimateContext {
pub(super) model: String,
pub(super) prompt_tokens: u64,
}
pub(super) fn from_vv_llm_response(
response: vv_llm::ChatResponse,
estimate: Option<UsageEstimateContext>,
) -> LLMResponse {
let mut raw = Metadata::new();
raw.insert("id".to_string(), Value::String(response.id));
raw.insert("model".to_string(), Value::String(response.model));
if let Some(reasoning_content) = response
.reasoning_content
.as_deref()
.filter(|value| !value.is_empty())
{
raw.insert(
"reasoning_content".to_string(),
Value::String(reasoning_content.to_string()),
);
}
let token_usage = response.usage.map(from_vv_llm_usage).unwrap_or_else(|| {
estimate_missing_usage(&response.content, &response.tool_calls, estimate)
});
raw.insert(
"usage".to_string(),
Value::Object(token_usage.provider_usage.clone()),
);
LLMResponse {
content: response.content,
tool_calls: response
.tool_calls
.into_iter()
.enumerate()
.map(|(index, tool_call)| from_vv_llm_tool_call(tool_call, index))
.collect(),
raw,
token_usage,
}
}
pub(super) fn from_vv_llm_tool_call(tool_call: vv_llm::ToolCall, index: usize) -> ToolCall {
let mut normalized = ToolCall::from_raw_arguments(
normalize_tool_call_id(&tool_call.id, index),
normalize_tool_call_name(&tool_call.name),
Value::String(tool_call.arguments),
);
if let Some(extra_content) = tool_call.extra_content {
normalized.extra_content = Some(match normalized.extra_content.take() {
Some(existing) => merge_tool_call_extra_content(existing, extra_content),
None => extra_content,
});
}
normalized
}
pub(super) fn merge_tool_call_extra_content(existing: Value, extra_content: Value) -> Value {
match (existing, extra_content) {
(Value::Object(mut existing), Value::Object(extra)) => {
existing.extend(extra);
Value::Object(existing)
}
(existing, extra_content) => {
serde_json::json!({"parse_error": existing, "provider_extra_content": extra_content})
}
}
}
fn normalize_tool_call_id(id: &str, index: usize) -> String {
let id = id.trim();
if id.is_empty() {
format!("call_generated_{index}")
} else {
id.to_string()
}
}
fn normalize_tool_call_name(name: &str) -> String {
name.replace(' ', "")
}
pub(super) fn from_vv_llm_usage(usage: vv_llm::ChatUsage) -> TokenUsage {
let provider_usage = usage
.raw_usage
.as_ref()
.and_then(Value::as_object)
.cloned()
.unwrap_or_else(|| {
serde_json::to_value(&usage)
.ok()
.and_then(|value| value.as_object().cloned())
.unwrap_or_default()
});
let mut normalization_input = Value::Object(provider_usage.clone());
overlay_typed_counts(&mut normalization_input, &usage);
overlay_typed_cache_usage(&mut normalization_input, &usage);
let mut normalized = normalize_token_usage_with_hints(
&normalization_input,
Some(UsageSource::ProviderReported),
None,
);
normalized.provider_usage = provider_usage;
normalized
}
fn overlay_typed_counts(raw: &mut Value, usage: &vv_llm::ChatUsage) {
let Some(raw) = raw.as_object_mut() else {
return;
};
let native_input_shape = raw.contains_key("input_tokens") && !raw.contains_key("prompt_tokens");
let fields = if native_input_shape {
vec![
("input_tokens", usage.input_tokens),
("output_tokens", usage.output_tokens),
]
} else {
vec![
("prompt_tokens", usage.prompt_tokens),
("completion_tokens", usage.completion_tokens),
("total_tokens", usage.total_tokens),
("input_tokens", usage.input_tokens),
("output_tokens", usage.output_tokens),
]
};
for (key, value) in fields {
if let Some(value) = value {
raw.entry(key.to_string())
.or_insert_with(|| Value::from(value));
}
}
}
fn overlay_typed_cache_usage(raw: &mut Value, usage: &vv_llm::ChatUsage) {
let Some(raw) = raw.as_object_mut() else {
return;
};
let openai_usage = raw.contains_key("prompt_tokens")
|| raw.contains_key("completion_tokens")
|| raw.contains_key("prompt_tokens_details");
if let Some(cache_read) = usage.cache_read_input_tokens {
if openai_usage {
insert_token_detail(raw, "cached_tokens", cache_read);
} else {
raw.insert(
"cache_read_input_tokens".to_string(),
Value::from(cache_read),
);
}
}
if let Some(cache_creation) = usage.cache_creation_input_tokens {
if openai_usage {
insert_token_detail(raw, "cache_creation_tokens", cache_creation);
} else {
raw.insert(
"cache_creation_input_tokens".to_string(),
Value::from(cache_creation),
);
}
}
}
fn insert_token_detail(raw: &mut serde_json::Map<String, Value>, key: &str, value: u32) {
let details = raw
.entry("prompt_tokens_details".to_string())
.or_insert_with(|| Value::Object(serde_json::Map::new()));
if !details.is_object() {
*details = Value::Object(serde_json::Map::new());
}
details
.as_object_mut()
.expect("token details were normalized to an object")
.insert(key.to_string(), Value::from(value));
}
pub(super) fn estimate_missing_usage(
content: &str,
tool_calls: &[vv_llm::ToolCall],
estimate: Option<UsageEstimateContext>,
) -> TokenUsage {
let Some(estimate) = estimate else {
return TokenUsage::default();
};
let completion_payload = completion_payload_for_usage(content, tool_calls);
let completion_tokens = count_tokens(&completion_payload, &estimate.model);
let total_tokens = estimate.prompt_tokens + completion_tokens;
TokenUsage {
input_tokens: Some(estimate.prompt_tokens),
output_tokens: Some(completion_tokens),
total_tokens: Some(total_tokens),
usage_source: UsageSource::Estimated,
..TokenUsage::default()
}
}
pub(super) fn completion_payload_for_usage(
content: &str,
tool_calls: &[vv_llm::ToolCall],
) -> String {
if tool_calls.is_empty() {
return content.to_string();
}
let mut payload = content.to_string();
if let Ok(tool_payload) = serde_json::to_string(tool_calls) {
if !payload.is_empty() {
payload.push('\n');
}
payload.push_str(&tool_payload);
}
payload
}
#[cfg(test)]
mod tests {
use super::from_vv_llm_usage;
use crate::types::{CacheUsageStatus, UsageSource};
#[test]
fn typed_openai_cache_zero_is_accounted_without_rewriting_raw_usage() {
let raw = serde_json::json!({
"prompt_tokens": 11,
"completion_tokens": 7,
"total_tokens": 18
});
let normalized = from_vv_llm_usage(vv_llm::ChatUsage {
prompt_tokens: Some(11),
completion_tokens: Some(7),
total_tokens: Some(18),
input_tokens: Some(11),
output_tokens: Some(7),
cache_read_input_tokens: Some(0),
cache_creation_input_tokens: None,
raw_usage: Some(raw.clone()),
});
assert_eq!(normalized.usage_source, UsageSource::ProviderReported);
assert_eq!(
normalized.cache_usage.status,
CacheUsageStatus::ProviderReported
);
assert_eq!(normalized.cache_usage.read_input_tokens, Some(0));
assert_eq!(normalized.cache_usage.uncached_input_tokens, Some(11));
assert_eq!(normalized.provider_usage, raw.as_object().unwrap().clone());
}
#[test]
fn typed_anthropic_cache_produces_total_and_uncached_input() {
let raw = serde_json::json!({
"input_tokens": 11,
"output_tokens": 7
});
let normalized = from_vv_llm_usage(vv_llm::ChatUsage {
prompt_tokens: Some(11),
completion_tokens: Some(7),
total_tokens: Some(18),
input_tokens: Some(11),
output_tokens: Some(7),
cache_read_input_tokens: Some(5),
cache_creation_input_tokens: Some(3),
raw_usage: Some(raw.clone()),
});
assert_eq!(normalized.input_tokens, Some(19));
assert_eq!(normalized.total_tokens, Some(26));
assert_eq!(normalized.cache_usage.read_input_tokens, Some(5));
assert_eq!(normalized.cache_usage.write_input_tokens, Some(3));
assert_eq!(normalized.cache_usage.uncached_input_tokens, Some(14));
assert_eq!(normalized.provider_usage, raw.as_object().unwrap().clone());
}
#[test]
fn typed_anthropic_cache_precedes_invalid_raw_without_mutating_it() {
let raw = serde_json::json!({
"input_tokens": 11,
"output_tokens": 7,
"cache_read_input_tokens": "invalid",
"cache_creation_input_tokens": null
});
let normalized = from_vv_llm_usage(vv_llm::ChatUsage {
prompt_tokens: Some(11),
completion_tokens: Some(7),
total_tokens: Some(18),
input_tokens: Some(11),
output_tokens: Some(7),
cache_read_input_tokens: Some(5),
cache_creation_input_tokens: Some(3),
raw_usage: Some(raw.clone()),
});
assert_eq!(normalized.input_tokens, Some(19));
assert_eq!(normalized.cache_usage.read_input_tokens, Some(5));
assert_eq!(normalized.cache_usage.write_input_tokens, Some(3));
assert_eq!(normalized.cache_usage.uncached_input_tokens, Some(14));
assert_eq!(normalized.provider_usage, raw.as_object().unwrap().clone());
}
#[test]
fn typed_cache_value_precedes_invalid_raw_detail_without_mutating_it() {
let raw = serde_json::json!({
"prompt_tokens": 11,
"completion_tokens": 7,
"total_tokens": 18,
"prompt_tokens_details": {"cached_tokens": "invalid"}
});
let normalized = from_vv_llm_usage(vv_llm::ChatUsage {
prompt_tokens: Some(11),
completion_tokens: Some(7),
total_tokens: Some(18),
input_tokens: Some(11),
output_tokens: Some(7),
cache_read_input_tokens: Some(5),
cache_creation_input_tokens: None,
raw_usage: Some(raw.clone()),
});
assert_eq!(normalized.cache_usage.read_input_tokens, Some(5));
assert_eq!(normalized.cache_usage.uncached_input_tokens, Some(6));
assert_eq!(normalized.provider_usage, raw.as_object().unwrap().clone());
}
}