Update usage_tracker.py

usage_tracker.py

@@ -2,130 +2,115 @@ import json
 import os
 import datetime
 import threading
-from collections import defaultdict
+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):
-        """Loads usage data from the JSON file."""
+    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)
-                    # Ensure all necessary keys exist, initialize if not
-                    data.setdefault('total_requests', 0)
-                    data.setdefault('models', {})
-                    data.setdefault('api_endpoints', {})
-                    data.setdefault('recent_daily_usage', {})
-
-                    # Convert loaded daily usage dictionaries back to defaultdicts for internal use
-                    loaded_daily_usage = {}
-                    for date, entities in data['recent_daily_usage'].items():
-                        loaded_daily_usage[date] = defaultdict(int, entities)
-                    data['recent_daily_usage'] = loaded_daily_usage
-
-                    return data
-            except json.JSONDecodeError:
-                print(f"Warning: Could not decode JSON from {self.data_file}. Starting with empty data.")
-                return self._initialize_empty_data()
+                    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):
-        """Initializes an empty data structure for usage tracking."""
-        return {
-            'total_requests': 0,
-            'models': {},
-            'api_endpoints': {},
-            'recent_daily_usage': {} # This will hold defaultdicts for each date
-        }
-
-    def _convert_defaultdicts_to_dicts(self, obj):
-        """Recursively converts defaultdicts to dicts for JSON serialization."""
-        if isinstance(obj, defaultdict):
-            return {k: self._convert_defaultdicts_to_dicts(v) for k, v in obj.items()}
-        elif isinstance(obj, dict):
-            return {k: self._convert_defaultdicts_to_dicts(v) for k, v in obj.items()}
-        elif isinstance(obj, list):
-            return [self._convert_defaultdicts_to_dicts(elem) for elem in obj]
-        return obj
+    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."""
+        """Saves current usage data to the JSON file periodically."""
         with self.lock:
-            # Convert defaultdicts to regular dicts before saving
-            data_to_save = self._convert_defaultdicts_to_dicts(self.data)
             try:
                 with open(self.data_file, 'w') as f:
-                    json.dump(data_to_save, f, indent=4)
+                    json.dump(self.data, f, indent=4)
             except IOError as e:
                 print(f"Error saving usage data to {self.data_file}: {e}")
 
-    def record_request(self, model: str = "unknown", endpoint: str = "unknown"):
-        """Records a single API request, updating model, endpoint, and daily usage."""
+    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()
-            current_date = now.strftime("%Y-%m-%d")
-            current_time = now.strftime("%Y-%m-%d %I:%M:%S %p")
-
-            # Update total requests
-            self.data['total_requests'] += 1
-
-            # Update model usage
-            if model not in self.data['models']:
-                self.data['models'][model] = {
-                    'total_requests': 0,
-                    'first_used': current_time,
-                    'last_used': current_time
-                }
-            self.data['models'][model]['total_requests'] += 1
-            self.data['models'][model]['last_used'] = current_time
-
-            # Update API endpoint usage
-            if endpoint not in self.data['api_endpoints']:
-                self.data['api_endpoints'][endpoint] = {
-                    'total_requests': 0,
-                    'first_used': current_time,
-                    'last_used': current_time
-                }
-            self.data['api_endpoints'][endpoint]['total_requests'] += 1
-            self.data['api_endpoints'][endpoint]['last_used'] = current_time
-
-            # Update daily usage
-            # Ensure the inner dictionary for the current_date is a defaultdict
-            if current_date not in self.data['recent_daily_usage']:
-                self.data['recent_daily_usage'][current_date] = defaultdict(int)
-            elif not isinstance(self.data['recent_daily_usage'][current_date], defaultdict):
-                # This handles cases where data was loaded as a plain dict
-                self.data['recent_daily_usage'][current_date] = defaultdict(int, self.data['recent_daily_usage'][current_date])
-
-            self.data['recent_daily_usage'][current_date][model] += 1
-            self.data['recent_daily_usage'][current_date][endpoint] += 1
-
-            # Removed the line that converts defaultdict back to dict here.
-            # Conversion will now happen only during save_data.
-
-
-    def get_usage_summary(self, days: int = 7):
-        """Generates a summary of usage data for the last 'days'."""
+            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': self.data['total_requests'],
-                'models': self.data['models'],
-                'api_endpoints': self.data['api_endpoints'],
-                'recent_daily_usage': {}
+                'total_requests': 0,
+                'model_usage': defaultdict(int),
+                'endpoint_usage': defaultdict(int),
+                'daily_usage': defaultdict(lambda: defaultdict(int)),
+                'unique_ips': set(),
+                'recent_requests': []
             }
 
-            # Filter daily usage for the last 'days' and ensure it's a plain dict for the summary
-            today = datetime.date.today()
-            for i in range(days):
-                date = (today - datetime.timedelta(days=i)).strftime("%Y-%m-%d")
-                if date in self.data['recent_daily_usage']:
-                    # Convert the defaultdict for this date to a plain dict for the summary
-                    summary['recent_daily_usage'][date] = dict(self.data['recent_daily_usage'][date])
+            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