From fa7e06319889cf4b371afa5db7d336ec95f95418 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 31 Oct 2023 22:25:43 -0700 Subject: [PATCH] (feat) add bedrock.cohere streaming --- litellm/llms/bedrock.py | 2 ++ litellm/utils.py | 8 ++++++++ 2 files changed, 10 insertions(+) diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index 0d2a75faf0f..27db0bbb7b5 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index 1f2145e52ef..19c6705a3d7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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"]