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

fix(utils.py): fix deepinfra streaming

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-01 18:46:50 -08:00
parent 0072d796f6
commit 241f0aad5e
2 changed files with 14 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

1
litellm/tests/test_streaming.py
View file

@ -465,6 +465,7 @@ def test_completion_mistral_api_stream():
def test_completion_deep_infra_stream(): def test_completion_deep_infra_stream():
# deep infra currently includes role in the 2nd chunk # deep infra currently includes role in the 2nd chunk
# waiting for them to make a fix on this # waiting for them to make a fix on this
litellm.set_verbose = True
try: try:
messages = [ messages = [
{"role": "system", "content": "You are a helpful assistant."}, {"role": "system", "content": "You are a helpful assistant."},

13
litellm/utils.py
View file

@ -8048,6 +8048,7 @@ class CustomStreamWrapper:
if len(original_chunk.choices) > 0: if len(original_chunk.choices) > 0:
try: try:
delta = dict(original_chunk.choices[0].delta) delta = dict(original_chunk.choices[0].delta)
print_verbose(f"original delta: {delta}")
model_response.choices[0].delta = Delta(**delta) model_response.choices[0].delta = Delta(**delta)
except Exception as e: except Exception as e:
model_response.choices[0].delta = Delta() model_response.choices[0].delta = Delta()
@ -8056,9 +8057,21 @@ class CustomStreamWrapper:
model_response.system_fingerprint = ( model_response.system_fingerprint = (
original_chunk.system_fingerprint original_chunk.system_fingerprint
) )
print_verbose(f"self.sent_first_chunk: {self.sent_first_chunk}")
if self.sent_first_chunk == False: if self.sent_first_chunk == False:
model_response.choices[0].delta["role"] = "assistant" model_response.choices[0].delta["role"] = "assistant"
self.sent_first_chunk = True self.sent_first_chunk = True
elif self.sent_first_chunk == True and hasattr(
model_response.choices[0].delta, "role"
):
_initial_delta = model_response.choices[
0
].delta.model_dump()
_initial_delta.pop("role", None)
model_response.choices[0].delta = Delta(**_initial_delta)
print_verbose(
f"model_response.choices[0].delta: {model_response.choices[0].delta}"
)
else: else:
## else ## else
completion_obj["content"] = model_response_str completion_obj["content"] = model_response_str