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: