From 366c7ce3b82ef1217ed88ab6411f2045a17ad880 Mon Sep 17 00:00:00 2001 From: Drew Ripberger Date: Fri, 7 Jun 2024 12:56:51 -0400 Subject: [PATCH 1/3] fix: OpenAI API handle optional usage field --- plugins/openai_plugin.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/plugins/openai_plugin.py b/plugins/openai_plugin.py index 87efdcf5..33c98d08 100644 --- a/plugins/openai_plugin.py +++ b/plugins/openai_plugin.py @@ -224,14 +224,15 @@ def streaming_request_http(self, query: dict, user_id: int, test_end_time: float # Last token comes with finish_reason set. if message.get("choices", [])[0].get("finish_reason", None): - result.output_tokens = message["usage"]["completion_tokens"] - result.input_tokens = message["usage"]["prompt_tokens"] - result.stop_reason = message["choices"][0]["finish_reason"] + if message.get("usage"): + result.output_tokens = message["usage"]["completion_tokens"] + result.input_tokens = message["usage"]["prompt_tokens"] + # If test duration timeout didn't happen before the last token is received, + # total tokens before the timeout will be equal to the total tokens in the response. + if not result.output_tokens_before_timeout: + result.output_tokens_before_timeout = result.output_tokens - # If test duration timeout didn't happen before the last token is received, - # total tokens before the timeout will be equal to the total tokens in the response. - if not result.output_tokens_before_timeout: - result.output_tokens_before_timeout = result.output_tokens + result.stop_reason = message["choices"][0]["finish_reason"] except KeyError: logging.exception("KeyError, unexpected response format in line: %s", line) From 1c90aa795348cc8cc2392bd5238eb3c730687b29 Mon Sep 17 00:00:00 2001 From: Drew Ripberger Date: Mon, 10 Jun 2024 13:32:45 -0400 Subject: [PATCH 2/3] default value for tokens before timeout when no usage --- plugins/openai_plugin.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/plugins/openai_plugin.py b/plugins/openai_plugin.py index 33c98d08..5b1a2285 100644 --- a/plugins/openai_plugin.py +++ b/plugins/openai_plugin.py @@ -227,10 +227,6 @@ def streaming_request_http(self, query: dict, user_id: int, test_end_time: float if message.get("usage"): result.output_tokens = message["usage"]["completion_tokens"] result.input_tokens = message["usage"]["prompt_tokens"] - # If test duration timeout didn't happen before the last token is received, - # total tokens before the timeout will be equal to the total tokens in the response. - if not result.output_tokens_before_timeout: - result.output_tokens_before_timeout = result.output_tokens result.stop_reason = message["choices"][0]["finish_reason"] @@ -249,5 +245,10 @@ def streaming_request_http(self, query: dict, user_id: int, test_end_time: float logger.warning("Output token count not found in response, length of token list") result.output_tokens = len(tokens) + # If test duration timeout didn't happen before the last token is received, + # total tokens before the timeout will be equal to the total tokens in the response. + if not result.output_tokens_before_timeout: + result.output_tokens_before_timeout = result.output_tokens + result.calculate_results() return result From 5c267bd19cbe553dca17acaf2312a4b9e76d0d35 Mon Sep 17 00:00:00 2001 From: Drew Ripberger Date: Mon, 10 Jun 2024 14:53:30 -0400 Subject: [PATCH 3/3] safer access to the choices field --- plugins/openai_plugin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plugins/openai_plugin.py b/plugins/openai_plugin.py index 5b1a2285..108c6392 100644 --- a/plugins/openai_plugin.py +++ b/plugins/openai_plugin.py @@ -223,12 +223,12 @@ def streaming_request_http(self, query: dict, user_id: int, test_end_time: float tokens.append(token) # Last token comes with finish_reason set. - if message.get("choices", [])[0].get("finish_reason", None): + if message.get("choices", [{}])[0].get("finish_reason", None): if message.get("usage"): result.output_tokens = message["usage"]["completion_tokens"] result.input_tokens = message["usage"]["prompt_tokens"] - result.stop_reason = message["choices"][0]["finish_reason"] + result.stop_reason = message.get("choices", [{}])[0].get("finish_reason", None) except KeyError: logging.exception("KeyError, unexpected response format in line: %s", line)