Mirrowel commited on
Commit
1199a4e
·
1 Parent(s): 9ee0808

fix(proxy): handle None values in streaming tool-call aggregation

Browse files

- Prevent KeyError by ensuring the 'function' key exists before populating it.
- Guard against appending None when streaming tool-call name or arguments.
- Defensively handle None function_call values.

Files changed (1) hide show
  1. src/proxy_app/main.py +12 -5
src/proxy_app/main.py CHANGED
@@ -256,22 +256,29 @@ async def streaming_response_wrapper(
256
  for tc_chunk in value:
257
  index = tc_chunk["index"]
258
  if index not in aggregated_tool_calls:
259
- aggregated_tool_calls[index] = {"id": None, "type": "function", "function": {"name": "", "arguments": ""}}
 
 
 
260
  if tc_chunk.get("id"):
261
  aggregated_tool_calls[index]["id"] = tc_chunk["id"]
262
  if "function" in tc_chunk:
263
  if "name" in tc_chunk["function"]:
264
- aggregated_tool_calls[index]["function"]["name"] += tc_chunk["function"]["name"]
 
265
  if "arguments" in tc_chunk["function"]:
266
- aggregated_tool_calls[index]["function"]["arguments"] += tc_chunk["function"]["arguments"]
 
267
 
268
  elif key == "function_call":
269
  if "function_call" not in final_message:
270
  final_message["function_call"] = {"name": "", "arguments": ""}
271
  if "name" in value:
272
- final_message["function_call"]["name"] += value["name"]
 
273
  if "arguments" in value:
274
- final_message["function_call"]["arguments"] += value["arguments"]
 
275
 
276
  else: # Generic key handling for other data like 'reasoning'
277
  if key not in final_message:
 
256
  for tc_chunk in value:
257
  index = tc_chunk["index"]
258
  if index not in aggregated_tool_calls:
259
+ aggregated_tool_calls[index] = {"function": {"name": "", "arguments": ""}} # Initialize with minimal required keys
260
+ # Ensure 'function' key exists for this index before accessing its sub-keys
261
+ if "function" not in aggregated_tool_calls[index]:
262
+ aggregated_tool_calls[index]["function"] = {"name": "", "arguments": ""}
263
  if tc_chunk.get("id"):
264
  aggregated_tool_calls[index]["id"] = tc_chunk["id"]
265
  if "function" in tc_chunk:
266
  if "name" in tc_chunk["function"]:
267
+ if tc_chunk["function"]["name"] is not None:
268
+ aggregated_tool_calls[index]["function"]["name"] += tc_chunk["function"]["name"]
269
  if "arguments" in tc_chunk["function"]:
270
+ if tc_chunk["function"]["arguments"] is not None:
271
+ aggregated_tool_calls[index]["function"]["arguments"] += tc_chunk["function"]["arguments"]
272
 
273
  elif key == "function_call":
274
  if "function_call" not in final_message:
275
  final_message["function_call"] = {"name": "", "arguments": ""}
276
  if "name" in value:
277
+ if value["name"] is not None:
278
+ final_message["function_call"]["name"] += value["name"]
279
  if "arguments" in value:
280
+ if value["arguments"] is not None:
281
+ final_message["function_call"]["arguments"] += value["arguments"]
282
 
283
  else: # Generic key handling for other data like 'reasoning'
284
  if key not in final_message: