fix: handle optional usage field in OpenAI API (#47)

* fix: OpenAI API handle optional usage field * default value for tokens before timeout when no usage * safer access to the choices field
openshift-psap · Jun 26, 2024 · 8d64d85 · 8d64d85
1 parent 3d76871
commit 8d64d85
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/plugins/openai_plugin.py b/plugins/openai_plugin.py
@@ -223,15 +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):
-                    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("choices", [{}])[0].get("finish_reason", None):
+                    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.get("choices", [{}])[0].get("finish_reason", None)
 
             except KeyError:
                 logging.exception("KeyError, unexpected response format in line: %s", line)
@@ -248,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