vv_agent/runtime/
token_usage.rs

1use serde_json::Value;
2
3use crate::types::{CycleRecord, TaskTokenUsage, TokenUsage};
4
5pub fn normalize_token_usage(raw_usage: &Value) -> TokenUsage {
6    let Some(raw) = raw_usage.as_object() else {
7        return TokenUsage::default();
8    };
9
10    let prompt_tokens = read_int(raw.get("prompt_tokens"));
11    let completion_tokens = read_int(raw.get("completion_tokens"));
12    let input_tokens = read_int(raw.get("input_tokens")).or(prompt_tokens);
13    let output_tokens = read_int(raw.get("output_tokens")).or(completion_tokens);
14    let total_tokens = read_int(raw.get("total_tokens")).unwrap_or_else(|| {
15        prompt_tokens.or(input_tokens).unwrap_or_default()
16            + completion_tokens.or(output_tokens).unwrap_or_default()
17    });
18    let cached_tokens = read_nested_int(
19        raw_usage,
20        &[
21            &["prompt_tokens_details", "cached_tokens"],
22            &["input_tokens_details", "cached_tokens"],
23            &["cache_read_input_tokens"],
24            &["cache_read_tokens"],
25        ],
26    )
27    .unwrap_or_default();
28    let reasoning_tokens = read_nested_int(
29        raw_usage,
30        &[
31            &["completion_tokens_details", "reasoning_tokens"],
32            &["output_tokens_details", "reasoning_tokens"],
33            &["reasoning_tokens"],
34        ],
35    )
36    .unwrap_or_default();
37    let cache_creation_tokens = read_nested_int(
38        raw_usage,
39        &[
40            &["input_tokens_details", "cache_creation_tokens"],
41            &["prompt_tokens_details", "cache_creation_tokens"],
42            &["cache_creation_input_tokens"],
43            &["cache_creation_tokens"],
44        ],
45    )
46    .unwrap_or_default();
47
48    TokenUsage {
49        prompt_tokens: prompt_tokens.or(input_tokens).unwrap_or_default(),
50        completion_tokens: completion_tokens.or(output_tokens).unwrap_or_default(),
51        total_tokens,
52        cached_tokens,
53        reasoning_tokens,
54        input_tokens: input_tokens.unwrap_or_default(),
55        output_tokens: output_tokens.unwrap_or_default(),
56        cache_creation_tokens,
57        raw: raw_usage.clone(),
58    }
59}
60
61pub fn summarize_task_token_usage(cycles: &[CycleRecord]) -> TaskTokenUsage {
62    let mut summary = TaskTokenUsage::default();
63    for cycle in cycles {
64        summary.add_cycle(cycle.index, cycle.token_usage.clone());
65    }
66    summary
67}
68
69fn read_nested_int(source: &Value, path_options: &[&[&str]]) -> Option<u64> {
70    for path in path_options {
71        let mut current = source;
72        let mut matched = true;
73        for key in *path {
74            let Some(next) = current.as_object().and_then(|object| object.get(*key)) else {
75                matched = false;
76                break;
77            };
78            current = next;
79        }
80        if matched {
81            if let Some(value) = read_int(Some(current)) {
82                return Some(value);
83            }
84        }
85    }
86    None
87}
88
89fn read_int(value: Option<&Value>) -> Option<u64> {
90    match value? {
91        Value::Bool(_) => None,
92        Value::Number(number) => number
93            .as_u64()
94            .or_else(|| {
95                number
96                    .as_i64()
97                    .and_then(|value| (value >= 0).then_some(value as u64))
98            })
99            .or_else(|| number.as_f64().and_then(float_to_u64)),
100        Value::String(value) => value.trim().parse::<u64>().ok(),
101        _ => None,
102    }
103}
104
105fn float_to_u64(value: f64) -> Option<u64> {
106    value
107        .is_finite()
108        .then_some(value.trunc())
109        .and_then(|value| {
110            if value >= 0.0 {
111                Some(value as u64)
112            } else {
113                None
114            }
115        })
116}
vv_agent/runtime/token_usage.rs

vv_agent/runtime/
token_usage.rs
