diff --git a/litellm/__init__.py b/litellm/__init__.py index fe0cbe49e..8b8a50e46 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -179,6 +179,11 @@ aleph_alpha_models = [ baseten_models = ["qvv0xeq", "q841o8w", "31dxrj3"] # FALCON 7B # WizardLM # Mosaic ML +bedrock_models = [ + "amazon.titan-tg1-large", + "ai21.j2-grande-instruct" +] + model_list = ( open_ai_chat_completion_models + open_ai_text_completion_models @@ -212,6 +217,7 @@ provider_list = [ "bedrock", "vllm", "nlp_cloud", + "bedrock", "custom", # custom apis ] @@ -226,6 +232,7 @@ models_by_provider = { "openrouter": openrouter_models, "vertex_ai": vertex_chat_models + vertex_text_models, "ai21": ai21_models, + "bedrock": bedrock_models, } ####### EMBEDDING MODELS ################### diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index 0310dd325..bc5d0e721 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -15,6 +15,32 @@ class BedrockError(Exception): self.message ) # Call the base class constructor with the parameters it needs +def init_bedrock_client(boto3, region_name): + import subprocess + try: + client = boto3.client( + service_name="bedrock", + region_name=region_name, + endpoint_url=f'https://bedrock.{region_name}.amazonaws.com' + ) + except: + try: + command1 = "python3 -m pip install https://github.com/BerriAI/litellm/raw/main/cookbook/bedrock_resources/boto3-1.28.21-py3-none-any.whl" + subprocess.run(command1, shell=True, check=True) + # Command 2: Install boto3 from URL + command2 = "python3 -m pip install https://github.com/BerriAI/litellm/raw/main/cookbook/bedrock_resources/botocore-1.31.21-py3-none-any.whl" + subprocess.run(command2, shell=True, check=True) + + import boto3 + client = boto3.client( + service_name="bedrock", + region_name=region_name, + endpoint_url=f'https://bedrock.{region_name}.amazonaws.com' + ) + except Exception as e: + raise e + return client + """ BEDROCK AUTH Keys/Vars os.environ['AWS_ACCESS_KEY_ID'] = "" @@ -36,17 +62,17 @@ def completion( ): import sys if 'boto3' not in sys.modules: - import boto3 + try: + import boto3 + except: + raise Exception("Please Install boto3 to use bedrock with LiteLLM, run 'pip install boto3'") region_name = ( get_secret("AWS_REGION_NAME") or "us-west-2" # default to us-west-2 ) - client = boto3.client( - service_name="bedrock", - region_name=region_name - ) + client = init_bedrock_client(boto3, region_name) model = model prompt = "" diff --git a/litellm/main.py b/litellm/main.py index 26d54cf67..8e1913a2c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -748,6 +748,29 @@ def completion( ) return response + ## RESPONSE OBJECT + response = model_response + elif custom_llm_provider == "bedrock": + # boto3 reads keys from .env + model_response = bedrock.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=encoding, + logging_obj=logging + ) + + if "stream" in optional_params and optional_params["stream"] == True: ## [BETA] + # don't try to access stream object, + response = CustomStreamWrapper( + iter(model_response), model, custom_llm_provider="bedrock", logging_obj=logging + ) + return response + ## RESPONSE OBJECT response = model_response elif custom_llm_provider == "vllm":