diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 0bfa87f2..2922a862 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -25,7 +25,7 @@ from databricks import sdk from langchain_core.callbacks import AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun from langchain_core.language_models import BaseChatModel -from langchain_core.language_models.base import LanguageModelInput +from langchain_core.language_models.base import LangSmithParams, LanguageModelInput from langchain_core.messages import ( AIMessage, AIMessageChunk, @@ -66,7 +66,7 @@ from openai.types.chat import ChatCompletion, ChatCompletionChunk from openai.types.completion_usage import CompletionUsage from openai.types.responses import Response, ResponseStreamEvent, ResponseUsage -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator from typing_extensions import override from databricks_langchain.utils import get_async_openai_client, get_openai_client @@ -272,10 +272,29 @@ class GetPopulation(BaseModel): model_config = ConfigDict(populate_by_name=True) + @classmethod + def is_lc_serializable(cls) -> bool: + """Return whether this model can be serialized by LangChain.""" + return True + + @classmethod + def get_lc_namespace(cls) -> List[str]: + """Get the namespace of the LangChain object.""" + return ["databricks_langchain", "chat_models"] + + @property + def lc_secrets(self) -> Dict[str, str]: + """Map the Databricks token field to its environment variable.""" + return {"databricks_token": "DATABRICKS_TOKEN"} + model: str = Field(alias="endpoint") """Name of Databricks Model Serving endpoint to query.""" target_uri: Optional[str] = None """The target MLflow deployment URI to use. Deprecated: use workspace_client instead.""" + databricks_host: Optional[str] = None + """Databricks workspace URL. If omitted, the SDK uses its default authentication.""" + databricks_token: Optional[SecretStr] = Field(default=None, repr=False) + """Databricks personal access token. Serialized as the ``DATABRICKS_TOKEN`` secret.""" workspace_client: Optional[sdk.WorkspaceClient] = Field(default=None, exclude=True) """Optional WorkspaceClient instance to use for authentication. If not provided, uses default authentication.""" temperature: Optional[float] = None @@ -366,6 +385,18 @@ def async_client(self) -> AsyncOpenAI: def __init__(self, **kwargs: Any): super().__init__(**kwargs) + if self.workspace_client is None and ( + self.databricks_host is not None or self.databricks_token is not None + ): + self.workspace_client = sdk.WorkspaceClient( + host=self.databricks_host, + token=( + self.databricks_token.get_secret_value() + if self.databricks_token is not None + else None + ), + ) + # Handle deprecated target_uri parameter if self.target_uri: warnings.warn( @@ -1392,14 +1423,40 @@ class AnswerWithJustification(BaseModel): def _identifying_params(self) -> Dict[str, Any]: return self._default_params + def _get_databricks_host(self) -> Optional[str]: + """Get the workspace host without exposing authentication state.""" + if self.databricks_host: + return self.databricks_host + if self.workspace_client is not None: + host = getattr(self.workspace_client.config, "host", None) + if isinstance(host, str): + return host + return None + def _get_invocation_params( self, stop: Optional[List[str]] = None, **kwargs: Any ) -> Dict[str, Any]: - """Get the parameters used to invoke the model FOR THE CALLBACKS.""" - return { - **self._default_params, - **super()._get_invocation_params(stop=stop, **kwargs), + """Get safe model and workspace configuration for callbacks and tracing.""" + params = super()._get_invocation_params(stop=stop, **kwargs) + optional_params = { + "databricks_host": self._get_databricks_host(), + "timeout": self.timeout, + "max_retries": self.max_retries, } + params.update({key: value for key, value in optional_params.items() if value is not None}) + if self.use_ai_gateway: + params["use_ai_gateway"] = True + if self.use_ai_gateway_native_api: + params["use_ai_gateway_native_api"] = True + if self.use_responses_api: + params["use_responses_api"] = True + return params + + def _get_ls_params(self, stop: Optional[List[str]] = None, **kwargs: Any) -> LangSmithParams: + """Get standard LangSmith trace metadata for ChatDatabricks.""" + params = super()._get_ls_params(stop=stop, **kwargs) + params["ls_provider"] = "databricks" + return params @property def _llm_type(self) -> str: diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index 5e85ff7e..dd3424f2 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -9,6 +9,7 @@ import mlflow # noqa: F401 import pytest from databricks import sdk +from langchain_core.load import load from langchain_core.messages import ( AIMessage, AIMessageChunk, @@ -71,6 +72,99 @@ def test_dict(llm: ChatDatabricks) -> None: assert "workspace_client" not in d +def test_langchain_serialization_redacts_explicit_credentials() -> None: + with patch("databricks_langchain.chat_models.sdk.WorkspaceClient") as workspace_client: + llm = ChatDatabricks( + model="test-model", + databricks_host="https://example.cloud.databricks.com", + databricks_token="test-token", + ) + + serialized = llm.to_json() + assert serialized["type"] == "constructor" + assert serialized["id"] == [ + "databricks_langchain", + "chat_models", + "ChatDatabricks", + ] + assert serialized["kwargs"]["databricks_host"] == ("https://example.cloud.databricks.com") + assert serialized["kwargs"]["databricks_token"] == { + "lc": 1, + "type": "secret", + "id": ["DATABRICKS_TOKEN"], + } + assert "workspace_client" not in serialized["kwargs"] + assert "test-token" not in json.dumps(serialized) + workspace_client.assert_called_once_with( + host="https://example.cloud.databricks.com", token="test-token" + ) + + +def test_langchain_serialization_round_trip() -> None: + with patch("databricks_langchain.chat_models.sdk.WorkspaceClient"): + llm = ChatDatabricks( + model="test-model", + databricks_host="https://example.cloud.databricks.com", + databricks_token="test-token", + ) + restored = load( + llm.to_json(), + secrets_map={"DATABRICKS_TOKEN": "restored-token"}, + valid_namespaces=["databricks_langchain"], + allowed_objects=[ChatDatabricks], + ) + + assert isinstance(restored, ChatDatabricks) + assert restored.model == "test-model" + assert restored.databricks_host == "https://example.cloud.databricks.com" + assert restored.databricks_token is not None + assert restored.databricks_token.get_secret_value() == "restored-token" + + +def test_invocation_params_include_safe_playground_configuration() -> None: + workspace_client = Mock(spec=sdk.WorkspaceClient) + workspace_client.config.host = "https://example.cloud.databricks.com" + llm = ChatDatabricks( + model="system.ai.claude-sonnet-5", + workspace_client=workspace_client, + temperature=0.2, + max_tokens=512, + timeout=30, + max_retries=4, + use_ai_gateway=True, + ) + + params = llm._get_invocation_params(stop=["DONE"]) + + assert params["model"] == "system.ai.claude-sonnet-5" + assert params["databricks_host"] == "https://example.cloud.databricks.com" + assert params["temperature"] == 0.2 + assert params["max_tokens"] == 512 + assert params["timeout"] == 30 + assert params["max_retries"] == 4 + assert params["use_ai_gateway"] is True + assert "use_responses_api" not in params + assert params["stop"] == ["DONE"] + assert "databricks_token" not in params + assert "workspace_client" not in params + + ls_params = llm._get_ls_params(stop=["DONE"]) + assert ls_params["ls_provider"] == "databricks" + assert ls_params["ls_model_name"] == "system.ai.claude-sonnet-5" + assert ls_params["ls_temperature"] == 0.2 + assert ls_params["ls_max_tokens"] == 512 + assert ls_params["ls_stop"] == ["DONE"] + + +@pytest.mark.parametrize("mode", ["use_ai_gateway_native_api", "use_responses_api"]) +def test_invocation_params_include_enabled_api_mode(mode: str) -> None: + llm = ChatDatabricks(model="test-model", **{mode: True}) + + params = llm._get_invocation_params() + + assert params[mode] is True + + def test_workspace_client_parameter() -> None: """Test the workspace_client parameter works correctly.""" from unittest.mock import Mock, patch