diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 33658d49063..69e3cc43322 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -94,6 +94,15 @@ def handle_messages_with_content_list_to_str_conversion( return messages +def strip_name_from_message(message: AllMessageValues, allowed_name_roles: List[str] = ["user"]) -> AllMessageValues: + """ + Removes 'name' from message + """ + msg_copy = message.copy() + if msg_copy.get("role") not in allowed_name_roles: + msg_copy.pop("name", None) # type: ignore + return msg_copy + def strip_name_from_messages( messages: List[AllMessageValues], allowed_name_roles: List[str] = ["user"] ) -> List[AllMessageValues]: diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index dd136c54264..9ee6cd25d20 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -26,7 +26,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( - strip_name_from_messages, + strip_name_from_message ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.anthropic import AllAnthropicToolsValues @@ -332,8 +332,11 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): _message = message.model_dump(exclude_none=True) else: _message = message + _message = strip_name_from_message(_message, allowed_name_roles=["user"]) + # Move message-level cache_control into a content block when content is a string. + if "cache_control" in _message and isinstance(_message.get("content"), str): + _message = self._move_cache_control_into_string_content_block(_message) new_messages.append(_message) - new_messages = strip_name_from_messages(new_messages) if is_async: return super()._transform_messages( @@ -344,6 +347,32 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): messages=new_messages, model=model, is_async=cast(Literal[False], False) ) + def _move_cache_control_into_string_content_block(self, message: AllMessageValues) -> AllMessageValues: + """ + Moves message-level cache_control into a content block when content is a string. + + Transforms: + {"role": "user", "content": "text", "cache_control": {...}} + Into: + {"role": "user", "content": [{"type": "text", "text": "text", "cache_control": {...}}]} + + This is required for Anthropic's prompt caching API when cache_control is specified + at the message level but content is a simple string (not already an array of content blocks). + """ + content = message.get("content") + # Create new message with cache_control moved into content block + transformed_message = cast(dict[str, Any], message.copy()) + cache_control = transformed_message.pop("cache_control") + transformed_message["content"] = [ + { + "type": "text", + "text": content, + "cache_control": cache_control, + } + ] + return cast(AllMessageValues, transformed_message) + + @staticmethod def extract_content_str( content: Optional[AllDatabricksContentValues], diff --git a/tests/llm_translation/test_databricks.py b/tests/llm_translation/test_databricks.py index f0667b25afd..40fc712f2b7 100644 --- a/tests/llm_translation/test_databricks.py +++ b/tests/llm_translation/test_databricks.py @@ -54,7 +54,7 @@ def mock_chat_response() -> Dict[str, Any]: } -def mock_chat_response_claude_prompt_caching() -> Dict[str, Any]: +def mock_chat_response_anthropic_prompt_caching() -> Dict[str, Any]: return { "id": "msg_01234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ", "object": "chat.completion", @@ -65,7 +65,7 @@ def mock_chat_response_claude_prompt_caching() -> Dict[str, Any]: "index": 0, "message": { "role": "assistant", - "content": "The text you've provided consists entirely of the phrase \"example text\" repeated many times without any variation or additional content. There is no specific information, narrative, argument, or structured content to explain. This appears to be placeholder or filler text that would typically be replaced with actual content in a final document.", + "content": "I notice that you've provided a repetitive text that simply repeats \"example text\" many times rather than actual content to summarize. \n\nTo provide you with a meaningful summary, I would need:\n- Actual substantive text with real information, arguments, or narrative\n- Content that has key points, themes, or conclusions to extract\n- Material with varying ideas or concepts to synthesize\n\nCould you please share the actual text you'd like me to summarize? I'm ready to help once you provide content with real information to work with.", "refusal": None, "function_call": None, "tool_calls": None, @@ -76,20 +76,26 @@ def mock_chat_response_claude_prompt_caching() -> Dict[str, Any]: "logprobs": None, } ], - "usage": { - "prompt_tokens": 1556, - "completion_tokens": 65, - "total_tokens": 1621, + "usage": { + "completion_tokens": 117, + "prompt_tokens": 1549, + "total_tokens": 1666, "completion_tokens_details": None, - "prompt_tokens_details": None, + "prompt_tokens_details": { + "audio_tokens": None, + "cached_tokens": 0, + "text_tokens": None, + "image_tokens": None, + "cache_creation_tokens": 1545 + }, "cache_read_input_tokens": 0, - "cache_creation_input_tokens": 1552, + "cache_creation_input_tokens": 1545 }, "service_tier": None, "system_fingerprint": None, } -def mock_chat_response_claude_prompt_caching_repeat() -> Dict[str, Any]: +def mock_chat_response_anthropic_prompt_caching_not_enough_tokens() -> Dict[str, Any]: return { "id": "msg_01234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ", "object": "chat.completion", @@ -100,7 +106,7 @@ def mock_chat_response_claude_prompt_caching_repeat() -> Dict[str, Any]: "index": 0, "message": { "role": "assistant", - "content": "The text you've provided consists entirely of the phrase \"example text\" repeated many times without any variation or additional content. There is no specific information, narrative, argument, or structured content to explain. This appears to be placeholder or filler text that would typically be replaced with actual content in a final document.", + "content": "I notice that you've provided a repetitive text that simply repeats \"example text\" many times rather than actual content to summarize. \n\nTo provide you with a meaningful summary, I would need:\n- Actual substantive text with real information, arguments, or narrative\n- Content that has key points, themes, or conclusions to extract\n- Material with varying ideas or concepts to synthesize\n\nCould you please share the actual text you'd like me to summarize? I'm ready to help once you provide content with real information to work with.", "refusal": None, "function_call": None, "tool_calls": None, @@ -111,26 +117,73 @@ def mock_chat_response_claude_prompt_caching_repeat() -> Dict[str, Any]: "logprobs": None, } ], - "usage": { - "prompt_tokens": 1556, - "completion_tokens": 65, - "total_tokens": 1621, + "usage": { + "completion_tokens": 117, + "prompt_tokens": 1549, + "total_tokens": 1666, "completion_tokens_details": None, - "prompt_tokens_details": None, - "cache_read_input_tokens": 1552, - "cache_creation_input_tokens": 0, + "prompt_tokens_details": { + "audio_tokens": None, + "cached_tokens": 0, + "text_tokens": None, + "image_tokens": None, + "cache_creation_tokens": 0 + }, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0 + }, + "service_tier": None, + "system_fingerprint": None, + } + +def mock_chat_response_anthropic_prompt_caching_repeat() -> Dict[str, Any]: + return { + "id": "msg_01234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ", + "object": "chat.completion", + "created": 1761118943, + "model": "claude-3-7-sonnet", # Mock model name for testing + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "I notice that you've provided a repetitive text that simply repeats \"example text\" many times rather than actual content to summarize. \n\nTo provide you with a meaningful summary, I would need:\n- Actual substantive text with real information, arguments, or narrative\n- Content that has key points, themes, or conclusions to extract\n- Material with varying ideas or concepts to synthesize\n\nCould you please share the actual text you'd like me to summarize? I'm ready to help once you provide content with real information to work with.", + "refusal": None, + "function_call": None, + "tool_calls": None, + "annotations": None, + "audio": None, + }, + "finish_reason": "stop", + "logprobs": None, + } + ], + "usage": { + "completion_tokens": 117, + "prompt_tokens": 1549, + "total_tokens": 1666, + "completion_tokens_details": None, + "prompt_tokens_details": { + "audio_tokens": None, + "cached_tokens": 0, + "text_tokens": None, + "image_tokens": None, + "cache_creation_tokens": 1545 + }, + "cache_read_input_tokens": 1545, + "cache_creation_input_tokens": 0 }, "service_tier": None, "system_fingerprint": None, } -def mock_chat_response_nonclaude_prompt_caching() -> Dict[str, Any]: +def mock_chat_response_nonanthropic_prompt_caching() -> Dict[str, Any]: return { "id": "msg_01234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ", "object": "chat.completion", "created": 1761119150, - "model": "gpt-oss-20b", + "model": "gpt-oss-20b", # Mock model nama for testing "choices": [ { "index": 0, @@ -881,7 +934,7 @@ async def test_databricks_embeddings(sync_mode, monkeypatch): ) -def test_completion_with_prompt_caching_claude_model(monkeypatch): +def test_completion_with_prompt_caching_anthropic_model(monkeypatch): base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints" api_key = "dapimykey" monkeypatch.setenv("DATABRICKS_API_BASE", base_url) @@ -890,7 +943,7 @@ def test_completion_with_prompt_caching_claude_model(monkeypatch): sync_handler = HTTPHandler() mock_response = Mock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.json.return_value = mock_chat_response_claude_prompt_caching() + mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching() mock_text = 'example text' * 512 messages = [ @@ -936,13 +989,13 @@ def test_completion_with_prompt_caching_claude_model(monkeypatch): # Check the response object returned from litellm.completion() assert 'claude-3-7-sonnet' in response['model'] assert response['usage']['cache_read_input_tokens'] == 0 - assert response['usage']['cache_creation_input_tokens'] == 1552 - assert response['usage']['prompt_tokens'] == 1556 - assert response['usage']['completion_tokens'] == 65 - assert response['usage']['total_tokens'] == 1621 + assert response['usage']['cache_creation_input_tokens'] == 1545 + assert response['usage']['prompt_tokens'] == 1549 + assert response['usage']['completion_tokens'] == 117 + assert response['usage']['total_tokens'] == 1666 -def test_completion_with_prompt_caching_claude_model_repeat(monkeypatch): +def test_completion_with_prompt_caching_anthropic_model_repeat(monkeypatch): base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints" api_key = "dapimykey" monkeypatch.setenv("DATABRICKS_API_BASE", base_url) @@ -951,7 +1004,7 @@ def test_completion_with_prompt_caching_claude_model_repeat(monkeypatch): sync_handler = HTTPHandler() mock_response = Mock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.json.return_value = mock_chat_response_claude_prompt_caching_repeat() + mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching_repeat() mock_text = 'example text' * 512 messages = [ @@ -998,14 +1051,14 @@ def test_completion_with_prompt_caching_claude_model_repeat(monkeypatch): # TODO: add test for entire expected output schema in the future # Check the response object returned from litellm.completion() assert 'claude-3-7-sonnet' in response['model'] - assert response['usage']['cache_read_input_tokens'] == 1552 + assert response['usage']['cache_read_input_tokens'] == 1545 assert response['usage']['cache_creation_input_tokens'] == 0 - assert response['usage']['prompt_tokens'] == 1556 - assert response['usage']['completion_tokens'] == 65 - assert response['usage']['total_tokens'] == 1621 + assert response['usage']['prompt_tokens'] == 1549 + assert response['usage']['completion_tokens'] == 117 + assert response['usage']['total_tokens'] == 1666 -def test_completion_with_prompt_caching_nonclaude_model(monkeypatch): +def test_completion_with_prompt_caching_nonanthropic_model(monkeypatch): base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints" api_key = "dapimykey" monkeypatch.setenv("DATABRICKS_API_BASE", base_url) @@ -1014,7 +1067,7 @@ def test_completion_with_prompt_caching_nonclaude_model(monkeypatch): sync_handler = HTTPHandler() mock_response = Mock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.json.return_value = mock_chat_response_nonclaude_prompt_caching() + mock_response.json.return_value = mock_chat_response_nonanthropic_prompt_caching() mock_text = 'example text' * 512 messages = [ @@ -1148,3 +1201,177 @@ def test_databricks_anthropic_function_call_with_no_schema(model, monkeypatch): assert len(response.choices[0].message.tool_calls) == 1 assert response.choices[0].message.tool_calls[0].function.name == "get_current_weather" + +def test_databricks_anthropic_user_string_content_cache_injection(monkeypatch): + base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints" + api_key = "dapimykey" + monkeypatch.setenv("DATABRICKS_API_BASE", base_url) + monkeypatch.setenv("DATABRICKS_API_KEY", api_key) + + sync_handler = HTTPHandler() + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching() + + mock_text = 'example text' * 512 + messages = [ + { + "role": "system", + "content": "You are an expert summarizer." + }, + { + "role": "user", + "content": mock_text + } + ] + cache_control_injection_points = [ + { + "location": "message", + "role": "user" + } + ] + + with patch.object(HTTPHandler, "post", return_value=mock_response) as mock_post: + response = litellm.completion( + model="databricks/databricks-claude-3-7-sonnet", + messages=messages, + client=sync_handler, + temperature=0.5, + cache_control_injection_points=cache_control_injection_points, + extraparam="testpassingextraparam", + ) + assert ( + mock_post.call_args.kwargs["headers"]["Content-Type"] == "application/json" + ) + assert ( + mock_post.call_args.kwargs["headers"]["Authorization"] + == f"Bearer {api_key}" + ) + assert mock_post.call_args.kwargs["url"] == f"{base_url}/chat/completions" + assert mock_post.call_args.kwargs["stream"] == False + + # TODO: add test for entire expected output schema in the future + # Check the response object returned from litellm.completion() + assert 'claude-3-7-sonnet' in response['model'] + assert response['usage']['cache_read_input_tokens'] == 0 + assert response['usage']['cache_creation_input_tokens'] == 1545 + assert response['usage']['prompt_tokens'] == 1549 + assert response['usage']['completion_tokens'] == 117 + assert response['usage']['total_tokens'] == 1666 + + +def test_databricks_anthropic_system_string_content_cache_injection(monkeypatch): + base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints" + api_key = "dapimykey" + monkeypatch.setenv("DATABRICKS_API_BASE", base_url) + monkeypatch.setenv("DATABRICKS_API_KEY", api_key) + + sync_handler = HTTPHandler() + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching() + + mock_text = 'example text' * 512 + messages = [ + { + "role": "system", + "content": mock_text + }, + { + "role": "user", + "content": "You are an expert summarizer." + } + ] + cache_control_injection_points = [ + { + "location": "message", + "role": "system" + } + ] + + with patch.object(HTTPHandler, "post", return_value=mock_response) as mock_post: + response = litellm.completion( + model="databricks/databricks-claude-3-7-sonnet", + messages=messages, + client=sync_handler, + temperature=0.5, + cache_control_injection_points=cache_control_injection_points, + extraparam="testpassingextraparam", + ) + assert ( + mock_post.call_args.kwargs["headers"]["Content-Type"] == "application/json" + ) + assert ( + mock_post.call_args.kwargs["headers"]["Authorization"] + == f"Bearer {api_key}" + ) + assert mock_post.call_args.kwargs["url"] == f"{base_url}/chat/completions" + assert mock_post.call_args.kwargs["stream"] == False + + # TODO: add test for entire expected output schema in the future + # Check the response object returned from litellm.completion() + assert 'claude-3-7-sonnet' in response['model'] + assert response['usage']['cache_read_input_tokens'] == 0 + assert response['usage']['cache_creation_input_tokens'] == 1545 + assert response['usage']['prompt_tokens'] == 1549 + assert response['usage']['completion_tokens'] == 117 + assert response['usage']['total_tokens'] == 1666 + + + +def test_databricks_anthropic_system_string_content_cache_injection_not_enough_tokens(monkeypatch): + base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints" + api_key = "dapimykey" + monkeypatch.setenv("DATABRICKS_API_BASE", base_url) + monkeypatch.setenv("DATABRICKS_API_KEY", api_key) + + sync_handler = HTTPHandler() + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching_not_enough_tokens() + + mock_text = 'example text' * 512 + messages = [ + { + "role": "system", + "content": "You are a helpful assistant that explains the content of the given text." + }, + { + "role": "user", + "content": mock_text + } + ] + cache_control_injection_points = [ + { + "location": "message", + "role": "system" + } + ] + + with patch.object(HTTPHandler, "post", return_value=mock_response) as mock_post: + response = litellm.completion( + model="databricks/databricks-claude-3-7-sonnet", + messages=messages, + client=sync_handler, + temperature=0.5, + cache_control_injection_points=cache_control_injection_points, + extraparam="testpassingextraparam", + ) + assert ( + mock_post.call_args.kwargs["headers"]["Content-Type"] == "application/json" + ) + assert ( + mock_post.call_args.kwargs["headers"]["Authorization"] + == f"Bearer {api_key}" + ) + assert mock_post.call_args.kwargs["url"] == f"{base_url}/chat/completions" + assert mock_post.call_args.kwargs["stream"] == False + + # TODO: add test for entire expected output schema in the future + # Check the response object returned from litellm.completion() + assert 'claude-3-7-sonnet' in response['model'] + assert response['usage']['cache_read_input_tokens'] == 0 + assert response['usage']['cache_creation_input_tokens'] == 0 + assert response['usage']['prompt_tokens'] == 1549 + assert response['usage']['completion_tokens'] == 117 + assert response['usage']['total_tokens'] == 1666 \ No newline at end of file