From 3cfff030c6e9aabd2028361bedf7e3c081dafc74 Mon Sep 17 00:00:00 2001 From: Jose Moreno Ortega Date: Thu, 13 Mar 2025 15:08:02 +0100 Subject: [PATCH 1/6] Enhance ChatResult output with model details and token usage --- .../langchain/src/databricks_langchain/chat_models.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 0704b523..3dd93fec 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -308,7 +308,12 @@ def _convert_response_to_chat_result(self, response: Mapping[str, Any]) -> ChatR for choice in response["choices"] ] usage = response.get("usage", {}) - return ChatResult(generations=generations, llm_output=usage) + llm_output = { + "token_usage": response.get("usage", {}), + "model_name": response.get("model", self.model), + "system_fingerprint": response.get("system_fingerprint", ""), + } + return ChatResult(generations=generations, llm_output=llm_output) def _stream( self, From 2345c5a8270e03dc771079f21b84ac0659e9beca Mon Sep 17 00:00:00 2001 From: Jose Moreno Ortega <30633897+jmoreno11@users.noreply.github.com> Date: Thu, 13 Mar 2025 15:23:46 +0100 Subject: [PATCH 2/6] Update chat_models.py --- integrations/langchain/src/databricks_langchain/chat_models.py | 1 - 1 file changed, 1 deletion(-) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 3dd93fec..b4559c17 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -307,7 +307,6 @@ def _convert_response_to_chat_result(self, response: Mapping[str, Any]) -> ChatR ) for choice in response["choices"] ] - usage = response.get("usage", {}) llm_output = { "token_usage": response.get("usage", {}), "model_name": response.get("model", self.model), From bb949fd9b1589b6090df55e1890dd449ce5537f2 Mon Sep 17 00:00:00 2001 From: Jose Moreno Ortega Date: Wed, 19 Mar 2025 19:18:04 +0100 Subject: [PATCH 3/6] Dump metadata in llm_output and create new unit test --- .../src/databricks_langchain/chat_models.py | 7 ++++--- .../tests/unit_tests/test_chat_models.py | 20 +++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index b4559c17..fbc29fdc 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -308,10 +308,11 @@ def _convert_response_to_chat_result(self, response: Mapping[str, Any]) -> ChatR for choice in response["choices"] ] llm_output = { - "token_usage": response.get("usage", {}), - "model_name": response.get("model", self.model), - "system_fingerprint": response.get("system_fingerprint", ""), + k: v for k, v in response.items() if k not in ("choices", "content", "role", "type") } + if "model" in llm_output and "model_name" not in llm_output: + llm_output["model_name"] = llm_output["model"] + return ChatResult(generations=generations, llm_output=llm_output) def _stream( diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index df281d51..7f69d378 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -345,3 +345,23 @@ def test_convert_tool_message_chunk() -> None: def test_convert_message_to_dict_function() -> None: with pytest.raises(ValueError, match="Function messages are not supported"): _convert_message_to_dict(FunctionMessage(content="", name="name")) + + +def test_convert_response_to_chat_result_llm_output(llm: ChatDatabricks) -> None: + """Test that _convert_response_to_chat_result correctly sets llm_output.""" + + result = llm._convert_response_to_chat_result(_MOCK_CHAT_RESPONSE) + + # Verify that llm_output contains the full response metadata + assert "model_name" in result.llm_output + assert "usage" in result.llm_output + assert result.llm_output["model_name"] == _MOCK_CHAT_RESPONSE["model"] + + # Verify that usage information is included directly in llm_output + assert result.llm_output["usage"] == _MOCK_CHAT_RESPONSE["usage"] + + # Verify that choices, content, role, and type are excluded from llm_output + assert "choices" not in result.llm_output + assert "content" not in result.llm_output + assert "role" not in result.llm_output + assert "type" not in result.llm_output From dd38d3b7567bbd5540bed6942bdfdf1ee73061c4 Mon Sep 17 00:00:00 2001 From: Jose Moreno Ortega Date: Wed, 19 Mar 2025 19:56:51 +0100 Subject: [PATCH 4/6] Fix linting error --- .../src/databricks_langchain/chat_models.py | 3 +-- .../langchain/tests/unit_tests/test_chat_models.py | 12 ++++++------ 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index fbc29fdc..9328e604 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -18,6 +18,7 @@ Union, ) +from databricks_langchain.utils import get_deployment_client from langchain_core.callbacks import CallbackManagerForLLMRun from langchain_core.language_models import BaseChatModel from langchain_core.language_models.base import LanguageModelInput @@ -54,8 +55,6 @@ from mlflow.deployments import BaseDeploymentClient # type: ignore from pydantic import BaseModel, ConfigDict, Field -from databricks_langchain.utils import get_deployment_client - logger = logging.getLogger(__name__) diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index 7f69d378..2b4fe444 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -4,6 +4,12 @@ import mlflow # type: ignore # noqa: F401 import pytest +from databricks_langchain.chat_models import ( + ChatDatabricks, + _convert_dict_to_message, + _convert_dict_to_message_chunk, + _convert_message_to_dict, +) from langchain_core.messages import ( AIMessage, AIMessageChunk, @@ -21,12 +27,6 @@ from langchain_core.runnables import RunnableMap from pydantic import BaseModel, Field -from databricks_langchain.chat_models import ( - ChatDatabricks, - _convert_dict_to_message, - _convert_dict_to_message_chunk, - _convert_message_to_dict, -) from tests.utils.chat_models import ( # noqa: F401 _MOCK_CHAT_RESPONSE, _MOCK_STREAM_RESPONSE, From cac1d70f68e93e75e059775d6927cbcb50dd2c38 Mon Sep 17 00:00:00 2001 From: Jose Moreno Ortega Date: Wed, 19 Mar 2025 21:58:05 +0100 Subject: [PATCH 5/6] revert changes --- .../src/databricks_langchain/chat_models.py | 3 ++- .../langchain/tests/unit_tests/test_chat_models.py | 12 ++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 9328e604..fbc29fdc 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -18,7 +18,6 @@ Union, ) -from databricks_langchain.utils import get_deployment_client from langchain_core.callbacks import CallbackManagerForLLMRun from langchain_core.language_models import BaseChatModel from langchain_core.language_models.base import LanguageModelInput @@ -55,6 +54,8 @@ from mlflow.deployments import BaseDeploymentClient # type: ignore from pydantic import BaseModel, ConfigDict, Field +from databricks_langchain.utils import get_deployment_client + logger = logging.getLogger(__name__) diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index 2b4fe444..7f69d378 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -4,12 +4,6 @@ import mlflow # type: ignore # noqa: F401 import pytest -from databricks_langchain.chat_models import ( - ChatDatabricks, - _convert_dict_to_message, - _convert_dict_to_message_chunk, - _convert_message_to_dict, -) from langchain_core.messages import ( AIMessage, AIMessageChunk, @@ -27,6 +21,12 @@ from langchain_core.runnables import RunnableMap from pydantic import BaseModel, Field +from databricks_langchain.chat_models import ( + ChatDatabricks, + _convert_dict_to_message, + _convert_dict_to_message_chunk, + _convert_message_to_dict, +) from tests.utils.chat_models import ( # noqa: F401 _MOCK_CHAT_RESPONSE, _MOCK_STREAM_RESPONSE, From 60a1dc296c0225af9faaa464b5f4a176905f8679 Mon Sep 17 00:00:00 2001 From: Jose Moreno Ortega <30633897+jmoreno11@users.noreply.github.com> Date: Wed, 19 Mar 2025 22:13:14 +0100 Subject: [PATCH 6/6] Fix linting errors (#2) --- .github/workflows/main.yml | 1 + .../langchain/src/databricks_langchain/chat_models.py | 2 +- .../langchain/tests/unit_tests/test_chat_models.py | 8 ++++---- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index d40a4ad4..90c88c4d 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -10,6 +10,7 @@ on: - synchronize - reopened - ready_for_review + workflow_dispatch: jobs: lint: diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index fbc29fdc..c714686d 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -312,7 +312,7 @@ def _convert_response_to_chat_result(self, response: Mapping[str, Any]) -> ChatR } if "model" in llm_output and "model_name" not in llm_output: llm_output["model_name"] = llm_output["model"] - + return ChatResult(generations=generations, llm_output=llm_output) def _stream( diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index 7f69d378..451a845a 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -349,17 +349,17 @@ def test_convert_message_to_dict_function() -> None: def test_convert_response_to_chat_result_llm_output(llm: ChatDatabricks) -> None: """Test that _convert_response_to_chat_result correctly sets llm_output.""" - + result = llm._convert_response_to_chat_result(_MOCK_CHAT_RESPONSE) - + # Verify that llm_output contains the full response metadata assert "model_name" in result.llm_output assert "usage" in result.llm_output assert result.llm_output["model_name"] == _MOCK_CHAT_RESPONSE["model"] - + # Verify that usage information is included directly in llm_output assert result.llm_output["usage"] == _MOCK_CHAT_RESPONSE["usage"] - + # Verify that choices, content, role, and type are excluded from llm_output assert "choices" not in result.llm_output assert "content" not in result.llm_output