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

lc_core/token_counter/
tracker.rs
