add bedrock to main and init

This commit is contained in:
ishaan-jaff 2023-09-14 13:42:21 -07:00
parent 46984bb7c6
commit 38bdb9335c
3 changed files with 61 additions and 5 deletions

View file

@ -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 ###################

View file

@ -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:
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 = ""

View file

@ -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":