Skip to content

Commit

Permalink
default value for tokens before timeout when no usage
Browse files Browse the repository at this point in the history
  • Loading branch information
drewrip committed Jun 10, 2024
1 parent 366c7ce commit 1c90aa7
Showing 1 changed file with 5 additions and 4 deletions.
9 changes: 5 additions & 4 deletions plugins/openai_plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand All @@ -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

0 comments on commit 1c90aa7

Please sign in to comment.