fixed linting and tests

Author: cemde

maseval/core/model.py

@@ -100,7 +100,7 @@ class ChatResponse:
     content: Optional[str] = None
     tool_calls: Optional[List[Dict[str, Any]]] = None
     role: str = "assistant"
-    usage: Optional[Dict[str, int]] = None
+    usage: Optional[Dict[str, Any]] = None
     model: Optional[str] = None
     stop_reason: Optional[str] = None
 

tests/conftest.py

@@ -48,7 +48,7 @@ def __init__(
         model_id: str = "test-model",
         responses: Optional[List[Optional[str]]] = None,
         tool_calls: Optional[List[Optional[List[Dict[str, Any]]]]] = None,
-        usage: Optional[Dict[str, int]] = None,
+        usage: Optional[Dict[str, Any]] = None,
         stop_reason: Optional[str] = None,
         seed: Optional[int] = None,
     ):

tests/test_core/test_benchmark/test_usage_collection.py

@@ -6,7 +6,6 @@
 
 import pytest
 from maseval import TaskQueue
-from maseval.core.usage import StaticPricingCalculator
 
 
 @pytest.mark.core
@@ -71,10 +70,12 @@ def test_usage_persists_across_task_repetitions(self):
         """Benchmark.usage accumulates across multiple tasks."""
         from conftest import DummyBenchmark
 
-        tasks = TaskQueue.from_list([
-            {"query": "Task 1", "environment_data": {}},
-            {"query": "Task 2", "environment_data": {}},
-        ])
+        tasks = TaskQueue.from_list(
+            [
+                {"query": "Task 1", "environment_data": {}},
+                {"query": "Task 2", "environment_data": {}},
+            ]
+        )
         benchmark = DummyBenchmark()
         benchmark.run(tasks, agent_data={"model": "test"})
 

tests/test_core/test_registry.py

@@ -13,6 +13,7 @@
 from maseval.core.registry import ComponentRegistry
 from maseval.core.tracing import TraceableMixin
 from maseval.core.config import ConfigurableMixin
+from maseval.core.usage import UsageTrackableMixin
 
 
 # ==================== Test Components ====================
@@ -276,6 +277,7 @@ def gather_traces(self) -> Dict[str, Any]:
 
     def gather_usage(self):
         from maseval.core.usage import TokenUsage
+
         return TokenUsage(
             cost=self._cost,
             input_tokens=self._input_tokens,
@@ -297,10 +299,6 @@ def gather_usage(self):
         raise RuntimeError("Usage collection failed")
 
 
-# Ensure MockUsageComponent also inherits UsageTrackableMixin
-from maseval.core.usage import UsageTrackableMixin
-
-
 class UsageAwareComponent(TraceableMixin, UsageTrackableMixin):
     """Component with both tracing and usage tracking."""
 
@@ -315,6 +313,7 @@ def gather_traces(self) -> Dict[str, Any]:
 
     def gather_usage(self):
         from maseval.core.usage import TokenUsage
+
         return TokenUsage(
             cost=self._cost,
             input_tokens=self._input_tokens,
@@ -362,7 +361,6 @@ def test_non_usage_component_not_in_usage_registry(self):
 
     def test_collect_usage_basic(self):
         """collect_usage returns structured dict with usage from registered components."""
-        from maseval.core.usage import TokenUsage
 
         registry = ComponentRegistry()
         model = UsageAwareComponent(cost=0.10, input_tokens=500, output_tokens=200)
@@ -445,7 +443,10 @@ def test_usage_by_component_accumulates(self):
         by_comp = registry.usage_by_component
         assert "models:main_model" in by_comp
 
+        from maseval.core.usage import TokenUsage
+
         total = by_comp["models:main_model"]
+        assert isinstance(total, TokenUsage)
         assert total.input_tokens == 300
         assert total.output_tokens == 150
         assert total.cost == pytest.approx(0.30)