From c8a1f8e3cc5779125d4cb39b806da68785b6da40 Mon Sep 17 00:00:00 2001 From: Drew Ripberger Date: Mon, 10 Jun 2024 13:32:45 -0400 Subject: [PATCH] default value for tokens before timeout when no usage --- plugins/openai_plugin.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/plugins/openai_plugin.py b/plugins/openai_plugin.py index 33c98d08..f05ec2e1 100644 --- a/plugins/openai_plugin.py +++ b/plugins/openai_plugin.py @@ -178,6 +178,8 @@ def streaming_request_http(self, query: dict, user_id: int, test_end_time: float if found and data != b"[DONE]": try: message = json.loads(data) + print("Original message received:") + print(message) logger.debug("Message: %s", message) if "/v1/chat/completions" in self.host and not message["choices"][0]['delta'].get('content'): message["choices"][0]['delta']['content']="" @@ -227,10 +229,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 +247,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