Fix bug with import_list and add unit tests to test

Author: bitterpanda63

aikido_zen/storage/ai_statistics.py

@@ -41,9 +41,9 @@ def import_list(self, ai_stats_list):
                 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"]
+            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()

aikido_zen/storage/ai_statistics_test.py

@@ -242,3 +242,191 @@ def test_get_stats_consistency_after_multiple_calls(stats):
             },
         },
     ]
+
+
+def test_import_list_with_empty_stats(stats):
+    # Test importing a list into empty statistics
+    stats.import_list(
+        [
+            {
+                "provider": "openai",
+                "model": "gpt-4",
+                "calls": 1,
+                "tokens": {
+                    "input": 100,
+                    "output": 50,
+                    "total": 150,
+                },
+            },
+            {
+                "provider": "anthropic",
+                "model": "claude-3",
+                "calls": 1,
+                "tokens": {
+                    "input": 120,
+                    "output": 60,
+                    "total": 180,
+                },
+            },
+        ]
+    )
+
+    result = stats.get_stats()
+    assert len(result) == 2
+    assert result[0] == {
+        "provider": "openai",
+        "model": "gpt-4",
+        "calls": 1,
+        "tokens": {
+            "input": 100,
+            "output": 50,
+            "total": 150,
+        },
+    }
+    assert result[1] == {
+        "provider": "anthropic",
+        "model": "claude-3",
+        "calls": 1,
+        "tokens": {
+            "input": 120,
+            "output": 60,
+            "total": 180,
+        },
+    }
+
+
+def test_import_list_with_existing_stats(stats):
+    # Add some initial statistics
+    stats.on_ai_call(
+        provider="openai", model="gpt-4", input_tokens=100, output_tokens=50
+    )
+
+    # Import a list that includes an existing provider:model combination
+    stats.import_list(
+        [
+            {
+                "provider": "openai",
+                "model": "gpt-4",
+                "calls": 1,
+                "tokens": {
+                    "input": 200,
+                    "output": 100,
+                    "total": 300,
+                },
+            },
+            {
+                "provider": "anthropic",
+                "model": "claude-3",
+                "calls": 1,
+                "tokens": {
+                    "input": 120,
+                    "output": 60,
+                    "total": 180,
+                },
+            },
+        ]
+    )
+
+    result = stats.get_stats()
+    assert len(result) == 2
+    assert result[0] == {
+        "provider": "openai",
+        "model": "gpt-4",
+        "calls": 2,  # Initial call + imported call
+        "tokens": {
+            "input": 300,  # 100 + 200
+            "output": 150,  # 50 + 100
+            "total": 450,  # 150 + 300
+        },
+    }
+    assert result[1] == {
+        "provider": "anthropic",
+        "model": "claude-3",
+        "calls": 1,
+        "tokens": {
+            "input": 120,
+            "output": 60,
+            "total": 180,
+        },
+    }
+
+
+def test_import_list_with_overlapping_and_new_entries(stats):
+    # Add some initial statistics
+    stats.on_ai_call(
+        provider="openai", model="gpt-4", input_tokens=100, output_tokens=50
+    )
+    stats.on_ai_call(
+        provider="anthropic", model="claude-3", input_tokens=120, output_tokens=60
+    )
+
+    # Import a list that includes both existing and new provider:model combinations
+    stats.import_list(
+        [
+            {
+                "provider": "openai",
+                "model": "gpt-4",
+                "calls": 1,
+                "tokens": {
+                    "input": 200,
+                    "output": 100,
+                    "total": 300,
+                },
+            },
+            {
+                "provider": "anthropic",
+                "model": "claude-3",
+                "calls": 1,
+                "tokens": {
+                    "input": 120,
+                    "output": 60,
+                    "total": 180,
+                },
+            },
+            {
+                "provider": "mistral",
+                "model": "mistral-7b",
+                "calls": 1,
+                "tokens": {
+                    "input": 150,
+                    "output": 75,
+                    "total": 225,
+                },
+            },
+        ]
+    )
+
+    result = stats.get_stats()
+    assert len(result) == 3
+    result.sort(key=lambda x: f"{x['provider']}:{x['model']}")
+
+    assert result[0] == {
+        "provider": "anthropic",
+        "model": "claude-3",
+        "calls": 2,  # Initial call + imported call
+        "tokens": {
+            "input": 240,  # 120 + 120
+            "output": 120,  # 60 + 60
+            "total": 360,  # 180 + 180
+        },
+    }
+    assert result[1] == {
+        "provider": "mistral",
+        "model": "mistral-7b",
+        "calls": 1,
+        "tokens": {
+            "input": 150,
+            "output": 75,
+            "total": 225,
+        },
+    }
+    assert result[2] == {
+        "provider": "openai",
+        "model": "gpt-4",
+        "calls": 2,  # Initial call + imported call
+        "tokens": {
+            "input": 300,  # 100 + 200
+            "output": 150,  # 50 + 100
+            "total": 450,  # 150 + 300
+        },
+    }