diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 8ae02bd65ed..a4a1d3a1136 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -227,11 +227,18 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): if description is not None: function_params["description"] = cast(Union[dict, str], description) - return DatabricksTool( + databricks_tool = DatabricksTool( type="function", function=function_params, ) + if tool.get("cache_control"): + databricks_tool["cache_control"] = cast(dict, tool.get("cache_control")) + + return databricks_tool + + + def _map_openai_to_dbrx_tool(self, model: str, tools: List) -> List[DatabricksTool]: # if not claude, send as is if "claude" not in model: diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index c862bf0e4ac..dcf35be7e63 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -57,9 +57,10 @@ class DatabricksFunction(TypedDict, total=False): strict: bool -class DatabricksTool(TypedDict): +class DatabricksTool(TypedDict, total=False): function: DatabricksFunction type: Literal["function"] + cache_control: Optional[dict] class DatabricksMessage(TypedDict, total=False): 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 f9b5b5fe29c..466355bc02c 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 @@ -123,6 +123,12 @@ def test_convert_anthropic_tool_to_databricks_tool_without_description(): assert databricks_tool is not None assert databricks_tool["type"] == "function" assert databricks_tool["function"].get("description") is None + assert "cache_control" not in databricks_tool + + anthropic_tool["cache_control"] = {"type": "ephemeral"} + databricks_tool = config.convert_anthropic_tool_to_databricks_tool(anthropic_tool) + assert databricks_tool is not None + assert databricks_tool.get("cache_control") == {"type": "ephemeral"} def test_transform_choices_with_citations():