agents.base.agent_with_token_tracking¶
Agent base class with integrated token usage tracking.
This module provides an enhanced Agent base class that automatically tracks token usage for all LLM interactions, providing cost analysis and capacity monitoring capabilities.
Classes¶
Agent base class with automatic token usage tracking. |
Module Contents¶
- class agents.base.agent_with_token_tracking.TokenTrackingAgent¶
Bases:
haive.agents.base.agent.AgentAgent base class with automatic token usage tracking.
This enhanced agent automatically tracks token usage for all LLM interactions, providing detailed metrics on token consumption, costs, and capacity usage. It uses MessagesStateWithTokenUsage as the default state schema.
Additional features: - Automatic token extraction from LLM responses - Cost calculation based on provider pricing - Capacity percentage monitoring - Token usage history tracking - Conversation cost analysis
Example
class MyAgent(TokenTrackingAgent): def build_graph(self): # Your graph logic pass agent = MyAgent( name="cost_aware_agent", engine=llm_engine, track_costs=True, input_cost_per_1k=0.003, output_cost_per_1k=0.015 ) # After running result = agent.invoke({"query": "Hello"}) usage = agent.get_token_usage_summary() print(f"Total tokens: {usage['total_tokens']}") print(f"Total cost: ${usage['total_cost']:.4f}")
- calculate_conversation_costs()¶
Calculate costs for the current conversation.
- Return type:
None