From 0c35252d769eee594185635426650a1ffd67fbc1 Mon Sep 17 00:00:00 2001 From: Jake Steelman Date: Tue, 7 Oct 2025 12:38:41 -0500 Subject: [PATCH] add a check for non-delta non-choices usage chunk --- .../src/databricks_langchain/chat_models.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 79a54148..f66ffa21 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -688,6 +688,18 @@ def _stream( logprobs=generation_info.get("logprobs"), ) yield generation_chunk + elif chunk.usage and stream_usage: + # Some models send a final chunk that does not have + # a delta or choices, but does have usage info + if not usage_chunk_emitted: + input_tokens = getattr(chunk.usage, "prompt_tokens", None) + output_tokens = getattr(chunk.usage, "completion_tokens", None) + if input_tokens is not None and output_tokens is not None: + final_usage = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } # Emit special usage chunk at end of stream if stream_usage and final_usage and not usage_chunk_emitted: