From 4cd06cda2c7c787b31516bd19dc729bc93a693ea Mon Sep 17 00:00:00 2001 From: Venky Iyer Date: Thu, 25 Jan 2024 08:48:31 -0800 Subject: [PATCH] Be robust to null values while updating token counts (#359) --- instructor/patch.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/instructor/patch.py b/instructor/patch.py index 7f2cb72..d8d143a 100644 --- a/instructor/patch.py +++ b/instructor/patch.py @@ -237,9 +237,9 @@ async def retry_async( response: ChatCompletion = await func(*args, **kwargs) stream = kwargs.get("stream", False) if isinstance(response, ChatCompletion) and response.usage is not None: - total_usage.completion_tokens += response.usage.completion_tokens - total_usage.prompt_tokens += response.usage.prompt_tokens - total_usage.total_tokens += response.usage.total_tokens + total_usage.completion_tokens += response.usage.completion_tokens or 0 + total_usage.prompt_tokens += response.usage.prompt_tokens or 0 + total_usage.total_tokens += response.usage.total_tokens or 0 response.usage = ( total_usage # Replace each response usage with the total usage ) @@ -300,9 +300,9 @@ def retry_sync( response = func(*args, **kwargs) stream = kwargs.get("stream", False) if isinstance(response, ChatCompletion) and response.usage is not None: - total_usage.completion_tokens += response.usage.completion_tokens - total_usage.prompt_tokens += response.usage.prompt_tokens - total_usage.total_tokens += response.usage.total_tokens + total_usage.completion_tokens += response.usage.completion_tokens or 0 + total_usage.prompt_tokens += response.usage.prompt_tokens or 0 + total_usage.total_tokens += response.usage.total_tokens or 0 response.usage = ( total_usage # Replace each response usage with the total usage )