aphilippov/python-server-api
0
1import logging2import os3import pathlib4import re5import subprocess6import sys7import time8from concurrent.futures import ThreadPoolExecutor, TimeoutError9from hashlib import md510from typing import Callable, Dict, List, Optional, Tuple, Union11 12from autogen import oai13 14try:15 import docker16except ImportError:17 docker = None18 19DEFAULT_MODEL = "gpt-4"20FAST_MODEL = "gpt-3.5-turbo"21# Regular expression for finding a code block22# ```[ \t]*(\w+)?[ \t]*\r?\n(.*?)[ \t]*\r?\n``` Matches multi-line code blocks.23# The [ \t]* matches the potential spaces before language name.24# The (\w+)? matches the language, where the ? indicates it is optional.25# The [ \t]* matches the potential spaces (not newlines) after language name.26# The \r?\n makes sure there is a linebreak after ```.27# The (.*?) matches the code itself (non-greedy).28# The \r?\n makes sure there is a linebreak before ```.29# The [ \t]* matches the potential spaces before closing ``` (the spec allows indentation).30CODE_BLOCK_PATTERN = r"```[ \t]*(\w+)?[ \t]*\r?\n(.*?)\r?\n[ \t]*```"31WORKING_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extensions")32UNKNOWN = "unknown"33TIMEOUT_MSG = "Timeout"34DEFAULT_TIMEOUT = 60035WIN32 = sys.platform == "win32"36PATH_SEPARATOR = WIN32 and "\\" or "/"37 38logger = logging.getLogger(__name__)39 40 41def content_str(content: Union[str, List, None]) -> str:42 """Converts `content` into a string format.43 44 This function processes content that may be a string, a list of mixed text and image URLs, or None,45 and converts it into a string. Text is directly appended to the result string, while image URLs are46 represented by a placeholder image token. If the content is None, an empty string is returned.47 48 Args:49 - content (Union[str, List, None]): The content to be processed. Can be a string, a list of dictionaries50 representing text and image URLs, or None.51 52 Returns:53 str: A string representation of the input content. Image URLs are replaced with an image token.54 55 Note:56 - The function expects each dictionary in the list to have a "type" key that is either "text" or "image_url".57 For "text" type, the "text" key's value is appended to the result. For "image_url", an image token is appended.58 - This function is useful for handling content that may include both text and image references, especially59 in contexts where images need to be represented as placeholders.60 """61 if content is None:62 return ""63 if isinstance(content, str):64 return content65 if not isinstance(content, list):66 raise TypeError(f"content must be None, str, or list, but got {type(content)}")67 68 rst = ""69 for item in content:70 if not isinstance(item, dict):71 raise TypeError("Wrong content format: every element should be dict if the content is a list.")72 assert "type" in item, "Wrong content format. Missing 'type' key in content's dict."73 if item["type"] == "text":74 rst += item["text"]75 elif item["type"] == "image_url":76 rst += "<image>"77 else:78 raise ValueError(f"Wrong content format: unknown type {item['type']} within the content")79 return rst80 81 82def infer_lang(code):83 """infer the language for the code.84 TODO: make it robust.85 """86 if code.startswith("python ") or code.startswith("pip") or code.startswith("python3 "):87 return "sh"88 89 # check if code is a valid python code90 try:91 compile(code, "test", "exec")92 return "python"93 except SyntaxError:94 # not a valid python code95 return UNKNOWN96 97 98# TODO: In the future move, to better support https://spec.commonmark.org/0.30/#fenced-code-blocks99# perhaps by using a full Markdown parser.100def extract_code(101 text: Union[str, List], pattern: str = CODE_BLOCK_PATTERN, detect_single_line_code: bool = False102) -> List[Tuple[str, str]]:103 """Extract code from a text.104 105 Args:106 text (str or List): The content to extract code from. The content can be107 a string or a list, as returned by standard GPT or multimodal GPT.108 pattern (str, optional): The regular expression pattern for finding the109 code block. Defaults to CODE_BLOCK_PATTERN.110 detect_single_line_code (bool, optional): Enable the new feature for111 extracting single line code. Defaults to False.112 113 Returns:114 list: A list of tuples, each containing the language and the code.115 If there is no code block in the input text, the language would be "unknown".116 If there is code block but the language is not specified, the language would be "".117 """118 text = content_str(text)119 if not detect_single_line_code:120 match = re.findall(pattern, text, flags=re.DOTALL)121 return match if match else [(UNKNOWN, text)]122 123 # Extract both multi-line and single-line code block, separated by the | operator124 # `([^`]+)`: Matches inline code.125 code_pattern = re.compile(CODE_BLOCK_PATTERN + r"|`([^`]+)`")126 code_blocks = code_pattern.findall(text)127 128 # Extract the individual code blocks and languages from the matched groups129 extracted = []130 for lang, group1, group2 in code_blocks:131 if group1:132 extracted.append((lang.strip(), group1.strip()))133 elif group2:134 extracted.append(("", group2.strip()))135 136 return extracted137 138 139def generate_code(pattern: str = CODE_BLOCK_PATTERN, **config) -> Tuple[str, float]:140 """(openai<1) Generate code.141 142 Args:143 pattern (Optional, str): The regular expression pattern for finding the code block.144 The default pattern is for finding a code block in a markdown file.145 config (Optional, dict): The configuration for the API call.146 147 Returns:148 str: The generated code.149 float: The cost of the generation.150 """151 response = oai.Completion.create(**config)152 return extract_code(oai.Completion.extract_text(response)[0], pattern), response["cost"]153 154 155_IMPROVE_FUNCTION_CONFIG = {156 "prompt": """Improve the function '{func_name}' to achieve the objective '{objective}'.157The current implementation of the function is as follows:158{file_string}""",159 "model": DEFAULT_MODEL,160 "request_timeout": 600,161}162 163 164def improve_function(file_name, func_name, objective, **config):165 """(openai<1) Improve the function to achieve the objective."""166 params = {**_IMPROVE_FUNCTION_CONFIG, **config}167 # read the entire file into a str168 with open(file_name, "r") as f:169 file_string = f.read()170 response = oai.Completion.create(171 {"func_name": func_name, "objective": objective, "file_string": file_string}, **params172 )173 return oai.Completion.extract_text(response)[0], response["cost"]174 175 176_IMPROVE_CODE_CONFIG = {177 "prompt": """Analyze the code in the following files and return a list of suggestions for improvement{followup}, to achieve the objective of '{objective}'.178{code}179""",180 "model": DEFAULT_MODEL,181 "request_timeout": 900,182}183 184 185def improve_code(files, objective, suggest_only=True, **config):186 """(openai<1) Improve the code to achieve a given objective.187 188 Args:189 files (list): A list of file names containing the source code.190 objective (str): The objective to achieve.191 suggest_only (bool): Whether to return only the suggestions or the improved code.192 config (Optional, dict): The configuration for the API call.193 194 Returns:195 str: The improved code if suggest_only=False; a list of suggestions if suggest_only=True (default).196 float: The cost of the generation.197 """198 code = ""199 for file_name in files:200 # read the entire file into a string201 with open(file_name, "r") as f:202 file_string = f.read()203 code += f"""{file_name}:204{file_string}205 206"""207 params = {**_IMPROVE_CODE_CONFIG, **config}208 followup = "" if suggest_only else " followed by the improved code"209 response = oai.Completion.create({"objective": objective, "code": code, "followup": followup}, **params)210 return oai.Completion.extract_text(response)[0], response["cost"]211 212 213def timeout_handler(signum, frame):214 raise TimeoutError("Timed out!")215 216 217def _cmd(lang):218 if lang.startswith("python") or lang in ["bash", "sh", "powershell"]:219 return lang220 if lang in ["shell"]:221 return "sh"222 if lang in ["ps1"]:223 return "powershell"224 raise NotImplementedError(f"{lang} not recognized in code execution")225 226 227def execute_code(228 code: Optional[str] = None,229 timeout: Optional[int] = None,230 filename: Optional[str] = None,231 work_dir: Optional[str] = None,232 use_docker: Optional[Union[List[str], str, bool]] = None,233 lang: Optional[str] = "python",234) -> Tuple[int, str, str]:235 """Execute code in a docker container.236 This function is not tested on MacOS.237 238 Args:239 code (Optional, str): The code to execute.240 If None, the code from the file specified by filename will be executed.241 Either code or filename must be provided.242 timeout (Optional, int): The maximum execution time in seconds.243 If None, a default timeout will be used. The default timeout is 600 seconds. On Windows, the timeout is not enforced when use_docker=False.244 filename (Optional, str): The file name to save the code or where the code is stored when `code` is None.245 If None, a file with a randomly generated name will be created.246 The randomly generated file will be deleted after execution.247 The file name must be a relative path. Relative paths are relative to the working directory.248 work_dir (Optional, str): The working directory for the code execution.249 If None, a default working directory will be used.250 The default working directory is the "extensions" directory under251 "path_to_autogen".252 use_docker (Optional, list, str or bool): The docker image to use for code execution.253 If a list or a str of image name(s) is provided, the code will be executed in a docker container254 with the first image successfully pulled.255 If None, False or empty, the code will be executed in the current environment.256 Default is None, which will be converted into an empty list when docker package is available.257 Expected behaviour:258 - If `use_docker` is explicitly set to True and the docker package is available, the code will run in a Docker container.259 - If `use_docker` is explicitly set to True but the Docker package is missing, an error will be raised.260 - If `use_docker` is not set (i.e., left default to None) and the Docker package is not available, a warning will be displayed, but the code will run natively.261 If the code is executed in the current environment,262 the code must be trusted.263 lang (Optional, str): The language of the code. Default is "python".264 265 Returns:266 int: 0 if the code executes successfully.267 str: The error message if the code fails to execute; the stdout otherwise.268 image: The docker image name after container run when docker is used.269 """270 if all((code is None, filename is None)):271 error_msg = f"Either {code=} or {filename=} must be provided."272 logger.error(error_msg)273 raise AssertionError(error_msg)274 275 if use_docker and docker is None:276 error_msg = "Cannot use docker because the python docker package is not available."277 logger.error(error_msg)278 raise AssertionError(error_msg)279 280 # Warn if use_docker was unspecified (or None), and cannot be provided (the default).281 # In this case the current behavior is to fall back to run natively, but this behavior282 # is subject to change.283 if use_docker is None:284 if docker is None:285 use_docker = False286 logger.warning(287 "execute_code was called without specifying a value for use_docker. Since the python docker package is not available, code will be run natively. Note: this fallback behavior is subject to change"288 )289 else:290 # Default to true291 use_docker = True292 293 timeout = timeout or DEFAULT_TIMEOUT294 original_filename = filename295 if WIN32 and lang in ["sh", "shell"] and (not use_docker):296 lang = "ps1"297 if filename is None:298 code_hash = md5(code.encode()).hexdigest()299 # create a file with a automatically generated name300 filename = f"tmp_code_{code_hash}.{'py' if lang.startswith('python') else lang}"301 if work_dir is None:302 work_dir = WORKING_DIR303 filepath = os.path.join(work_dir, filename)304 file_dir = os.path.dirname(filepath)305 os.makedirs(file_dir, exist_ok=True)306 if code is not None:307 with open(filepath, "w", encoding="utf-8") as fout:308 fout.write(code)309 # check if already running in a docker container310 in_docker_container = os.path.exists("/.dockerenv")311 if not use_docker or in_docker_container:312 # already running in a docker container313 cmd = [314 sys.executable if lang.startswith("python") else _cmd(lang),315 f".\\{filename}" if WIN32 else filename,316 ]317 if WIN32:318 logger.warning("SIGALRM is not supported on Windows. No timeout will be enforced.")319 result = subprocess.run(320 cmd,321 cwd=work_dir,322 capture_output=True,323 text=True,324 )325 else:326 with ThreadPoolExecutor(max_workers=1) as executor:327 future = executor.submit(328 subprocess.run,329 cmd,330 cwd=work_dir,331 capture_output=True,332 text=True,333 )334 try:335 result = future.result(timeout=timeout)336 except TimeoutError:337 if original_filename is None:338 os.remove(filepath)339 return 1, TIMEOUT_MSG, None340 if original_filename is None:341 os.remove(filepath)342 if result.returncode:343 logs = result.stderr344 if original_filename is None:345 abs_path = str(pathlib.Path(filepath).absolute())346 logs = logs.replace(str(abs_path), "").replace(filename, "")347 else:348 abs_path = str(pathlib.Path(work_dir).absolute()) + PATH_SEPARATOR349 logs = logs.replace(str(abs_path), "")350 else:351 logs = result.stdout352 return result.returncode, logs, None353 354 # create a docker client355 client = docker.from_env()356 image_list = (357 ["python:3-alpine", "python:3", "python:3-windowsservercore"]358 if use_docker is True359 else [use_docker]360 if isinstance(use_docker, str)361 else use_docker362 )363 for image in image_list:364 # check if the image exists365 try:366 client.images.get(image)367 break368 except docker.errors.ImageNotFound:369 # pull the image370 print("Pulling image", image)371 try:372 client.images.pull(image)373 break374 except docker.errors.DockerException:375 print("Failed to pull image", image)376 # get a randomized str based on current time to wrap the exit code377 exit_code_str = f"exitcode{time.time()}"378 abs_path = pathlib.Path(work_dir).absolute()379 cmd = [380 "sh",381 "-c",382 f"{_cmd(lang)} {filename}; exit_code=$?; echo -n {exit_code_str}; echo -n $exit_code; echo {exit_code_str}",383 ]384 # create a docker container385 container = client.containers.run(386 image,387 command=cmd,388 working_dir="/workspace",389 detach=True,390 # get absolute path to the working directory391 volumes={abs_path: {"bind": "/workspace", "mode": "rw"}},392 )393 start_time = time.time()394 while container.status != "exited" and time.time() - start_time < timeout:395 # Reload the container object396 container.reload()397 if container.status != "exited":398 container.stop()399 container.remove()400 if original_filename is None:401 os.remove(filepath)402 return 1, TIMEOUT_MSG, image403 # get the container logs404 logs = container.logs().decode("utf-8").rstrip()405 # commit the image406 tag = filename.replace("/", "")407 container.commit(repository="python", tag=tag)408 # remove the container409 container.remove()410 # check if the code executed successfully411 exit_code = container.attrs["State"]["ExitCode"]412 if exit_code == 0:413 # extract the exit code from the logs414 pattern = re.compile(f"{exit_code_str}(\\d+){exit_code_str}")415 match = pattern.search(logs)416 exit_code = 1 if match is None else int(match.group(1))417 # remove the exit code from the logs418 logs = logs if match is None else pattern.sub("", logs)419 420 if original_filename is None:421 os.remove(filepath)422 if exit_code:423 logs = logs.replace(f"/workspace/{filename if original_filename is None else ''}", "")424 # return the exit code, logs and image425 return exit_code, logs, f"python:{tag}"426 427 428_GENERATE_ASSERTIONS_CONFIG = {429 "prompt": """Given the signature and docstring, write the exactly same number of assertion(s) for the provided example(s) in the docstring, without assertion messages.430 431func signature:432{definition}433assertions:""",434 "model": FAST_MODEL,435 "max_tokens": 256,436 "stop": "\n\n",437}438 439 440def generate_assertions(definition: str, **config) -> Tuple[str, float]:441 """(openai<1) Generate assertions for a function.442 443 Args:444 definition (str): The function definition, including the signature and docstr.445 config (Optional, dict): The configuration for the API call.446 447 Returns:448 str: The generated assertions.449 float: The cost of the generation.450 """451 params = {**_GENERATE_ASSERTIONS_CONFIG, **config}452 response = oai.Completion.create(453 {"definition": definition},454 **params,455 )456 assertions = oai.Completion.extract_text(response)[0]457 return assertions, response["cost"]458 459 460def _remove_check(response):461 """Remove the check function from the response."""462 # find the position of the check function463 pos = response.find("def check(")464 if pos == -1:465 return response466 return response[:pos]467 468 469def eval_function_completions(470 responses: List[str],471 definition: str,472 test: Optional[str] = None,473 entry_point: Optional[str] = None,474 assertions: Optional[Union[str, Callable[[str], Tuple[str, float]]]] = None,475 timeout: Optional[float] = 3,476 use_docker: Optional[bool] = True,477) -> Dict:478 """(openai<1) Select a response from a list of responses for the function completion task (using generated assertions), and/or evaluate if the task is successful using a gold test.479 480 Args:481 responses (list): The list of responses.482 definition (str): The input definition.483 test (Optional, str): The test code.484 entry_point (Optional, str): The name of the function.485 assertions (Optional, str or Callable): The assertion code which serves as a filter of the responses, or an assertion generator.486 When provided, only the responses that pass the assertions will be considered for the actual test (if provided).487 timeout (Optional, float): The timeout for executing the code.488 489 Returns:490 dict: The success metrics.491 """492 n = len(responses)493 if assertions is None:494 # no assertion filter495 success_list = []496 for i in range(n):497 response = _remove_check(responses[i])498 code = (499 f"{response}\n{test}\ncheck({entry_point})"500 if response.startswith("def")501 else f"{definition}{response}\n{test}\ncheck({entry_point})"502 )503 success = execute_code(code, timeout=timeout, use_docker=use_docker)[0] == 0504 success_list.append(success)505 return {506 "expected_success": 1 - pow(1 - sum(success_list) / n, n),507 "success": any(s for s in success_list),508 }509 if callable(assertions) and n > 1:510 # assertion generator511 assertions, gen_cost = assertions(definition)512 else:513 assertions, gen_cost = None, 0514 if n > 1 or test is None:515 for i in range(n):516 response = responses[i] = _remove_check(responses[i])517 code = (518 f"{response}\n{assertions}" if response.startswith("def") else f"{definition}{response}\n{assertions}"519 )520 succeed_assertions = execute_code(code, timeout=timeout, use_docker=use_docker)[0] == 0521 if succeed_assertions:522 break523 else:524 # just test, no need to check assertions525 succeed_assertions = False526 i, response = 0, responses[0]527 if test is None:528 # no test code529 return {530 "index_selected": i,531 "succeed_assertions": succeed_assertions,532 "gen_cost": gen_cost,533 "assertions": assertions,534 }535 code_test = (536 f"{response}\n{test}\ncheck({entry_point})"537 if response.startswith("def")538 else f"{definition}{response}\n{test}\ncheck({entry_point})"539 )540 success = execute_code(code_test, timeout=timeout, use_docker=use_docker)[0] == 0541 return {542 "index_selected": i,543 "succeed_assertions": succeed_assertions,544 "success": success,545 "gen_cost": gen_cost,546 "assertions": assertions,547 }548 549 550_FUNC_COMPLETION_PROMPT = "# Python 3{definition}"551_FUNC_COMPLETION_STOP = ["\nclass", "\ndef", "\nif", "\nprint"]552_IMPLEMENT_CONFIGS = [553 {"model": FAST_MODEL, "prompt": _FUNC_COMPLETION_PROMPT, "temperature": 0, "cache_seed": 0},554 {"model": FAST_MODEL, "prompt": _FUNC_COMPLETION_PROMPT, "stop": _FUNC_COMPLETION_STOP, "n": 7, "cache_seed": 0},555 {"model": DEFAULT_MODEL, "prompt": _FUNC_COMPLETION_PROMPT, "temperature": 0, "cache_seed": 1},556 {"model": DEFAULT_MODEL, "prompt": _FUNC_COMPLETION_PROMPT, "stop": _FUNC_COMPLETION_STOP, "n": 2, "cache_seed": 2},557 {"model": DEFAULT_MODEL, "prompt": _FUNC_COMPLETION_PROMPT, "stop": _FUNC_COMPLETION_STOP, "n": 1, "cache_seed": 2},558]559 560 561class PassAssertionFilter:562 def __init__(self, assertions):563 self._assertions = assertions564 self.cost = 0565 self.metrics = self.responses = None566 567 def pass_assertions(self, context, response, **_):568 """(openai<1) Check if the response passes the assertions."""569 responses = oai.Completion.extract_text(response)570 metrics = eval_function_completions(responses, context["definition"], assertions=self._assertions)571 self._assertions = metrics["assertions"]572 self.cost += metrics["gen_cost"]573 self.metrics = metrics574 self.responses = responses575 return metrics["succeed_assertions"]576 577 578def implement(579 definition: str,580 configs: Optional[List[Dict]] = None,581 assertions: Optional[Union[str, Callable[[str], Tuple[str, float]]]] = generate_assertions,582) -> Tuple[str, float]:583 """(openai<1) Implement a function from a definition.584 585 Args:586 definition (str): The function definition, including the signature and docstr.587 configs (list): The list of configurations for completion.588 assertions (Optional, str or Callable): The assertion code which serves as a filter of the responses, or an assertion generator.589 590 Returns:591 str: The implementation.592 float: The cost of the implementation.593 int: The index of the configuration which generates the implementation.594 """595 cost = 0596 configs = configs or _IMPLEMENT_CONFIGS597 if len(configs) > 1 and callable(assertions):598 assertions, cost = assertions(definition)599 assertion_filter = PassAssertionFilter(assertions)600 response = oai.Completion.create(601 {"definition": definition}, config_list=configs, filter_func=assertion_filter.pass_assertions602 )603 cost += assertion_filter.cost + response["cost"]604 return assertion_filter.responses[assertion_filter.metrics["index_selected"]], cost, response["config_id"]605 606 # for i, config in enumerate(configs):607 # response = oai.Completion.create({"definition": definition}, **config)608 # cost += oai.Completion.cost(response)609 # responses = oai.Completion.extract_text(response)610 # metrics = eval_function_completions(responses, definition, assertions=assertions)611 # assertions = metrics["assertions"]612 # cost += metrics["gen_cost"]613 # if metrics["succeed_assertions"] or i == len(configs) - 1:614 # return responses[metrics["index_selected"]], cost, i615