81 lines
3.8 KiB
Python
81 lines
3.8 KiB
Python
"""Normalize local bridge batches before persistence and per-tool execution."""
|
|
|
|
from copy import deepcopy
|
|
import json
|
|
|
|
from agent.message_sanitization import coalesce_tool_call_id, deterministic_call_id, uniquify_tool_call_ids
|
|
from agent.transports.types import ToolCall
|
|
from tools.connectors import is_connector_name
|
|
from tools.connectors.gateway.config import MAX_CALLS_PER_DISPATCH
|
|
from tools.tool_search_catalog import TOOL_CALL_NAME
|
|
from tools.tool_search_validation import normalize_tool_call_entries
|
|
|
|
|
|
def expand_local_tool_batches(tool_calls: list, *, provider_data: dict | None = None) -> list:
|
|
"""Give every local entry the ordinary agent execution and persistence path.
|
|
|
|
Registry-only batch dispatch loses the live agent's inline tools, guardrails,
|
|
and per-server concurrency policy. Split the *new* assistant turn instead,
|
|
before it is persisted; never rewrite an earlier conversation message.
|
|
Connector-only batches keep their existing dispatcher and wire shape.
|
|
"""
|
|
expanded, replacements = [], {}
|
|
for call in tool_calls:
|
|
if call.function.name != TOOL_CALL_NAME:
|
|
expanded.append(call)
|
|
continue
|
|
try:
|
|
args = call.function.arguments
|
|
entries, error = normalize_tool_call_entries(json.loads(args) if isinstance(args, str) else args)
|
|
except (TypeError, ValueError, AttributeError):
|
|
entries, error = [], True
|
|
if (error or not 1 < len(entries) <= MAX_CALLS_PER_DISPATCH
|
|
or all(is_connector_name(entry["name"]) for entry in entries)):
|
|
# Leave malformed/oversized requests intact for the dispatcher's error.
|
|
expanded.append(call)
|
|
continue
|
|
parent_id = coalesce_tool_call_id(call)
|
|
children = []
|
|
for index, entry in enumerate(entries):
|
|
arguments = json.dumps({"calls": [entry]}, ensure_ascii=False)
|
|
if index == 0:
|
|
child = deepcopy(call)
|
|
child.function.arguments = arguments
|
|
else:
|
|
# A new entry must not inherit the parent's Responses item id or
|
|
# provider signature. IDs are deterministic for stable replay.
|
|
child = ToolCall(
|
|
id=deterministic_call_id(TOOL_CALL_NAME, parent_id + arguments, index),
|
|
name=TOOL_CALL_NAME, arguments=arguments,
|
|
)
|
|
expanded.append(child)
|
|
children.append(child)
|
|
replacements[parent_id] = children
|
|
uniquify_tool_call_ids(expanded)
|
|
# Native providers replay ordered blocks instead of rebuilding them from
|
|
# tool_calls. Replace the matching tool block in place, preserving signed
|
|
# thinking/text blocks and their order exactly.
|
|
if provider_data and replacements:
|
|
for key, replace in (("anthropic_content_blocks", _anthropic_children),
|
|
("bedrock_content_blocks", _bedrock_children)):
|
|
if blocks := provider_data.get(key):
|
|
provider_data[key] = [child for block in blocks for child in replace(block, replacements)]
|
|
return expanded
|
|
|
|
|
|
def _anthropic_children(block, replacements):
|
|
children = replacements.get(block.get("id")) if block.get("type") == "tool_use" else None
|
|
if children is None:
|
|
return [block]
|
|
return [{**block, "id": coalesce_tool_call_id(child),
|
|
"input": json.loads(child.function.arguments)} for child in children]
|
|
|
|
|
|
def _bedrock_children(block, replacements):
|
|
tool = block.get("toolUse")
|
|
children = replacements.get(tool.get("toolUseId")) if isinstance(tool, dict) else None
|
|
if children is None:
|
|
return [block]
|
|
return [{**block, "toolUse": {**tool, "toolUseId": coalesce_tool_call_id(child),
|
|
"input": json.loads(child.function.arguments)}} for child in children]
|