lc_core/token_counter/mod.rs
1//! Token counter and cost tracking
2//!
3//! Provides token counting (tiktoken), usage statistics, cost estimation,
4//! and a `TokenTrackingLLM` wrapper.
5
6pub mod counter;
7pub mod tiktoken;
8pub mod tracker;
9
10pub use counter::{CharRatioCounter, TokenCounter, TrackerTokenUsage};
11pub use tiktoken::TiktokenCounter;
12pub use tracker::{ModelPricing, TokenTrackingLLM};
13
14use std::sync::LazyLock;
15use tiktoken_rs::CoreBPE;
16
17/// Error type for token counting.
18#[derive(Debug, Clone, thiserror::Error)]
19#[non_exhaustive]
20pub enum TokenCounterError {
21 /// Failed to load the tiktoken encoder (e.g. the vendored BPE file is missing).
22 #[error("failed to load tiktoken encoder: {0}")]
23 EncoderLoad(String),
24}
25
26/// Global cached tiktoken encoder (cl100k_base, for GPT-3.5/4/4o).
27///
28/// Holds a `Result` so encoder-load failure surfaces as a `count_tokens`
29/// error instead of a process-wide `expect` panic on first use (Q9).
30static GLOBAL_ENCODER: LazyLock<Result<CoreBPE, TokenCounterError>> = LazyLock::new(|| {
31 tiktoken_rs::cl100k_base().map_err(|e| TokenCounterError::EncoderLoad(e.to_string()))
32});
33
34/// Count tokens in text using the global tiktoken encoder.
35///
36/// This is a convenience function that uses a lazily-initialized
37/// cl100k_base encoder (suitable for GPT-3.5/4/4o models).
38///
39/// Returns an error if the tiktoken encoder could not be loaded (e.g. the
40/// vendored BPE file is missing), rather than panicking.
41///
42/// # Examples
43/// ```no_run
44/// use lc_core::token_counter::count_tokens;
45///
46/// let n = count_tokens("Hello, world!").unwrap();
47/// assert!(n > 0);
48/// ```
49pub fn count_tokens(text: &str) -> Result<usize, TokenCounterError> {
50 let encoder = GLOBAL_ENCODER.as_ref().map_err(Clone::clone)?;
51 Ok(encoder.encode_with_special_tokens(text).len())
52}