add litellm cost tracking feature to Usage objects through track_cost model setting

Author: habema

src/agents/extensions/models/litellm_model.py

@@ -124,6 +124,15 @@ async def get_response(
 
             if hasattr(response, "usage"):
                 response_usage = response.usage
+
+                # Extract cost from LiteLLM's hidden params if cost tracking is enabled.
+                cost = None
+                if model_settings.track_cost:
+                    if hasattr(response, "_hidden_params") and isinstance(
+                        response._hidden_params, dict
+                    ):
+                        cost = response._hidden_params.get("response_cost")
+
                 usage = (
                     Usage(
                         requests=1,
@@ -142,6 +151,7 @@ async def get_response(
                             )
                             or 0
                         ),
+                        cost=cost,
                     )
                     if response.usage
                     else Usage()
@@ -201,10 +211,67 @@ async def stream_response(
 
             final_response: Response | None = None
             async for chunk in ChatCmplStreamHandler.handle_stream(response, stream):
-                yield chunk
-
+                # Intercept the response.completed event to calculate and attach cost.
                 if chunk.type == "response.completed":
                     final_response = chunk.response
+                    # Calculate cost using LiteLLM's completion_cost function if enabled.
+                    # Streaming responses don't include cost in _hidden_params, so we
+                    # calculate it from the final token counts.
+                    if model_settings.track_cost and final_response.usage:
+                        try:
+                            # Create a mock ModelResponse for cost calculation.
+                            # Include token details (cached, reasoning) for accurate pricing.
+                            from litellm.types.utils import (
+                                Choices as LitellmChoices,
+                                CompletionTokensDetailsWrapper,
+                                Message as LitellmMessage,
+                                ModelResponse as LitellmModelResponse,
+                                PromptTokensDetailsWrapper,
+                                Usage as LitellmUsage,
+                            )
+
+                            # Extract token details for accurate cost calculation.
+                            cached_tokens = (
+                                final_response.usage.input_tokens_details.cached_tokens
+                                if final_response.usage.input_tokens_details
+                                else 0
+                            )
+                            reasoning_tokens = (
+                                final_response.usage.output_tokens_details.reasoning_tokens
+                                if final_response.usage.output_tokens_details
+                                else 0
+                            )
+
+                            mock_response = LitellmModelResponse(
+                                choices=[
+                                    LitellmChoices(
+                                        index=0,
+                                        message=LitellmMessage(role="assistant", content=""),
+                                    )
+                                ],
+                                usage=LitellmUsage(
+                                    prompt_tokens=final_response.usage.input_tokens,
+                                    completion_tokens=final_response.usage.output_tokens,
+                                    total_tokens=final_response.usage.total_tokens,
+                                    prompt_tokens_details=PromptTokensDetailsWrapper(
+                                        cached_tokens=cached_tokens
+                                    ),
+                                    completion_tokens_details=CompletionTokensDetailsWrapper(
+                                        reasoning_tokens=reasoning_tokens
+                                    ),
+                                ),
+                                model=self.model,
+                            )
+                            cost = litellm.completion_cost(completion_response=mock_response)
+                            # Attach cost as a custom attribute on the Response object so
+                            # run.py can access it when creating the Usage object.
+                            final_response._litellm_cost = cost
+                        except Exception:
+                            # If cost calculation fails (e.g., unknown model), continue
+                            # without cost.
+                            pass
+
+                yield chunk
 
             if tracing.include_data() and final_response:
                 span_generation.span_data.output = [final_response.model_dump()]

src/agents/model_settings.py

@@ -120,6 +120,12 @@ class ModelSettings:
     """Whether to include usage chunk.
     Only available for Chat Completions API."""
 
+    track_cost: bool = False
+    """Whether to track and calculate cost for model calls.
+    When enabled, the SDK will populate `Usage.cost` with cost estimates.
+    Currently only supported for LiteLLM models. For other providers, cost will remain None.
+    Defaults to False."""
+
     # TODO: revisit ResponseIncludable | str if ResponseIncludable covers more cases
     # We've added str to support missing ones like
     # "web_search_call.action.sources" etc.

src/agents/run.py

@@ -1104,6 +1104,9 @@ async def _run_single_turn_streamed(
             prompt=prompt_config,
         ):
             if isinstance(event, ResponseCompletedEvent):
+                # Extract cost if it was attached by LiteLLM model.
+                cost = getattr(event.response, "_litellm_cost", None)
+                
                 usage = (
                     Usage(
                         requests=1,
@@ -1112,6 +1115,7 @@ async def _run_single_turn_streamed(
                         total_tokens=event.response.usage.total_tokens,
                         input_tokens_details=event.response.usage.input_tokens_details,
                         output_tokens_details=event.response.usage.output_tokens_details,
+                        cost=cost,
                     )
                     if event.response.usage
                     else Usage()

src/agents/usage.py

@@ -27,6 +27,10 @@ class Usage:
     total_tokens: int = 0
     """Total tokens sent and received, across all requests."""
 
+    cost: float | None = None
+    """Total cost in USD for the requests. Only available for certain model providers
+    (e.g., LiteLLM). Will be None for models that don't provide cost information."""
+
     def add(self, other: "Usage") -> None:
         self.requests += other.requests if other.requests else 0
         self.input_tokens += other.input_tokens if other.input_tokens else 0
@@ -41,3 +45,7 @@ def add(self, other: "Usage") -> None:
             reasoning_tokens=self.output_tokens_details.reasoning_tokens
             + other.output_tokens_details.reasoning_tokens
         )
+
+        # Aggregate cost if either has a value.
+        if self.cost is not None or other.cost is not None:
+            self.cost = (self.cost or 0.0) + (other.cost or 0.0)

tests/test_usage.py

@@ -50,3 +50,52 @@ def test_usage_add_aggregates_with_none_values():
     assert u1.total_tokens == 15
     assert u1.input_tokens_details.cached_tokens == 4
     assert u1.output_tokens_details.reasoning_tokens == 6
+
+
+def test_usage_cost_defaults_to_none():
+    """Test that cost field defaults to None."""
+    usage = Usage()
+    assert usage.cost is None
+
+
+def test_usage_add_with_cost():
+    """Test that cost is aggregated correctly when both usages have cost."""
+    u1 = Usage(requests=1, input_tokens=10, output_tokens=20, total_tokens=30, cost=0.001)
+    u2 = Usage(requests=1, input_tokens=15, output_tokens=25, total_tokens=40, cost=0.002)
+
+    u1.add(u2)
+
+    assert u1.cost == 0.003
+    assert u1.requests == 2
+    assert u1.total_tokens == 70
+
+
+def test_usage_add_with_partial_cost():
+    """Test that cost is preserved when only one usage has cost."""
+    u1 = Usage(requests=1, input_tokens=10, output_tokens=20, total_tokens=30, cost=0.001)
+    u2 = Usage(requests=1, input_tokens=15, output_tokens=25, total_tokens=40)  # no cost
+
+    u1.add(u2)
+
+    assert u1.cost == 0.001
+    assert u1.requests == 2
+
+
+def test_usage_add_with_cost_none_plus_value():
+    """Test that cost aggregation works when first usage has no cost."""
+    u1 = Usage(requests=1, input_tokens=10, output_tokens=20, total_tokens=30)
+    u2 = Usage(requests=1, input_tokens=15, output_tokens=25, total_tokens=40, cost=0.002)
+
+    u1.add(u2)
+
+    assert u1.cost == 0.002
+
+
+def test_usage_add_with_both_cost_none():
+    """Test that cost remains None when neither usage has cost."""
+    u1 = Usage(requests=1, input_tokens=10, output_tokens=20, total_tokens=30)
+    u2 = Usage(requests=1, input_tokens=15, output_tokens=25, total_tokens=40)
+
+    u1.add(u2)
+
+    assert u1.cost is None