fix(utils.py): support bedrock mistral streaming

This commit is contained in:
Krrish Dholakia 2024-03-29 07:56:10 -07:00
parent cdb940d504
commit cd53291b62
2 changed files with 48 additions and 0 deletions

View file

@ -974,6 +974,43 @@ def test_completion_bedrock_ai21_stream():
# test_completion_bedrock_ai21_stream() # test_completion_bedrock_ai21_stream()
def test_completion_bedrock_mistral_stream():
try:
litellm.set_verbose = False
response = completion(
model="bedrock/mistral.mixtral-8x7b-instruct-v0:1",
messages=[
{
"role": "user",
"content": "Be as verbose as possible and give as many details as possible, how does a court case get to the Supreme Court?",
}
],
temperature=1,
max_tokens=20,
stream=True,
)
print(response)
complete_response = ""
has_finish_reason = False
# Add any assertions here to check the response
for idx, chunk in enumerate(response):
# print
chunk, finished = streaming_format_tests(idx, chunk)
has_finish_reason = finished
complete_response += chunk
if finished:
break
if has_finish_reason is False:
raise Exception("finish reason not set for last chunk")
if complete_response.strip() == "":
raise Exception("Empty response received")
print(f"completion_response: {complete_response}")
except RateLimitError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_sagemaker_weird_response(): def test_sagemaker_weird_response():
""" """
When the stream ends, flush any remaining holding chunks. When the stream ends, flush any remaining holding chunks.

View file

@ -9093,6 +9093,17 @@ class CustomStreamWrapper:
if stop_reason != None: if stop_reason != None:
is_finished = True is_finished = True
finish_reason = stop_reason finish_reason = stop_reason
######## bedrock.mistral mappings ###############
elif "outputs" in chunk_data:
if (
len(chunk_data["outputs"]) == 1
and chunk_data["outputs"][0].get("text", None) is not None
):
text = chunk_data["outputs"][0]["text"]
stop_reason = chunk_data.get("stop_reason", None)
if stop_reason != None:
is_finished = True
finish_reason = stop_reason
######## bedrock.cohere mappings ############### ######## bedrock.cohere mappings ###############
# meta mapping # meta mapping
elif "generation" in chunk_data: elif "generation" in chunk_data: