File size: 4,129 Bytes
01523b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from __future__ import annotations

import os
import subprocess
import multiprocessing
from typing import TYPE_CHECKING, Any, List, Tuple

from agentverse.logging import get_logger
from agentverse.agents import ExecutorAgent
from agentverse.message import ExecutorMessage, SolverMessage
from agentverse.logging import logger

from . import BaseExecutor, executor_registry


def execute_command(command: str, result_list) -> str:
    # TODO: make it more secure
    result = subprocess.run(command, capture_output=True, shell=True, encoding="utf-8")
    result_list.append(f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}")
    # return f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"


@executor_registry.register("code-test")
class CodeTestExecutor(BaseExecutor):
    has_test: dict = {}
    timeout: int = 10

    async def astep(
        self,
        agent: ExecutorAgent,
        task_description: str,
        solution: List[SolverMessage],
        *args,
        **kwargs,
    ) -> Any:
        solution = solution[0].content
        os.makedirs("tmp", exist_ok=True)
        self.write_to_file("tmp/main.py", solution)
        manager = multiprocessing.Manager()
        result = manager.list()
        if task_description not in self.has_test:
            response = (await agent.astep(task_description, solution)).content
            self.write_to_file(response["file_path"], response["code"])
            self.has_test[task_description] = f"python {response['file_path']}"
            p = multiprocessing.Process(
                target=execute_command, args=(f"python {response['file_path']}", result)
            )
            p.start()
            p.join(timeout=self.timeout + 1)
            if p.is_alive():
                p.kill()
            # result = execute_command(f"python {response['file_path']}")
        else:
            # result = execute_command(self.has_test[task_description])
            p = multiprocessing.Process(
                target=execute_command, args=(self.has_test[task_description], result)
            )
            p.start()
            p.join(timeout=self.timeout + 1)
            if p.is_alive():
                p.kill()
        if not result:
            result.append("Execution timed out.")
        return [ExecutorMessage(content=result[0], sender="Code Tester")]

    def step(
        self,
        agent: ExecutorAgent,
        task_description: str,
        solution: List[SolverMessage],
        *args,
        **kwargs,
    ) -> Any:
        solution = solution[0].content
        os.makedirs("tmp", exist_ok=True)
        self.write_to_file("tmp/main.py", solution)
        manager = multiprocessing.Manager()
        result = manager.list()
        if task_description not in self.has_test:
            response = agent.step(task_description, solution).content
            self.write_to_file(response["file_path"], response["code"])
            self.has_test[task_description] = f"python {response['file_path']}"
            p = multiprocessing.Process(
                target=execute_command, args=(f"python {response['file_path']}", result)
            )
            p.start()
            p.join(timeout=self.timeout + 1)
            if p.is_alive():
                p.kill()
            # result = execute_command(f"python {response['file_path']}")
        else:
            # result = execute_command(self.has_test[task_description])
            p = multiprocessing.Process(
                target=execute_command, args=(self.has_test[task_description], result)
            )
            p.start()
            p.join(timeout=self.timeout + 1)
            if p.is_alive():
                p.kill()
        if not result:
            result.append("Execution timed out.")
        return [ExecutorMessage(content=result[0], sender="Code Tester")]

    def write_to_file(self, file_name, file_content):
        # TODO: generalize this method to a common tool
        try:
            with open(file_name, "w") as f:
                f.write(file_content)
                f.flush()
        except:
            logger.error(f"Failed to write to {file_name}")