fix(proxy_server.py): drop none values in streaming response

This commit is contained in:
Krrish Dholakia 2024-02-27 14:37:29 -08:00
parent 05e096ce25
commit 57998c28dc
2 changed files with 15 additions and 11 deletions

View file

@ -2115,10 +2115,9 @@ async def async_data_generator(response, user_api_key_dict):
try: try:
start_time = time.time() start_time = time.time()
async for chunk in response: async for chunk in response:
verbose_proxy_logger.debug(f"returned chunk: {chunk}") chunk = chunk.model_dump_json(exclude_none=True)
assert isinstance(chunk, litellm.ModelResponse)
try: try:
yield f"data: {json.dumps(chunk.model_dump(exclude_none=True))}\n\n" yield f"data: {chunk}\n\n"
except Exception as e: except Exception as e:
yield f"data: {str(e)}\n\n" yield f"data: {str(e)}\n\n"

View file

@ -205,18 +205,18 @@ def map_finish_reason(
class FunctionCall(OpenAIObject): class FunctionCall(OpenAIObject):
arguments: str arguments: str
name: str name: Optional[str] = None
class Function(OpenAIObject): class Function(OpenAIObject):
arguments: str arguments: str
name: str name: Optional[str] = None
class ChatCompletionDeltaToolCall(OpenAIObject): class ChatCompletionDeltaToolCall(OpenAIObject):
id: str id: Optional[str] = None
function: Function function: Function
type: str type: Optional[str] = None
index: int index: int
@ -275,8 +275,11 @@ class Delta(OpenAIObject):
super(Delta, self).__init__(**params) super(Delta, self).__init__(**params)
self.content = content self.content = content
self.role = role self.role = role
self.function_call = function_call if function_call is not None and isinstance(function_call, dict):
if tool_calls is not None and isinstance(tool_calls, dict): self.function_call = FunctionCall(**function_call)
else:
self.function_call = function_call
if tool_calls is not None and isinstance(tool_calls, list):
self.tool_calls = [] self.tool_calls = []
for tool_call in tool_calls: for tool_call in tool_calls:
if tool_call.get("index", None) is None: if tool_call.get("index", None) is None:
@ -8727,7 +8730,7 @@ class CustomStreamWrapper:
or original_chunk.choices[0].delta.tool_calls is not None or original_chunk.choices[0].delta.tool_calls is not None
): ):
try: try:
delta = dict(original_chunk.choices[0].delta) delta = original_chunk.choices[0].delta
model_response.system_fingerprint = ( model_response.system_fingerprint = (
original_chunk.system_fingerprint original_chunk.system_fingerprint
) )
@ -8762,7 +8765,9 @@ class CustomStreamWrapper:
is None is None
): ):
t.function.arguments = "" t.function.arguments = ""
model_response.choices[0].delta = Delta(**delta) _json_delta = delta.model_dump()
print_verbose(f"_json_delta: {_json_delta}")
model_response.choices[0].delta = Delta(**_json_delta)
except Exception as e: except Exception as e:
traceback.print_exc() traceback.print_exc()
model_response.choices[0].delta = Delta() model_response.choices[0].delta = Delta()