From f9cf23a2fcfb8c4a4977683ea442c67c62716bae Mon Sep 17 00:00:00 2001 From: gziz Date: Thu, 19 Dec 2024 21:57:54 -0800 Subject: [PATCH 1/4] Fix openai client token usage (#4769) --- .../models/openai/_openai_client.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 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 b7276888411f..67f901b4fe82 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 @@ -285,13 +285,6 @@ def calculate_vision_tokens(image: Image, detail: str = "auto") -> int: return total_tokens -def _add_usage(usage1: RequestUsage, usage2: RequestUsage) -> RequestUsage: - return RequestUsage( - prompt_tokens=usage1.prompt_tokens + usage2.prompt_tokens, - completion_tokens=usage1.completion_tokens + usage2.completion_tokens, - ) - - def convert_tools( tools: Sequence[Tool | ToolSchema], ) -> List[ChatCompletionToolParam]: @@ -561,8 +554,7 @@ async def create( logprobs=logprobs, ) - _add_usage(self._actual_usage, usage) - _add_usage(self._total_usage, usage) + self.add_usage(usage) # TODO - why is this cast needed? return response @@ -778,8 +770,7 @@ async def create_stream( logprobs=logprobs, ) - _add_usage(self._actual_usage, usage) - _add_usage(self._total_usage, usage) + self.add_usage(usage) yield result @@ -886,6 +877,12 @@ def remaining_tokens(self, messages: Sequence[LLMMessage], tools: Sequence[Tool def capabilities(self) -> ModelCapabilities: return self._model_capabilities + def add_usage(self, usage: RequestUsage) -> None: + self._total_usage = RequestUsage( + prompt_tokens=self._total_usage.prompt_tokens + usage.prompt_tokens, + completion_tokens=self._total_usage.completion_tokens + usage.completion_tokens, + ) + class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient): """Chat completion client for OpenAI hosted models. From 0136055ad1355d948a7fbcd2d410366dc83181d8 Mon Sep 17 00:00:00 2001 From: gziz Date: Thu, 26 Dec 2024 09:39:33 -0600 Subject: [PATCH 2/4] Include actual_usage in add_usage function --- .../src/autogen_ext/models/openai/_openai_client.py | 9 +++++---- 1 file changed, 5 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 67f901b4fe82..c151d7be19ae 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 @@ -878,10 +878,11 @@ def capabilities(self) -> ModelCapabilities: return self._model_capabilities def add_usage(self, usage: RequestUsage) -> None: - self._total_usage = RequestUsage( - prompt_tokens=self._total_usage.prompt_tokens + usage.prompt_tokens, - completion_tokens=self._total_usage.completion_tokens + usage.completion_tokens, - ) + self._total_usage.prompt_tokens += usage.prompt_tokens + self._total_usage.completion_tokens += usage.completion_tokens + + self._actual_usage.prompt_tokens += usage.prompt_tokens + self._actual_usage.completion_tokens += usage.completion_tokens class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient): From a940f7305d46c82f4a3c12f77796ebb8e7ab6590 Mon Sep 17 00:00:00 2001 From: Jack Gerrits Date: Mon, 30 Dec 2024 15:29:15 -0500 Subject: [PATCH 3/4] Make sure result is assigned --- .../src/autogen_ext/models/openai/_openai_client.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 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 1eea3e309deb..5484c2a1dca7 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 @@ -269,6 +269,13 @@ def calculate_vision_tokens(image: Image, detail: str = "auto") -> int: return total_tokens +def _add_usage(usage1: RequestUsage, usage2: RequestUsage) -> RequestUsage: + return RequestUsage( + prompt_tokens=usage1.prompt_tokens + usage2.prompt_tokens, + completion_tokens=usage1.completion_tokens + usage2.completion_tokens, + ) + + def convert_tools( tools: Sequence[Tool | ToolSchema], ) -> List[ChatCompletionToolParam]: @@ -551,7 +558,8 @@ async def create( logprobs=logprobs, ) - self.add_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 @@ -767,7 +775,8 @@ async def create_stream( logprobs=logprobs, ) - self.add_usage(usage) + self._total_usage = _add_usage(self._total_usage, usage) + self._actual_usage = _add_usage(self._actual_usage, usage) yield result From ef95eceae7ad1b6f1a403dbe9fbeba2b4e34a5f9 Mon Sep 17 00:00:00 2001 From: Jack Gerrits Date: Mon, 30 Dec 2024 15:29:41 -0500 Subject: [PATCH 4/4] remove method --- .../src/autogen_ext/models/openai/_openai_client.py | 7 ------- 1 file changed, 7 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 5484c2a1dca7..6f032a3393b3 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 @@ -888,13 +888,6 @@ def capabilities(self) -> ModelCapabilities: # type: ignore def model_info(self) -> ModelInfo: return self._model_info - def add_usage(self, usage: RequestUsage) -> None: - self._total_usage.prompt_tokens += usage.prompt_tokens - self._total_usage.completion_tokens += usage.completion_tokens - - self._actual_usage.prompt_tokens += usage.prompt_tokens - self._actual_usage.completion_tokens += usage.completion_tokens - class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component[OpenAIClientConfigurationConfigModel]): """Chat completion client for OpenAI hosted models.