lc_core/token_counter/
tracker.rs

1//! Token 追踪 LLM 包装器与成本估算
2
3use std::sync::Arc;
4
5use crate::language_models::LLMResult;
6use crate::{BaseChatModel, RunnableConfig};
7use lc_schema::Message;
8use tokio::sync::Mutex;
9
10use super::counter::{TokenCounter, TrackerTokenUsage};
11use super::tiktoken::TiktokenCounter;
12use super::TokenCounterError;
13
14/// 带 Token 统计的 LLM 包装器
15///
16/// 包装任意 `BaseChatModel`,自动累计 prompt / completion token 用量,
17/// 优先使用 LLM 返回的真实 usage,无则用 tiktoken 估算。
18pub struct TokenTrackingLLM<L: BaseChatModel> {
19    llm: L,
20    counter: Arc<dyn TokenCounter>,
21    usage: Arc<Mutex<TrackerTokenUsage>>,
22}
23
24impl<L: BaseChatModel> TokenTrackingLLM<L> {
25    /// 用自定义计数器包装 LLM。
26    pub fn new(llm: L, counter: Arc<dyn TokenCounter>) -> Self {
27        Self {
28            llm,
29            counter,
30            usage: Arc::new(Mutex::new(TrackerTokenUsage::new())),
31        }
32    }
33
34    /// 用 Tiktoken(cl100k_base)计数器包装
35    pub fn for_openai(llm: L) -> Result<Self, TokenCounterError> {
36        let counter = TiktokenCounter::new()?;
37        Ok(Self::new(llm, Arc::new(counter)))
38    }
39
40    /// 调用 LLM 并统计 token
41    pub async fn chat(
42        &self,
43        messages: Vec<Message>,
44        config: Option<RunnableConfig>,
45    ) -> Result<LLMResult, L::Error> {
46        let estimated_prompt = self.counter.count_messages(&messages);
47        let result = self.llm.chat(messages, config).await?;
48
49        // 优先用 LLM 返回的真实 usage,否则用估算。
50        // `TrackerTokenUsage` 与 `language_models::TokenUsage` 同为 usize，
51        // 无需精度损失转换（Q6）。
52        let (prompt, completion) = result
53            .token_usage
54            .as_ref()
55            .map(|u| (u.prompt_tokens, u.completion_tokens))
56            .unwrap_or((
57                estimated_prompt as usize,
58                self.counter.count_tokens(&result.content) as usize,
59            ));
60
61        self.usage.lock().await.add(prompt, completion);
62        Ok(result)
63    }
64
65    /// 获取累计用量
66    pub async fn get_usage(&self) -> TrackerTokenUsage {
67        self.usage.lock().await.clone()
68    }
69
70    /// 重置统计
71    pub async fn reset(&self) {
72        self.usage.lock().await.reset();
73    }
74
75    /// 估算成本(美元)
76    pub async fn estimate_cost(&self, pricing: &ModelPricing) -> f64 {
77        let usage = self.get_usage().await;
78        pricing.calculate(usage.prompt_tokens, usage.completion_tokens)
79    }
80}
81
82/// 模型定价(每 1K token 价格,美元)
83pub struct ModelPricing {
84    /// 每 1K prompt token 价格(美元)
85    pub prompt_price_per_1k: f64,
86    /// 每 1K completion token 价格(美元)
87    pub completion_price_per_1k: f64,
88}
89
90impl ModelPricing {
91    /// 创建自定义模型定价。
92    pub fn new(prompt: f64, completion: f64) -> Self {
93        Self {
94            prompt_price_per_1k: prompt,
95            completion_price_per_1k: completion,
96        }
97    }
98
99    /// gpt-4o-mini 定价(美元/1K token)
100    pub fn gpt4o_mini() -> Self {
101        Self::new(0.15, 0.60)
102    }
103
104    /// gpt-4o 定价(美元/1K token)
105    pub fn gpt4o() -> Self {
106        Self::new(2.50, 10.00)
107    }
108
109    /// 计算成本
110    pub fn calculate(&self, prompt: usize, completion: usize) -> f64 {
111        (prompt as f64 / 1000.0) * self.prompt_price_per_1k
112            + (completion as f64 / 1000.0) * self.completion_price_per_1k
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    #[test]
121    fn test_model_pricing_gpt4o_mini() {
122        let p = ModelPricing::gpt4o_mini();
123        // 1000 prompt * 0.15/1k + 1000 completion * 0.60/1k = 0.75
124        let cost = p.calculate(1000, 1000);
125        assert!((cost - 0.75).abs() < 0.001);
126    }
127
128    #[test]
129    fn test_model_pricing_zero() {
130        let p = ModelPricing::gpt4o_mini();
131        assert_eq!(p.calculate(0, 0), 0.0);
132    }
133
134    #[test]
135    fn test_model_pricing_custom() {
136        let p = ModelPricing::new(1.0, 2.0);
137        // 500 * 1.0/1k + 250 * 2.0/1k = 0.5 + 0.5 = 1.0
138        let cost = p.calculate(500, 250);
139        assert!((cost - 1.0).abs() < 0.001);
140    }
141
142    // NOTE: Tests that require OpenAIChat live in the lc-providers crate
143    // because lc-core cannot depend on lc-providers (circular dependency).
144    // The TokenTrackingLLM integration is tested there instead.
145}
lc_core/token_counter/tracker.rs

lc_core/token_counter/
tracker.rs
