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` 是本地追踪累计用量（字段为 `u32`），两者职责不同。
90#[derive(Debug, Clone, Default, PartialEq)]
91pub struct TrackerTokenUsage {
92    pub prompt_tokens: u32,
93    pub completion_tokens: u32,
94    pub total_tokens: u32,
95}
96
97impl TrackerTokenUsage {
98    pub fn new() -> Self {
99        Self::default()
100    }
101
102    /// 累加用量
103    pub fn add(&mut self, prompt: u32, completion: u32) {
104        self.prompt_tokens += prompt;
105        self.completion_tokens += completion;
106        self.total_tokens = self.prompt_tokens + self.completion_tokens;
107    }
108
109    pub fn reset(&mut self) {
110        *self = Self::default();
111    }
112}
113
114// Re-export as TokenUsage for backward compatibility within this module
115pub use TrackerTokenUsage as TokenUsage;
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_usage_add() {
123        let mut u = TrackerTokenUsage::new();
124        u.add(10, 20);
125        assert_eq!(u.prompt_tokens, 10);
126        assert_eq!(u.completion_tokens, 20);
127        assert_eq!(u.total_tokens, 30);
128    }
129
130    #[test]
131    fn test_usage_accumulate() {
132        let mut u = TrackerTokenUsage::new();
133        u.add(10, 20);
134        u.add(5, 5);
135        assert_eq!(u.prompt_tokens, 15);
136        assert_eq!(u.total_tokens, 40);
137    }
138
139    #[test]
140    fn test_usage_reset() {
141        let mut u = TrackerTokenUsage::new();
142        u.add(10, 20);
143        u.reset();
144        assert_eq!(u, TrackerTokenUsage::new());
145    }
146}
lc_core/token_counter/counter.rs

lc_core/token_counter/
counter.rs
