diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 0704b523..fe8f544e 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -303,7 +303,10 @@ def _convert_response_to_chat_result(self, response: Mapping[str, Any]) -> ChatR generations = [ ChatGeneration( message=_convert_dict_to_message(choice["message"]), - generation_info=choice.get("usage", {}), + generation_info={ + **choice.get("usage", {}), + "logprobs": choice.get("logprobs"), + }, ) for choice in response["choices"] ]