katu_core/
usage.rs

1//! # katu_core::usage
2//!
3//! ## 职责
4//! 定义 LLM 请求的 token 用量与费用计量类型。
5//!
6//! ## 依赖
7//! 无（本模块是 katu-core 的底层类型）
8//!
9//! ## 对外接口
10//! - `Usage` — token 用量统计
11//! - `Cost` — 费用明细（美元）
12//!
13//! ## 调用者
14//! - `katu_core::message` — AssistantMessage 持有 Usage
15//! - 上层 crate 通过 `katu_core::usage::*` 使用
16
17use serde::{Deserialize, Serialize};
18
19// ---------------------------------------------------------------------------
20// Cost
21// ---------------------------------------------------------------------------
22
23/// 费用明细（单位：美元）。
24///
25/// 按 token 类别拆分，`total` 为各项之和。
26/// 仅在上层持有定价表时可计算，因此 `Usage` 中为 `Option<Cost>`。
27#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
28pub struct Cost {
29    /// 输入 token 费用
30    pub input: f64,
31    /// 输出 token 费用
32    pub output: f64,
33    /// 缓存读取 token 费用
34    pub cache_read: f64,
35    /// 缓存写入 token 费用
36    pub cache_write: f64,
37    /// 总费用（各项之和）
38    pub total: f64,
39}
40
41// ---------------------------------------------------------------------------
42// Usage
43// ---------------------------------------------------------------------------
44
45/// LLM 请求的 token 用量统计。
46///
47/// 字段语义遵循"包含式总量 + 非重叠分解"惯例
48/// （与 OpenAI / Anthropic / AI SDK 对齐）：
49///
50/// - `input_tokens` — 总输入 token，**包含**缓存读/写
51/// - `output_tokens` — 总输出 token，**包含**推理 token
52/// - `total_tokens` — provider 报告的总量，或 `input + output`
53///
54/// 非重叠分解：
55/// - `cache_read_tokens` + `cache_write_tokens` + 非缓存部分 = `input_tokens`
56/// - `reasoning_tokens` ≤ `output_tokens`
57///
58/// 部分 provider 不报告细分字段，此时为 `0` 或由上层推断。
59#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
60pub struct Usage {
61    /// 总输入 token（含缓存读/写）
62    pub input_tokens: u32,
63    /// 总输出 token（含推理）
64    pub output_tokens: u32,
65    /// 从缓存读取的输入 token
66    pub cache_read_tokens: u32,
67    /// 写入缓存的输入 token
68    pub cache_write_tokens: u32,
69    /// 推理/思考 token（output 的子集），`None` 表示 provider 未报告
70    pub reasoning_tokens: Option<u32>,
71    /// 总 token（provider 报告值或 input + output）
72    pub total_tokens: u32,
73    /// 费用明细，仅在上层持有定价表时可计算
74    pub cost: Option<Cost>,
75}
76
77impl Usage {
78    /// 可见输出 token = output - reasoning，至少为 0。
79    pub fn visible_output_tokens(&self) -> u32 {
80        self.output_tokens
81            .saturating_sub(self.reasoning_tokens.unwrap_or(0))
82    }
83
84    /// 非缓存输入 token = input - cache_read - cache_write，至少为 0。
85    pub fn non_cached_input_tokens(&self) -> u32 {
86        self.input_tokens
87            .saturating_sub(self.cache_read_tokens)
88            .saturating_sub(self.cache_write_tokens)
89    }
90}
91
92
93// ---------------------------------------------------------------------------
94// Tests
95// ---------------------------------------------------------------------------
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_usage_default_is_zero() {
103        let u = Usage::default();
104        assert_eq!(u.input_tokens, 0);
105        assert_eq!(u.output_tokens, 0);
106        assert_eq!(u.total_tokens, 0);
107        assert!(u.cost.is_none());
108    }
109
110    #[test]
111    fn test_visible_output_tokens() {
112        let u = Usage {
113            output_tokens: 100,
114            reasoning_tokens: Some(40),
115            ..Default::default()
116        };
117        assert_eq!(u.visible_output_tokens(), 60);
118    }
119
120    #[test]
121    fn test_visible_output_tokens_no_reasoning() {
122        let u = Usage {
123            output_tokens: 100,
124            reasoning_tokens: None,
125            ..Default::default()
126        };
127        assert_eq!(u.visible_output_tokens(), 100);
128    }
129
130    #[test]
131    fn test_visible_output_tokens_saturating() {
132        let u = Usage {
133            output_tokens: 10,
134            reasoning_tokens: Some(999),
135            ..Default::default()
136        };
137        assert_eq!(u.visible_output_tokens(), 0);
138    }
139
140    #[test]
141    fn test_non_cached_input_tokens() {
142        let u = Usage {
143            input_tokens: 1000,
144            cache_read_tokens: 600,
145            cache_write_tokens: 200,
146            ..Default::default()
147        };
148        assert_eq!(u.non_cached_input_tokens(), 200);
149    }
150
151    #[test]
152    fn test_non_cached_input_tokens_saturating() {
153        let u = Usage {
154            input_tokens: 100,
155            cache_read_tokens: 80,
156            cache_write_tokens: 80,
157            ..Default::default()
158        };
159        assert_eq!(u.non_cached_input_tokens(), 0);
160    }
161
162    #[test]
163    fn test_usage_serde_roundtrip() {
164        let u = Usage {
165            input_tokens: 500,
166            output_tokens: 200,
167            cache_read_tokens: 100,
168            cache_write_tokens: 50,
169            reasoning_tokens: Some(30),
170            total_tokens: 700,
171            cost: Some(Cost {
172                input: 0.005,
173                output: 0.010,
174                cache_read: 0.001,
175                cache_write: 0.002,
176                total: 0.018,
177            }),
178        };
179        let json = serde_json::to_string(&u).unwrap();
180        let restored: Usage = serde_json::from_str(&json).unwrap();
181        assert_eq!(u, restored);
182    }
183
184    #[test]
185    fn test_cost_serde_roundtrip() {
186        let c = Cost {
187            input: 0.01,
188            output: 0.03,
189            cache_read: 0.001,
190            cache_write: 0.002,
191            total: 0.043,
192        };
193        let json = serde_json::to_string(&c).unwrap();
194        let restored: Cost = serde_json::from_str(&json).unwrap();
195        assert_eq!(c, restored);
196    }
197}
katu_core/usage.rs

katu_core/
usage.rs
