From bd5a8d582b3ce4dac24ea313b4d3b33b79329b6d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Albert=20=C3=96rwall?= Date: Fri, 4 Apr 2025 05:19:21 +0200 Subject: [PATCH] Fix prompt caching for Anthropic tool calls (#9706) * Add prompt cache support to Anhtropic tool calls * Fix linting issue and add test --- .../prompt_templates/factory.py | 27 ++++++-- litellm/types/llms/anthropic.py | 1 + tests/llm_translation/test_prompt_factory.py | 66 +++++++++++++++++-- 3 files changed, 84 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0673307929..15c8cc275c 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1254,6 +1254,7 @@ def convert_function_to_anthropic_tool_invoke( id=str(uuid.uuid4()), name=_name, input=json.loads(_arguments) if _arguments else {}, + cache_control=None, ) ] return anthropic_tool_invoke @@ -1300,8 +1301,13 @@ def convert_to_anthropic_tool_invoke( ] } """ - anthropic_tool_invoke = [ - AnthropicMessagesToolUseParam( + anthropic_tool_invoke = [] + + for tool in tool_calls: + if not get_attribute_or_key(tool, "type") == "function": + continue + + _anthropic_tool_use_param = AnthropicMessagesToolUseParam( type="tool_use", id=get_attribute_or_key(tool, "id"), name=get_attribute_or_key(get_attribute_or_key(tool, "function"), "name"), @@ -1310,10 +1316,20 @@ def convert_to_anthropic_tool_invoke( get_attribute_or_key(tool, "function"), "arguments" ) ), + cache_control=None, ) - for tool in tool_calls - if get_attribute_or_key(tool, "type") == "function" - ] + + _content_element = add_cache_control_to_content( + anthropic_content_element=_anthropic_tool_use_param, + orignal_content_element=dict(tool), + ) + + if "cache_control" in _content_element: + _anthropic_tool_use_param["cache_control"] = _content_element[ + "cache_control" + ] + + anthropic_tool_invoke.append(_anthropic_tool_use_param) return anthropic_tool_invoke @@ -1324,6 +1340,7 @@ def add_cache_control_to_content( AnthropicMessagesImageParam, AnthropicMessagesTextParam, AnthropicMessagesDocumentParam, + AnthropicMessagesToolUseParam, ChatCompletionThinkingBlock, ], orignal_content_element: Union[dict, AllMessageValues], diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 367b2421fd..003c0bc62d 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -57,6 +57,7 @@ class AnthropicMessagesToolUseParam(TypedDict): id: str name: str input: dict + cache_control: Optional[Union[dict, ChatCompletionCachedContent]] AnthropicMessagesAssistantMessageValues = Union[ diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 314aceac7d..5ccc2e3aac 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -300,6 +300,60 @@ def test_anthropic_cache_controls_pt(): print("translated_messages: ", translated_messages) +def test_anthropic_cache_controls_tool_calls_pt(): + """ + Tests that cache_control is properly set in tool_calls when converting messages + for the Anthropic API. + """ + messages = [ + { + "role": "user", + "content": "Can you help me get the weather?", + }, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "weather-tool-id-123", + "function": { + "arguments": '{"location": "San Francisco"}', + "name": "get_weather", + }, + "type": "function", + } + ], + "cache_control": {"type": "ephemeral"}, + }, + { + "role": "function", + "content": '{"temperature": 72, "unit": "fahrenheit", "description": "sunny"}', + "name": "get_weather", + "tool_call_id": "weather-tool-id-123", + "cache_control": {"type": "ephemeral"}, + }, + ] + + translated_messages = anthropic_messages_pt( + messages, model="claude-3-sonnet-20240229", llm_provider="anthropic" + ) + + print("Translated tool call messages:", translated_messages) + + assert translated_messages[0]["role"] == "user" + + assert translated_messages[1]["role"] == "assistant" + for content_item in translated_messages[1]["content"]: + if content_item["type"] == "tool_use": + assert content_item["cache_control"] is None + assert content_item["name"] == "get_weather" + + assert translated_messages[2]["role"] == "user" + for content_item in translated_messages[2]["content"]: + if content_item["type"] == "tool_result": + assert content_item["cache_control"] == {"type": "ephemeral"} + + @pytest.mark.parametrize("provider", ["bedrock", "anthropic"]) def test_bedrock_parallel_tool_calling_pt(provider): """ @@ -701,7 +755,7 @@ def test_hf_chat_template(): "add_eos_token": False, "bos_token": { "__type": "AddedToken", - "content": "<｜begin▁of▁sentence｜>", + "content": "", "lstrip": False, "normalized": True, "rstrip": False, @@ -710,7 +764,7 @@ def test_hf_chat_template(): "clean_up_tokenization_spaces": False, "eos_token": { "__type": "AddedToken", - "content": "<｜end▁of▁sentence｜>", + "content": "", "lstrip": False, "normalized": True, "rstrip": False, @@ -720,7 +774,7 @@ def test_hf_chat_template(): "model_max_length": 16384, "pad_token": { "__type": "AddedToken", - "content": "<｜end▁of▁sentence｜>", + "content": "", "lstrip": False, "normalized": True, "rstrip": False, @@ -729,7 +783,7 @@ def test_hf_chat_template(): "sp_model_kwargs": {}, "unk_token": None, "tokenizer_class": "LlamaTokenizerFast", - "chat_template": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='') %}{%- for message in messages %}{%- if message['role'] == 'system' %}{% set ns.system_prompt = message['content'] %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{'<｜User｜>' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{%- set ns.is_first = true -%}{%- else %}{{'\\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{% set content = message['content'] %}{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}{{'<｜Assistant｜>' + content + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<｜Assistant｜>\\n'}}{% endif %}", + "chat_template": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='') %}{%- for message in messages %}{%- if message['role'] == 'system' %}{% set ns.system_prompt = message['content'] %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{' ' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{' ' + tool['type'] + ' ' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + ' '}}{%- set ns.is_first = true -%}{%- else %}{{' ' + tool['type'] + ' ' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + ' '}}{{' '}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{' ' + message['content'] + ' '}}{%- set ns.is_tool = false -%}{%- else %}{% set content = message['content'] %}{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}{{' ' + content + ' '}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{' ' + message['content'] + ' '}}{%- set ns.is_output_first = false %}{%- else %}{{' ' + message['content'] + ' '}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{' '}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{' '}}{% endif %}", }, ) @@ -741,7 +795,9 @@ def test_hf_chat_template(): print(chat_template) assert ( chat_template.rstrip() - == """<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>What is the weather in Copenhagen?<｜Assistant｜>""" + == """You are a helpful assistant. +What is the weather in Copenhagen? +""" )