diff --git a/openai_server/autogen_utils.py b/openai_server/autogen_utils.py index b49dc03df..97423d6b6 100644 --- a/openai_server/autogen_utils.py +++ b/openai_server/autogen_utils.py @@ -1,17 +1,32 @@ import logging import os import re +import subprocess import sys import typing -from typing import List +from hashlib import md5 +from pathlib import Path +from typing import Any, Callable, ClassVar, Dict, List, Optional, Union +from types import SimpleNamespace import uuid +from autogen.code_utils import PYTHON_VARIANTS, WIN32, _cmd, TIMEOUT_MSG from autogen.coding import LocalCommandLineCodeExecutor, CodeBlock from autogen.coding.base import CommandLineCodeResult from autogen import ConversableAgent from autogen import GroupChatManager import backoff +from autogen.coding.func_with_reqs import ( + FunctionWithRequirements, + FunctionWithRequirementsStr, +) +from autogen.coding.utils import silence_pip, _get_file_name_from_content + +from typing_extensions import ParamSpec + +A = ParamSpec("A") + from openai_server.autogen_streaming import iostream_generator from openai_server.backend_utils import convert_gen_kwargs from openai_server.agent_utils import in_pycharm, set_python_path @@ -23,6 +38,19 @@ class H2OLocalCommandLineCodeExecutor(LocalCommandLineCodeExecutor): + def __init__( + self, + timeout: int = 60, + virtual_env_context: Optional[SimpleNamespace] = None, + work_dir: Union[Path, str] = Path("."), + functions: List[Union[FunctionWithRequirements[Any, A], Callable[..., Any], FunctionWithRequirementsStr]] = [], + functions_module: str = "functions", + execution_policies: Optional[Dict[str, bool]] = None, + autogen_code_restrictions_level: int = 2, + ): + super().__init__(timeout, virtual_env_context, work_dir, functions, functions_module, execution_policies) + self.autogen_code_restrictions_level = autogen_code_restrictions_level + @staticmethod def remove_comments_strings(code: str, lang: str) -> str: if verbose: @@ -145,6 +173,87 @@ def sanitize_command(lang: str, code: str) -> None: if match.group(f"pat{i}"): raise ValueError(f"{danger_mark}: {patterns[pattern]}\n\n{cleaned_code}") + def __execute_code_dont_check_setup(self, code_blocks: List[CodeBlock]) -> CommandLineCodeResult: + # nearly identical to parent, but with control over guardrails via self.sanitize_command + logs_all = "" + file_names = [] + exitcode = -2 + for code_block in code_blocks: + lang, code = code_block.language, code_block.code + lang = lang.lower() + + if self.autogen_code_restrictions_level >= 2: + self.sanitize_command(lang, code) + elif self.autogen_code_restrictions_level == 1: + LocalCommandLineCodeExecutor.sanitize_command(lang, code) + code = silence_pip(code, lang) + + if lang in PYTHON_VARIANTS: + lang = "python" + + if WIN32 and lang in ["sh", "shell"]: + lang = "ps1" + + if lang not in self.SUPPORTED_LANGUAGES: + # In case the language is not supported, we return an error message. + exitcode = 1 + logs_all += "\n" + f"unknown language {lang}" + break + + execute_code = self.execution_policies.get(lang, False) + try: + # Check if there is a filename comment + filename = _get_file_name_from_content(code, self._work_dir) + except ValueError: + return CommandLineCodeResult(exit_code=1, output="Filename is not in the workspace") + + if filename is None: + # create a file with an automatically generated name + code_hash = md5(code.encode()).hexdigest() + filename = f"tmp_code_{code_hash}.{'py' if lang.startswith('python') else lang}" + written_file = (self._work_dir / filename).resolve() + with written_file.open("w", encoding="utf-8") as f: + f.write(code) + file_names.append(written_file) + + if not execute_code: + # Just return a message that the file is saved. + logs_all += f"Code saved to {str(written_file)}\n" + exitcode = 0 + continue + + program = _cmd(lang) + cmd = [program, str(written_file.absolute())] + env = os.environ.copy() + + if self._virtual_env_context: + virtual_env_abs_path = os.path.abspath(self._virtual_env_context.bin_path) + path_with_virtualenv = rf"{virtual_env_abs_path}{os.pathsep}{env['PATH']}" + env["PATH"] = path_with_virtualenv + if WIN32: + activation_script = os.path.join(virtual_env_abs_path, "activate.bat") + cmd = [activation_script, "&&", *cmd] + + try: + result = subprocess.run( + cmd, cwd=self._work_dir, capture_output=True, text=True, timeout=float(self._timeout), env=env + ) + except subprocess.TimeoutExpired: + logs_all += "\n" + TIMEOUT_MSG + # Same exit code as the timeout command on linux. + exitcode = 124 + break + + logs_all += result.stderr + logs_all += result.stdout + exitcode = result.returncode + + if exitcode != 0: + break + + code_file = str(file_names[0]) if len(file_names) > 0 else None + return CommandLineCodeResult(exit_code=exitcode, output=logs_all, code_file=code_file) + def _execute_code_dont_check_setup(self, code_blocks: List[CodeBlock]) -> CommandLineCodeResult: try: # skip code blocks with # execution: false @@ -164,15 +273,16 @@ def _execute_code_dont_check_setup(self, code_blocks: List[CodeBlock]) -> Comman os.environ['TERM'] = 'dumb' """ + code_block.code - ret = super()._execute_code_dont_check_setup(code_blocks) - except Exception as e: - if 'exitcode' in str(e) and 'local variable' in str(e): - return CommandLineCodeResult(exit_code=0, + ret = self.__execute_code_dont_check_setup(code_blocks) + + if ret.exit_code == -2 and len(code_blocks) > 0: + ret = CommandLineCodeResult(exit_code=0, output='Code block present, but no code executed (execution tag was false or not present for all code blocks). This is expected if you had code blocks but they were not meant for python or shell execution. For example, you may have shown code for demonstration purposes. If this is expected, then move on normally without concern.') + except Exception as e: if danger_mark in str(e): print(f"Code Danger Error: {e}\n\n{code_blocks}", file=sys.stderr) # dont' fail, just return the error so LLM can adjust - return CommandLineCodeResult(exit_code=1, output=str(e)) + ret = CommandLineCodeResult(exit_code=1, output=str(e)) else: raise try: @@ -181,7 +291,7 @@ def _execute_code_dont_check_setup(self, code_blocks: List[CodeBlock]) -> Comman if bad_output_mark in str(e): print(f"Code Output Danger Error: {e}\n\n{code_blocks}\n\n{ret}", file=sys.stderr) # dont' fail, just return the error so LLM can adjust - return CommandLineCodeResult(exit_code=1, output=str(e)) + ret = CommandLineCodeResult(exit_code=1, output=str(e)) else: raise ret = self.truncate_output(ret) @@ -290,10 +400,6 @@ def truncate_output(ret: CommandLineCodeResult) -> CommandLineCodeResult: return ret -# override the original method with the new one -# required because original class does not use super() but references its own method -LocalCommandLineCodeExecutor.sanitize_command = H2OLocalCommandLineCodeExecutor.sanitize_command - error_patterns = [ r"Rate limit reached", r"Connection timeout", diff --git a/src/version.py b/src/version.py index 269f37101..cd8e47fc9 100644 --- a/src/version.py +++ b/src/version.py @@ -1 +1 @@ -__version__ = "122332ef576358589f3dff64301e7ea0622870f8" +__version__ = "34d2ee55eee948203c15f6a3cd355d347b0029ae"