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/// Token 用量统计（计数器模块内部类型）
14///
15/// 注意：`language_models::TokenUsage` 是 LLM API 返回的用量（字段为 `usize`），
16/// 此 `TrackerTokenUsage` 是本地追踪累计用量（字段为 `u32`），两者职责不同。
17#[derive(Debug, Clone, Default, PartialEq)]
18pub struct TrackerTokenUsage {
19    pub prompt_tokens: u32,
20    pub completion_tokens: u32,
21    pub total_tokens: u32,
22}
23
24impl TrackerTokenUsage {
25    pub fn new() -> Self {
26        Self::default()
27    }
28
29    /// 累加用量
30    pub fn add(&mut self, prompt: u32, completion: u32) {
31        self.prompt_tokens += prompt;
32        self.completion_tokens += completion;
33        self.total_tokens = self.prompt_tokens + self.completion_tokens;
34    }
35
36    pub fn reset(&mut self) {
37        *self = Self::default();
38    }
39}
40
41// Re-export as TokenUsage for backward compatibility within this module
42pub use TrackerTokenUsage as TokenUsage;
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47
48    #[test]
49    fn test_usage_add() {
50        let mut u = TrackerTokenUsage::new();
51        u.add(10, 20);
52        assert_eq!(u.prompt_tokens, 10);
53        assert_eq!(u.completion_tokens, 20);
54        assert_eq!(u.total_tokens, 30);
55    }
56
57    #[test]
58    fn test_usage_accumulate() {
59        let mut u = TrackerTokenUsage::new();
60        u.add(10, 20);
61        u.add(5, 5);
62        assert_eq!(u.prompt_tokens, 15);
63        assert_eq!(u.total_tokens, 40);
64    }
65
66    #[test]
67    fn test_usage_reset() {
68        let mut u = TrackerTokenUsage::new();
69        u.add(10, 20);
70        u.reset();
71        assert_eq!(u, TrackerTokenUsage::new());
72    }
73}
lc_core/token_counter/counter.rs

lc_core/token_counter/
counter.rs
