Update usage_tracker.py

usage_tracker.py

@@ -1,106 +1,108 @@
-import os
 import json
-from datetime import datetime, timedelta
+import os
+import datetime
+import threading
+from collections import defaultdict
 
 class UsageTracker:
-    def __init__(self, filename='usage_data.json'):
-        self.filename = filename
-        self.usage_data = self.load_usage_data()
+    def __init__(self, data_file="usage_data.json"):
+        self.data_file = data_file
+        self.lock = threading.Lock()
+        self.data = self._load_data()
 
-    def load_usage_data(self):
-        """Load existing usage data from file or create a new structure"""
-        try:
-            if os.path.exists(self.filename):
-                with open(self.filename, 'r') as f:
-                    return json.load(f)
-        except (json.JSONDecodeError, FileNotFoundError):
-            pass
-        
-        # Initialize usage data structure
+    def _load_data(self):
+        """Loads usage data from the JSON file."""
+        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', {})
+                    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()
+        return self._initialize_empty_data()
+
+    def _initialize_empty_data(self):
+        """Initializes an empty data structure for usage tracking."""
         return {
-            "total_requests": 0,
-            "models": {},
-            "daily_usage": {},
-            "api_endpoints": {}
+            'total_requests': 0,
+            'models': {},
+            'api_endpoints': {},
+            'recent_daily_usage': {}
         }
 
-    def record_request(self, model=None, endpoint=None):
-        """Record a request for a specific model and/or endpoint"""
-        # Get current date
-        today = datetime.now().strftime("%Y-%m-%d")
-        
-        # Update total requests
-        self.usage_data["total_requests"] += 1
-        
-        # Track model if provided
-        if model:
-            if model not in self.usage_data["models"]:
-                self.usage_data["models"][model] = {
-                    "total_requests": 0,
-                    "first_used": datetime.now().isoformat(),
-                    "last_used": datetime.now().isoformat()
+    def save_data(self):
+        """Saves current usage data to the JSON file."""
+        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 record_request(self, model: str = "unknown", endpoint: str = "unknown"):
+        """Records a single API request, updating model, endpoint, and daily usage."""
+        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
                 }
-            
-            model_data = self.usage_data["models"][model]
-            model_data["total_requests"] += 1
-            model_data["last_used"] = datetime.now().isoformat()
-        
-        # Track endpoint if provided
-        if endpoint:
-            if endpoint not in self.usage_data["api_endpoints"]:
-                self.usage_data["api_endpoints"][endpoint] = {
-                    "total_requests": 0,
-                    "first_used": datetime.now().isoformat(),
-                    "last_used": datetime.now().isoformat()
+            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
+            if current_date not in self.data['recent_daily_usage']:
+                self.data['recent_daily_usage'][current_date] = defaultdict(int)
             
-            endpoint_data = self.usage_data["api_endpoints"][endpoint]
-            endpoint_data["total_requests"] += 1
-            endpoint_data["last_used"] = datetime.now().isoformat()
-        
-        # Update daily usage
-        if today not in self.usage_data["daily_usage"]:
-            self.usage_data["daily_usage"][today] = {}
-        
-        # Track daily usage for model
-        if model:
-            if model not in self.usage_data["daily_usage"][today]:
-                self.usage_data["daily_usage"][today][model] = 0
-            self.usage_data["daily_usage"][today][model] += 1
-        
-        # Track daily usage for endpoint
-        if endpoint:
-            if endpoint not in self.usage_data["daily_usage"][today]:
-                self.usage_data["daily_usage"][today][endpoint] = 0
-            self.usage_data["daily_usage"][today][endpoint] += 1
-        
-        # Save updated usage data
-        self.save_usage_data()
+            # Convert defaultdict back to dict for JSON serialization
+            self.data['recent_daily_usage'][current_date][model] += 1
+            self.data['recent_daily_usage'][current_date][endpoint] += 1
+            
+            # Ensure the daily usage is stored as a regular dict
+            self.data['recent_daily_usage'][current_date] = dict(self.data['recent_daily_usage'][current_date])
 
-    def save_usage_data(self):
-        """Save usage data to file"""
-        try:
-            with open(self.filename, 'w') as f:
-                json.dump(self.usage_data, f, indent=4)
-        except Exception as e:
-            print(f"Error saving usage data: {e}")
 
-    def get_usage_summary(self, days=7):
-        """Get usage summary for the last specified number of days"""
-        # Calculate the date threshold
-        threshold_date = (datetime.now() - timedelta(days=days)).strftime("%Y-%m-%d")
-        
-        # Prepare summary
-        summary = {
-            "total_requests": self.usage_data["total_requests"],
-            "models": self.usage_data["models"],
-            "api_endpoints": self.usage_data["api_endpoints"],
-            "recent_daily_usage": {}
-        }
-        
-        # Collect recent daily usage
-        for date, date_data in sorted(self.usage_data["daily_usage"].items(), reverse=True):
-            if date >= threshold_date:
-                summary["recent_daily_usage"][date] = date_data
-        
-        return summary
+    def get_usage_summary(self, days: int = 7):
+        """Generates a summary of usage data for the last 'days'."""
+        with self.lock:
+            summary = {
+                'total_requests': self.data['total_requests'],
+                'models': self.data['models'],
+                'api_endpoints': self.data['api_endpoints'],
+                'recent_daily_usage': {}
+            }
+
+            # Filter daily usage for the last 'days'
+            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']:
+                    summary['recent_daily_usage'][date] = self.data['recent_daily_usage'][date]
+            
+            return summary
+