lc_core/token_counter/
counter.rs

1//! Token 计数器 trait 与用量统计
2
3use lc_schema::Message;
4
5/// Token 计数器 trait
6pub trait TokenCounter: Send + Sync {
7    /// 计算文本 token 数
8    fn count_tokens(&self, text: &str) -> u32;
9    /// 计算消息列表 token 数
10    fn count_messages(&self, messages: &[Message]) -> u32;
11}
12
13/// 字符比估算计数器(零依赖快路径)
14///
15/// `count_tokens = len / ratio`,不依赖 tiktoken BPE 模型,离线/测试环境可用。
16/// 默认 `ratio = 4` 对齐旧的 `len/4` 粗略估算;`count_messages` 沿用
17/// `TiktokenCounter` 的开销结构(每条消息 4 + 名称 + 图片 1000/张,结尾 2),
18/// 保证与 BPE 口径的预算语义一致。
19#[derive(Debug, Clone)]
20pub struct CharRatioCounter {
21    ratio: u32,
22}
23
24impl CharRatioCounter {
25    /// 创建字符比计数器。
26    ///
27    /// `ratio` 为每 token 的字符数,至少为 1。
28    pub fn new(ratio: u32) -> Self {
29        Self {
30            ratio: ratio.max(1),
31        }
32    }
33}
34
35impl TokenCounter for CharRatioCounter {
36    fn count_tokens(&self, text: &str) -> u32 {
37        text.len() as u32 / self.ratio
38    }
39
40    fn count_messages(&self, messages: &[Message]) -> u32 {
41        let mut total = 0u32;
42        for msg in messages {
43            total += 4; // OpenAI 消息格式开销
44            total += self.count_tokens(&msg.content);
45            if let Some(name) = &msg.name {
46                total += self.count_tokens(name);
47            }
48            // 图片内容粗略计为 1000 token/图
49            for _ in &msg.images {
50                total += 1000;
51            }
52        }
53        total += 2; // 对话边界标记
54        total
55    }
56}
57
58#[cfg(test)]
59mod char_ratio_tests {
60    use super::*;
61
62    #[test]
63    fn test_char_ratio_counts_bytes() {
64        let counter = CharRatioCounter::new(4);
65        assert_eq!(counter.count_tokens(""), 0);
66        assert_eq!(counter.count_tokens("Hello"), 1); // 5 / 4 = 1
67        assert_eq!(counter.count_tokens("Hello World"), 2); // 11 / 4 = 2
68    }
69
70    #[test]
71    fn test_char_ratio_ratio_at_least_one() {
72        let counter = CharRatioCounter::new(0);
73        assert!(counter.count_tokens("x") >= 1);
74    }
75
76    #[test]
77    fn test_char_ratio_count_messages_matches_structure() {
78        let counter = CharRatioCounter::new(4);
79        let msgs = vec![Message::system("You are helpful."), Message::human("Hi")];
80        let n = counter.count_messages(&msgs);
81        // 至少含 2*4 开销 + 2 边界
82        assert!(n >= 10);
83    }
84}
85
86/// Token 用量统计（计数器模块内部类型）
87///
88/// 注意：`language_models::TokenUsage` 是 LLM API 返回的用量（字段为 `usize`），
89/// 此 `TrackerTokenUsage` 是本地追踪累计用量，**字段同为 `usize`**，两者可互转
90/// 而无精度损失（Q6：统一底层类型，不再存在 `usize` vs `u32` 的命名冲突）。
91/// 这里不使用 `TokenUsage` 别名，避免与 `language_models::TokenUsage` 混淆。
92#[derive(Debug, Clone, Default, PartialEq)]
93pub struct TrackerTokenUsage {
94    /// 累计的 prompt token 数
95    pub prompt_tokens: usize,
96    /// 累计的 completion token 数
97    pub completion_tokens: usize,
98    /// 累计的 total token 数(prompt + completion)
99    pub total_tokens: usize,
100}
101
102impl TrackerTokenUsage {
103    /// 创建空的用量统计(全零)。
104    pub fn new() -> Self {
105        Self::default()
106    }
107
108    /// 累加用量
109    pub fn add(&mut self, prompt: usize, completion: usize) {
110        self.prompt_tokens += prompt;
111        self.completion_tokens += completion;
112        self.total_tokens = self.prompt_tokens + self.completion_tokens;
113    }
114
115    /// 重置用量统计为全零。
116    pub fn reset(&mut self) {
117        *self = Self::default();
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn test_usage_add() {
127        let mut u = TrackerTokenUsage::new();
128        u.add(10, 20);
129        assert_eq!(u.prompt_tokens, 10);
130        assert_eq!(u.completion_tokens, 20);
131        assert_eq!(u.total_tokens, 30);
132    }
133
134    #[test]
135    fn test_usage_accumulate() {
136        let mut u = TrackerTokenUsage::new();
137        u.add(10, 20);
138        u.add(5, 5);
139        assert_eq!(u.prompt_tokens, 15);
140        assert_eq!(u.total_tokens, 40);
141    }
142
143    #[test]
144    fn test_usage_reset() {
145        let mut u = TrackerTokenUsage::new();
146        u.add(10, 20);
147        u.reset();
148        assert_eq!(u, TrackerTokenUsage::new());
149    }
150}
lc_core/token_counter/counter.rs

lc_core/token_counter/
counter.rs
