Spaces:
Sleeping
Sleeping
File size: 8,544 Bytes
5fad09e |
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 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 |
import sys
import ast
import re
import subprocess
from typing import Any, Dict, Tuple, List
from io import StringIO
from aworld.logs.util import logger
from aworld.config.conf import ToolConfig
from examples.tools.tool_action import PythonToolAction
from aworld.core.common import ActionModel, Observation, ActionResult
from aworld.core.tool.base import Tool, AgentInput, ToolFactory
from aworld.utils import import_package
from aworld.tools.utils import build_observation
@ToolFactory.register(name="python_execute",
desc="python interpreter tool",
supported_action=PythonToolAction,
conf_file_name=f'python_execute_tool.yaml')
class PythonTool(Tool):
def __init__(self,
conf: ToolConfig,
**kwargs) -> None:
"""
Initialize the PythonExecutor
Args:
conf: tool config
**kwargs: -
Return:
None
"""
super(PythonTool, self).__init__(conf, **kwargs)
self.type = "function"
self.local_namespace = {}
self.global_namespace = {}
self.original_stdout = sys.stdout
self.output_buffer = StringIO()
self.installed_packages = set()
import_package('langchain_experimental')
from langchain_experimental.utilities.python import PythonREPL
self.python_repl = PythonREPL()
def extract_imports(self, code: str) -> set:
"""
Extract import statements
Args:
code: python code
Returns:
set: import statements
"""
imports = set()
try:
tree = ast.parse(code)
for node in ast.walk(tree):
if isinstance(node, ast.Import):
# deal import xxx or import xxx as yyy
for name in node.names:
package_name = name.name.split('.')[0]
imports.add(package_name)
elif isinstance(node, ast.ImportFrom):
# deal from xxx import yyy or from xxx.yyy import zzz
if node.module:
package_name = node.module.split('.')[0]
imports.add(package_name)
except SyntaxError:
import_pattern = r'^import\s+([\w\s,]+)|from\s+(\w+)'
for line in code.split('\n'):
line = line.strip()
match = re.match(import_pattern, line)
if match:
if match.group(1):
packages = [p.strip() for p in match.group(1).split(',')]
for package in packages:
if package:
package_name = package.split()[0]
imports.add(package_name)
elif match.group(2):
imports.add(match.group(2))
return imports
def install_dependencies(self,
packages: set) -> None:
"""
Install dependency packages
Args:
packages: python third packages
Returns:
None
"""
for package in packages:
try:
__import__(package)
except ImportError:
try:
subprocess.check_call([sys.executable, "-m", "pip", "install", package])
self.installed_packages.add(package)
except subprocess.CalledProcessError as e:
logger.warning(f"Failed to install {package}: {str(e)}")
def uninstall_dependencies(self) -> None:
"""
Uninstall dependency packages
Args:
-
Returns:
None
"""
try:
for package in self.installed_packages:
try:
subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", package])
except subprocess.CalledProcessError as e:
logger.warning(f"Failed to uninstall {package}: {str(e)}")
self.installed_packages.clear()
except Exception as e:
logger.warning(f"Failed to uninstall dependencies: {repr(e)}")
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.close()
self.local_namespace = {}
self.global_namespace = {}
self._finished = False
self.installed_packages.clear()
return build_observation(observer=self.name(),
ability=PythonToolAction.EXECUTE.value.name), {}
def close(self) -> None:
"""
Close the executor
Returns:
None
"""
try:
self.uninstall_dependencies()
sys.stdout = self.original_stdout
self.output_buffer.close()
self.local_namespace.clear()
self.global_namespace.clear()
except:
pass
finally:
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.step_finished = False
reward = 0
fail_error = ""
observation = build_observation(observer=self.name(),
ability=PythonToolAction.EXECUTE.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:
code = action.params.get("code", "")
if not code:
logger.warning(f"{action} no code to execute.")
continue
try:
_, output, error = self.execute(code)
observation.content = output
except Exception as e:
error = str(e)
output = error
observation.action_result.append(
ActionResult(is_done=True,
success=False if error else True,
content=f"{output}",
error=f"{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, code, timeout=300):
"""
Execute the code
Args:
code: python code
timeout: timeout seconds
Returns:
result, output, error
"""
required_packages = self.extract_imports(code)
self.install_dependencies(required_packages)
self.python_repl.globals = self.global_namespace
self.python_repl.locals = self.local_namespace
error = None
try:
output = self.python_repl.run(code, timeout)
except Exception as e:
error = f'{repr(e)}'
finally:
self.uninstall_dependencies()
return '', output, error
def get_execute_result(self):
"""
Get the execute result
Returns:
output, error
"""
output = None
error = ''
try:
output = self.output_buffer.getvalue()
self.output_buffer.truncate(0)
self.output_buffer.seek(0)
sys.stdout = self.original_stdout
except Exception as e:
error = f'{repr(e)}'
logger.warning(f"Failed to get output, {repr(e)}")
return output, error
|