fixed cost tracking of agent frameworks

Author: cemde

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Usage and cost tracking as a first-class collection axis alongside tracing and configuration. `Usage` and `TokenUsage` data classes record billable resource consumption (tokens, API calls, custom units). `UsageTrackableMixin` enables automatic collection via `gather_usage()`. `ModelAdapter` tracks token usage automatically after each `chat()` call with no changes required from benchmark implementers. (PR: #45)
 - Pluggable cost calculation via `CostCalculator` protocol. `StaticPricingCalculator` computes cost from user-supplied per-token rates (supports USD, EUR, credits, or any unit). Pass a `cost_calculator` to any `ModelAdapter` to fill in `Usage.cost` when the provider doesn't report it. Provider-reported cost always takes precedence. (PR: #45)
 - `LiteLLMCostCalculator` in `maseval.interface.usage` for automatic pricing via LiteLLM's bundled model database. Supports `custom_pricing` overrides and `model_id_map` for remapping adapter model IDs to LiteLLM's naming convention. Requires `litellm`. (PR: #45)
+- Cost calculation for agent adapters. `AgentAdapter` now accepts `cost_calculator` and `model_id` parameters. For smolagents, CAMEL, and LlamaIndex, both the model ID and cost calculator are auto-detected (model ID from the framework's agent object, calculator via `LiteLLMCostCalculator` if litellm is installed). For LangGraph, `model_id` must be passed explicitly since graphs can contain multiple models. Explicit `cost_calculator` and `model_id` always override auto-detection. (PR: #45)
 - `UsageReporter` post-hoc analysis utility for slicing usage data from benchmark reports by task, component, or model. Create via `UsageReporter.from_reports(benchmark.reports)`. (PR: #45)
 - Live usage totals accessible during benchmark execution via `benchmark.usage` (grand total) and `benchmark.usage_by_component` (per-component breakdowns). Totals persist across task repetitions. (PR: #45)
 - `ComponentRegistry` gains usage collection: `collect_usage()`, `total_usage`, and `usage_by_component` properties, parallel to existing trace and config collection. (PR: #45)

docs/guides/usage-tracking.md

@@ -16,7 +16,7 @@ MASEval tracks how much each benchmark run consumes (tokens, API calls, dollars)
 
 **Model adapters** track every `chat()` call: input tokens, output tokens, cached tokens, reasoning tokens. No setup needed.
 
-**Agent adapters** aggregate token usage from the underlying framework's execution. Each framework adapter (`SmolAgentAdapter`, `CamelAgentAdapter`, `LangGraphAgentAdapter`, `LlamaIndexAgentAdapter`) extracts usage from its framework's native data structures (memory steps, response metadata, message annotations, or execution logs respectively).
+**Agent adapters** aggregate token usage from the underlying framework's execution. Each framework adapter (`SmolAgentAdapter`, `CamelAgentAdapter`, `LangGraphAgentAdapter`, `LlamaIndexAgentAdapter`) extracts usage from its framework's native data structures (memory steps, response metadata, message annotations, or execution logs respectively). Cost is computed automatically when litellm is installed (see [Agent Cost Tracking](#agent-cost-tracking) below).
 
 **Benchmarks** collect usage from all registered components after each task and include it in reports.
 
@@ -55,6 +55,51 @@ print(f"{usage.input_tokens} in, {usage.output_tokens} out")
 
 This works across all supported frameworks (smolagents, CAMEL, LangGraph, and LlamaIndex). The adapter handles the framework-specific extraction; you always call `gather_usage()`.
 
+### Agent Cost Tracking
+
+Agent adapters auto-detect cost when possible. For smolagents, CAMEL, and LlamaIndex, the adapter reads the model ID from the framework's agent object and uses `LiteLLMCostCalculator` if litellm is installed. No configuration needed:
+
+```python
+# Cost tracking works automatically if litellm is installed
+adapter = SmolAgentAdapter(agent, name="researcher")
+adapter.run("What's the capital of France?")
+print(f"Cost: ${adapter.gather_usage().cost:.4f}")
+```
+
+For **LangGraph**, the model ID cannot be auto-detected because a graph can contain multiple models across its nodes. Pass `model_id` explicitly:
+
+```python
+adapter = LangGraphAgentAdapter(
+    compiled_graph, "agent",
+    model_id="gpt-4o-mini",  # Required for cost tracking
+)
+```
+
+To override auto-detection or use custom pricing, pass `cost_calculator` and/or `model_id`:
+
+```python
+from maseval import StaticPricingCalculator
+
+calculator = StaticPricingCalculator({
+    "my-model": {"input": 0.001, "output": 0.002},
+})
+
+adapter = SmolAgentAdapter(
+    agent, name="researcher",
+    cost_calculator=calculator,
+    model_id="my-model",
+)
+```
+
+| Framework | Model ID | Cost Calculator |
+|-----------|----------|-----------------|
+| smolagents | Auto (`agent.model.model_id`) | Auto (`LiteLLMCostCalculator`) |
+| CAMEL | Auto (`agent.model_backend.model_type`) | Auto (`LiteLLMCostCalculator`) |
+| LlamaIndex | Auto (`agent.llm.metadata.model_name`) | Auto (`LiteLLMCostCalculator`) |
+| LangGraph | **Manual** (`model_id=...`) | Auto (`LiteLLMCostCalculator`) |
+
+If litellm is not installed, auto-creation of the calculator is skipped and cost stays at `0.0`. Tokens are always tracked regardless.
+
 ### In Benchmarks
 
 Usage is collected automatically alongside traces and configs after each task. Each report includes a `"usage"` key:

examples/five_a_day_benchmark/five_a_day_benchmark.py

@@ -967,7 +967,7 @@ def load_benchmark_data(
     def _fmt_usage(usage):
         parts = [f"cost=${usage.cost:.6f}"]
         if isinstance(usage, TokenUsage):
-            parts.append(f"in={usage.input_tokens}  out={usage.output_tokens}")
+            parts.append(f"in={usage.input_tokens:>10,}  out={usage.output_tokens:>10,}")
         if usage.units:
             parts.append(f"units={dict(usage.units)}")
         return "  ".join(parts)

maseval/core/agent.py

@@ -5,25 +5,56 @@
 from .history import MessageHistory
 from .tracing import TraceableMixin
 from .config import ConfigurableMixin
-from .usage import UsageTrackableMixin
+from .usage import Usage, TokenUsage, UsageTrackableMixin, CostCalculator
 
 
 class AgentAdapter(ABC, TraceableMixin, ConfigurableMixin, UsageTrackableMixin):
     """Wraps an agent from any framework to provide a standard interface.
 
     This Adapter provides:
+
     - Unified execution interface via `run()`
     - Callback hooks for monitoring
     - Message history management via getter/setter
     - Framework-agnostic tracing
+    - Automatic cost calculation from token usage (when a cost calculator is available)
+
+    Cost Tracking:
+        Agent adapters track token usage from the underlying framework. To also
+        compute cost, you can pass a ``cost_calculator`` and optionally a ``model_id``.
+
+        Most framework adapters auto-detect both the model ID (from the framework's
+        agent object) and the cost calculator (using ``LiteLLMCostCalculator`` if
+        litellm is installed). This means cost tracking often works with zero
+        configuration.
+
+        To override or disable auto-detection, pass explicit values::
+
+            adapter = SmolAgentAdapter(
+                agent, name="researcher",
+                cost_calculator=StaticPricingCalculator({...}),
+                model_id="my-custom-model",
+            )
+
+        Pass ``cost_calculator=None`` explicitly to disable cost calculation
+        even when auto-detection would otherwise enable it.
     """
 
-    def __init__(self, agent_instance: Any, name: str, callbacks: Optional[List[AgentCallback]] = None):
+    def __init__(
+        self,
+        agent_instance: Any,
+        name: str,
+        callbacks: Optional[List[AgentCallback]] = None,
+        cost_calculator: Optional[CostCalculator] = None,
+        model_id: Optional[str] = None,
+    ):
         self.agent = agent_instance
         self.name = name
         self.callbacks = callbacks or []
         self.messages: Optional[MessageHistory] = None
         self.logs: List[Dict[str, Any]] = []
+        self._cost_calculator = cost_calculator
+        self._model_id = model_id
 
     def run(self, query: str) -> Any:
         """Executes the agent and returns the result."""
@@ -105,6 +136,70 @@ def get_messages(self) -> MessageHistory:
         """
         return self.messages if self.messages is not None else MessageHistory()
 
+    def gather_usage(self) -> Usage:
+        """Gather usage with automatic cost calculation.
+
+        Calls ``_gather_usage()`` for raw token counts, then applies
+        the cost calculator if one is available and cost is still ``0.0``.
+
+        The ``model_id`` used for cost calculation is resolved in order:
+
+        1. Explicit ``model_id`` passed to ``__init__``
+        2. Auto-detected from the framework agent via ``_resolve_model_id()``
+
+        Subclasses should override ``_gather_usage()`` (not this method)
+        to provide framework-specific token extraction.
+
+        Returns:
+            Usage (or TokenUsage) with cost filled in when possible.
+        """
+        usage = self._gather_usage()
+        if isinstance(usage, TokenUsage) and usage.cost == 0.0:
+            calculator = self._resolve_cost_calculator()
+            if calculator is not None:
+                mid = self._model_id or self._resolve_model_id()
+                if mid:
+                    cost = calculator.calculate_cost(usage, mid)
+                    if cost is not None:
+                        usage.cost = cost
+        return usage
+
+    def _gather_usage(self) -> Usage:
+        """Gather raw token usage from the framework.
+
+        Override this in subclasses to extract token counts from the
+        framework's native data structures.
+
+        Returns:
+            Usage or TokenUsage with token counts (cost may be 0.0).
+        """
+        return Usage()
+
+    def _resolve_model_id(self) -> Optional[str]:
+        """Auto-detect the model ID from the framework agent.
+
+        Override in subclasses to extract the model identifier from
+        the framework's agent object (e.g., ``self.agent.model.model_id``
+        for smolagents).
+
+        Returns:
+            Model ID string, or ``None`` if not detectable.
+        """
+        return None
+
+    def _resolve_cost_calculator(self) -> Optional[CostCalculator]:
+        """Resolve the cost calculator to use.
+
+        Returns the explicit calculator if one was provided, otherwise
+        returns ``None``. Framework-specific subclasses can override this
+        to auto-create a calculator (e.g., ``LiteLLMCostCalculator``)
+        when the required dependencies are available.
+
+        Returns:
+            A CostCalculator, or ``None`` if cost calculation is not available.
+        """
+        return self._cost_calculator
+
     def gather_traces(self) -> Dict[str, Any]:
         """Gather execution traces from this agent.
 

maseval/interface/agents/camel.py

@@ -175,7 +175,9 @@ class CamelAgentAdapter(AgentAdapter):
         camel-ai to be installed: `pip install maseval[camel]`
     """
 
-    def __init__(self, agent_instance: Any, name: str, callbacks: Optional[List[Any]] = None):
+    def __init__(
+        self, agent_instance: Any, name: str, callbacks: Optional[List[Any]] = None, cost_calculator: Any = None, model_id: Optional[str] = None
+    ):
         """Initialize the CAMEL adapter.
 
         Note: We don't call super().__init__() to avoid initializing self.logs as a list,
@@ -185,11 +187,19 @@ def __init__(self, agent_instance: Any, name: str, callbacks: Optional[List[Any]
             agent_instance: CAMEL ChatAgent instance
             name: Agent name for identification
             callbacks: Optional list of AgentCallback instances
+            cost_calculator: Optional cost calculator. If not provided, a
+                ``LiteLLMCostCalculator`` is created automatically when litellm
+                is available.
+            model_id: Optional model ID for cost calculation. If not provided,
+                auto-detected from ``agent.model_backend.model_type``.
         """
         self.agent = agent_instance
         self.name = name
         self.callbacks = callbacks or []
         self.messages = None
+        self._cost_calculator = cost_calculator
+        self._model_id = model_id
+        self._auto_calculator = None  # Lazy-initialized
         # Store responses from each step() call
         self._responses: List[Any] = []
         # Store errors that occur during execution (for comprehensive logging)
@@ -419,7 +429,36 @@ def gather_traces(self) -> Dict[str, Any]:
 
         return base_traces
 
-    def gather_usage(self) -> Usage:
+    def _resolve_model_id(self):
+        """Auto-detect model ID from CAMEL agent.
+
+        CAMEL's ChatAgent stores the model backend in ``model_backend``
+        (or ``model`` for older versions). The backend has a ``model_type``
+        enum whose ``.value`` is the model ID string.
+        """
+        try:
+            backend = getattr(self.agent, "model_backend", None) or getattr(self.agent, "model", None)
+            if backend is not None and hasattr(backend, "model_type"):
+                model_type = backend.model_type
+                return model_type.value if hasattr(model_type, "value") else str(model_type)
+        except Exception:
+            pass
+        return None
+
+    def _resolve_cost_calculator(self):
+        """Return the cost calculator, auto-creating one if litellm is available."""
+        if self._cost_calculator is not None:
+            return self._cost_calculator
+        if self._auto_calculator is None:
+            try:
+                from maseval.interface.usage import LiteLLMCostCalculator
+
+                self._auto_calculator = LiteLLMCostCalculator()
+            except (ImportError, Exception):
+                self._auto_calculator = False
+        return self._auto_calculator if self._auto_calculator is not False else None
+
+    def _gather_usage(self) -> Usage:
         """Gather aggregated token usage across all CAMEL agent responses.
 
         Walks stored ``ChatAgentResponse`` objects and sums their

maseval/interface/agents/langgraph.py

@@ -117,19 +117,39 @@ def chatbot(state: MessagesState):
         langgraph to be installed: `pip install maseval[langgraph]`
     """
 
-    def __init__(self, agent_instance: Any, name: str, callbacks: Optional[List[Any]] = None, config: Optional[Dict[str, Any]] = None):
+    def __init__(
+        self,
+        agent_instance: Any,
+        name: str,
+        callbacks: Optional[List[Any]] = None,
+        config: Optional[Dict[str, Any]] = None,
+        cost_calculator: Any = None,
+        model_id: Optional[str] = None,
+    ):
         """Initialize the LangGraph adapter.
 
         Args:
             agent_instance: Compiled LangGraph graph
             name: Agent name
             callbacks: Optional list of callbacks
-            config: Optional LangGraph config dict (for stateful graphs with checkpointer)
-                   Should include 'configurable': {'thread_id': '...'} for persistent state
+            config: Optional LangGraph config dict (for stateful graphs with checkpointer).
+                Should include ``'configurable': {'thread_id': '...'}`` for persistent state.
+            cost_calculator: Optional cost calculator. If not provided, a
+                ``LiteLLMCostCalculator`` is created automatically when litellm
+                is available.
+            model_id: Model ID for cost calculation. LangGraph graphs can contain
+                multiple models across nodes, so the model ID cannot be auto-detected.
+                Pass the primary model's ID here to enable cost tracking::
+
+                    LangGraphAgentAdapter(
+                        graph, "agent",
+                        model_id="gpt-4o-mini",
+                    )
         """
-        super().__init__(agent_instance, name, callbacks)
+        super().__init__(agent_instance, name, callbacks, cost_calculator=cost_calculator, model_id=model_id)
         self._langgraph_config = config
         self._last_result = None
+        self._auto_calculator = None  # Lazy-initialized
 
     def get_messages(self) -> MessageHistory:
         """Get message history from LangGraph.
@@ -214,7 +234,20 @@ def gather_config(self) -> dict[str, Any]:
 
         return base_config
 
-    def gather_usage(self) -> Usage:
+    def _resolve_cost_calculator(self):
+        """Return the cost calculator, auto-creating one if litellm is available."""
+        if self._cost_calculator is not None:
+            return self._cost_calculator
+        if self._auto_calculator is None:
+            try:
+                from maseval.interface.usage import LiteLLMCostCalculator
+
+                self._auto_calculator = LiteLLMCostCalculator()
+            except (ImportError, Exception):
+                self._auto_calculator = False
+        return self._auto_calculator if self._auto_calculator is not False else None
+
+    def _gather_usage(self) -> Usage:
         """Gather aggregated token usage from LangGraph message metadata.
 
         Walks messages from the last graph execution (or persistent state)