feat: Add support for serialization and tracing metadata params

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 @@
 
     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 __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(
@@ -547,7 +578,7 @@
                 content_blocks.append(item_dict)
 
                 try:
                     args = json.loads(item.arguments, strict=False)  # ty:ignore[unresolved-attribute, invalid-argument-type]: astral-sh/ty#1479 should fix this
                     error = None
                 except json.JSONDecodeError as e:
                     error = str(e)
@@ -575,8 +606,8 @@
                 content_blocks.append(
                     {
                         "role": "tool",
                         "content": item.output,  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
                         "tool_call_id": item.call_id,  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
                     }
                 )
             elif item.type in (
@@ -1392,14 +1423,40 @@
     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:
@@ -1733,15 +1790,15 @@
         item = chunk.item  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
         if item.type == "function_call_output":
             lc_chunk = ToolMessageChunk(
                 content=item.output,  # ty:ignore[unresolved-attribute, invalid-argument-type]: astral-sh/ty#1479 should fix this
                 tool_call_id=item.call_id,  # ty:ignore[unresolved-attribute, invalid-argument-type]: astral-sh/ty#1479 should fix this
             )
         elif item.type == "function_call":
             id = item.call_id  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
             tool_call_chunks.append(
                 tool_call_chunk(
                     name=item.name,  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
                     args=item.arguments,  # ty:ignore[unresolved-attribute, invalid-argument-type]: astral-sh/ty#1479 should fix this
                     id=item.call_id,  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
                 )
             )

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 @@
     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
@@ -738,8 +832,8 @@
         id=ID,
         tool_calls=[
             {
                 "name": tool_calls[0]["function"]["name"],  # type: ignore[index]
                 "args": json.loads(tool_calls[0]["function"]["arguments"]),  # type: ignore[index]
                 "id": ID,
                 "type": "tool_call",
             }