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

(feat) add bedrock.cohere streaming

Browse source
This commit is contained in:
ishaan-jaff 2023-10-31 22:25:43 -07:00
parent 8a401076de
commit c6229b7113
2 changed files with 10 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/bedrock.py
View file

@ -326,6 +326,8 @@ def completion(
for k, v in config.items():
if k not in inference_params: # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in
inference_params[k] = v
if optional_params.get("stream", False) == True:
inference_params["stream"] = True # cohere requires stream = True in inference params
data = json.dumps({
"prompt": prompt,
**inference_params

8
litellm/utils.py
View file

@ -3763,6 +3763,14 @@ class CustomStreamWrapper:
if stop_reason != None:
is_finished = True
finish_reason = stop_reason
######## bedrock.cohere mappings ###############
# cohere mapping
elif "text" in chunk_data:
text = chunk_data["text"] # bedrock.cohere
# cohere mapping for finish reason
elif "finish_reason" in chunk_data:
finish_reason = chunk_data["finish_reason"]
is_finished = True
elif chunk_data.get("completionReason", None):
is_finished = True
finish_reason = chunk_data["completionReason"]