From 0035799ba9e4c8b050c60282cdd86b796d809d1f Mon Sep 17 00:00:00 2001 From: Samuel Monson Date: Tue, 24 Sep 2024 13:32:39 -0400 Subject: [PATCH] fix: flatten usage conditional --- plugins/openai_plugin.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/plugins/openai_plugin.py b/plugins/openai_plugin.py index 3796092b..422463f0 100644 --- a/plugins/openai_plugin.py +++ b/plugins/openai_plugin.py @@ -243,17 +243,16 @@ def streaming_request_http(self, query: dict, user_id: int, test_end_time: float # Check for usage statistics message = self._process_resp(resps[-1]['data']) - if message: - # If stream_options.include_usage == True then the final - # message contains only token stats - if not message.get("choices") and message.get('usage'): - result.output_tokens = deepget(message, "usage", "completion_tokens") - result.input_tokens = deepget(message, "usage", "prompt_tokens") - # We don't want to record this message - resps.pop() - else: - # TODO This signals that the request is faulty - logger.warn("Usage token missing") + # If stream_options.include_usage == True then the final + # message contains only token stats + if message and not message.get("choices") and message.get('usage'): + result.output_tokens = deepget(message, "usage", "completion_tokens") + result.input_tokens = deepget(message, "usage", "prompt_tokens") + # We don't want to record this message + resps.pop() + else: + # TODO This signals that the request is faulty + logger.warn("Usage token missing") # Iterate through all responses tokens = []