phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-28 04:04:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

bedrock.py fixes

Browse source
This commit is contained in:
ishaan-jaff 2023-09-14 14:15:39 -07:00
parent 7740fe7dda
commit ccb0b7fc78
1 changed files with 26 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

47
litellm/llms/bedrock.py
View file

@ -15,7 +15,9 @@ class BedrockError(Exception):
self.message self.message
) # Call the base class constructor with the parameters it needs ) # Call the base class constructor with the parameters it needs
def init_bedrock_client(boto3, region_name): def init_bedrock_client(region_name):
import sys
import boto3
import subprocess import subprocess
try: try:
client = boto3.client( client = boto3.client(
@ -23,7 +25,7 @@ def init_bedrock_client(boto3, region_name):
region_name=region_name, region_name=region_name,
endpoint_url=f'https://bedrock.{region_name}.amazonaws.com' endpoint_url=f'https://bedrock.{region_name}.amazonaws.com'
) )
except: except Exception as e:
try: try:
command1 = "python3 -m pip install https://github.com/BerriAI/litellm/raw/main/cookbook/bedrock_resources/boto3-1.28.21-py3-none-any.whl" 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) subprocess.run(command1, shell=True, check=True)
@ -60,21 +62,16 @@ def completion(
litellm_params=None, litellm_params=None,
logger_fn=None, logger_fn=None,
): ):
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 = ( region_name = (
get_secret("AWS_REGION_NAME") or get_secret("AWS_REGION_NAME") or
"us-west-2" # default to us-west-2 "us-west-2" # default to us-west-2
) )
client = init_bedrock_client(boto3, region_name) client = init_bedrock_client(region_name)
model = model model = model
provider = model.split(".")[0]
prompt = "" prompt = ""
for message in messages: for message in messages:
if "role" in message: if "role" in message:
@ -88,18 +85,23 @@ def completion(
) )
else: else:
prompt += f"{message['content']}" prompt += f"{message['content']}"
if provider == "ai21":
data = json.dumps({ data = json.dumps({
"inputText": prompt, "prompt": prompt,
"textGenerationConfig":{
"maxTokenCount":4096,
"stopSequences":[],
"temperature":0,
"topP":0.9
}
}) })
else: # amazon titan
data = json.dumps({
"inputText": prompt,
"textGenerationConfig":{
"maxTokenCount":4096,
"stopSequences":[],
"temperature":0,
"topP":0.9
}
})
## LOGGING ## LOGGING
logging_obj.pre_call( logging_obj.pre_call(
input=prompt, input=prompt,
@ -130,8 +132,11 @@ def completion(
) )
print_verbose(f"raw model_response: {response}") print_verbose(f"raw model_response: {response}")
## RESPONSE OBJECT ## RESPONSE OBJECT
outputText = response_body.get('results')[0].get('outputText') outputText = "default"
print(outputText) if provider == "ai21":
outputText = response_body.get('completions')[0].get('data').get('text')
else: # amazon titan
outputText = response_body.get('results')[0].get('outputText')
if "error" in outputText: if "error" in outputText:
raise BedrockError( raise BedrockError(
message=outputText["error"], message=outputText["error"],