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

fix(bedrock.py): support ai21 / bedrock streaming

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-29 16:34:48 -08:00
parent 3b89cff65e
commit ab76daa90b
4 changed files with 97 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

64
litellm/llms/bedrock.py
View file

@ -393,29 +393,55 @@ def completion(
accept = 'application/json'
contentType = 'application/json'
if stream == True:
## LOGGING
request_str = f"""
response = client.invoke_model_with_response_stream(
body={data},
modelId={model},
accept=accept,
contentType=contentType
)
"""
logging_obj.pre_call(
if provider == "ai21":
## LOGGING
request_str = f"""
response = client.invoke_model(
body={data},
modelId={model},
accept=accept,
contentType=contentType
)
"""
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={"complete_input_dict": data, "request_str": request_str},
)
response = client.invoke_model_with_response_stream(
body=data,
modelId=model,
accept=accept,
contentType=contentType
)
response = response.get('body')
return response
)
response = client.invoke_model(
body=data,
modelId=model,
accept=accept,
contentType=contentType
)
response = response.get('body').read()
return response
else:
## LOGGING
request_str = f"""
response = client.invoke_model_with_response_stream(
body={data},
modelId={model},
accept=accept,
contentType=contentType
)
"""
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={"complete_input_dict": data, "request_str": request_str},
)
response = client.invoke_model_with_response_stream(
body=data,
modelId=model,
accept=accept,
contentType=contentType
)
response = response.get('body')
return response
try:
## LOGGING
request_str = f"""