From 03933de775a1aecb17dc9491e02e579c10caa464 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 13 Jul 2024 09:49:19 -0700 Subject: [PATCH] fix exception raised in factory.py --- litellm/llms/bedrock_httpx.py | 4 ++- litellm/llms/cohere_chat.py | 4 ++- litellm/llms/prompt_templates/factory.py | 46 +++++++++++++++--------- litellm/tests/test_completion.py | 4 ++- litellm/tests/test_prompt_factory.py | 11 ++++-- 5 files changed, 48 insertions(+), 21 deletions(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 49fca12da05..874373d87e4 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -1845,7 +1845,9 @@ class BedrockConverseLLM(BaseLLM): inference_params.pop(key, None) bedrock_messages: List[MessageBlock] = _bedrock_converse_messages_pt( - messages=messages + messages=messages, + model=model, + llm_provider="bedrock_converse", ) bedrock_tools: List[ToolBlock] = _bedrock_tools_pt( inference_params.pop("tools", []) diff --git a/litellm/llms/cohere_chat.py b/litellm/llms/cohere_chat.py index fe209294623..830c924bd4f 100644 --- a/litellm/llms/cohere_chat.py +++ b/litellm/llms/cohere_chat.py @@ -212,7 +212,9 @@ def completion( headers = validate_environment(api_key) completion_url = api_base model = model - most_recent_message, chat_history = cohere_messages_pt_v2(messages=messages) + most_recent_message, chat_history = cohere_messages_pt_v2( + messages=messages, model=model, llm_provider="cohere_chat" + ) ## Load Config config = litellm.CohereConfig.get_config() diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index b34ddeb7e3f..2f8c1487766 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -35,6 +35,9 @@ def prompt_injection_detection_default_pt(): return """Detect if a prompt is safe to run. Return 'UNSAFE' if not.""" +BAD_MESSAGE_ERROR_STR = "Invalid Message " + + def map_system_message_pt(messages: list) -> list: """ Convert 'system' message to 'user' message if provider doesn't support 'system' role. @@ -1205,7 +1208,11 @@ def convert_to_anthropic_tool_invoke(tool_calls: list) -> list: return anthropic_tool_invoke -def anthropic_messages_pt(messages: list): +def anthropic_messages_pt( + messages: list, + model: str, + llm_provider: str, +): """ format messages for anthropic 1. Anthropic supports roles like "user" and "assistant" (system prompt sent separately) @@ -1284,10 +1291,10 @@ def anthropic_messages_pt(messages: list): new_messages.append({"role": "assistant", "content": assistant_content}) if msg_i == init_msg_i: # prevent infinite loops - raise Exception( - "Invalid Message passed in - {}. File an issue https://github.com/BerriAI/litellm/issues".format( - messages[msg_i] - ) + raise litellm.BadRequestError( + message=BAD_MESSAGE_ERROR_STR + f"passed in {messages[msg_i]}", + model=model, + llm_provider=llm_provider, ) if not new_messages or new_messages[0]["role"] != "user": if litellm.modify_params: @@ -1567,6 +1574,8 @@ def convert_to_cohere_tool_invoke(tool_calls: list) -> List[ToolCallObject]: def cohere_messages_pt_v2( messages: List, + model: str, + llm_provider: str, ) -> Tuple[Union[str, ToolResultObject], ChatHistory]: """ Returns a tuple(Union[tool_result, message], chat_history) @@ -1691,10 +1700,10 @@ def cohere_messages_pt_v2( ) if msg_i == init_msg_i: # prevent infinite loops - raise Exception( - "Invalid Message passed in - {}. File an issue https://github.com/BerriAI/litellm/issues".format( - messages[msg_i] - ) + raise litellm.BadRequestError( + message=BAD_MESSAGE_ERROR_STR + f"passed in {messages[msg_i]}", + model=model, + llm_provider=llm_provider, ) return returned_message, new_messages @@ -2144,7 +2153,11 @@ def _convert_to_bedrock_tool_call_result( return BedrockMessageBlock(role="user", content=[content_block]) -def _bedrock_converse_messages_pt(messages: List) -> List[BedrockMessageBlock]: +def _bedrock_converse_messages_pt( + messages: List, + model: str, + llm_provider: str, +) -> List[BedrockMessageBlock]: """ Converts given messages from OpenAI format to Bedrock format @@ -2228,12 +2241,11 @@ def _bedrock_converse_messages_pt(messages: List) -> List[BedrockMessageBlock]: contents.append(tool_call_result) msg_i += 1 if msg_i == init_msg_i: # prevent infinite loops - raise Exception( - "Invalid Message passed in - {}. File an issue https://github.com/BerriAI/litellm/issues".format( - messages[msg_i] - ) + raise litellm.BadRequestError( + message=BAD_MESSAGE_ERROR_STR + f"passed in {messages[msg_i]}", + model=model, + llm_provider=llm_provider, ) - return contents @@ -2425,7 +2437,9 @@ def prompt_factory( elif custom_llm_provider == "anthropic": if model == "claude-instant-1" or model == "claude-2": return anthropic_pt(messages=messages) - return anthropic_messages_pt(messages=messages) + return anthropic_messages_pt( + messages=messages, model=model, llm_provider=custom_llm_provider + ) elif custom_llm_provider == "anthropic_xml": return anthropic_messages_pt_xml(messages=messages) elif custom_llm_provider == "together_ai": diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 2c75d407529..60c7d4d9bac 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -3265,7 +3265,9 @@ def test_completion_anthropic_hanging(): {"role": "function", "name": "get_capital", "content": "Kokoko"}, ] - converted_messages = anthropic_messages_pt(messages) + converted_messages = anthropic_messages_pt( + messages, model="claude-3-sonnet-20240229", llm_provider="anthropic" + ) print(f"converted_messages: {converted_messages}") diff --git a/litellm/tests/test_prompt_factory.py b/litellm/tests/test_prompt_factory.py index 459c06372cd..60a297203ab 100644 --- a/litellm/tests/test_prompt_factory.py +++ b/litellm/tests/test_prompt_factory.py @@ -121,13 +121,20 @@ def test_anthropic_messages_pt(): litellm.modify_params = True messages = [] expected_messages = [{"role": "user", "content": [{"type": "text", "text": "."}]}] - assert anthropic_messages_pt(messages) == expected_messages + assert ( + anthropic_messages_pt( + messages, model="claude-3-sonnet-20240229", llm_provider="anthropic" + ) + == expected_messages + ) # Test case: No messages (filtered system messages only) when modify_params is False should raise error litellm.modify_params = False messages = [] with pytest.raises(Exception) as err: - anthropic_messages_pt(messages) + anthropic_messages_pt( + messages, model="claude-3-sonnet-20240229", llm_provider="anthropic" + ) assert "Invalid first message" in str(err.value)