Compare commits
1 Commits
fix/749-v2
...
fix/795
| Author | SHA1 | Date | |
|---|---|---|---|
|
45679eef8a |
288
agent/gemma4_tool_hardening.py
Normal file
288
agent/gemma4_tool_hardening.py
Normal file
@@ -0,0 +1,288 @@
|
||||
"""Gemma 4 tool calling hardening — parse, validate, benchmark.
|
||||
|
||||
Gemma 4 has native multimodal function calling but its output format
|
||||
may differ from OpenAI/Claude. This module provides:
|
||||
|
||||
1. Gemma4ToolParser — robust parsing for Gemma 4's tool call format
|
||||
2. Parallel tool call detection and splitting
|
||||
3. Tool call success rate tracking and benchmarking
|
||||
4. Fallback parsing strategies for malformed output
|
||||
|
||||
Usage:
|
||||
from agent.gemma4_tool_hardening import Gemma4ToolParser
|
||||
parser = Gemma4ToolParser()
|
||||
tool_calls = parser.parse(response_text)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCallAttempt:
|
||||
"""Record of a single tool call parsing attempt."""
|
||||
raw_text: str
|
||||
parsed: bool
|
||||
tool_name: str
|
||||
arguments: dict
|
||||
error: str
|
||||
strategy: str # "native", "json_block", "regex", "fallback"
|
||||
timestamp: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class Gemma4BenchmarkResult:
|
||||
"""Result of a tool calling benchmark run."""
|
||||
total_calls: int = 0
|
||||
successful_parses: int = 0
|
||||
parallel_calls: int = 0
|
||||
strategies_used: Dict[str, int] = field(default_factory=dict)
|
||||
avg_parse_time_ms: float = 0.0
|
||||
success_rate: float = 0.0
|
||||
errors: List[str] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"total_calls": self.total_calls,
|
||||
"successful_parses": self.successful_parses,
|
||||
"parallel_calls": self.parallel_calls,
|
||||
"success_rate": round(self.success_rate, 3),
|
||||
"strategies_used": self.strategies_used,
|
||||
"avg_parse_time_ms": round(self.avg_parse_time_ms, 2),
|
||||
"error_count": len(self.errors),
|
||||
"errors": self.errors[:10],
|
||||
}
|
||||
|
||||
|
||||
class Gemma4ToolParser:
|
||||
"""Robust tool call parser for Gemma 4 output format.
|
||||
|
||||
Tries multiple parsing strategies in order:
|
||||
1. Native OpenAI format (standard tool_calls)
|
||||
2. JSON code blocks (```json ... ```)
|
||||
3. Regex extraction (function_name + arguments patterns)
|
||||
4. Heuristic fallback (best-effort extraction)
|
||||
"""
|
||||
|
||||
# Patterns for Gemma 4 tool call formats
|
||||
_JSON_BLOCK_PATTERN = re.compile(
|
||||
r'```(?:json)?\s*\n?(.*?)\n?```',
|
||||
re.DOTALL | re.IGNORECASE,
|
||||
)
|
||||
_FUNCTION_CALL_PATTERN = re.compile(
|
||||
r'(?:function|tool|call)[:\s]*(\w+)\s*\(\s*({.*?})\s*\)',
|
||||
re.DOTALL | re.IGNORECASE,
|
||||
)
|
||||
_GEMMA_INLINE_PATTERN = re.compile(
|
||||
r'\[(?:tool_call|function_call)\]\s*(\w+)\s*:\s*({.*?})',
|
||||
re.DOTALL | re.IGNORECASE,
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
self._attempts: List[ToolCallAttempt] = []
|
||||
self._benchmark = Gemma4BenchmarkResult()
|
||||
|
||||
@property
|
||||
def benchmark(self) -> Gemma4BenchmarkResult:
|
||||
return self._benchmark
|
||||
|
||||
def parse(self, response_text: str, expected_tools: List[str] = None) -> List[Dict[str, Any]]:
|
||||
"""Parse tool calls from model response using multiple strategies.
|
||||
|
||||
Returns list of tool call dicts in OpenAI format:
|
||||
[{"id": "...", "type": "function", "function": {"name": "...", "arguments": "..."}}]
|
||||
"""
|
||||
t0 = time.monotonic()
|
||||
self._benchmark.total_calls += 1
|
||||
|
||||
# Strategy 1: Native OpenAI format
|
||||
result = self._try_native_parse(response_text)
|
||||
if result:
|
||||
self._record_attempt(response_text, True, result, "native")
|
||||
self._benchmark.successful_parses += 1
|
||||
if len(result) > 1:
|
||||
self._benchmark.parallel_calls += 1
|
||||
self._benchmark.strategies_used["native"] = self._benchmark.strategies_used.get("native", 0) + 1
|
||||
self._update_timing(t0)
|
||||
return result
|
||||
|
||||
# Strategy 2: JSON code blocks
|
||||
result = self._try_json_block_parse(response_text, expected_tools)
|
||||
if result:
|
||||
self._record_attempt(response_text, True, result, "json_block")
|
||||
self._benchmark.successful_parses += 1
|
||||
if len(result) > 1:
|
||||
self._benchmark.parallel_calls += 1
|
||||
self._benchmark.strategies_used["json_block"] = self._benchmark.strategies_used.get("json_block", 0) + 1
|
||||
self._update_timing(t0)
|
||||
return result
|
||||
|
||||
# Strategy 3: Regex extraction
|
||||
result = self._try_regex_parse(response_text)
|
||||
if result:
|
||||
self._record_attempt(response_text, True, result, "regex")
|
||||
self._benchmark.successful_parses += 1
|
||||
self._benchmark.strategies_used["regex"] = self._benchmark.strategies_used.get("regex", 0) + 1
|
||||
self._update_timing(t0)
|
||||
return result
|
||||
|
||||
# Strategy 4: Heuristic fallback
|
||||
result = self._try_heuristic_parse(response_text, expected_tools)
|
||||
if result:
|
||||
self._record_attempt(response_text, True, result, "fallback")
|
||||
self._benchmark.successful_parses += 1
|
||||
self._benchmark.strategies_used["fallback"] = self._benchmark.strategies_used.get("fallback", 0) + 1
|
||||
self._update_timing(t0)
|
||||
return result
|
||||
|
||||
# All strategies failed
|
||||
self._record_attempt(response_text, False, [], "none")
|
||||
self._benchmark.errors.append(f"Failed to parse: {response_text[:200]}")
|
||||
self._update_timing(t0)
|
||||
return []
|
||||
|
||||
def _try_native_parse(self, text: str) -> List[Dict[str, Any]]:
|
||||
"""Try parsing standard OpenAI tool_calls JSON."""
|
||||
try:
|
||||
data = json.loads(text)
|
||||
if isinstance(data, dict) and "tool_calls" in data:
|
||||
return data["tool_calls"]
|
||||
if isinstance(data, list):
|
||||
if all(isinstance(item, dict) and "function" in item for item in data):
|
||||
return data
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
return []
|
||||
|
||||
def _try_json_block_parse(self, text: str, expected_tools: List[str] = None) -> List[Dict[str, Any]]:
|
||||
"""Extract tool calls from JSON code blocks."""
|
||||
matches = self._JSON_BLOCK_PATTERN.findall(text)
|
||||
calls = []
|
||||
for match in matches:
|
||||
try:
|
||||
data = json.loads(match.strip())
|
||||
if isinstance(data, dict):
|
||||
if "name" in data and "arguments" in data:
|
||||
calls.append(self._to_openai_format(data["name"], data["arguments"]))
|
||||
elif "function" in data and "arguments" in data:
|
||||
calls.append(self._to_openai_format(data["function"], data["arguments"]))
|
||||
elif isinstance(data, list):
|
||||
for item in data:
|
||||
if isinstance(item, dict) and "name" in item:
|
||||
args = item.get("arguments", item.get("args", {}))
|
||||
calls.append(self._to_openai_format(item["name"], args))
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
return calls
|
||||
|
||||
def _try_regex_parse(self, text: str) -> List[Dict[str, Any]]:
|
||||
"""Extract tool calls using regex patterns."""
|
||||
calls = []
|
||||
|
||||
# Pattern: function_name({...})
|
||||
for match in self._FUNCTION_CALL_PATTERN.finditer(text):
|
||||
name = match.group(1)
|
||||
args_str = match.group(2)
|
||||
try:
|
||||
args = json.loads(args_str)
|
||||
calls.append(self._to_openai_format(name, args))
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
# Pattern: [tool_call] name: {...}
|
||||
for match in self._GEMMA_INLINE_PATTERN.finditer(text):
|
||||
name = match.group(1)
|
||||
args_str = match.group(2)
|
||||
try:
|
||||
args = json.loads(args_str)
|
||||
calls.append(self._to_openai_format(name, args))
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
return calls
|
||||
|
||||
def _try_heuristic_parse(self, text: str, expected_tools: List[str] = None) -> List[Dict[str, Any]]:
|
||||
"""Best-effort heuristic extraction."""
|
||||
if not expected_tools:
|
||||
return []
|
||||
|
||||
calls = []
|
||||
for tool_name in expected_tools:
|
||||
# Look for tool name near JSON-like content
|
||||
pattern = re.compile(
|
||||
rf'{re.escape(tool_name)}\s*[\(:]\s*({{[^}}]+}})',
|
||||
re.IGNORECASE,
|
||||
)
|
||||
match = pattern.search(text)
|
||||
if match:
|
||||
try:
|
||||
args = json.loads(match.group(1))
|
||||
calls.append(self._to_openai_format(tool_name, args))
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
return calls
|
||||
|
||||
def _to_openai_format(self, name: str, arguments: Any) -> Dict[str, Any]:
|
||||
"""Convert to OpenAI tool call format."""
|
||||
import uuid
|
||||
args_str = json.dumps(arguments) if isinstance(arguments, dict) else str(arguments)
|
||||
return {
|
||||
"id": f"call_{uuid.uuid4().hex[:24]}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"arguments": args_str,
|
||||
},
|
||||
}
|
||||
|
||||
def _record_attempt(self, text: str, success: bool, result: list, strategy: str):
|
||||
self._attempts.append(ToolCallAttempt(
|
||||
raw_text=text[:500],
|
||||
parsed=success,
|
||||
tool_name=result[0]["function"]["name"] if result else "",
|
||||
arguments={},
|
||||
error="" if success else "parse failed",
|
||||
strategy=strategy,
|
||||
timestamp=time.time(),
|
||||
))
|
||||
|
||||
def _update_timing(self, t0: float):
|
||||
elapsed = (time.monotonic() - t0) * 1000
|
||||
n = self._benchmark.total_calls
|
||||
self._benchmark.avg_parse_time_ms = (
|
||||
(self._benchmark.avg_parse_time_ms * (n - 1) + elapsed) / n
|
||||
)
|
||||
self._benchmark.success_rate = (
|
||||
self._benchmark.successful_parses / n if n > 0 else 0
|
||||
)
|
||||
|
||||
def format_report(self) -> str:
|
||||
"""Format benchmark report."""
|
||||
b = self._benchmark
|
||||
lines = [
|
||||
"Gemma 4 Tool Calling Benchmark",
|
||||
"=" * 40,
|
||||
f"Total attempts: {b.total_calls}",
|
||||
f"Successful parses: {b.successful_parses}",
|
||||
f"Success rate: {b.success_rate:.1%}",
|
||||
f"Parallel calls: {b.parallel_calls}",
|
||||
f"Avg parse time: {b.avg_parse_time_ms:.2f}ms",
|
||||
"",
|
||||
"Strategies used:",
|
||||
]
|
||||
for strategy, count in sorted(b.strategies_used.items(), key=lambda x: -x[1]):
|
||||
lines.append(f" {strategy}: {count}")
|
||||
|
||||
if b.errors:
|
||||
lines.append("")
|
||||
lines.append(f"Errors ({len(b.errors)}):")
|
||||
for err in b.errors[:5]:
|
||||
lines.append(f" {err[:100]}")
|
||||
|
||||
return "\n".join(lines)
|
||||
@@ -1,136 +0,0 @@
|
||||
"""Tests for batch tool execution — Issue #749."""
|
||||
import asyncio
|
||||
import sys
|
||||
from pathlib import Path
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from tools.batch_executor import (
|
||||
ToolSafety, ToolCall, BatchResult,
|
||||
classify_tool_safety, classify_calls,
|
||||
execute_batch_sync, get_tool_safety_report
|
||||
)
|
||||
|
||||
|
||||
class TestClassification:
|
||||
def test_parallel_safe_read(self):
|
||||
assert classify_tool_safety("file_read") == ToolSafety.PARALLEL_SAFE
|
||||
|
||||
def test_sequential_write(self):
|
||||
assert classify_tool_safety("file_write") == ToolSafety.SEQUENTIAL
|
||||
|
||||
def test_destructive_terminal(self):
|
||||
assert classify_tool_safety("terminal") == ToolSafety.DESTRUCTIVE
|
||||
|
||||
def test_unknown_defaults_sequential(self):
|
||||
assert classify_tool_safety("unknown_tool") == ToolSafety.SEQUENTIAL
|
||||
|
||||
def test_prefix_match(self):
|
||||
assert classify_tool_safety("file_read_special") == ToolSafety.PARALLEL_SAFE
|
||||
|
||||
|
||||
class TestClassifyCalls:
|
||||
def test_classifies_multiple(self):
|
||||
calls = [
|
||||
{"name": "file_read", "arguments": "{}"},
|
||||
{"name": "file_write", "arguments": "{}"},
|
||||
{"name": "terminal", "arguments": "{}"},
|
||||
]
|
||||
result = classify_calls(calls)
|
||||
assert len(result) == 3
|
||||
assert result[0].safety == ToolSafety.PARALLEL_SAFE
|
||||
assert result[1].safety == ToolSafety.SEQUENTIAL
|
||||
assert result[2].safety == ToolSafety.DESTRUCTIVE
|
||||
|
||||
|
||||
class TestBatchExecution:
|
||||
def test_parallel_execution(self):
|
||||
"""Parallel-safe calls should execute faster than sequential."""
|
||||
import time
|
||||
|
||||
def slow_executor(name, args):
|
||||
time.sleep(0.1)
|
||||
return f"result_{name}"
|
||||
|
||||
calls = [
|
||||
{"name": "file_read", "arguments": "{}"},
|
||||
{"name": "file_search", "arguments": "{}"},
|
||||
{"name": "web_search", "arguments": "{}"},
|
||||
]
|
||||
|
||||
start = time.time()
|
||||
result = execute_batch_sync(calls, slow_executor)
|
||||
duration = time.time() - start
|
||||
|
||||
# Should be faster than 0.3s (3 * 0.1) since parallel
|
||||
assert duration < 0.25
|
||||
assert result.parallel_count == 3
|
||||
assert len(result.errors) == 0
|
||||
|
||||
def test_sequential_execution(self):
|
||||
"""Sequential calls should execute one at a time."""
|
||||
import time
|
||||
|
||||
def slow_executor(name, args):
|
||||
time.sleep(0.05)
|
||||
return f"result_{name}"
|
||||
|
||||
calls = [
|
||||
{"name": "file_write", "arguments": "{}"},
|
||||
{"name": "file_patch", "arguments": "{}"},
|
||||
]
|
||||
|
||||
start = time.time()
|
||||
result = execute_batch_sync(calls, slow_executor)
|
||||
duration = time.time() - start
|
||||
|
||||
# Should take at least 0.1s (2 * 0.05) since sequential
|
||||
assert duration >= 0.1
|
||||
assert result.sequential_count == 2
|
||||
|
||||
def test_mixed_execution(self):
|
||||
"""Mixed calls: parallel first, then sequential."""
|
||||
calls = [
|
||||
{"name": "file_read", "arguments": "{}"},
|
||||
{"name": "file_write", "arguments": "{}"},
|
||||
{"name": "web_search", "arguments": "{}"},
|
||||
]
|
||||
|
||||
def executor(name, args):
|
||||
return f"result_{name}"
|
||||
|
||||
result = execute_batch_sync(calls, executor)
|
||||
assert result.parallel_count == 2
|
||||
assert result.sequential_count == 1
|
||||
assert len(result.errors) == 0
|
||||
|
||||
def test_error_handling(self):
|
||||
"""Errors in one call shouldn't stop others."""
|
||||
def failing_executor(name, args):
|
||||
if name == "file_write":
|
||||
raise Exception("Write failed")
|
||||
return "ok"
|
||||
|
||||
calls = [
|
||||
{"name": "file_read", "arguments": "{}"},
|
||||
{"name": "file_write", "arguments": "{}"},
|
||||
]
|
||||
|
||||
result = execute_batch_sync(calls, failing_executor)
|
||||
assert len(result.errors) == 1
|
||||
assert "file_write" in result.errors[0]
|
||||
|
||||
|
||||
class TestSafetyReport:
|
||||
def test_report_format(self):
|
||||
calls = [
|
||||
ToolCall(name="file_read", args={}, safety=ToolSafety.PARALLEL_SAFE, duration=0.1),
|
||||
ToolCall(name="file_write", args={}, safety=ToolSafety.SEQUENTIAL, duration=0.2),
|
||||
]
|
||||
report = get_tool_safety_report(calls)
|
||||
assert "Parallel-safe: 1" in report
|
||||
assert "Sequential: 1" in report
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import pytest
|
||||
pytest.main([__file__, "-v"])
|
||||
94
tests/test_gemma4_tool_hardening.py
Normal file
94
tests/test_gemma4_tool_hardening.py
Normal file
@@ -0,0 +1,94 @@
|
||||
"""Tests for Gemma 4 tool calling hardening."""
|
||||
|
||||
import json
|
||||
import pytest
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from agent.gemma4_tool_hardening import Gemma4ToolParser, Gemma4BenchmarkResult
|
||||
|
||||
|
||||
class TestNativeParse:
|
||||
def test_standard_tool_calls(self):
|
||||
parser = Gemma4ToolParser()
|
||||
text = json.dumps({"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read_file", "arguments": '{"path": "test.py"}'}}]})
|
||||
result = parser.parse(text)
|
||||
assert len(result) == 1
|
||||
assert result[0]["function"]["name"] == "read_file"
|
||||
|
||||
def test_list_format(self):
|
||||
parser = Gemma4ToolParser()
|
||||
text = json.dumps([{"id": "c1", "type": "function", "function": {"name": "terminal", "arguments": '{"command": "ls"}'}}])
|
||||
result = parser.parse(text)
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
class TestJsonBlockParse:
|
||||
def test_json_code_block(self):
|
||||
parser = Gemma4ToolParser()
|
||||
text = 'Here is the tool call:\n```json\n{"name": "read_file", "arguments": {"path": "test.py"}}\n```'
|
||||
result = parser.parse(text)
|
||||
assert len(result) == 1
|
||||
assert result[0]["function"]["name"] == "read_file"
|
||||
|
||||
def test_multiple_json_blocks(self):
|
||||
parser = Gemma4ToolParser()
|
||||
text = '```json\n{"name": "read_file", "arguments": {"path": "a.py"}}\n```\n```json\n{"name": "read_file", "arguments": {"path": "b.py"}}\n```'
|
||||
result = parser.parse(text)
|
||||
assert len(result) == 2
|
||||
|
||||
def test_list_in_json_block(self):
|
||||
parser = Gemma4ToolParser()
|
||||
text = '```json\n[{"name": "terminal", "arguments": {"command": "ls"}}]\n```'
|
||||
result = parser.parse(text)
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
class TestRegexParse:
|
||||
def test_function_call_pattern(self):
|
||||
parser = Gemma4ToolParser()
|
||||
text = 'I will call read_file({"path": "test.py"}) now.'
|
||||
result = parser.parse(text)
|
||||
assert len(result) == 1
|
||||
assert result[0]["function"]["name"] == "read_file"
|
||||
|
||||
def test_gemma_inline_pattern(self):
|
||||
parser = Gemma4ToolParser()
|
||||
text = '[tool_call] terminal: {"command": "pwd"}'
|
||||
result = parser.parse(text)
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
class TestHeuristicParse:
|
||||
def test_heuristic_with_expected_tools(self):
|
||||
parser = Gemma4ToolParser()
|
||||
text = 'Calling read_file({"path": "config.yaml"}) now'
|
||||
result = parser.parse(text, expected_tools=["read_file"])
|
||||
assert len(result) == 1
|
||||
|
||||
def test_heuristic_without_expected_tools(self):
|
||||
parser = Gemma4ToolParser()
|
||||
text = 'Some text with {"key": "value"} but no tool name'
|
||||
result = parser.parse(text)
|
||||
assert len(result) == 0
|
||||
|
||||
|
||||
class TestBenchmark:
|
||||
def test_benchmark_counts(self):
|
||||
parser = Gemma4ToolParser()
|
||||
parser.parse(json.dumps({"tool_calls": [{"id": "1", "type": "function", "function": {"name": "x", "arguments": "{}"}}]}))
|
||||
parser.parse('```json\n{"name": "y", "arguments": {}}\n```')
|
||||
parser.parse('no tool call here')
|
||||
b = parser.benchmark
|
||||
assert b.total_calls == 3
|
||||
assert b.successful_parses == 2
|
||||
assert abs(b.success_rate - 2/3) < 0.01
|
||||
|
||||
def test_report_format(self):
|
||||
parser = Gemma4ToolParser()
|
||||
parser.parse(json.dumps({"tool_calls": [{"id": "1", "type": "function", "function": {"name": "x", "arguments": "{}"}}]}))
|
||||
report = parser.format_report()
|
||||
assert "Gemma 4 Tool Calling Benchmark" in report
|
||||
assert "native" in report
|
||||
@@ -1,280 +0,0 @@
|
||||
"""Batch tool execution with parallel safety checks.
|
||||
|
||||
Classifies tool calls as parallel-safe vs sequential and executes
|
||||
parallel-safe calls concurrently while keeping destructive ops serialized.
|
||||
|
||||
Issue #749: feat: batch tool execution with parallel safety checks
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ToolSafety(Enum):
|
||||
"""Safety classification for tool calls."""
|
||||
PARALLEL_SAFE = "parallel_safe" # Can run concurrently
|
||||
SEQUENTIAL = "sequential" # Must run one at a time
|
||||
DESTRUCTIVE = "destructive" # Destructive, needs approval
|
||||
|
||||
|
||||
# Tool safety classifications
|
||||
_TOOL_SAFETY: Dict[str, ToolSafety] = {
|
||||
# Parallel-safe: reads, searches, non-destructive
|
||||
"file_read": ToolSafety.PARALLEL_SAFE,
|
||||
"file_search": ToolSafety.PARALLEL_SAFE,
|
||||
"web_search": ToolSafety.PARALLEL_SAFE,
|
||||
"web_extract": ToolSafety.PARALLEL_SAFE,
|
||||
"browser_snapshot": ToolSafety.PARALLEL_SAFE,
|
||||
"browser_vision": ToolSafety.PARALLEL_SAFE,
|
||||
"browser_get_images": ToolSafety.PARALLEL_SAFE,
|
||||
"skill_view": ToolSafety.PARALLEL_SAFE,
|
||||
"memory_search": ToolSafety.PARALLEL_SAFE,
|
||||
"memory_recall": ToolSafety.PARALLEL_SAFE,
|
||||
"session_search": ToolSafety.PARALLEL_SAFE,
|
||||
|
||||
# Sequential: writes, edits, state changes
|
||||
"file_write": ToolSafety.SEQUENTIAL,
|
||||
"file_patch": ToolSafety.SEQUENTIAL,
|
||||
"file_append": ToolSafety.SEQUENTIAL,
|
||||
"browser_navigate": ToolSafety.SEQUENTIAL,
|
||||
"browser_click": ToolSafety.SEQUENTIAL,
|
||||
"browser_type": ToolSafety.SEQUENTIAL,
|
||||
"browser_scroll": ToolSafety.SEQUENTIAL,
|
||||
"memory_store": ToolSafety.SEQUENTIAL,
|
||||
"memory_update": ToolSafety.SEQUENTIAL,
|
||||
"cronjob": ToolSafety.SEQUENTIAL,
|
||||
"send_message": ToolSafety.SEQUENTIAL,
|
||||
|
||||
# Destructive: needs approval
|
||||
"terminal": ToolSafety.DESTRUCTIVE,
|
||||
"execute_code": ToolSafety.DESTRUCTIVE,
|
||||
"browser_execute_js": ToolSafety.DESTRUCTIVE,
|
||||
"delegate_task": ToolSafety.DESTRUCTIVE,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""A single tool call with metadata."""
|
||||
name: str
|
||||
args: Dict[str, Any]
|
||||
call_id: str = ""
|
||||
safety: ToolSafety = ToolSafety.SEQUENTIAL
|
||||
result: Optional[Any] = None
|
||||
error: Optional[str] = None
|
||||
duration: float = 0.0
|
||||
started_at: float = 0.0
|
||||
completed_at: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchResult:
|
||||
"""Result of batch tool execution."""
|
||||
calls: List[ToolCall] = field(default_factory=list)
|
||||
parallel_count: int = 0
|
||||
sequential_count: int = 0
|
||||
total_duration: float = 0.0
|
||||
errors: List[str] = field(default_factory=list)
|
||||
|
||||
|
||||
def classify_tool_safety(tool_name: str) -> ToolSafety:
|
||||
"""Classify a tool call's safety level."""
|
||||
# Check exact match first
|
||||
if tool_name in _TOOL_SAFETY:
|
||||
return _TOOL_SAFETY[tool_name]
|
||||
|
||||
# Check prefix matches
|
||||
for pattern, safety in _TOOL_SAFETY.items():
|
||||
if tool_name.startswith(pattern):
|
||||
return safety
|
||||
|
||||
# Default to sequential for unknown tools
|
||||
return ToolSafety.SEQUENTIAL
|
||||
|
||||
|
||||
def classify_calls(tool_calls: List[Dict[str, Any]]) -> List[ToolCall]:
|
||||
"""Classify a list of tool calls by safety level."""
|
||||
calls = []
|
||||
for i, tc in enumerate(tool_calls):
|
||||
name = tc.get("name", tc.get("function", {}).get("name", ""))
|
||||
args = tc.get("arguments", tc.get("function", {}).get("arguments", {}))
|
||||
if isinstance(args, str):
|
||||
import json
|
||||
try:
|
||||
args = json.loads(args)
|
||||
except Exception:
|
||||
args = {}
|
||||
|
||||
call_id = tc.get("id", f"call_{i}")
|
||||
safety = classify_tool_safety(name)
|
||||
|
||||
calls.append(ToolCall(
|
||||
name=name,
|
||||
args=args,
|
||||
call_id=call_id,
|
||||
safety=safety,
|
||||
))
|
||||
|
||||
return calls
|
||||
|
||||
|
||||
async def execute_parallel(
|
||||
calls: List[ToolCall],
|
||||
executor: Callable[[str, Dict[str, Any]], Any],
|
||||
) -> List[ToolCall]:
|
||||
"""Execute parallel-safe calls concurrently."""
|
||||
async def run_call(call: ToolCall) -> ToolCall:
|
||||
call.started_at = time.time()
|
||||
try:
|
||||
# Run in thread pool to avoid blocking
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: executor(call.name, call.args),
|
||||
)
|
||||
call.result = result
|
||||
except Exception as e:
|
||||
call.error = str(e)
|
||||
logger.error(f"Parallel call {call.name} failed: {e}")
|
||||
finally:
|
||||
call.completed_at = time.time()
|
||||
call.duration = call.completed_at - call.started_at
|
||||
return call
|
||||
|
||||
# Execute all parallel-safe calls concurrently
|
||||
tasks = [run_call(call) for call in calls]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# Handle exceptions from gather
|
||||
processed = []
|
||||
for i, result in enumerate(results):
|
||||
if isinstance(result, Exception):
|
||||
calls[i].error = str(result)
|
||||
calls[i].completed_at = time.time()
|
||||
calls[i].duration = calls[i].completed_at - calls[i].started_at
|
||||
processed.append(calls[i])
|
||||
else:
|
||||
processed.append(result)
|
||||
|
||||
return processed
|
||||
|
||||
|
||||
async def execute_sequential(
|
||||
calls: List[ToolCall],
|
||||
executor: Callable[[str, Dict[str, Any]], Any],
|
||||
) -> List[ToolCall]:
|
||||
"""Execute sequential/destructive calls one at a time."""
|
||||
for call in calls:
|
||||
call.started_at = time.time()
|
||||
try:
|
||||
result = executor(call.name, call.args)
|
||||
call.result = result
|
||||
except Exception as e:
|
||||
call.error = str(e)
|
||||
logger.error(f"Sequential call {call.name} failed: {e}")
|
||||
finally:
|
||||
call.completed_at = time.time()
|
||||
call.duration = call.completed_at - call.started_at
|
||||
|
||||
return calls
|
||||
|
||||
|
||||
async def execute_batch(
|
||||
tool_calls: List[Dict[str, Any]],
|
||||
executor: Callable[[str, Dict[str, Any]], Any],
|
||||
max_parallel: int = 5,
|
||||
) -> BatchResult:
|
||||
"""Execute a batch of tool calls with parallel safety checks.
|
||||
|
||||
Args:
|
||||
tool_calls: List of tool call dicts (OpenAI format)
|
||||
executor: Function to execute a single tool call (name, args) -> result
|
||||
max_parallel: Maximum concurrent parallel calls
|
||||
|
||||
Returns:
|
||||
BatchResult with all call results and timing info
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Classify all calls
|
||||
calls = classify_calls(tool_calls)
|
||||
|
||||
# Split by safety level
|
||||
parallel_calls = [c for c in calls if c.safety == ToolSafety.PARALLEL_SAFE]
|
||||
sequential_calls = [c for c in calls if c.safety != ToolSafety.PARALLEL_SAFE]
|
||||
|
||||
result = BatchResult(
|
||||
calls=calls,
|
||||
parallel_count=len(parallel_calls),
|
||||
sequential_count=len(sequential_calls),
|
||||
)
|
||||
|
||||
# Execute parallel calls concurrently
|
||||
if parallel_calls:
|
||||
logger.info(f"Executing {len(parallel_calls)} parallel-safe calls concurrently")
|
||||
|
||||
# Batch into chunks of max_parallel
|
||||
for i in range(0, len(parallel_calls), max_parallel):
|
||||
chunk = parallel_calls[i:i + max_parallel]
|
||||
await execute_parallel(chunk, executor)
|
||||
|
||||
# Execute sequential calls one at a time
|
||||
if sequential_calls:
|
||||
logger.info(f"Executing {len(sequential_calls)} sequential calls")
|
||||
await execute_sequential(sequential_calls, executor)
|
||||
|
||||
# Collect errors
|
||||
for call in calls:
|
||||
if call.error:
|
||||
result.errors.append(f"{call.name}: {call.error}")
|
||||
|
||||
result.total_duration = time.time() - start_time
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def execute_batch_sync(
|
||||
tool_calls: List[Dict[str, Any]],
|
||||
executor: Callable[[str, Dict[str, Any]], Any],
|
||||
max_parallel: int = 5,
|
||||
) -> BatchResult:
|
||||
"""Synchronous wrapper for execute_batch."""
|
||||
return asyncio.run(execute_batch(tool_calls, executor, max_parallel))
|
||||
|
||||
|
||||
def get_tool_safety_report(calls: List[ToolCall]) -> str:
|
||||
"""Generate a human-readable safety report."""
|
||||
parallel = [c for c in calls if c.safety == ToolSafety.PARALLEL_SAFE]
|
||||
sequential = [c for c in calls if c.safety == ToolSafety.SEQUENTIAL]
|
||||
destructive = [c for c in calls if c.safety == ToolSafety.DESTRUCTIVE]
|
||||
|
||||
lines = ["Tool Safety Report:"]
|
||||
lines.append(f" Parallel-safe: {len(parallel)}")
|
||||
lines.append(f" Sequential: {len(sequential)}")
|
||||
lines.append(f" Destructive: {len(destructive)}")
|
||||
|
||||
if parallel:
|
||||
lines.append("\nParallel-safe calls:")
|
||||
for c in parallel:
|
||||
status = "✓" if not c.error else "✗"
|
||||
lines.append(f" {status} {c.name} ({c.duration:.2f}s)")
|
||||
|
||||
if sequential:
|
||||
lines.append("\nSequential calls:")
|
||||
for c in sequential:
|
||||
status = "✓" if not c.error else "✗"
|
||||
lines.append(f" {status} {c.name} ({c.duration:.2f}s)")
|
||||
|
||||
if destructive:
|
||||
lines.append("\nDestructive calls:")
|
||||
for c in destructive:
|
||||
status = "✓" if not c.error else "✗"
|
||||
lines.append(f" {status} {c.name} ({c.duration:.2f}s)")
|
||||
|
||||
return "\n".join(lines)
|
||||
Reference in New Issue
Block a user