fix cost (#2292)

microsoft · Apr 5, 2024 · 4a44093 · 4a44093
1 parent a4c1d58
commit 4a44093
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/autogen/oai/client.py b/autogen/oai/client.py
@@ -289,6 +289,8 @@ def cost(self, response: Union[ChatCompletion, Completion]) -> float:
 
         n_input_tokens = response.usage.prompt_tokens if response.usage is not None else 0  # type: ignore [union-attr]
         n_output_tokens = response.usage.completion_tokens if response.usage is not None else 0  # type: ignore [union-attr]
+        if n_output_tokens is None:
+            n_output_tokens = 0
         tmp_price1K = OAI_PRICE1K[model]
         # First value is input token rate, second value is output token rate
         if isinstance(tmp_price1K, tuple):
@@ -805,6 +807,8 @@ def update_usage(usage_summary, response_usage):
             cost = response_usage["cost"]
             prompt_tokens = response_usage["prompt_tokens"]
             completion_tokens = response_usage["completion_tokens"]
+            if completion_tokens is None:
+                completion_tokens = 0
             total_tokens = response_usage["total_tokens"]
 
             if usage_summary is None: