forked from phoenix/litellm-mirror
add bedrock to main and init
This commit is contained in:
parent
46984bb7c6
commit
38bdb9335c
3 changed files with 61 additions and 5 deletions
|
@ -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 ###################
|
||||
|
|
|
@ -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 = ""
|
||||
|
|
|
@ -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":
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue