default value for tokens before timeout when no usage

openshift-psap · Jun 10, 2024 · c8a1f8e · c8a1f8e
1 parent 366c7ce
commit c8a1f8e
Showing 1 changed file with 7 additions and 4 deletions.
diff --git 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