File size: 4,430 Bytes
67befa7 |
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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
"""Statistics tracking for API usage and token counts."""
import time
from collections import defaultdict, deque
from threading import Lock
from typing import Dict, Any
from dataclasses import dataclass, field
@dataclass
class RequestStats:
"""Statistics for a single request."""
timestamp: float
prompt_tokens: int
completion_tokens: int
total_tokens: int
model: str
finish_reason: str
@dataclass
class AggregateStats:
"""Aggregate statistics."""
total_requests: int = 0
total_prompt_tokens: int = 0
total_completion_tokens: int = 0
total_tokens: int = 0
requests_by_model: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
tokens_by_model: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
finish_reasons: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
recent_requests: deque = field(default_factory=lambda: deque(maxlen=100)) # Keep last 100 requests
class StatsTracker:
"""Thread-safe statistics tracker."""
def __init__(self):
self._lock = Lock()
self._stats = AggregateStats()
self._start_time = time.time()
def record_request(self, stats: RequestStats):
"""Record a request's statistics."""
with self._lock:
self._stats.total_requests += 1
self._stats.total_prompt_tokens += stats.prompt_tokens
self._stats.total_completion_tokens += stats.completion_tokens
self._stats.total_tokens += stats.total_tokens
self._stats.requests_by_model[stats.model] += 1
self._stats.tokens_by_model[stats.model] += stats.total_tokens
self._stats.finish_reasons[stats.finish_reason] += 1
self._stats.recent_requests.append(stats)
def get_stats(self) -> Dict[str, Any]:
"""Get current statistics."""
with self._lock:
uptime_seconds = time.time() - self._start_time
uptime_hours = uptime_seconds / 3600
# Calculate averages
avg_prompt_tokens = (
self._stats.total_prompt_tokens / self._stats.total_requests
if self._stats.total_requests > 0 else 0
)
avg_completion_tokens = (
self._stats.total_completion_tokens / self._stats.total_requests
if self._stats.total_requests > 0 else 0
)
avg_total_tokens = (
self._stats.total_tokens / self._stats.total_requests
if self._stats.total_requests > 0 else 0
)
# Calculate requests per hour
requests_per_hour = (
self._stats.total_requests / uptime_hours
if uptime_hours > 0 else 0
)
# Calculate tokens per hour
tokens_per_hour = (
self._stats.total_tokens / uptime_hours
if uptime_hours > 0 else 0
)
return {
"uptime_seconds": int(uptime_seconds),
"uptime_hours": round(uptime_hours, 2),
"total_requests": self._stats.total_requests,
"total_prompt_tokens": self._stats.total_prompt_tokens,
"total_completion_tokens": self._stats.total_completion_tokens,
"total_tokens": self._stats.total_tokens,
"average_prompt_tokens": round(avg_prompt_tokens, 2),
"average_completion_tokens": round(avg_completion_tokens, 2),
"average_total_tokens": round(avg_total_tokens, 2),
"requests_per_hour": round(requests_per_hour, 2),
"tokens_per_hour": round(tokens_per_hour, 2),
"requests_by_model": dict(self._stats.requests_by_model),
"tokens_by_model": dict(self._stats.tokens_by_model),
"finish_reasons": dict(self._stats.finish_reasons),
"recent_requests_count": len(self._stats.recent_requests),
}
def reset(self):
"""Reset all statistics."""
with self._lock:
self._stats = AggregateStats()
self._start_time = time.time()
# Global stats tracker instance
_stats_tracker = StatsTracker()
def get_stats_tracker() -> StatsTracker:
"""Get the global stats tracker instance."""
return _stats_tracker
|