389 lines
13 KiB
Python
389 lines
13 KiB
Python
import json
|
|
import logging
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError, as_completed
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any, Callable, Dict, List, Optional, Set, Tuple
|
|
|
|
from rp.core.debug import debug_trace
|
|
|
|
logger = logging.getLogger("rp")
|
|
|
|
|
|
class ToolPriority(Enum):
|
|
CRITICAL = 1
|
|
HIGH = 2
|
|
NORMAL = 3
|
|
LOW = 4
|
|
|
|
|
|
@dataclass
|
|
class ToolCall:
|
|
tool_id: str
|
|
function_name: str
|
|
arguments: Dict[str, Any]
|
|
priority: ToolPriority = ToolPriority.NORMAL
|
|
timeout: float = 30.0
|
|
depends_on: Set[str] = field(default_factory=set)
|
|
retries: int = 3
|
|
retry_delay: float = 1.0
|
|
|
|
|
|
@dataclass
|
|
class ToolResult:
|
|
tool_id: str
|
|
function_name: str
|
|
success: bool
|
|
result: Any
|
|
error: Optional[str] = None
|
|
duration: float = 0.0
|
|
retries_used: int = 0
|
|
|
|
|
|
class ToolExecutor:
|
|
|
|
def __init__(
|
|
self,
|
|
max_workers: int = 10,
|
|
default_timeout: float = 30.0,
|
|
max_retries: int = 3,
|
|
retry_delay: float = 1.0
|
|
):
|
|
self.max_workers = max_workers
|
|
self.default_timeout = default_timeout
|
|
self.max_retries = max_retries
|
|
self.retry_delay = retry_delay
|
|
self._tool_registry: Dict[str, Callable] = {}
|
|
self._execution_stats: Dict[str, Dict[str, Any]] = {}
|
|
|
|
@debug_trace
|
|
def register_tool(self, name: str, func: Callable):
|
|
self._tool_registry[name] = func
|
|
|
|
@debug_trace
|
|
def register_tools(self, tools: Dict[str, Callable]):
|
|
self._tool_registry.update(tools)
|
|
|
|
@debug_trace
|
|
def _execute_single_tool(
|
|
self,
|
|
tool_call: ToolCall,
|
|
context: Optional[Dict[str, Any]] = None
|
|
) -> ToolResult:
|
|
start_time = time.time()
|
|
retries_used = 0
|
|
last_error = None
|
|
|
|
for attempt in range(tool_call.retries + 1):
|
|
try:
|
|
if tool_call.function_name not in self._tool_registry:
|
|
return ToolResult(
|
|
tool_id=tool_call.tool_id,
|
|
function_name=tool_call.function_name,
|
|
success=False,
|
|
result=None,
|
|
error=f"Unknown tool: {tool_call.function_name}",
|
|
duration=time.time() - start_time
|
|
)
|
|
|
|
func = self._tool_registry[tool_call.function_name]
|
|
|
|
if context:
|
|
result = func(**tool_call.arguments, **context)
|
|
else:
|
|
result = func(**tool_call.arguments)
|
|
|
|
duration = time.time() - start_time
|
|
self._update_stats(tool_call.function_name, duration, True)
|
|
|
|
return ToolResult(
|
|
tool_id=tool_call.tool_id,
|
|
function_name=tool_call.function_name,
|
|
success=True,
|
|
result=result,
|
|
duration=duration,
|
|
retries_used=retries_used
|
|
)
|
|
|
|
except Exception as e:
|
|
last_error = str(e)
|
|
retries_used = attempt + 1
|
|
logger.warning(
|
|
f"Tool {tool_call.function_name} failed (attempt {attempt + 1}): {last_error}"
|
|
)
|
|
if attempt < tool_call.retries:
|
|
time.sleep(tool_call.retry_delay * (attempt + 1))
|
|
|
|
duration = time.time() - start_time
|
|
self._update_stats(tool_call.function_name, duration, False)
|
|
|
|
return ToolResult(
|
|
tool_id=tool_call.tool_id,
|
|
function_name=tool_call.function_name,
|
|
success=False,
|
|
result=None,
|
|
error=last_error,
|
|
duration=duration,
|
|
retries_used=retries_used
|
|
)
|
|
|
|
def _update_stats(self, tool_name: str, duration: float, success: bool):
|
|
if tool_name not in self._execution_stats:
|
|
self._execution_stats[tool_name] = {
|
|
"total_calls": 0,
|
|
"successful_calls": 0,
|
|
"failed_calls": 0,
|
|
"total_duration": 0.0,
|
|
"avg_duration": 0.0
|
|
}
|
|
|
|
stats = self._execution_stats[tool_name]
|
|
stats["total_calls"] += 1
|
|
stats["total_duration"] += duration
|
|
stats["avg_duration"] = stats["total_duration"] / stats["total_calls"]
|
|
|
|
if success:
|
|
stats["successful_calls"] += 1
|
|
else:
|
|
stats["failed_calls"] += 1
|
|
|
|
@debug_trace
|
|
def execute_parallel(
|
|
self,
|
|
tool_calls: List[ToolCall],
|
|
context: Optional[Dict[str, Any]] = None
|
|
) -> List[ToolResult]:
|
|
if not tool_calls:
|
|
return []
|
|
|
|
dependency_graph = self._build_dependency_graph(tool_calls)
|
|
execution_order = self._topological_sort(dependency_graph)
|
|
|
|
results: Dict[str, ToolResult] = {}
|
|
|
|
for batch in execution_order:
|
|
batch_calls = [tc for tc in tool_calls if tc.tool_id in batch]
|
|
batch_results = self._execute_batch(batch_calls, context)
|
|
|
|
for result in batch_results:
|
|
results[result.tool_id] = result
|
|
if not result.success:
|
|
failed_dependents = self._get_dependents(result.tool_id, tool_calls)
|
|
for dep_id in failed_dependents:
|
|
if dep_id not in results:
|
|
results[dep_id] = ToolResult(
|
|
tool_id=dep_id,
|
|
function_name=next(
|
|
tc.function_name for tc in tool_calls if tc.tool_id == dep_id
|
|
),
|
|
success=False,
|
|
result=None,
|
|
error=f"Dependency {result.tool_id} failed"
|
|
)
|
|
|
|
return [results[tc.tool_id] for tc in tool_calls if tc.tool_id in results]
|
|
|
|
def _execute_batch(
|
|
self,
|
|
tool_calls: List[ToolCall],
|
|
context: Optional[Dict[str, Any]] = None
|
|
) -> List[ToolResult]:
|
|
results = []
|
|
|
|
sorted_calls = sorted(tool_calls, key=lambda x: x.priority.value)
|
|
|
|
with ThreadPoolExecutor(max_workers=min(len(sorted_calls), self.max_workers)) as executor:
|
|
future_to_call = {}
|
|
|
|
for tool_call in sorted_calls:
|
|
future = executor.submit(
|
|
self._execute_with_timeout,
|
|
tool_call,
|
|
context
|
|
)
|
|
future_to_call[future] = tool_call
|
|
|
|
for future in as_completed(future_to_call):
|
|
tool_call = future_to_call[future]
|
|
try:
|
|
result = future.result()
|
|
results.append(result)
|
|
except Exception as e:
|
|
results.append(ToolResult(
|
|
tool_id=tool_call.tool_id,
|
|
function_name=tool_call.function_name,
|
|
success=False,
|
|
result=None,
|
|
error=str(e)
|
|
))
|
|
|
|
return results
|
|
|
|
def _execute_with_timeout(
|
|
self,
|
|
tool_call: ToolCall,
|
|
context: Optional[Dict[str, Any]] = None
|
|
) -> ToolResult:
|
|
timeout = tool_call.timeout or self.default_timeout
|
|
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
future = executor.submit(self._execute_single_tool, tool_call, context)
|
|
try:
|
|
return future.result(timeout=timeout)
|
|
except FuturesTimeoutError:
|
|
return ToolResult(
|
|
tool_id=tool_call.tool_id,
|
|
function_name=tool_call.function_name,
|
|
success=False,
|
|
result=None,
|
|
error=f"Tool execution timed out after {timeout}s"
|
|
)
|
|
|
|
def _build_dependency_graph(
|
|
self,
|
|
tool_calls: List[ToolCall]
|
|
) -> Dict[str, Set[str]]:
|
|
graph = {tc.tool_id: tc.depends_on.copy() for tc in tool_calls}
|
|
return graph
|
|
|
|
def _topological_sort(
|
|
self,
|
|
graph: Dict[str, Set[str]]
|
|
) -> List[Set[str]]:
|
|
in_degree = {node: 0 for node in graph}
|
|
for node in graph:
|
|
for dep in graph[node]:
|
|
if dep in in_degree:
|
|
in_degree[node] += 1
|
|
|
|
batches = []
|
|
remaining = set(graph.keys())
|
|
|
|
while remaining:
|
|
batch = {
|
|
node for node in remaining
|
|
if all(dep not in remaining for dep in graph[node])
|
|
}
|
|
|
|
if not batch:
|
|
batch = {min(remaining, key=lambda x: in_degree.get(x, 0))}
|
|
|
|
batches.append(batch)
|
|
remaining -= batch
|
|
|
|
return batches
|
|
|
|
def _get_dependents(
|
|
self,
|
|
tool_id: str,
|
|
tool_calls: List[ToolCall]
|
|
) -> Set[str]:
|
|
dependents = set()
|
|
for tc in tool_calls:
|
|
if tool_id in tc.depends_on:
|
|
dependents.add(tc.tool_id)
|
|
dependents.update(self._get_dependents(tc.tool_id, tool_calls))
|
|
return dependents
|
|
|
|
def execute_sequential(
|
|
self,
|
|
tool_calls: List[ToolCall],
|
|
context: Optional[Dict[str, Any]] = None
|
|
) -> List[ToolResult]:
|
|
results = []
|
|
for tool_call in tool_calls:
|
|
result = self._execute_with_timeout(tool_call, context)
|
|
results.append(result)
|
|
return results
|
|
|
|
def get_statistics(self) -> Dict[str, Any]:
|
|
return {
|
|
"tool_stats": self._execution_stats.copy(),
|
|
"registered_tools": list(self._tool_registry.keys()),
|
|
"total_tools": len(self._tool_registry)
|
|
}
|
|
|
|
def clear_statistics(self):
|
|
self._execution_stats.clear()
|
|
|
|
|
|
def create_tool_executor_from_assistant(assistant) -> ToolExecutor:
|
|
from rp.tools.command import kill_process, run_command, tail_process
|
|
from rp.tools.database import db_get, db_query, db_set
|
|
from rp.tools.filesystem import (
|
|
chdir, getpwd, index_source_directory, list_directory,
|
|
mkdir, read_file, search_replace, write_file
|
|
)
|
|
from rp.tools.interactive_control import (
|
|
close_interactive_session, list_active_sessions,
|
|
read_session_output, send_input_to_session, start_interactive_session
|
|
)
|
|
from rp.tools.memory import (
|
|
add_knowledge_entry, delete_knowledge_entry, get_knowledge_by_category,
|
|
get_knowledge_entry, get_knowledge_statistics, search_knowledge,
|
|
update_knowledge_importance
|
|
)
|
|
from rp.tools.patch import apply_patch, create_diff, display_file_diff
|
|
from rp.tools.python_exec import python_exec
|
|
from rp.tools.web import http_fetch, web_search, web_search_news
|
|
from rp.tools.agents import (
|
|
collaborate_agents, create_agent, execute_agent_task, list_agents, remove_agent
|
|
)
|
|
from rp.tools.filesystem import (
|
|
clear_edit_tracker, display_edit_summary, display_edit_timeline
|
|
)
|
|
|
|
executor = ToolExecutor(
|
|
max_workers=10,
|
|
default_timeout=30.0,
|
|
max_retries=3
|
|
)
|
|
|
|
tools = {
|
|
"http_fetch": http_fetch,
|
|
"run_command": run_command,
|
|
"tail_process": tail_process,
|
|
"kill_process": kill_process,
|
|
"start_interactive_session": start_interactive_session,
|
|
"send_input_to_session": send_input_to_session,
|
|
"read_session_output": read_session_output,
|
|
"close_interactive_session": close_interactive_session,
|
|
"list_active_sessions": list_active_sessions,
|
|
"read_file": lambda **kw: read_file(**kw, db_conn=assistant.db_conn),
|
|
"write_file": lambda **kw: write_file(**kw, db_conn=assistant.db_conn),
|
|
"list_directory": list_directory,
|
|
"mkdir": mkdir,
|
|
"chdir": chdir,
|
|
"getpwd": getpwd,
|
|
"db_set": lambda **kw: db_set(**kw, db_conn=assistant.db_conn),
|
|
"db_get": lambda **kw: db_get(**kw, db_conn=assistant.db_conn),
|
|
"db_query": lambda **kw: db_query(**kw, db_conn=assistant.db_conn),
|
|
"web_search": web_search,
|
|
"web_search_news": web_search_news,
|
|
"python_exec": lambda **kw: python_exec(**kw, python_globals=assistant.python_globals),
|
|
"index_source_directory": index_source_directory,
|
|
"search_replace": lambda **kw: search_replace(**kw, db_conn=assistant.db_conn),
|
|
"create_diff": create_diff,
|
|
"apply_patch": lambda **kw: apply_patch(**kw, db_conn=assistant.db_conn),
|
|
"display_file_diff": display_file_diff,
|
|
"display_edit_summary": display_edit_summary,
|
|
"display_edit_timeline": display_edit_timeline,
|
|
"clear_edit_tracker": clear_edit_tracker,
|
|
"create_agent": create_agent,
|
|
"list_agents": list_agents,
|
|
"execute_agent_task": execute_agent_task,
|
|
"remove_agent": remove_agent,
|
|
"collaborate_agents": collaborate_agents,
|
|
"add_knowledge_entry": add_knowledge_entry,
|
|
"get_knowledge_entry": get_knowledge_entry,
|
|
"search_knowledge": search_knowledge,
|
|
"get_knowledge_by_category": get_knowledge_by_category,
|
|
"update_knowledge_importance": update_knowledge_importance,
|
|
"delete_knowledge_entry": delete_knowledge_entry,
|
|
"get_knowledge_statistics": get_knowledge_statistics,
|
|
}
|
|
|
|
executor.register_tools(tools)
|
|
return executor
|