From 530335a0016a8a6a16d07b5ee0227db7a8413150 Mon Sep 17 00:00:00 2001 From: Rohan Thacker Date: Tue, 31 Dec 2024 15:07:23 +0530 Subject: [PATCH] Added: Operator overloads for `Request Usage` --- .../src/autogen_core/models/_types.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/python/packages/autogen-core/src/autogen_core/models/_types.py b/python/packages/autogen-core/src/autogen_core/models/_types.py index fb118562e4d3..fa133b48315a 100644 --- a/python/packages/autogen-core/src/autogen_core/models/_types.py +++ b/python/packages/autogen-core/src/autogen_core/models/_types.py @@ -51,6 +51,25 @@ class RequestUsage: prompt_tokens: int completion_tokens: int + def __add__(self, other: "RequestUsage") -> "RequestUsage": + # Runtime type check ensures robustness, even though static analysis indicates unreachable code. + if not isinstance(other, RequestUsage): + raise TypeError( + f"Unsupported operand type(s) for +: 'RequestUsage' and '{type(other).__name__}'" + ) + return RequestUsage( + prompt_tokens=self.prompt_tokens + other.prompt_tokens, + completion_tokens=self.completion_tokens + other.completion_tokens, + ) + + def __iadd__(self, other: "RequestUsage") -> "RequestUsage": + # Runtime type check ensures robustness, even though static analysis indicates unreachable code. + if not isinstance(other, RequestUsage): + return NotImplemented + self.prompt_tokens += other.prompt_tokens + self.completion_tokens += other.completion_tokens + return self + FinishReasons = Literal["stop", "length", "function_calls", "content_filter"]