From 78a2013e5136e87b443522c17be8111568802fc7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 14 Aug 2024 17:03:10 -0700 Subject: [PATCH] add test for large context in system message for anthropic --- litellm/llms/anthropic.py | 36 ++++--- .../tests/test_anthropic_prompt_caching.py | 99 +++++++++++++++++++ litellm/types/llms/anthropic.py | 8 +- litellm/types/llms/openai.py | 2 +- 4 files changed, 128 insertions(+), 17 deletions(-) diff --git a/litellm/llms/anthropic.py b/litellm/llms/anthropic.py index 19fca056bdf..cf58163461a 100644 --- a/litellm/llms/anthropic.py +++ b/litellm/llms/anthropic.py @@ -35,6 +35,7 @@ from litellm.types.llms.anthropic import ( AnthropicResponseContentBlockText, AnthropicResponseContentBlockToolUse, AnthropicResponseUsageBlock, + AnthropicSystemMessageContent, ContentBlockDelta, ContentBlockStart, ContentBlockStop, @@ -907,7 +908,7 @@ class AnthropicChatCompletion(BaseLLM): # Separate system prompt from rest of message system_prompt_indices = [] system_prompt = "" - system_prompt_dict = None + anthropic_system_message_list = None for idx, message in enumerate(messages): if message["role"] == "system": valid_content: bool = False @@ -915,19 +916,24 @@ class AnthropicChatCompletion(BaseLLM): system_prompt += message["content"] valid_content = True elif isinstance(message["content"], list): - for content in message["content"]: - system_prompt += content.get("text", "") - valid_content = True + for _content in message["content"]: + anthropic_system_message_content = ( + AnthropicSystemMessageContent( + type=_content.get("type"), + text=_content.get("text"), + ) + ) + if "cache_control" in _content: + anthropic_system_message_content["cache_control"] = ( + _content["cache_control"] + ) - # Handle Anthropic API context caching - if "cache_control" in message: - system_prompt_dict = [ - { - "cache_control": message["cache_control"], - "text": system_prompt, - "type": "text", - } - ] + if anthropic_system_message_list is None: + anthropic_system_message_list = [] + anthropic_system_message_list.append( + anthropic_system_message_content + ) + valid_content = True if valid_content: system_prompt_indices.append(idx) @@ -938,8 +944,8 @@ class AnthropicChatCompletion(BaseLLM): optional_params["system"] = system_prompt # Handling anthropic API Prompt Caching - if system_prompt_dict is not None: - optional_params["system"] = system_prompt_dict + if anthropic_system_message_list is not None: + optional_params["system"] = anthropic_system_message_list # Format rest of message according to anthropic guidelines try: messages = prompt_factory( diff --git a/litellm/tests/test_anthropic_prompt_caching.py b/litellm/tests/test_anthropic_prompt_caching.py index 8f57e96065f..87bfc23f841 100644 --- a/litellm/tests/test_anthropic_prompt_caching.py +++ b/litellm/tests/test_anthropic_prompt_caching.py @@ -220,3 +220,102 @@ async def test_anthropic_api_prompt_caching_basic(): assert (response.usage.cache_read_input_tokens > 0) or ( response.usage.cache_creation_input_tokens > 0 ) + + +@pytest.mark.asyncio +async def test_litellm_anthropic_prompt_caching_system(): + # https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching#prompt-caching-examples + # LArge Context Caching Example + mock_response = AsyncMock() + + def return_val(): + return { + "id": "msg_01XFDUDYJgAACzvnptvVoYEL", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Hello!"}], + "model": "claude-3-5-sonnet-20240620", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + } + + mock_response.json = return_val + + litellm.set_verbose = True + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=mock_response, + ) as mock_post: + # Act: Call the litellm.acompletion function + response = await litellm.acompletion( + api_key="mock_api_key", + model="anthropic/claude-3-5-sonnet-20240620", + messages=[ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are an AI assistant tasked with analyzing legal documents.", + }, + { + "type": "text", + "text": "Here is the full text of a complex legal agreement", + "cache_control": {"type": "ephemeral"}, + }, + ], + }, + { + "role": "user", + "content": "what are the key terms and conditions in this agreement?", + }, + ], + extra_headers={ + "anthropic-version": "2023-06-01", + "anthropic-beta": "prompt-caching-2024-07-31", + }, + ) + + # Print what was called on the mock + print("call args=", mock_post.call_args) + + expected_url = "https://api.anthropic.com/v1/messages" + expected_headers = { + "accept": "application/json", + "content-type": "application/json", + "anthropic-version": "2023-06-01", + "anthropic-beta": "prompt-caching-2024-07-31", + "x-api-key": "mock_api_key", + } + + expected_json = { + "system": [ + { + "type": "text", + "text": "You are an AI assistant tasked with analyzing legal documents.", + }, + { + "type": "text", + "text": "Here is the full text of a complex legal agreement", + "cache_control": {"type": "ephemeral"}, + }, + ], + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "what are the key terms and conditions in this agreement?", + } + ], + } + ], + "max_tokens": 4096, + "model": "claude-3-5-sonnet-20240620", + } + + mock_post.assert_called_once_with( + expected_url, json=expected_json, headers=expected_headers, timeout=600.0 + ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 2eb2aef5495..f14aa20c733 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -94,6 +94,12 @@ class AnthropicMetadata(TypedDict, total=False): user_id: str +class AnthropicSystemMessageContent(TypedDict, total=False): + type: str + text: str + cache_control: Optional[dict] + + class AnthropicMessagesRequest(TypedDict, total=False): model: Required[str] messages: Required[ @@ -108,7 +114,7 @@ class AnthropicMessagesRequest(TypedDict, total=False): metadata: AnthropicMetadata stop_sequences: List[str] stream: bool - system: str + system: Union[str, List] temperature: float tool_choice: AnthropicMessagesToolChoice tools: List[AnthropicMessagesTool] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 0d67d5d602c..5d2c416f9cd 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -361,7 +361,7 @@ class ChatCompletionToolMessage(TypedDict): class ChatCompletionSystemMessage(TypedDict, total=False): role: Required[Literal["system"]] - content: Required[str] + content: Required[Union[str, List]] name: str