Skip to content

Commit 74816b1

Browse files
authored
Fix usage is None case (#114)
During testing fib with SGLang found that `data["usage"]` my be a `None`. Fix it.
1 parent a328f78 commit 74816b1

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

src/flexible_inference_benchmark/engine/backend_functions.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -531,14 +531,14 @@ async def async_request_openai_chat_completions(
531531
generated_text += reasoning_content
532532
most_recent_timestamp = timestamp
533533

534-
if "usage" in data:
535-
if data["usage"]["completion_tokens"]:
534+
if "usage" in data and data["usage"] is not None:
535+
if data["usage"].get("completion_tokens"):
536536
output.output_len = int(data["usage"]["completion_tokens"])
537537
if process_span:
538538
process_span.set_attribute(
539539
"fib.completion_tokens", output.output_len
540540
)
541-
if data["usage"]["prompt_tokens"]:
541+
if data["usage"].get("prompt_tokens"):
542542
output.prompt_len = int(data["usage"]["prompt_tokens"])
543543
if process_span:
544544
process_span.set_attribute("fib.prompt_tokens", output.prompt_len)

0 commit comments

Comments
 (0)