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

(feat) bedrock add finish_reason to streaming responses

Browse source
This commit is contained in:
ishaan-jaff 2023-10-12 16:22:34 -07:00
parent 0937d07d6d
commit e01d83cea6
2 changed files with 11 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/tests/test_streaming.py
View file

@ -656,14 +656,14 @@ def test_completion_replicate_stream_bad_key():
# test_completion_replicate_stream_bad_key() # test_completion_replicate_stream_bad_key()
def test_completion_bedrock_ai21_stream(): def test_completion_bedrock_claude_stream():
try: try:
litellm.set_verbose=True litellm.set_verbose=False
response = completion( response = completion(
model="bedrock/anthropic.claude-instant-v1", model="bedrock/anthropic.claude-instant-v1",
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?"}], 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, temperature=1,
max_tokens=4096, max_tokens=20,
stream=True, stream=True,
) )
print(response) print(response)
@ -687,7 +687,7 @@ def test_completion_bedrock_ai21_stream():
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
test_completion_bedrock_ai21_stream() # test_completion_bedrock_claude_stream()
def test_completion_sagemaker_stream(): def test_completion_sagemaker_stream():

9
litellm/utils.py
View file

@ -3236,9 +3236,14 @@ class CustomStreamWrapper:
finish_reason = "" finish_reason = ""
if "outputText" in chunk_data: if "outputText" in chunk_data:
text = chunk_data['outputText'] text = chunk_data['outputText']
if "completion" in chunk_data: # anthropic mapping
elif "completion" in chunk_data:
text = chunk_data['completion'] # bedrock.anthropic text = chunk_data['completion'] # bedrock.anthropic
if chunk_data.get("completionReason", None): stop_reason = chunk_data.get("stop_reason", None)
if stop_reason != None:
is_finished = True
finish_reason = stop_reason
elif chunk_data.get("completionReason", None):
is_finished = True is_finished = True
finish_reason = chunk_data["completionReason"] finish_reason = chunk_data["completionReason"]
elif chunk.get("error", None): elif chunk.get("error", None):