|
|
""" |
|
|
Monitoring utilities for tracking execution timing and token usage. |
|
|
""" |
|
|
|
|
|
from dataclasses import dataclass, field |
|
|
|
|
|
__all__ = ["Timing", "TokenUsage"] |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class TokenUsage: |
|
|
""" |
|
|
Contains the token usage information for a given step or run. |
|
|
""" |
|
|
|
|
|
input_tokens: int |
|
|
output_tokens: int |
|
|
total_tokens: int = field(init=False) |
|
|
|
|
|
def __post_init__(self): |
|
|
self.total_tokens = self.input_tokens + self.output_tokens |
|
|
|
|
|
def dict(self): |
|
|
return { |
|
|
"input_tokens": self.input_tokens, |
|
|
"output_tokens": self.output_tokens, |
|
|
"total_tokens": self.total_tokens, |
|
|
} |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class Timing: |
|
|
""" |
|
|
Contains the timing information for a given step or run. |
|
|
""" |
|
|
|
|
|
start_time: float |
|
|
end_time: float | None = None |
|
|
|
|
|
@property |
|
|
def duration(self): |
|
|
return None if self.end_time is None else self.end_time - self.start_time |
|
|
|
|
|
def dict(self): |
|
|
return { |
|
|
"start_time": self.start_time, |
|
|
"end_time": self.end_time, |
|
|
"duration": self.duration, |
|
|
} |
|
|
|
|
|
def __repr__(self) -> str: |
|
|
return f"Timing(start_time={self.start_time}, end_time={self.end_time}, duration={self.duration})" |
|
|
|
