From e504141c0741e7550418386fb93c41d5e794b9f6 Mon Sep 17 00:00:00 2001 From: Debanjum Date: Mon, 18 Aug 2025 18:53:30 -0700 Subject: [PATCH] Fix to calculate usage from openai api streaming completion During streaming chunk.chunk contains usage data. This regression must have appeared while tuning openai stream processors --- src/khoj/processor/conversation/openai/utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/khoj/processor/conversation/openai/utils.py b/src/khoj/processor/conversation/openai/utils.py index 3b706df5..51d71cc0 100644 --- a/src/khoj/processor/conversation/openai/utils.py +++ b/src/khoj/processor/conversation/openai/utils.py @@ -219,6 +219,10 @@ def completion_with_backoff( # Json dump tool calls into aggregated response aggregated_response = json.dumps([tool_call.__dict__ for tool_call in tool_calls]) + # Align chunk definition with non-streaming mode for post stream completion usage + if hasattr(chunk, "chunk"): + chunk = chunk.chunk + # Calculate cost of chat input_tokens = chunk.usage.prompt_tokens if hasattr(chunk, "usage") and chunk.usage else 0 output_tokens = chunk.usage.completion_tokens if hasattr(chunk, "usage") and chunk.usage else 0