diff --git a/docs/usage.md b/docs/usage.md index bedae99b3..84008832f 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -54,8 +54,8 @@ The SDK automatically tracks usage for each API request in `request_usage_entrie ```python result = await Runner.run(agent, "What's the weather in Tokyo?") -for request in enumerate(result.context_wrapper.usage.request_usage_entries): - print(f"Request {i + 1}: {request.input_tokens} in, {request.output_tokens} out") +for i, request in enumerate(result.context_wrapper.usage.request_usage_entries): + print(f"Request {i + 1}: Input={request.input_tokens} tokens, Output={request.output_tokens} tokens, metadata={request.metadata}") ``` ## Accessing usage with sessions diff --git a/src/agents/models/interface.py b/src/agents/models/interface.py index f25934780..4df4d0a97 100644 --- a/src/agents/models/interface.py +++ b/src/agents/models/interface.py @@ -36,6 +36,9 @@ def include_data(self) -> bool: class Model(abc.ABC): """The base interface for calling an LLM.""" + # The model name. Subclasses can set this in __init__. + model: str = "" + @abc.abstractmethod async def get_response( self, diff --git a/src/agents/run.py b/src/agents/run.py index fce7b4840..1f64becb6 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -1402,7 +1402,14 @@ async def _run_single_turn_streamed( usage=usage, response_id=event.response.id, ) - context_wrapper.usage.add(usage) + context_wrapper.usage.add( + usage, + metadata={ + "model_name": model.model, + "agent_name": agent.name, + "response_id": event.response.id, + }, + ) if isinstance(event, ResponseOutputItemDoneEvent): output_item = event.item @@ -1819,7 +1826,14 @@ async def _get_new_response( prompt=prompt_config, ) - context_wrapper.usage.add(new_response.usage) + context_wrapper.usage.add( + new_response.usage, + metadata={ + "model_name": model.model, + "agent_name": agent.name, + "response_id": new_response.response_id, + }, + ) # If we have run hooks, or if the agent has hooks, we need to call them after the LLM call await asyncio.gather( diff --git a/src/agents/usage.py b/src/agents/usage.py index a10778123..dbc0076e9 100644 --- a/src/agents/usage.py +++ b/src/agents/usage.py @@ -1,4 +1,7 @@ +from __future__ import annotations + from dataclasses import field +from typing import Any from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails from pydantic.dataclasses import dataclass @@ -23,6 +26,9 @@ class RequestUsage: output_tokens_details: OutputTokensDetails """Details about the output tokens for this individual request.""" + metadata: dict[str, Any] = field(default_factory=dict) + """Additional metadata for this request (e.g., model_name, agent_name, response_id).""" + @dataclass class Usage: @@ -70,13 +76,18 @@ def __post_init__(self) -> None: if self.output_tokens_details.reasoning_tokens is None: self.output_tokens_details = OutputTokensDetails(reasoning_tokens=0) - def add(self, other: "Usage") -> None: + def add( + self, + other: Usage, + metadata: dict[str, Any] | None = None, + ) -> None: """Add another Usage object to this one, aggregating all fields. This method automatically preserves request_usage_entries. Args: other: The Usage object to add to this one. + metadata: Additional metadata for this request """ self.requests += other.requests if other.requests else 0 self.input_tokens += other.input_tokens if other.input_tokens else 0 @@ -101,6 +112,7 @@ def add(self, other: "Usage") -> None: total_tokens=other.total_tokens, input_tokens_details=other.input_tokens_details, output_tokens_details=other.output_tokens_details, + metadata=metadata or {}, ) self.request_usage_entries.append(request_usage) elif other.request_usage_entries: diff --git a/tests/test_usage.py b/tests/test_usage.py index 9d89cc750..2d0dd086f 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -21,7 +21,14 @@ def test_usage_add_aggregates_all_fields(): total_tokens=15, ) - u1.add(u2) + u1.add( + u2, + metadata={ + "model_name": "gpt-5", + "agent_name": "test-agent", + "response_id": "resp-1", + }, + ) assert u1.requests == 3 assert u1.input_tokens == 17 @@ -42,7 +49,14 @@ def test_usage_add_aggregates_with_none_values(): total_tokens=15, ) - u1.add(u2) + u1.add( + u2, + metadata={ + "model_name": "gpt-5", + "agent_name": "test-agent", + "response_id": "resp-1", + }, + ) assert u1.requests == 2 assert u1.input_tokens == 7 @@ -60,6 +74,11 @@ def test_request_usage_creation(): total_tokens=300, input_tokens_details=InputTokensDetails(cached_tokens=10), output_tokens_details=OutputTokensDetails(reasoning_tokens=20), + metadata={ + "model_name": "gpt-5", + "agent_name": "test-agent", + "response_id": "resp-123", + }, ) assert request_usage.input_tokens == 100 @@ -67,6 +86,9 @@ def test_request_usage_creation(): assert request_usage.total_tokens == 300 assert request_usage.input_tokens_details.cached_tokens == 10 assert request_usage.output_tokens_details.reasoning_tokens == 20 + assert request_usage.metadata["model_name"] == "gpt-5" + assert request_usage.metadata["agent_name"] == "test-agent" + assert request_usage.metadata["response_id"] == "resp-123" def test_usage_add_preserves_single_request(): @@ -81,7 +103,14 @@ def test_usage_add_preserves_single_request(): total_tokens=300, ) - u1.add(u2) + u1.add( + u2, + metadata={ + "model_name": "gpt-5", + "agent_name": "test-agent", + "response_id": "resp-1", + }, + ) # Should preserve the request usage details assert len(u1.request_usage_entries) == 1 @@ -91,6 +120,9 @@ def test_usage_add_preserves_single_request(): assert request_usage.total_tokens == 300 assert request_usage.input_tokens_details.cached_tokens == 10 assert request_usage.output_tokens_details.reasoning_tokens == 20 + assert request_usage.metadata["model_name"] == "gpt-5" + assert request_usage.metadata["agent_name"] == "test-agent" + assert request_usage.metadata["response_id"] == "resp-1" def test_usage_add_ignores_zero_token_requests(): @@ -105,7 +137,14 @@ def test_usage_add_ignores_zero_token_requests(): total_tokens=0, ) - u1.add(u2) + u1.add( + u2, + metadata={ + "model_name": "gpt-5", + "agent_name": "test-agent", + "response_id": "resp-1", + }, + ) # Should not create a request_usage_entry for zero tokens assert len(u1.request_usage_entries) == 0 @@ -123,7 +162,14 @@ def test_usage_add_ignores_multi_request_usage(): total_tokens=300, ) - u1.add(u2) + u1.add( + u2, + metadata={ + "model_name": "gpt-5", + "agent_name": "test-agent", + "response_id": "resp-1", + }, + ) # Should not create a request usage entry for multi-request usage assert len(u1.request_usage_entries) == 0 @@ -141,7 +187,14 @@ def test_usage_add_merges_existing_request_usage_entries(): output_tokens_details=OutputTokensDetails(reasoning_tokens=20), total_tokens=300, ) - u1.add(u2) + u1.add( + u2, + metadata={ + "model_name": "gpt-5", + "agent_name": "agent-1", + "response_id": "resp-1", + }, + ) # Create second usage with request_usage_entries u3 = Usage( @@ -153,7 +206,14 @@ def test_usage_add_merges_existing_request_usage_entries(): total_tokens=125, ) - u1.add(u3) + u1.add( + u3, + metadata={ + "model_name": "gpt-5", + "agent_name": "agent-2", + "response_id": "resp-2", + }, + ) # Should have both request_usage_entries assert len(u1.request_usage_entries) == 2 @@ -163,12 +223,16 @@ def test_usage_add_merges_existing_request_usage_entries(): assert first.input_tokens == 100 assert first.output_tokens == 200 assert first.total_tokens == 300 + assert first.metadata["agent_name"] == "agent-1" + assert first.metadata["response_id"] == "resp-1" # Second request second = u1.request_usage_entries[1] assert second.input_tokens == 50 assert second.output_tokens == 75 assert second.total_tokens == 125 + assert second.metadata["agent_name"] == "agent-2" + assert second.metadata["response_id"] == "resp-2" def test_usage_add_with_pre_existing_request_usage_entries(): @@ -184,7 +248,14 @@ def test_usage_add_with_pre_existing_request_usage_entries(): output_tokens_details=OutputTokensDetails(reasoning_tokens=20), total_tokens=300, ) - u1.add(u2) + u1.add( + u2, + metadata={ + "model_name": "gpt-5", + "agent_name": "agent-1", + "response_id": "resp-1", + }, + ) # Create another usage with request_usage_entries u3 = Usage( @@ -197,7 +268,14 @@ def test_usage_add_with_pre_existing_request_usage_entries(): ) # Add u3 to u1 - u1.add(u3) + u1.add( + u3, + metadata={ + "model_name": "gpt-5", + "agent_name": "agent-2", + "response_id": "resp-2", + }, + ) # Should have both request_usage_entries assert len(u1.request_usage_entries) == 2 @@ -227,7 +305,14 @@ def test_anthropic_cost_calculation_scenario(): output_tokens_details=OutputTokensDetails(reasoning_tokens=0), total_tokens=150_000, ) - usage.add(req1) + usage.add( + req1, + metadata={ + "model_name": "gpt-5", + "agent_name": "test-agent", + "response_id": "resp-1", + }, + ) # Second request: 150K input tokens req2 = Usage( @@ -238,7 +323,14 @@ def test_anthropic_cost_calculation_scenario(): output_tokens_details=OutputTokensDetails(reasoning_tokens=0), total_tokens=225_000, ) - usage.add(req2) + usage.add( + req2, + metadata={ + "model_name": "gpt-5", + "agent_name": "test-agent", + "response_id": "resp-2", + }, + ) # Third request: 80K input tokens req3 = Usage( @@ -249,7 +341,14 @@ def test_anthropic_cost_calculation_scenario(): output_tokens_details=OutputTokensDetails(reasoning_tokens=0), total_tokens=120_000, ) - usage.add(req3) + usage.add( + req3, + metadata={ + "model_name": "gpt-5", + "agent_name": "test-agent", + "response_id": "resp-3", + }, + ) # Verify aggregated totals assert usage.requests == 3