Struct TokenTrackingLLM 

pub struct TokenTrackingLLM<L: BaseChatModel> { /* private fields */ }

LLM wrapper with token statistics

Wraps any BaseChatModel, accumulating prompt / completion token usage automatically, preferring the real usage returned by the LLM, falling back to tiktoken estimates.

v0.20.1: implements BaseChatModel itself, so a tracked LLM can be plugged directly into framework agents (e.g. FunctionCallingAgent::new); every chat/stream/bind_tools call inside the agent loop is counted, and get_usage/estimate_cost reflect the cumulative agent-run usage.

Implementations

impl<L: BaseChatModel> TokenTrackingLLM<L>

pub fn new(llm: L, counter: Arc<dyn TokenCounter>) -> Self

Wraps an LLM with a custom counter.

pub fn for_openai(llm: L) -> Result<Self, TokenCounterError>

Wraps with a Tiktoken (cl100k_base) counter

pub fn with_metrics_sink(self, sink: Arc<dyn MetricsSink>) -> Self

Attaches an observability sink: after each call that reports usage the wrapper exports a TokenUsage event (real when the provider reports it, otherwise the tiktoken estimate). The sink is shared across wrappers rebuilt by bind_tools/with_temperature/with_max_tokens.

pub fn with_cost_tracker(self, tracker: Arc<CostTracker>) -> Self

Attaches a CostTracker (B3): every counted call is priced under the model’s provider/id and aggregated on the shared tracker. Share one Arc<CostTracker> across models/runs to get per-run or per-session totals. The tracker keeps working when the model has no price entry (tokens/calls counted, cost 0).

pub fn with_provider(self, provider: impl Into<String>) -> Self

Declares the provider slug ("openai", "anthropic", …) used when looking up prices. Defaults to None (model-only table lookup).

pub async fn chat( &self, messages: Vec<Message>, config: Option<RunnableConfig>, ) -> Result<LLMResult, L::Error>

Calls the LLM and counts tokens.

Inherent method (kept for backward compatibility); delegates to Self::chat_tracked. When the model is reached through dyn BaseChatModel (e.g. inside an agent), the trait chat takes over — both count through the same helper, so the two paths never diverge.

pub async fn get_usage(&self) -> TrackerTokenUsage

Returns the cumulative usage

pub async fn reset(&self)

Resets the statistics

pub async fn estimate_cost(&self, pricing: &ModelPricing) -> f64

Estimates the cost (USD)

Trait Implementations

impl<L> BaseChatModel for TokenTrackingLLM<L>

where L: BaseChatModel + Send + Sync,

fn chat<'life0, 'async_trait>( &'life0 self, messages: Vec<Message>, config: Option<RunnableConfig>, ) -> Pin<Box<dyn Future<Output = Result<LLMResult, Self::Error>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait,

Chat with the model.

fn stream_chat<'life0, 'async_trait>( &'life0 self, messages: Vec<Message>, config: Option<RunnableConfig>, ) -> Pin<Box<dyn Future<Output = Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait,

Stream chat with the model.

fn bind_tools( &self, tools: Vec<ToolDefinition>, ) -> Option<Box<dyn BaseChatModel<Error = Self::Error> + Send + Sync>>

Bind tool definitions for function calling.

fn chat_with_system<'life0, 'async_trait>( &'life0 self, system: String, messages: Vec<Message>, ) -> Pin<Box<dyn Future<Output = Result<LLMResult, Self::Error>> + Send + 'async_trait>>

where Self: Sync + 'async_trait, 'life0: 'async_trait,

Chat with system prompt.

impl<L> BaseLanguageModel<Vec<Message>, LLMResult> for TokenTrackingLLM<L>

where L: BaseChatModel + Send + Sync,

fn model_name(&self) -> &str

Returns the model name.

fn get_num_tokens(&self, text: &str) -> usize

Calculates token count for text.

fn temperature(&self) -> Option<f32>

Returns the temperature parameter.

fn max_tokens(&self) -> Option<usize>

Returns the max tokens limit.

fn with_temperature(self, temp: f32) -> Self

where Self: Sized,

Sets the temperature parameter.

fn with_max_tokens(self, max: usize) -> Self

where Self: Sized,

Sets the max tokens limit.

impl<L> Runnable<Vec<Message>, LLMResult> for TokenTrackingLLM<L>

where L: BaseChatModel + Send + Sync,

type Error = <L as Runnable<Vec<Message>, LLMResult>>::Error

Error type.

fn invoke<'life0, 'async_trait>( &'life0 self, input: Vec<Message>, config: Option<RunnableConfig>, ) -> Pin<Box<dyn Future<Output = Result<LLMResult, Self::Error>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait,

Transforms single input to output.

fn stream<'life0, 'async_trait>( &'life0 self, input: Vec<Message>, config: Option<RunnableConfig>, ) -> Pin<Box<dyn Future<Output = Result<Pin<Box<dyn Stream<Item = Result<LLMResult, Self::Error>> + Send>>, Self::Error>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait,

Streaming output - for real-time responses (LLM, etc).

fn batch<'life0, 'async_trait>( &'life0 self, inputs: Vec<Input>, config: Option<RunnableConfig>, ) -> Pin<Box<dyn Future<Output = Result<Vec<Output>, Self::Error>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait,

Batch processing - transforms multiple inputs to outputs.

fn batch_as_completed<'life0, 'async_trait>( &'life0 self, inputs: Vec<Input>, config: Option<RunnableConfig>, ) -> Pin<Box<dyn Future<Output = Result<Vec<(usize, Output)>, Self::Error>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait,

Batch processing that returns results in completion order.

fn transform<'life0, 'async_trait>( &'life0 self, input: Pin<Box<dyn Stream<Item = Result<Input, Self::Error>> + Send>>, config: Option<RunnableConfig>, ) -> Pin<Box<dyn Future<Output = Result<Pin<Box<dyn Stream<Item = Result<Output, Self::Error>> + Send + '_>>, Self::Error>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait,

Stream-to-stream transformation - the core of LCEL streaming.