feat: rename assistant identifier to "rp" across config and runtime references
This commit is contained in:
@@ -0,0 +1,388 @@
|
||||
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
|
||||
Reference in New Issue
Block a user