key_token/
usage.rs

1use std::sync::atomic::{AtomicU32, Ordering};
2
3/// Token 累积器：streaming 过程中原子更新
4#[derive(Debug, Default)]
5pub struct UsageAccumulator {
6    input_tokens: AtomicU32,
7    output_tokens: AtomicU32,
8}
9
10impl UsageAccumulator {
11    pub fn new() -> Self {
12        Self::default()
13    }
14
15    /// 添加输出 token（流式响应中累积）
16    pub fn add_output(&self, tokens: u32) {
17        self.output_tokens.fetch_add(tokens, Ordering::Relaxed);
18    }
19
20    /// 设置输入 token（通常从 provider 报告获取）
21    pub fn set_input(&self, tokens: u32) {
22        self.input_tokens.store(tokens, Ordering::Relaxed);
23    }
24
25    /// 获取当前用量快照
26    pub fn snapshot(&self) -> (u32, u32) {
27        (
28            self.input_tokens.load(Ordering::Relaxed),
29            self.output_tokens.load(Ordering::Relaxed),
30        )
31    }
32
33    /// 获取总 token 数
34    pub fn total_tokens(&self) -> u32 {
35        let (input, output) = self.snapshot();
36        input + output
37    }
38}
39
40/// 从 snapshot 创建新的 UsageAccumulator
41///
42/// 注意：这会创建独立的 Accumulator，不会共享状态。
43/// 如需共享状态，请使用 Arc<UsageAccumulator>。
44impl From<(u32, u32)> for UsageAccumulator {
45    fn from((input, output): (u32, u32)) -> Self {
46        let acc = Self::new();
47        acc.set_input(input);
48        // 直接设置输出 token，避免循环
49        acc.output_tokens.store(output, Ordering::Relaxed);
50        acc
51    }
52}
53
54/// 最终用量记录
55#[derive(Debug, Clone, Copy)]
56pub struct UsageRecord {
57    pub input_tokens: u32,
58    pub output_tokens: u32,
59}
60
61impl UsageRecord {
62    pub fn total(&self) -> u32 {
63        self.input_tokens + self.output_tokens
64    }
65}
66
67impl From<(u32, u32)> for UsageRecord {
68    fn from((input, output): (u32, u32)) -> Self {
69        Self {
70            input_tokens: input,
71            output_tokens: output,
72        }
73    }
74}
key_token/usage.rs

key_token/
usage.rs
