Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 25 additions & 2 deletions src/agents/run_internal/run_loop.py
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,21 @@
]


def _get_model_name(model: Any) -> str | None:
"""Extract the string model name from a Model instance, if available.

Most built-in model implementations (``OpenAIResponsesModel``,
``OpenAIChatCompletionsModel``) expose a ``model`` attribute that contains
the underlying model name string. This helper retrieves it in a
forward-compatible, duck-typed way so that third-party model implementations
that may not have this attribute are handled gracefully.
"""
model_name = getattr(model, "model", None)
if isinstance(model_name, str):
return model_name
return None


async def _should_persist_stream_items(
*,
session: Session | None,
Expand Down Expand Up @@ -1401,7 +1416,11 @@ async def rewind_model_request() -> None:
)

if final_response is not None:
context_wrapper.usage.add(final_response.usage)
context_wrapper.usage.add(
final_response.usage,
agent_name=agent.name,
model_name=_get_model_name(model),
)
await asyncio.gather(
(
agent.hooks.on_llm_end(context_wrapper, agent, final_response)
Expand Down Expand Up @@ -1656,7 +1675,11 @@ async def rewind_model_request() -> None:
# new deltas.
server_conversation_tracker.mark_input_as_sent(filtered.input)

context_wrapper.usage.add(new_response.usage)
context_wrapper.usage.add(
new_response.usage,
agent_name=agent.name,
model_name=_get_model_name(model),
)

await asyncio.gather(
(
Expand Down
52 changes: 49 additions & 3 deletions src/agents/usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ def deserialize_usage(usage_data: Mapping[str, Any]) -> Usage:
entry.get("output_tokens_details") or {"reasoning_tokens": 0},
OutputTokensDetails(reasoning_tokens=0),
),
agent_name=entry.get("agent_name", None),
model_name=entry.get("model_name", None),
)
)

Expand Down Expand Up @@ -76,6 +78,20 @@ class RequestUsage:
output_tokens_details: OutputTokensDetails
"""Details about the output tokens for this individual request."""

agent_name: str | None = None
"""Name of the agent that made this request, if available.

Populated automatically when an agent makes a model call so that callers can attribute
token usage and costs to specific agents in multi-agent workflows.
"""

model_name: str | None = None
"""Name of the model used for this request, if available.

Populated automatically when an agent makes a model call so that callers can attribute
token usage and costs to specific models.
"""


def _normalize_input_tokens_details(
v: InputTokensDetails | PromptTokensDetails | None,
Expand Down Expand Up @@ -154,13 +170,23 @@ def __post_init__(self) -> None:
if output_details_none or output_reasoning_none:
self.output_tokens_details = OutputTokensDetails(reasoning_tokens=0)

def add(self, other: Usage) -> None:
def add(
self,
other: Usage,
*,
agent_name: str | None = None,
model_name: str | 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.
agent_name: Optional name of the agent making this request, used to annotate the
resulting ``RequestUsage`` entry for per-agent cost attribution.
model_name: Optional name of the model used for this request, used to annotate the
resulting ``RequestUsage`` entry for per-model cost attribution.
"""
self.requests += other.requests if other.requests else 0
self.input_tokens += other.input_tokens if other.input_tokens else 0
Expand Down Expand Up @@ -206,11 +232,26 @@ def add(self, other: Usage) -> None:
total_tokens=other.total_tokens,
input_tokens_details=input_details,
output_tokens_details=output_details,
agent_name=agent_name,
model_name=model_name,
)
self.request_usage_entries.append(request_usage)
elif other.request_usage_entries:
# If the other Usage already has individual request breakdowns, merge them.
self.request_usage_entries.extend(other.request_usage_entries)
# Apply agent_name/model_name to entries that don't already have them set,
# but copy each entry rather than mutating the original objects in place
# to avoid silent mis-attribution when the same Usage is added multiple times.
for entry in other.request_usage_entries:
annotated_entry = RequestUsage(
input_tokens=entry.input_tokens,
output_tokens=entry.output_tokens,
total_tokens=entry.total_tokens,
input_tokens_details=entry.input_tokens_details,
output_tokens_details=entry.output_tokens_details,
agent_name=agent_name if (agent_name is not None and entry.agent_name is None) else entry.agent_name,
model_name=model_name if (model_name is not None and entry.model_name is None) else entry.model_name,
)
self.request_usage_entries.append(annotated_entry)


def _serialize_usage_details(details: Any, default: dict[str, int]) -> dict[str, Any]:
Expand All @@ -228,7 +269,7 @@ def serialize_usage(usage: Usage) -> dict[str, Any]:
output_details = _serialize_usage_details(usage.output_tokens_details, {"reasoning_tokens": 0})

def _serialize_request_entry(entry: RequestUsage) -> dict[str, Any]:
return {
result: dict[str, Any] = {
"input_tokens": entry.input_tokens,
"output_tokens": entry.output_tokens,
"total_tokens": entry.total_tokens,
Expand All @@ -239,6 +280,11 @@ def _serialize_request_entry(entry: RequestUsage) -> dict[str, Any]:
entry.output_tokens_details, {"reasoning_tokens": 0}
),
}
if entry.agent_name is not None:
result["agent_name"] = entry.agent_name
if entry.model_name is not None:
result["model_name"] = entry.model_name
return result

return {
"requests": usage.requests,
Expand Down
Loading