53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
from typing import Any
|
|
|
|
from duck_core.tools.base import Tool, ToolResult
|
|
from duck_core.tools.coder import CoderTool
|
|
from duck_core.tools.file_read import FileReadTool
|
|
from duck_core.tools.file_write import FileWriteTool
|
|
from duck_core.tools.list_dir import ListDirTool
|
|
from duck_core.tools.search_files import SearchFilesTool
|
|
from duck_core.tools.shell_exec_safe import ShellExecSafeTool
|
|
|
|
|
|
class ToolGateway:
|
|
def __init__(self, tools: list[Tool]):
|
|
self.tools = {tool.name: tool for tool in tools}
|
|
|
|
@classmethod
|
|
def default(cls, workspace: str, model_client: Any = None) -> "ToolGateway":
|
|
return cls(
|
|
[
|
|
FileReadTool(workspace),
|
|
FileWriteTool(workspace),
|
|
ListDirTool(workspace),
|
|
SearchFilesTool(workspace),
|
|
ShellExecSafeTool(workspace),
|
|
CoderTool(model_client=model_client),
|
|
]
|
|
)
|
|
|
|
def with_model_client(self, model_client: Any) -> "ToolGateway":
|
|
"""Return a new gateway with model-dependent tools configured."""
|
|
new_tools = list(self.tools.values())
|
|
# Replace coder tool with one that has model_client
|
|
new_tools = [
|
|
t for t in new_tools if not isinstance(t, CoderTool)
|
|
] + [CoderTool(model_client=model_client)]
|
|
return self.__class__(new_tools)
|
|
|
|
async def run_action(
|
|
self, action: dict[str, Any], approved: bool = False, password: str | None = None
|
|
) -> ToolResult:
|
|
tool_name = str(action.get("tool", ""))
|
|
tool = self.tools.get(tool_name)
|
|
if tool is None:
|
|
return ToolResult(ok=False, error=f"Unknown tool: {tool_name}")
|
|
args = action.get("args") or {}
|
|
if not isinstance(args, dict):
|
|
return ToolResult(ok=False, error="Tool args must be an object")
|
|
if approved:
|
|
args = {**args, "_approved": True}
|
|
if password is not None:
|
|
args = {**args, "_password": password}
|
|
return await tool.run(args)
|