LZH-YS1998
78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
"""Tool planning helpers for Native Runtime V2."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from opc.layer4_tools.registry import ToolDefinition, ToolRegistry
|
|
|
|
|
|
_READ_ONLY_TOOL_NAMES = {
|
|
"file_read",
|
|
"file_search",
|
|
"list_dir",
|
|
"grep",
|
|
"glob",
|
|
"web_search",
|
|
"web_fetch",
|
|
"todo_read",
|
|
"agent_list",
|
|
"agent_wait",
|
|
}
|
|
|
|
_NON_CONCURRENT_TOOL_NAMES = {
|
|
"shell_exec",
|
|
"file_write",
|
|
"file_edit",
|
|
"apply_patch",
|
|
"python_exec",
|
|
"git_commit",
|
|
"agent_spawn",
|
|
"agent_wait",
|
|
"agent_send",
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class ToolBatch:
|
|
concurrency_safe: bool
|
|
calls: list[dict[str, Any]]
|
|
|
|
|
|
class ToolPlanner:
|
|
"""Determine tool execution ordering and concurrency."""
|
|
|
|
def __init__(self, registry: ToolRegistry, max_parallel_read_tools: int = 6) -> None:
|
|
self.registry = registry
|
|
self.max_parallel_read_tools = max(1, int(max_parallel_read_tools or 1))
|
|
|
|
def is_read_only(self, tool: ToolDefinition | None) -> bool:
|
|
if tool is None:
|
|
return False
|
|
if tool.read_only is not None:
|
|
return bool(tool.read_only)
|
|
if tool.name in _READ_ONLY_TOOL_NAMES:
|
|
return True
|
|
return tool.category in {"search", "read"} or tool.name.endswith("_read")
|
|
|
|
def is_concurrency_safe(self, tool: ToolDefinition | None) -> bool:
|
|
if tool is None:
|
|
return False
|
|
if tool.concurrency_safe is not None:
|
|
return bool(tool.concurrency_safe)
|
|
if tool.name in _NON_CONCURRENT_TOOL_NAMES:
|
|
return False
|
|
return self.is_read_only(tool)
|
|
|
|
def partition(self, tool_calls: list[dict[str, Any]]) -> list[ToolBatch]:
|
|
batches: list[ToolBatch] = []
|
|
for call in tool_calls:
|
|
tool = self.registry.get(str(call.get("function", "") or ""))
|
|
concurrency_safe = self.is_concurrency_safe(tool)
|
|
if concurrency_safe and batches and batches[-1].concurrency_safe:
|
|
batches[-1].calls.append(call)
|
|
else:
|
|
batches.append(ToolBatch(concurrency_safe=concurrency_safe, calls=[call]))
|
|
return batches
|