Skip to content

Commit 1dbe3a6

Browse files
liuruibinliuruibin
committed
fix: enhance tool call handling with new ID detection and smart argument merging
1 parent c8403f9 commit 1dbe3a6

File tree

1 file changed

+93
-6
lines changed

1 file changed

+93
-6
lines changed

‎apps/application/flow/tools.py‎

Lines changed: 93 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -443,9 +443,20 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_
443443
continue
444444
entry = _tool_fragments.setdefault(
445445
idx, {'id': '', 'name': '', 'arguments': ''})
446+
447+
# Detect if this is a NEW tool call on the same index (fragment reuse)
446448
raw_id = tc_chunk.get('id')
447449
if raw_id and str(raw_id).strip():
448-
entry['id'] = str(raw_id).strip()
450+
new_id = str(raw_id).strip()
451+
# If fragment is completed and ID changes, it's a new tool call - reset fragment
452+
if entry.get('completed') and entry.get('id') and entry['id'] != new_id:
453+
maxkb_logger.debug(
454+
f"Resetting completed fragment {idx}: old ID {entry['id']} -> new ID {new_id}")
455+
entry.clear()
456+
entry.update(
457+
{'id': '', 'name': '', 'arguments': ''})
458+
entry['id'] = new_id
459+
449460
func_name = tc_chunk.get('name')
450461
if func_name:
451462
entry['name'] = func_name
@@ -457,7 +468,25 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_
457468
except Exception:
458469
part_args = str(part_args) if part_args else ''
459470
if part_args:
460-
entry['arguments'] += part_args
471+
# Smart merging: if both existing and new args are valid JSON objects,
472+
# merge them as dicts instead of concatenating strings
473+
if entry['arguments']:
474+
try:
475+
existing_obj = json.loads(entry['arguments'])
476+
new_obj = json.loads(part_args)
477+
if isinstance(existing_obj, dict) and isinstance(new_obj, dict):
478+
# Merge the two objects
479+
merged = {**existing_obj, **new_obj}
480+
entry['arguments'] = json.dumps(
481+
merged, ensure_ascii=False)
482+
else:
483+
# Not both dicts, fall back to concatenation
484+
entry['arguments'] += part_args
485+
except (json.JSONDecodeError, ValueError):
486+
# Not valid JSON, just concatenate (streaming fragments)
487+
entry['arguments'] += part_args
488+
else:
489+
entry['arguments'] = part_args
461490

462491
# ----------------------------------------------------------------
463492
# 2. 兼容 additional_kwargs['tool_calls'] 方式（旧格式/非流式情况）
@@ -470,9 +499,20 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_
470499
continue
471500
entry = _tool_fragments.setdefault(
472501
idx, {'id': '', 'name': '', 'arguments': ''})
502+
503+
# Detect if this is a NEW tool call on the same index (fragment reuse)
473504
raw_id = tool_call.get('id')
474505
if raw_id and str(raw_id).strip():
475-
entry['id'] = str(raw_id).strip()
506+
new_id = str(raw_id).strip()
507+
# If fragment is completed and ID changes, it's a new tool call - reset fragment
508+
if entry.get('completed') and entry.get('id') and entry['id'] != new_id:
509+
maxkb_logger.debug(
510+
f"Resetting completed fragment {idx}: old ID {entry['id']} -> new ID {new_id}")
511+
entry.clear()
512+
entry.update(
513+
{'id': '', 'name': '', 'arguments': ''})
514+
entry['id'] = new_id
515+
476516
func = tool_call.get('function', {})
477517
if isinstance(func, dict):
478518
func_name = func.get('name')
@@ -488,7 +528,25 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_
488528
except Exception:
489529
part_args = str(part_args) if part_args else ''
490530
if part_args:
491-
entry['arguments'] += part_args
531+
# Smart merging: if both existing and new args are valid JSON objects,
532+
# merge them as dicts instead of concatenating strings
533+
if entry['arguments']:
534+
try:
535+
existing_obj = json.loads(entry['arguments'])
536+
new_obj = json.loads(part_args)
537+
if isinstance(existing_obj, dict) and isinstance(new_obj, dict):
538+
# Merge the two objects
539+
merged = {**existing_obj, **new_obj}
540+
entry['arguments'] = json.dumps(
541+
merged, ensure_ascii=False)
542+
else:
543+
# Not both dicts, fall back to concatenation
544+
entry['arguments'] += part_args
545+
except (json.JSONDecodeError, ValueError):
546+
# Not valid JSON, just concatenate (streaming fragments)
547+
entry['arguments'] += part_args
548+
else:
549+
entry['arguments'] = part_args
492550

493551
# ----------------------------------------------------------------
494552
# 3. 检测工具调用结束，更新 tool_calls_info
@@ -501,7 +559,22 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_
501559

502560
if is_finish_chunk:
503561
# 在 finish chunk 时，将所有未完成的 fragment 标记完成并更新 tool_calls_info
562+
maxkb_logger.debug(
563+
f"Processing finish chunk. Tool fragments: {_tool_fragments}")
504564
for idx, entry in _tool_fragments.items():
565+
if entry.get('completed'):
566+
maxkb_logger.debug(
567+
f"Skipping fragment {idx}: already completed")
568+
continue
569+
if not entry.get('id'):
570+
maxkb_logger.warning(
571+
f"Skipping fragment {idx}: missing id. Fragment: {entry}")
572+
continue
573+
if not entry.get('arguments'):
574+
maxkb_logger.warning(
575+
f"Skipping fragment {idx}: missing arguments. Fragment: {entry}")
576+
continue
577+
505578
if not entry.get('completed') and entry.get('id') and entry.get('arguments'):
506579
try:
507580
parsed_args = json.loads(entry['arguments'])
@@ -514,9 +587,20 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_
514587
'input': json.dumps(filtered_args, ensure_ascii=False)
515588
}
516589
entry['completed'] = True
590+
maxkb_logger.debug(
591+
f"Added tool call {entry['id']} to tool_calls_info")
517592
except (json.JSONDecodeError, ValueError) as e:
593+
# JSON parsing failed, but still add to tool_calls_info with raw arguments
594+
# to prevent "Tool ID not found" errors when ToolMessage arrives
518595
maxkb_logger.warning(
519-
f"Failed to parse tool arguments at finish: {entry['arguments']}, error: {e}")
596+
f"Failed to parse tool arguments at finish for tool {entry.get('id', 'unknown')}: "
597+
f"{entry['arguments']}, error: {e}. Using raw arguments.")
598+
tool_calls_info[entry['id']] = {
599+
'name': entry['name'],
600+
# Use raw arguments
601+
'input': entry['arguments']
602+
}
603+
entry['completed'] = True
520604

521605
# ----------------------------------------------------------------
522606
# 4. 修复 tool_call_chunks 中的空 id（回填已知 id）
@@ -592,7 +676,10 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_
592676
chunk[0].content = content
593677
else:
594678
maxkb_logger.warning(
595-
f"Tool ID {tool_id} not found in tool_calls_info. Available IDs: {list(tool_calls_info.keys())}")
679+
f"Tool ID {tool_id} not found in tool_calls_info. "
680+
f"Available IDs: {list(tool_calls_info.keys())}. "
681+
f"Tool fragments at this point: {_tool_fragments}"
682+
)
596683

597684
yield chunk[0]
598685

0 commit comments

Comments
 (0)