diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index d3df5bbf361..0f3530f85da 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -169,12 +169,17 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): if tool is None: return None + kwags = { + "name":tool["name"], + "parameters":cast(dict, tool.get("input_schema") or {}) + } + + if tool.get("description"): + kwags["description"] = tool.get("description") + return DatabricksTool( type="function", - function=DatabricksFunction( - name=tool["name"], - parameters=cast(dict, tool.get("input_schema") or {}), - ), + function=DatabricksFunction(**kwags), ) def _map_openai_to_dbrx_tool(self, model: str, tools: List) -> List[DatabricksTool]: diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index 112427c6b56..c362d065694 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -51,7 +51,7 @@ AllDatabricksContentValues = Union[str, List[AllDatabricksContentListValues]] class DatabricksFunction(TypedDict, total=False): name: Required[str] - description: dict + description: dict | str parameters: dict strict: bool diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 51a2e971c09..a14683fac17 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -94,6 +94,33 @@ def test_transform_choices_without_signature(): assert thinking_block["type"] == "thinking" assert thinking_block["thinking"] == "i'm thinking without signature." +def test_convert_anthropic_tool_to_databricks_tool_with_description(): + config = DatabricksConfig() + anthropic_tool = { + "name": "test_tool", + "description": "test description", + "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}} + } + + databricks_tool = config.convert_anthropic_tool_to_databricks_tool(anthropic_tool) + + assert databricks_tool is not None + assert databricks_tool["type"] == "function" + assert databricks_tool["function"]["description"] == "test description" + + +def test_convert_anthropic_tool_to_databricks_tool_without_description(): + config = DatabricksConfig() + anthropic_tool = { + "name": "test_tool", + "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}} + } + + databricks_tool = config.convert_anthropic_tool_to_databricks_tool(anthropic_tool) + + assert databricks_tool is not None + assert databricks_tool["type"] == "function" + assert databricks_tool["function"].get("description") is None def test_transform_choices_with_citations(): config = DatabricksConfig()