tokrs 0.2.3

A CLI tool for counting tokens
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! 项目通用结构体与枚举
//!
//! 定义了AppKind(app类型)枚举,列举支持的app,经 clap::ValueEnum 直接驱动 --app 参数
//! 定义了UsageEntry结构体保存用户数据
//! 定义了TokenTotals结构体用于从用户请求数据统计Token的数量
//!
use clap::ValueEnum;
use std::fmt;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, ValueEnum)]
pub enum AppKind {
    Claude,
    Codex,
    #[value(name = "opencode")]
    OpenCode,
    Gemini,
    Grok,
    Pi,
    Kimi,
}

impl AppKind {
    pub const ALL: [AppKind; 7] = [
        AppKind::Claude,
        AppKind::Codex,
        AppKind::OpenCode,
        AppKind::Gemini,
        AppKind::Grok,
        AppKind::Pi,
        AppKind::Kimi,
    ];

    pub fn as_str(&self) -> &'static str {
        match self {
            AppKind::Claude => "claude",
            AppKind::Codex => "codex",
            AppKind::OpenCode => "opencode",
            AppKind::Gemini => "gemini",
            AppKind::Grok => "grok",
            AppKind::Pi => "pi",
            AppKind::Kimi => "kimi",
        }
    }
}

impl fmt::Display for AppKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

#[derive(Debug, Clone)]
pub struct UsageEntry {
    pub app: AppKind,
    pub model: String,
    #[allow(dead_code)]
    pub session_id: Option<String>,
    pub created_at: i64,
    /// fresh input(与缓存无关的增量输入)
    ///
    /// codex/gemini/grok 上游 input 含缓存, 已在解析层扣除归一, 全 app 语义统一
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_read_tokens: u64,
    pub cache_creation_tokens: u64,
    /// 上游自报成本(USD), 仅当来源可信且 >0 时填充
    pub self_cost_usd: Option<f64>,
    /// 最终成本(USD): 自报优先, 否则由定价表估价, 均无则 None(计入 unpriced)
    pub cost_usd: Option<f64>,
}

impl UsageEntry {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        app: AppKind,
        model: String,
        session_id: Option<String>,
        created_at: i64,
        input_tokens: u64,
        output_tokens: u64,
        cache_read_tokens: u64,
        cache_creation_tokens: u64,
        self_cost_usd: Option<f64>,
    ) -> Self {
        Self {
            app,
            model,
            session_id,
            created_at,
            input_tokens,
            output_tokens,
            cache_read_tokens,
            cache_creation_tokens,
            self_cost_usd,
            cost_usd: None,
        }
    }

    #[cfg(test)]
    pub fn total_tokens(&self) -> u64 {
        self.input_tokens + self.output_tokens + self.cache_read_tokens + self.cache_creation_tokens
    }
}

#[derive(Debug, Default, Clone, Copy, PartialEq)]
pub struct TokenTotals {
    pub requests: u64,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_read_tokens: u64,
    pub cache_creation_tokens: u64,
    /// 已确定成本的请求成本合计(USD), 含自报与估价
    pub cost_usd: f64,
    /// 无任何成本来源的请求数, 成本合计不含这些请求
    pub unpriced: u64,
}

impl TokenTotals {
    pub fn add_entry(&mut self, entry: &UsageEntry) {
        self.requests += 1;
        self.input_tokens += entry.input_tokens;
        self.output_tokens += entry.output_tokens;
        self.cache_read_tokens += entry.cache_read_tokens;
        self.cache_creation_tokens += entry.cache_creation_tokens;
        match entry.cost_usd {
            Some(cost) => self.cost_usd += cost,
            None => self.unpriced += 1,
        }
    }

    pub fn total_tokens(&self) -> u64 {
        self.input_tokens + self.output_tokens + self.cache_read_tokens + self.cache_creation_tokens
    }
}