Spaces:
Running
Running
| import json | |
| import os | |
| import datetime | |
| import threading | |
| from collections import Counter, defaultdict | |
| from typing import List, Dict, Any, Optional | |
| from fastapi import Request | |
| class UsageTracker: | |
| def __init__(self, data_file="usage_data.json"): | |
| self.data_file = data_file | |
| self.lock = threading.Lock() | |
| self.data = self._load_data() | |
| self._schedule_save() | |
| def _load_data(self) -> Dict[str, Any]: | |
| """Loads usage data from the JSON file, ensuring data integrity.""" | |
| if os.path.exists(self.data_file): | |
| try: | |
| with open(self.data_file, 'r') as f: | |
| data = json.load(f) | |
| if isinstance(data, dict) and 'requests' in data: | |
| return data | |
| # If data is old format, try to convert it | |
| if isinstance(data.get('total_requests'), int): | |
| return self._convert_old_format(data) | |
| except (json.JSONDecodeError, TypeError): | |
| print(f"Warning: Could not decode JSON from {self.data_file}. Starting fresh.") | |
| return self._initialize_empty_data() | |
| def _initialize_empty_data(self) -> Dict[str, List]: | |
| """Initializes a new data structure for usage tracking.""" | |
| return {'requests': []} | |
| def _convert_old_format(self, old_data: Dict[str, Any]) -> Dict[str, List]: | |
| """Converts data from the old format to the new detailed format.""" | |
| print("Converting old usage data format to new format.") | |
| new_data = self._initialize_empty_data() | |
| # This is a simplification; a more robust conversion would be needed for full data recovery | |
| # For now, we are just starting fresh with the new structure to avoid complexity. | |
| return new_data | |
| def save_data(self): | |
| """Saves current usage data to the JSON file periodically.""" | |
| with self.lock: | |
| try: | |
| with open(self.data_file, 'w') as f: | |
| json.dump(self.data, f, indent=4) | |
| except IOError as e: | |
| print(f"Error saving usage data to {self.data_file}: {e}") | |
| def _schedule_save(self): | |
| """Schedules the data to be saved every 60 seconds.""" | |
| threading.Timer(60.0, self._schedule_save).start() | |
| self.save_data() | |
| def record_request(self, request: Optional[Request], model: str = "unknown", endpoint: str = "unknown"): | |
| """Records a single API request with detailed information.""" | |
| with self.lock: | |
| now = datetime.datetime.now(datetime.timezone.utc) | |
| ip_address = "N/A" | |
| user_agent = "N/A" | |
| if request: | |
| ip_address = request.client.host | |
| user_agent = request.headers.get("user-agent", "N/A") | |
| self.data['requests'].append({ | |
| 'timestamp': now.isoformat(), | |
| 'model': model, | |
| 'endpoint': endpoint, | |
| 'ip_address': ip_address, | |
| 'user_agent': user_agent, | |
| }) | |
| def get_usage_summary(self, days: int = 7) -> Dict[str, Any]: | |
| """Generates a comprehensive summary of usage data.""" | |
| with self.lock: | |
| summary = { | |
| 'total_requests': 0, | |
| 'model_usage': defaultdict(int), | |
| 'endpoint_usage': defaultdict(int), | |
| 'daily_usage': defaultdict(lambda: defaultdict(int)), | |
| 'unique_ips': set(), | |
| 'recent_requests': [] | |
| } | |
| cutoff_date = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=days) | |
| # Iterate backwards for recent requests | |
| for req in reversed(self.data['requests']): | |
| req_time = datetime.datetime.fromisoformat(req['timestamp']) | |
| # Update total requests (for all time) | |
| summary['total_requests'] += 1 | |
| summary['unique_ips'].add(req['ip_address']) | |
| if req_time >= cutoff_date: | |
| date_str = req_time.strftime("%Y-%m-%d") | |
| # Aggregate data for charts and tables | |
| summary['model_usage'][req['model']] += 1 | |
| summary['endpoint_usage'][req['endpoint']] += 1 | |
| summary['daily_usage'][date_str]['requests'] += 1 | |
| # Add to recent requests list | |
| if len(summary['recent_requests']) < 20: | |
| summary['recent_requests'].append(req) | |
| # Sort daily usage by date | |
| summary['daily_usage'] = dict(sorted(summary['daily_usage'].items())) | |
| summary['unique_ip_count'] = len(summary['unique_ips']) | |
| del summary['unique_ips'] # No need to send the whole set | |
| return summary | |