Sync AI stats and create a new merging function to facilitate

Author: bitterpanda63

aikido_zen/background_process/commands/sync_data.py

@@ -46,6 +46,9 @@ def process_sync_data(connection_manager, data, conn, queue=None):
     # Sync stats
     connection_manager.statistics.import_from_record(data.get("stats", {}))
 
+    # Sync ai stats
+    connection_manager.ai_stats.import_list(data.get("ai_stats", []))
+
     if connection_manager.conf.last_updated_at > 0:
         # Only report data if the config has been fetched.
         return {

aikido_zen/storage/ai_statistics.py

@@ -37,6 +37,16 @@ def on_ai_call(self, provider, model, input_tokens, output_tokens):
     def get_stats(self):
         return [copy.deepcopy(stats) for stats in self.calls.values()]
 
+    def import_list(self, ai_stats_list):
+        for new_entry in ai_stats_list:
+            existing_entry = self.ensure_provider_stats(
+                new_entry["provider"], new_entry["model"]
+            )
+            existing_entry["calls"] += new_entry["calls"]
+            existing_entry["tokens"]["input"] = new_entry["tokens"]["input"]
+            existing_entry["tokens"]["output"] = new_entry["tokens"]["output"]
+            existing_entry["tokens"]["total"] = new_entry["tokens"]["total"]
+
     def clear(self):
         self.calls.clear()