From d1cb4a195c825bb8886fa597bb889c819eded127 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 26 Jun 2024 16:19:05 -0700 Subject: [PATCH] fix(bedrock_httpx.py): Fix https://github.com/BerriAI/litellm/issues/4415 --- litellm/llms/bedrock.py | 5 ++ litellm/llms/bedrock_httpx.py | 30 +++++----- litellm/tests/test_bedrock_completion.py | 74 +++++++++++++++++++++--- 3 files changed, 88 insertions(+), 21 deletions(-) diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index d0d3bef6da..a8c47b3b91 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -1,3 +1,8 @@ +#################################### +######### DEPRECATED FILE ########## +#################################### +# logic moved to `bedrock_httpx.py` # + import copy import json import os diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 84ab10907c..14abec784f 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -261,20 +261,24 @@ class BedrockLLM(BaseLLM): # handle anthropic prompts and amazon titan prompts prompt = "" chat_history: Optional[list] = None + ## CUSTOM PROMPT + if model in custom_prompt_dict: + # check if the model has a registered custom prompt + model_prompt_details = custom_prompt_dict[model] + prompt = custom_prompt( + role_dict=model_prompt_details["roles"], + initial_prompt_value=model_prompt_details.get( + "initial_prompt_value", "" + ), + final_prompt_value=model_prompt_details.get("final_prompt_value", ""), + messages=messages, + ) + return prompt, None + ## ELSE if provider == "anthropic" or provider == "amazon": - if model in custom_prompt_dict: - # check if the model has a registered custom prompt - model_prompt_details = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, - ) - else: - prompt = prompt_factory( - model=model, messages=messages, custom_llm_provider="bedrock" - ) + prompt = prompt_factory( + model=model, messages=messages, custom_llm_provider="bedrock" + ) elif provider == "mistral": prompt = prompt_factory( model=model, messages=messages, custom_llm_provider="bedrock" diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index b953ca2a3a..24eefceeff 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -1,20 +1,31 @@ # @pytest.mark.skip(reason="AWS Suspended Account") -import sys, os +import os +import sys import traceback + from dotenv import load_dotenv load_dotenv() -import os, io +import io +import os sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path +from unittest.mock import AsyncMock, Mock, patch + import pytest + import litellm -from litellm import embedding, completion, completion_cost, Timeout, ModelResponse -from litellm import RateLimitError -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler -from unittest.mock import patch, AsyncMock, Mock +from litellm import ( + ModelResponse, + RateLimitError, + Timeout, + completion, + completion_cost, + embedding, +) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler # litellm.num_retries = 3 litellm.cache = None @@ -481,7 +492,10 @@ def test_completion_claude_3_base64(): def test_provisioned_throughput(): try: litellm.set_verbose = True - import botocore, json, io + import io + import json + + import botocore import botocore.session from botocore.stub import Stubber @@ -537,7 +551,6 @@ def test_completion_bedrock_mistral_completion_auth(): # aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] # aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"] # aws_region_name = os.environ["AWS_REGION_NAME"] - # os.environ.pop("AWS_ACCESS_KEY_ID", None) # os.environ.pop("AWS_SECRET_ACCESS_KEY", None) # os.environ.pop("AWS_REGION_NAME", None) @@ -624,3 +637,48 @@ async def test_bedrock_extra_headers(): assert "test" in mock_client_post.call_args.kwargs["headers"] assert mock_client_post.call_args.kwargs["headers"]["test"] == "hello world" mock_client_post.assert_called_once() + + +@pytest.mark.asyncio +async def test_bedrock_custom_prompt_template(): + """ + Check if custom prompt template used for bedrock models + + Reference: https://github.com/BerriAI/litellm/issues/4415 + """ + client = AsyncHTTPHandler() + + with patch.object(client, "post", new=AsyncMock()) as mock_client_post: + import json + + try: + response = await litellm.acompletion( + model="bedrock/mistral.OpenOrca", + messages=[{"role": "user", "content": "What's AWS?"}], + client=client, + roles={ + "system": { + "pre_message": "<|im_start|>system\n", + "post_message": "<|im_end|>", + }, + "assistant": { + "pre_message": "<|im_start|>assistant\n", + "post_message": "<|im_end|>", + }, + "user": { + "pre_message": "<|im_start|>user\n", + "post_message": "<|im_end|>", + }, + }, + bos_token="", + eos_token="<|im_end|>", + ) + except Exception as e: + pass + + print(f"mock_client_post.call_args: {mock_client_post.call_args}") + assert "prompt" in mock_client_post.call_args.kwargs["data"] + + prompt = json.loads(mock_client_post.call_args.kwargs["data"])["prompt"] + assert prompt == "<|im_start|>user\nWhat's AWS?<|im_end|>" + mock_client_post.assert_called_once()