Repository: geekan/MetaGPT
File: examples/build_customized_agent.py (Python)

Code before:
import re
import subprocess

from metagpt.actions import Action
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.logs import logger

class SimpleWriteCode(Action):

    PROMPT_TEMPLATE = """
    Write a python function that can {instruction} and provide two runnnable test cases.
    Return ```python your_code_here ``` with NO other texts,
    example: 
    ```python
    # function
    def add(a, b):
        return a + b
    # test cases
    print(add(1, 2))
    print(add(3, 4))
    ```
    your code:
    """

    def __init__(self, name="SimpleWriteCode", context=None, llm=None):
        super().__init__(name, context, llm)
    
    async def run(self, instruction: str):

        prompt = self.PROMPT_TEMPLATE.format(instruction=instruction)
        # logger.info(prompt)

        rsp = await self._aask(prompt)

        code_text = SimpleWriteCode.parse_code(rsp)

        return code_text
    
    @staticmethod
    def parse_code(rsp):
        pattern = r'```python(.*)```'
        match = re.search(pattern, rsp, re.DOTALL)
        code_text = match.group(1) if match else rsp
        return code_text

class SimpleRunCode(Action):
    def __init__(self, name="SimpleRunCode", context=None, llm=None):
        super().__init__(name, context, llm)
    
    async def run(self, code_text: str):
        result = subprocess.run(["python3", "-c", code_text], capture_output=True, text=True)
        code_result = result.stdout
        # exec(code_text)
        logger.info(f"{code_result=}")
        return code_result

class SimpleCoder(Role):
    def __init__(
        self,
        name: str = "Alice",
        profile: str = "SimpleCoder",
        **kwargs,
    ):
        super().__init__(name, profile, **kwargs)
        self._init_actions([SimpleWriteCode])

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: ready to {self._rc.todo}")
        todo = self._rc.todo
        msg = self._rc.memory.get()[-1]

        instruction = msg.content
        code_text = await SimpleWriteCode().run(instruction)
        msg = Message(content=code_text, role=self.profile, cause_by=todo)

        return msg

class RunnableCoder(Role):
    def __init__(
        self,
        name: str = "Alice",
        profile: str = "RunnableCoder",
        **kwargs,
    ):
        super().__init__(name, profile, **kwargs)
        self._init_actions([SimpleWriteCode, SimpleRunCode])

    async def _think(self) -> None:
        if self._rc.todo is None:
            self._set_state(0)
            return

        if self._rc.state + 1 < len(self._states):
            self._set_state(self._rc.state + 1)
        else:
            self._rc.todo = None

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: ready to {self._rc.todo}")
        todo = self._rc.todo
        msg = self._rc.memory.get()[-1]

        if isinstance(todo, SimpleWriteCode):
            instruction = msg.content
            code_text = await SimpleWriteCode().run(instruction)
            msg = Message(content=code_text, role=self.profile, cause_by=todo)
        
        elif isinstance(todo, SimpleRunCode):
            code_text = msg.content
            rsp = await SimpleRunCode().run(code_text)
            msg = Message(content=rsp, role=self.profile, cause_by=todo)

        self._rc.memory.add(msg)
        return msg

    async def _react(self) -> Message:
        while True:
            await self._think()
            if self._rc.todo is None:
                break
            await self._act()
        return Message(content="All job done", role=self.profile)

if __name__ == "__main__":
    import asyncio

    async def main():
        msg = "write a function that calculates the sum of a list"
        # role = SimpleCoder()
        role = RunnableCoder()
        logger.info(msg)
        result = await role.run(msg)
        logger.info(result)

    asyncio.run(main())

Review comment (suggestion):
it can be 
```
msg = "xxxx"
role = RunnableCoder()
result = asyncio.run(role.run(msg))
logger.info(result)
```

Code after:
            instruction = msg.content
            result = await SimpleWriteCode().run(instruction)

        elif isinstance(todo, SimpleRunCode):
            code_text = msg.content
            result = await SimpleRunCode().run(code_text)

        msg = Message(content=result, role=self.profile, cause_by=todo)
        self._rc.memory.add(msg)
        return msg

    async def _react(self) -> Message:
        while True:
            await self._think()
            if self._rc.todo is None:
                break
            await self._act()
        return Message(content="All job done", role=self.profile)

def main(msg="write a function that calculates the sum of a list"):
    # role = SimpleCoder()
    role = RunnableCoder()
    logger.info(msg)
    result = asyncio.run(role.run(msg))
    logger.info(result)

if __name__ == '__main__':
    fire.Fire(main)