Spaces:
Sleeping
Sleeping
File size: 7,144 Bytes
3a3f97e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 |
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
import subprocess
import os
import signal
import sys
from typing import Any, Dict, Tuple
from aworld.config.conf import ToolConfig
from examples.tools.tool_action import ShellAction
from aworld.core.common import ActionModel, Observation, ActionResult
from aworld.core.tool.base import Tool, AgentInput, ToolFactory
from aworld.logs.util import logger
from aworld.tools.utils import build_observation
@ToolFactory.register(name='shell',
desc="shell execute tool",
supported_action=ShellAction,
conf_file_name=f'shell_tool.yaml')
class ShellTool(Tool):
"""
used to execute shell commands, providing initialization, execution, and exit functions.
"""
def __init__(self, conf: ToolConfig, **kwargs) -> None:
"""
Initialize the ShellTool
Args:
conf: tool config
**kwargs: -
"""
super(ShellTool, self).__init__(conf, **kwargs)
self.type = "function"
self.working_dir = self.conf.get('working_dir')
self.env = self.conf.get('env') if self.conf.get('env') else os.environ.copy()
self.processes = []
def reset(self, *, seed: int | None = None, options: Dict[str, str] | None = None) -> Tuple[
AgentInput, dict[str, Any]]:
"""
Reset the executor
Args:
seed: -
options: -
Returns:
AgentInput, dict[str, Any]: -
"""
self.working_dir = None
self.env = os.environ.copy()
self.processes = []
self._finished = False
return build_observation(observer=self.name(),
ability=ShellAction.EXECUTE_SCRIPT.value.name), {}
def close(self) -> None:
"""
Close the executor
Returns:
None
"""
try:
for process in self.processes:
# Check whether the process is still running
if process.poll() is None:
try:
# Try to gracefully terminate the process
if sys.platform != "win32":
os.kill(process.pid, signal.SIGTERM)
else:
process.terminate()
except Exception as e:
logger.warning(f"An error occurred while terminating the process. e: {str(e)}")
except Exception as e:
logger.warning(f"Error while exiting Shell Executor. e: {str(e)}")
finally:
# Clear process list
self.processes = []
self._finished = True
def do_step(self,
actions: list[ActionModel],
**kwargs) -> Tuple[Observation, float, bool, bool, dict[str, Any]]:
"""
Step the executor
Args:
actions: actions
**kwargs: -
Returns:
Observation, float, bool, bool, dict[str, Any]: -
"""
self._finished = False
reward = 0
fail_error = ""
observation = build_observation(observer=self.name(),
ability=ShellAction.EXECUTE_SCRIPT.value.name)
try:
if not actions:
return (observation, reward,
kwargs.get("terminated",
False), kwargs.get("truncated", False), {
"exception": "actions is empty"
})
for action in actions:
cmd_string = action.params.get("command", "")
if not cmd_string:
continue
_, output, error = self.execute(cmd_string)
observation.content = output
observation.action_result.append(
ActionResult(is_done=True,
success=False if error else True,
content=output,
error=error,
keep=False))
reward = 1
except Exception as e:
fail_error = str(e)
finally:
self._finished = True
info = {"exception": fail_error}
info.update(kwargs)
return (observation,
reward,
kwargs.get("terminated", False),
kwargs.get("truncated", False),
info)
def execute(self, script: str, capture_output: bool = True, timeout: int = 5):
"""
exec shell script
Args:
script (str): shell script to execute
capture_output (bool): whether to capture the script output
timeout (int, optional): Command execution timeout (seconds)
Returns:
dict: action result
"""
try:
if capture_output:
process_ = subprocess.run(
script,
shell=True,
cwd=self.working_dir,
env=self.env,
timeout=timeout,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True
)
return {
'success': process_.returncode == 0,
'return_code': process_.returncode,
'stdout': process_.stdout,
'stderr': process_.stderr,
'script': script
}
else:
process_ = subprocess.Popen(
script,
shell=True,
cwd=self.working_dir,
env=self.env
)
self.processes.append(process_)
process_.wait(timeout=timeout)
return {
'success': process_.returncode == 0,
'return_code': process_.returncode,
'script': script
}
except subprocess.TimeoutExpired:
return {
'success': False,
'error': 'Timeout',
'script': script
}
except Exception as e:
return {
'success': False,
'error': str(e),
'script': script
}
def execute_async(self, script: str):
"""
Execute shell script asynchronously (no waiting)
Args:
script (str): The shell script to execute
Returns:
subprocess.Popen: Process object
"""
try:
process_ = subprocess.Popen(
script,
shell=True,
cwd=self.working_dir,
env=self.env
)
self.processes.append(process_)
return process_
except Exception as e:
logger.warning(f"An error occurred while executing the script asynchronously. e: {str(e)}")
return None
|