From 4a4409386c7a43d97fbc3ad8b0928bc38d1731c3 Mon Sep 17 00:00:00 2001
From: Yiran Wu <32823396+kevin666aa@users.noreply.github.com>
Date: Fri, 5 Apr 2024 17:12:56 -0400
Subject: [PATCH] fix cost (#2292)

---
 autogen/oai/client.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/autogen/oai/client.py b/autogen/oai/client.py
index a5f44d2bbb1f..de35e5c5273f 100644
--- 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: