From 0569689e6b70592e4997adc0172f1cc0efff85a3 Mon Sep 17 00:00:00 2001 From: Gerardo Moreno Date: Mon, 30 Dec 2024 14:33:06 -0600 Subject: [PATCH] Fix BaseOpenAIChatCompletionClient token usage (#4770) * Fix openai client token usage (#4769) * Include actual_usage in add_usage function * Make sure result is assigned * remove method --------- Co-authored-by: Jack Gerrits --- .../src/autogen_ext/models/openai/_openai_client.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index db5fd911302..6f032a3393b 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py @@ -558,8 +558,8 @@ async def create( logprobs=logprobs, ) - _add_usage(self._actual_usage, usage) - _add_usage(self._total_usage, usage) + self._total_usage = _add_usage(self._total_usage, usage) + self._actual_usage = _add_usage(self._actual_usage, usage) # TODO - why is this cast needed? return response @@ -775,8 +775,8 @@ async def create_stream( logprobs=logprobs, ) - _add_usage(self._actual_usage, usage) - _add_usage(self._total_usage, usage) + self._total_usage = _add_usage(self._total_usage, usage) + self._actual_usage = _add_usage(self._actual_usage, usage) yield result