File size: 2,818 Bytes
a2ec7b6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | """
Token Usage Tracking Utilities
===============================
"""
from typing import Dict, Any
class TokenUsageTracker:
"""Track token usage across agent executions."""
def __init__(self):
"""Initialize token usage tracker."""
self.reset()
def reset(self):
"""Reset all usage statistics."""
self._stats = {
"total_input_tokens": 0,
"total_output_tokens": 0,
"total_tokens": 0,
"total_turns": 0,
"total_execution_time": 0.0,
"successful_executions": 0,
"failed_executions": 0,
}
def update(self, success: bool, token_usage: Dict[str, int],
turn_count: int, execution_time: float):
"""
Update usage statistics.
Args:
success: Whether execution was successful
token_usage: Token usage dict with input_tokens, output_tokens, total_tokens
turn_count: Number of conversation turns
execution_time: Execution time in seconds
"""
if success:
self._stats["successful_executions"] += 1
else:
self._stats["failed_executions"] += 1
self._stats["total_input_tokens"] += token_usage.get("input_tokens", 0)
self._stats["total_output_tokens"] += token_usage.get("output_tokens", 0)
self._stats["total_tokens"] += token_usage.get("total_tokens", 0)
self._stats["total_turns"] += turn_count
self._stats["total_execution_time"] += execution_time
def get_stats(self) -> Dict[str, Any]:
"""
Get usage statistics with calculated averages.
Returns:
Dictionary containing usage statistics
"""
stats = self._stats.copy()
# Calculate averages
total_executions = stats["successful_executions"] + stats["failed_executions"]
if total_executions > 0:
stats["avg_input_tokens"] = stats["total_input_tokens"] / total_executions
stats["avg_output_tokens"] = stats["total_output_tokens"] / total_executions
stats["avg_total_tokens"] = stats["total_tokens"] / total_executions
stats["avg_turns"] = stats["total_turns"] / total_executions
stats["avg_execution_time"] = stats["total_execution_time"] / total_executions
stats["success_rate"] = (stats["successful_executions"] / total_executions * 100)
else:
stats.update({
"avg_input_tokens": 0.0,
"avg_output_tokens": 0.0,
"avg_total_tokens": 0.0,
"avg_turns": 0.0,
"avg_execution_time": 0.0,
"success_rate": 0.0,
})
return stats |