Struct TokenTrackingMiddleware

pub struct TokenTrackingMiddleware { /* private fields */ }

Token tracking middleware that wraps an LLM to monitor usage

Implementations

impl TokenTrackingMiddleware

pub fn new( config: TokenTrackingConfig, inner_model: Arc<dyn LanguageModel>, event_dispatcher: Option<Arc<EventDispatcher>>, ) -> Self

pub fn get_usage_stats(&self) -> Vec<TokenUsage>

Get accumulated usage statistics

pub fn get_total_usage(&self) -> TokenUsageSummary

Get total usage summary

pub fn clear_stats(&self)

Clear usage statistics

Trait Implementations

impl AgentMiddleware for TokenTrackingMiddleware

fn id(&self) -> &'static str

Unique identifier for logging and diagnostics.

fn modify_model_request<'life0, 'life1, 'life2, 'async_trait>( &'life0 self, _ctx: &'life1 mut MiddlewareContext<'life2>, ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait,

Apply middleware-specific mutations to the pending model request.

fn tools(&self) -> Vec<ToolBox>

Tools to expose when this middleware is active.

fn before_tool_execution<'life0, 'life1, 'life2, 'life3, 'async_trait>( &'life0 self, _tool_name: &'life1 str, _tool_args: &'life2 Value, _call_id: &'life3 str, ) -> Pin<Box<dyn Future<Output = Result<Option<AgentInterrupt>>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait,

Hook called before tool execution - can return an interrupt to pause execution.

impl LanguageModel for TokenTrackingMiddleware

fn generate<'life0, 'async_trait>( &'life0 self, request: LlmRequest, ) -> Pin<Box<dyn Future<Output = Result<LlmResponse>> + Send + 'async_trait>>

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

Generate a complete response (non-streaming)

fn generate_stream<'life0, 'async_trait>( &'life0 self, request: LlmRequest, ) -> Pin<Box<dyn Future<Output = Result<ChunkStream>> + Send + 'async_trait>>

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

Generate a streaming response Default implementation falls back to non-streaming generate()