Skip to main content

vv_agent/types/
token_usage.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
5pub struct TokenUsage {
6    pub prompt_tokens: u64,
7    pub completion_tokens: u64,
8    pub total_tokens: u64,
9    pub cached_tokens: u64,
10    pub reasoning_tokens: u64,
11    pub input_tokens: u64,
12    pub output_tokens: u64,
13    pub cache_creation_tokens: u64,
14    pub raw: Value,
15}
16
17impl TokenUsage {
18    pub fn has_usage(&self) -> bool {
19        self.prompt_tokens > 0
20            || self.completion_tokens > 0
21            || self.total_tokens > 0
22            || self.cached_tokens > 0
23            || self.reasoning_tokens > 0
24            || self.input_tokens > 0
25            || self.output_tokens > 0
26            || self.cache_creation_tokens > 0
27    }
28}
29
30#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
31pub struct CycleTokenUsage {
32    pub cycle_index: u32,
33    pub usage: TokenUsage,
34}
35
36#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
37pub struct TaskTokenUsage {
38    pub prompt_tokens: u64,
39    pub completion_tokens: u64,
40    pub total_tokens: u64,
41    pub cached_tokens: u64,
42    pub reasoning_tokens: u64,
43    pub input_tokens: u64,
44    pub output_tokens: u64,
45    pub cache_creation_tokens: u64,
46    pub cycles: Vec<CycleTokenUsage>,
47}
48
49impl TaskTokenUsage {
50    pub fn add_cycle(&mut self, cycle_index: u32, usage: TokenUsage) {
51        if !usage.has_usage() {
52            return;
53        }
54        self.prompt_tokens += usage.prompt_tokens;
55        self.completion_tokens += usage.completion_tokens;
56        self.total_tokens += usage.total_tokens;
57        self.cached_tokens += usage.cached_tokens;
58        self.reasoning_tokens += usage.reasoning_tokens;
59        self.input_tokens += usage.input_tokens;
60        self.output_tokens += usage.output_tokens;
61        self.cache_creation_tokens += usage.cache_creation_tokens;
62        self.cycles.push(CycleTokenUsage { cycle_index, usage });
63    }
64}