autoagents/environment.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2023/5/11 22:12
@Author  : alexanderwu
@File    : environment.py
@Modified From: https://github.com/geekan/MetaGPT/blob/main/metagpt/environment.py
"""
import asyncio
import re
import json
import datetime
import websockets
from common import MessageType, format_message, timestamp
from typing import Iterable

from pydantic import BaseModel, Field

from .roles import Role
from .actions import Requirement
from .roles import CustomRole, ActionObserver, Group, ROLES_LIST, ROLES_MAPPING

from .system.memory import Memory
from .system.schema import Message

class Environment(BaseModel):
    """环境，承载一批角色，角色可以向环境发布消息，可以被其他角色观察到"""

    roles: dict[str, Role] = Field(default_factory=dict)
    memory: Memory = Field(default_factory=Memory)
    history: str = Field(default='')
    new_roles_args: dict = Field(default_factory=dict)
    new_roles: dict[str, Role] = Field(default_factory=dict)
    steps: list = Field(default_factory=list)
    msg_json: list = Field(default_factory=list)
    json_log: str = Field(default='./logs/json_log.json')
    task_id: str = Field(default='')
    proxy: str = Field(default='')
    llm_api_key: str = Field(default='')
    serpapi_key: str = Field(default='')
    alg_msg_queue: object = Field(default=None)

    class Config:
        arbitrary_types_allowed = True


    def add_role(self, role: Role):
        """增加一个在当前环境的Role"""
        role.set_env(self)
        self.roles[role.profile] = role

    def add_roles(self, roles: Iterable[Role]):
        """增加一批在当前环境的Role"""
        for role in roles:
            self.add_role(role)

    def _parser_roles(self, text):
        """解析添加的Roles"""
        agents = re.findall('{[\s\S]*?}', text) # re.findall('{{.*}}', agents)
        agents_args = []
        for agent in agents:
            agent = json.loads(agent.strip())
            if len(agent.keys()) > 0:
                agents_args.append(agent)

        print('---------------Agents---------------')
        for i, agent in enumerate(agents_args):
            print('Role', i, agent)

        return agents_args
    
    def _parser_plan(self, context):
        """解析生成的计划Plan"""
        plan_context = re.findall('## Execution Plan([\s\S]*?)##', str(context))[0]
        steps = [v.split("\n")[0] for v in re.split("\n\d+\. ", plan_context)[1:]]
        print('---------------Steps---------------')
        for i, step in enumerate(steps):
            print('Step', i, step)
        
        steps.insert(0, '')
        return steps
    
    def create_roles(self, plan: list, args: dict):
        """创建Role""" 

        requirement_type = type('Requirement_Group', (Requirement,), {})
        self.add_role(Group(roles=args, steps=plan, watch_actions=[Requirement,requirement_type],  proxy=self.proxy, serpapi_api_key=self.serpapi_key, llm_api_key=self.llm_api_key))

        # existing_roles = dict()
        # for item in ROLES_LIST:
        #     existing_roles[item['name']] = item
                
        # init_actions, watch_actions = [], []
        # for role in args:
        #     class_name = role['name'].replace(' ', '_') + '_Requirement'
        #     requirement_type = type(class_name, (Requirement,), {})
        #     if role['name'] in existing_roles.keys():
        #         print('Add a predefiend role:', role['name'])
        #         role_object = ROLES_MAPPING[role['name']]
        #         if 'Engineer' in role['name']:
        #             _role = role_object(n_borg=2, use_code_review=True, proxy=self.proxy, llm_api_key=self.llm_api_key, serpapi_api_key=self.serpapi_key)
        #         else:
        #             _role = role_object(watch_actions=[requirement_type], proxy=self.proxy, llm_api_key=self.llm_api_key, serpapi_api_key=self.serpapi_key)
        #     else:
        #         print('Add a new role:', role['name'])
        #         _role = CustomRole(
        #             name=role['name'],
        #             profile=role['name'],
        #             goal=role['description'],
        #             role_prompt=role['prompt'],
        #             steps=role['steps'],
        #             tool=role['tools'],
        #             watch_actions=[requirement_type],
        #             proxy=self.proxy,
        #             llm_api_key=self.llm_api_key,
        #             serpapi_api_key=self.serpapi_key,
        #         )
                
        #     self.add_role(_role)
        #     watch_actions.append(requirement_type)
        #     init_actions.append(_role.init_actions)
            
        
        # init_actions.append(Requirement)
        # self.add_role(ActionObserver(steps=plan, watch_actions=init_actions, init_actions=watch_actions, proxy=self.proxy, llm_api_key=self.llm_api_key))

    async def publish_message(self, message: Message):
        """向当前环境发布信息"""
        # self.message_queue.put(message)
        self.memory.add(message)
        self.history += f"\n{message}"

        if 'Manager' in message.role:
            self.steps = self._parser_plan(message.content)
            self.new_roles_args = self._parser_roles(message.content)
            self.new_roles = self.create_roles(self.steps, self.new_roles_args)

        filename, file_content = None, None
        if hasattr(message.instruct_content, 'Type') and 'FILE' in message.instruct_content.Type:
            filename = message.instruct_content.Key
            file_type = re.findall('```(.*?)\n', str(message.content))[0]
            file_content = re.findall(f'```{file_type}([\s\S]*?)```', str(message.content))[0]
        
        if message.role and 'ActionObserver' != message.role:
            if hasattr(message.instruct_content, 'Response'):
                content = message.instruct_content.Response
            else:
                content = message.content

            msg = {   
                'timestamp': timestamp(),
                'role': message.role,
                'content': content,
                'file': {
                    'file_type': filename,
                    'file_data': file_content,
                }
            }

            if self.alg_msg_queue:
                self.alg_msg_queue.put_nowait(format_message(action=MessageType.RunTask.value, data={'task_id': self.task_id, 'task_message':msg}))
        
        if 'Agents Observer' in message.role:
            
            # send role list
            msg = {   
                'timestamp': timestamp(),
                'role': "Revised Role List",
                'content': self.new_roles_args,
                'file': {
                    'file_type': None,
                    'file_data': None,
                }
            }

            if self.alg_msg_queue:
                self.alg_msg_queue.put_nowait(format_message(action=MessageType.RunTask.value, data={'task_id': self.task_id, 'task_message':msg}))



    async def run(self, k=1):
        """处理一次所有Role的运行"""
        old_roles = []
        for _ in range(k):
            futures = []
            for key in self.roles.keys():
                old_roles.append(key)
                role = self.roles[key]
                future = role.run()
                futures.append(future)
            
            await asyncio.gather(*futures)

        if len(old_roles) < len(self.roles):
            while len(self.get_role(name='Group').steps) > 0:
                futures = []
                for key in self.roles.keys():
                    if key not in old_roles:
                        role = self.roles[key]
                        future = role.run()
                        futures.append(future)

                await asyncio.gather(*futures)

    def get_roles(self) -> dict[str, Role]:
        """获得环境内的所有Role"""
        return self.roles

    def get_role(self, name: str) -> Role:
        """获得环境内的指定Role"""
        return self.roles.get(name, None)