forked from phoenix/litellm-mirror
fix(utils.py): support bedrock mistral streaming
This commit is contained in:
parent
cdb940d504
commit
cd53291b62
2 changed files with 48 additions and 0 deletions
|
@ -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.
|
||||||
|
|
|
@ -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:
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue