Skip to content
Next Next commit
feat: implemented class using InputTokenDetails, will revert for comp…
…letions compatibility
  • Loading branch information
WJPBProjects committed May 20, 2025
commit a08fe82ac879dbd1856a03d6712e30e05fb20cac
19 changes: 19 additions & 0 deletions src/agents/usage.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,17 @@
from dataclasses import dataclass
from typing import TypeVar

from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails

T = TypeVar("T", bound="InputTokensDetails | OutputTokensDetails")


def add_numeric_fields(current: T, other: T) -> None:
for field in current.__dataclass_fields__:
v1 = getattr(current, field, 0)
v2 = getattr(other, field, 0)
if isinstance(v1, (int, float)) and isinstance(v2, (int, float)):
setattr(current, field, (v1 or 0) + (v2 or 0))


@dataclass
Expand All @@ -9,9 +22,13 @@ class Usage:
input_tokens: int = 0
"""Total input tokens sent, across all requests."""

input_tokens_details: InputTokensDetails = InputTokensDetails(cached_tokens=0)

output_tokens: int = 0
"""Total output tokens received, across all requests."""

output_tokens_details: OutputTokensDetails = OutputTokensDetails(reasoning_tokens=0)

total_tokens: int = 0
"""Total tokens sent and received, across all requests."""

Expand All @@ -20,3 +37,5 @@ def add(self, other: "Usage") -> None:
self.input_tokens += other.input_tokens if other.input_tokens else 0
self.output_tokens += other.output_tokens if other.output_tokens else 0
self.total_tokens += other.total_tokens if other.total_tokens else 0
add_numeric_fields(self.input_tokens_details, other.input_tokens_details)
add_numeric_fields(self.output_tokens_details, other.output_tokens_details)