Update usage_tracker.py

usage_tracker.py

@@ -21,6 +21,13 @@ class UsageTracker:
                     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.")
@@ -33,15 +40,27 @@ class UsageTracker:
             'total_requests': 0,
             'models': {},
             'api_endpoints': {},
-            'recent_daily_usage': {}
+            '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 save_data(self):
         """Saves current usage data to the JSON file."""
         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(self.data, f, indent=4)
+                    json.dump(data_to_save, f, indent=4)
             except IOError as e:
                 print(f"Error saving usage data to {self.data_file}: {e}")
 
@@ -76,15 +95,18 @@ class UsageTracker:
             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)
-            
-            # Convert defaultdict back to dict for JSON serialization
+            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
-            
-            # 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])
+
+            # 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):
@@ -97,12 +119,13 @@ class UsageTracker:
                 'recent_daily_usage': {}
             }
 
-            # Filter daily usage for the last 'days'
+            # 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']:
-                    summary['recent_daily_usage'][date] = self.data['recent_daily_usage'][date]
-            
+                    # 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])
+
             return summary