bump version to 0.1.8

agentverse/__init__.py

@@ -0,0 +1,24 @@
+from .output_parser import output_parser_registry
+from .environments import env_registry
+from .environments.simulation_env.rules.order import order_registry
+from .environments.simulation_env.rules.describer import describer_registry
+from .environments.simulation_env.rules.selector import selector_registry
+from .environments.simulation_env.rules.updater import updater_registry
+from .environments.simulation_env.rules.visibility import visibility_registry
+
+
+from .environments.tasksolving_env.rules.decision_maker import decision_maker_registry
+from .environments.tasksolving_env.rules.evaluator import evaluator_registry
+from .environments.tasksolving_env.rules.executor import executor_registry
+from .environments.tasksolving_env.rules.role_assigner import role_assigner_registry
+
+
+from .simulation import Simulation
+from .tasksolving import TaskSolving
+from .initialization import (
+    prepare_task_config,
+    load_agent,
+    load_environment,
+    load_llm,
+    load_memory,
+)

agentverse/agents/__init__.py

@@ -0,0 +1,20 @@
+# from .agent import Agent
+from agentverse.registry import Registry
+
+agent_registry = Registry(name="AgentRegistry")
+
+
+from .base import BaseAgent
+from agentverse.agents.simulation_agent.conversation import ConversationAgent
+from agentverse.agents.simulation_agent.tool import ToolAgent
+from agentverse.agents.simulation_agent.prisoner_dilemma import (
+    PoliceAgent,
+    PrisonerAgent,
+)
+
+from agentverse.agents.tasksolving_agent.role_assigner import RoleAssignerAgent
+from agentverse.agents.tasksolving_agent.critic import CriticAgent
+from agentverse.agents.tasksolving_agent.evaluator import EvaluatorAgent
+from agentverse.agents.tasksolving_agent.solver import SolverAgent
+from agentverse.agents.tasksolving_agent.manager import ManagerAgent
+from agentverse.agents.tasksolving_agent.executor import ExecutorAgent

agentverse/agents/base.py

@@ -0,0 +1,101 @@
+import logging
+from abc import abstractmethod
+from typing import List, NamedTuple, Set, Union
+from string import Template
+
+from pydantic import BaseModel, Field
+
+from agentverse.llms import BaseLLM
+from agentverse.memory import BaseMemory, ChatHistoryMemory
+from agentverse.message import Message
+from agentverse.output_parser import OutputParser
+from agentverse.memory_manipulator import BaseMemoryManipulator
+
+
+class BaseAgent(BaseModel):
+    name: str
+    llm: BaseLLM
+    output_parser: OutputParser
+    prepend_prompt_template: str = Field(default="")
+    append_prompt_template: str = Field(default="")
+    prompt_template: str = Field(default="")
+    role_description: str = Field(default="")
+    memory: BaseMemory = Field(default_factory=ChatHistoryMemory)
+    memory_manipulator: BaseMemoryManipulator = Field(
+        default_factory=BaseMemoryManipulator
+    )
+    max_retry: int = Field(default=3)
+    receiver: Set[str] = Field(default=set({"all"}))
+    async_mode: bool = Field(default=True)
+
+    @abstractmethod
+    def step(self, env_description: str = "") -> Message:
+        """Get one step response"""
+        pass
+
+    @abstractmethod
+    def astep(self, env_description: str = "") -> Message:
+        """Asynchronous version of step"""
+        pass
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Reset the agent"""
+        pass
+
+    @abstractmethod
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        """Add a message to the memory"""
+        pass
+
+    def get_spend(self) -> float:
+        return self.llm.get_spend()
+
+    def get_spend_formatted(self) -> str:
+        two_trailing = f"${self.get_spend():.2f}"
+        if two_trailing == "$0.00":
+            return f"${self.get_spend():.6f}"
+        return two_trailing
+
+    def get_all_prompts(self, **kwargs):
+        prepend_prompt = Template(self.prepend_prompt_template).safe_substitute(
+            **kwargs
+        )
+        append_prompt = Template(self.append_prompt_template).safe_substitute(**kwargs)
+        return prepend_prompt, append_prompt
+
+    def get_receiver(self) -> Set[str]:
+        return self.receiver
+
+    def set_receiver(self, receiver: Union[Set[str], str]) -> None:
+        if isinstance(receiver, str):
+            self.receiver = set({receiver})
+        elif isinstance(receiver, set):
+            self.receiver = receiver
+        else:
+            raise ValueError(
+                "input argument `receiver` must be a string or a set of string"
+            )
+
+    def add_receiver(self, receiver: Union[Set[str], str]) -> None:
+        if isinstance(receiver, str):
+            self.receiver.add(receiver)
+        elif isinstance(receiver, set):
+            self.receiver = self.receiver.union(receiver)
+        else:
+            raise ValueError(
+                "input argument `receiver` must be a string or a set of string"
+            )
+
+    def remove_receiver(self, receiver: Union[Set[str], str]) -> None:
+        if isinstance(receiver, str):
+            try:
+                self.receiver.remove(receiver)
+            except KeyError as e:
+                logging.warning(f"Receiver {receiver} not found.")
+        elif isinstance(receiver, set):
+            self.receiver = self.receiver.difference(receiver)
+        else:
+            raise ValueError(
+                "input argument `receiver` must be a string or a set of string"
+            )

agentverse/agents/simulation_agent/conversation.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+from colorama import Fore
+
+# import logging
+from agentverse.logging import get_logger
+import bdb
+from string import Template
+from typing import TYPE_CHECKING, List
+
+from agentverse.message import Message
+
+#from . import agent_registry
+#from .base import BaseAgent
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+
+logger = get_logger()
+
+
+@agent_registry.register("conversation")
+class ConversationAgent(BaseAgent):
+    def step(self, env_description: str = "") -> Message:
+        prompt = self._fill_prompt_template(env_description)
+
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                response = self.llm.generate_response(prompt)
+                parsed_response = self.output_parser.parse(response)
+                break
+            except KeyboardInterrupt:
+                raise
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+
+        message = Message(
+            content=""
+            if parsed_response is None
+            else parsed_response.return_values["output"],
+            sender=self.name,
+            receiver=self.get_receiver(),
+        )
+        return message
+
+    async def astep(self, env_description: str = "") -> Message:
+        """Asynchronous version of step"""
+        prompt = self._fill_prompt_template(env_description)
+
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                # if self.name == "Code Reviewer":
+                logger.debug(prompt, "Prompt", Fore.CYAN)
+                response = await self.llm.agenerate_response(prompt)
+
+                # logging.info(f"{self.name}'s request result:"
+                #              f" {response.content}")
+                parsed_response = self.output_parser.parse(response)
+                break
+            except (KeyboardInterrupt, bdb.BdbQuit):
+                raise
+            except Exception as e:
+                logger.error(e)
+                logger.warning("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+
+        message = Message(
+            content=""
+            if parsed_response is None
+            else parsed_response.return_values["output"],
+            sender=self.name,
+            receiver=self.get_receiver(),
+        )
+        return message
+
+    def _fill_prompt_template(self, env_description: str = "") -> str:
+        """Fill the placeholders in the prompt template
+
+        In the conversation agent, three placeholders are supported:
+        - ${agent_name}: the name of the agent
+        - ${env_description}: the description of the environment
+        - ${role_description}: the description of the role of the agent
+        - ${chat_history}: the chat history of the agent
+        """
+        input_arguments = {
+            "agent_name": self.name,
+            "env_description": env_description,
+            "role_description": self.role_description,
+            "chat_history": self.memory.to_string(add_sender_prefix=True),
+        }
+        return Template(self.prompt_template).safe_substitute(input_arguments)
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self) -> None:
+        """Reset the agent"""
+        self.memory.reset()
+        # TODO: reset receiver

agentverse/agents/simulation_agent/prisoner_dilemma.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import logging
+from string import Template
+from typing import TYPE_CHECKING, List
+
+from agentverse.message import Message
+
+# from . import agent_registry
+# from .base import BaseAgent
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+
+if TYPE_CHECKING:
+    from agentverse.environments.base import BaseEnvironment
+
+
+class PrisonerDilemaAgent(BaseAgent):
+    def step(
+        self,
+        environment: BaseEnvironment,
+        env_description: str = "",
+    ) -> Message:
+        prompt = self._fill_prompt_template(env_description)
+
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                response = self.llm.generate_response(prompt)
+                parsed_response = self.output_parser.parse(self, environment, response)
+                break
+            except Exception as e:
+                logging.error(e)
+                logging.warning("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logging.error(f"{self.name} failed to generate valid response.")
+
+        message = Message(
+            content=""
+            if parsed_response is None
+            else parsed_response.return_values["output"],
+            sender=self.name,
+            receiver=self.get_receiver(),
+        )
+        return message
+
+    async def astep(
+        self, environment: BaseEnvironment, env_description: str = ""
+    ) -> Message:
+        """Asynchronous version of step"""
+        prompt = self._fill_prompt_template(env_description)
+
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                response = await self.llm.agenerate_response(prompt)
+                parsed_response = self.output_parser.parse(self, environment, response)
+                break
+            except Exception as e:
+                logging.error(e)
+                logging.warning("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logging.error(f"{self.name} failed to generate valid response.")
+
+        message = Message(
+            content=""
+            if parsed_response is None
+            else parsed_response.return_values["output"],
+            sender=self.name,
+            receiver=self.get_receiver(),
+        )
+        return message
+
+    def _fill_prompt_template(self, env_description: str = "") -> str:
+        """Fill the placeholders in the prompt template
+
+        In the conversation agent, three placeholders are supported:
+        - ${agent_name}: the name of the agent
+        - ${env_description}: the description of the environment
+        - ${role_description}: the description of the role of the agent
+        - ${chat_history}: the chat history of the agent
+        """
+        input_arguments = {
+            "agent_name": self.name,
+            "env_description": env_description,
+            "role_description": self.role_description,
+            "chat_history": self.memory.to_string(add_sender_prefix=True),
+        }
+        return Template(self.prompt_template).safe_substitute(input_arguments)
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self) -> None:
+        """Reset the agent"""
+        self.memory.reset()
+        # TODO: reset receiver
+
+
+@agent_registry.register("police")
+class PoliceAgent(PrisonerDilemaAgent):
+    interrogating_form: str
+
+    def _fill_prompt_template(self, env_description: str = "") -> str:
+        """Fill the placeholders in the prompt template
+
+        In the conversation agent, three placeholders are supported:
+        - ${agent_name}: the name of the agent
+        - ${env_description}: the description of the environment
+        - ${role_description}: the description of the role of the agent
+        - ${chat_history}: the chat history of the agent
+        """
+        input_arguments = {
+            "agent_name": self.name,
+            "env_description": env_description,
+            "role_description": self.role_description,
+            "chat_history": self.memory.to_string(add_sender_prefix=True),
+        }
+
+        role_argument = {
+            "interrogating_form": self.interrogating_form,
+        }
+
+        role_description = Template(self.role_description).safe_substitute(
+            role_argument
+        )
+        input_arguments["role_description"] = role_description
+
+        return Template(self.prompt_template).safe_substitute(input_arguments)
+
+
+@agent_registry.register("prisoner")
+class PrisonerAgent(PrisonerDilemaAgent):
+    personality: str
+    relationship_with_another: str
+
+    def _fill_prompt_template(self, env_description: str = "") -> str:
+        """Fill the placeholders in the prompt template
+
+        In the conversation agent, three placeholders are supported:
+        - ${agent_name}: the name of the agent
+        - ${env_description}: the description of the environment
+        - ${role_description}: the description of the role of the agent
+        - ${chat_history}: the chat history of the agent
+        """
+        input_arguments = {
+            "agent_name": self.name,
+            "env_description": env_description,
+            "role_description": self.role_description,
+            "chat_history": self.memory.to_string(add_sender_prefix=True),
+        }
+
+        role_argument = {
+            "personality": self.personality,
+            "relationship_with_another": self.relationship_with_another,
+        }
+
+        role_description = Template(self.role_description).safe_substitute(
+            role_argument
+        )
+        input_arguments["role_description"] = role_description
+
+        return Template(self.prompt_template).safe_substitute(input_arguments)

agentverse/agents/simulation_agent/reflection.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+"""
+An agent based upon Observation-Planning-Reflection architecture.
+"""
+
+from logging import getLogger
+
+from abc import abstractmethod
+from typing import List, Set, Union, NamedTuple, TYPE_CHECKING
+
+from pydantic import BaseModel, Field, validator
+
+from agentverse.llms import BaseLLM
+from agentverse.memory import BaseMemory, ChatHistoryMemory
+from agentverse.message import Message
+from agentverse.output_parser import OutputParser
+
+from agentverse.message import Message
+from agentverse.agents.base import BaseAgent
+
+from datetime import datetime as dt
+import datetime
+
+#from . import agent_registry
+from string import Template
+
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+
+logger = getLogger(__file__)
+
+if TYPE_CHECKING:
+    from agentverse.environments.base import BaseEnvironment
+
+
+@agent_registry.register("reflection")
+class ReflectionAgent(BaseAgent):
+    async_mode: bool = (True,)
+    current_time: str = (None,)
+    environment: BaseEnvironment = None
+    step_cnt: int = 0
+
+    manipulated_memory: str = Field(
+        default="", description="one fragment used in prompt construction"
+    )
+
+    @validator("current_time")
+    def convert_str_to_dt(cls, current_time):
+        if not isinstance(current_time, str):
+            raise ValueError("current_time should be str")
+        return dt.strptime(current_time, "%Y-%m-%d %H:%M:%S")
+
+    def step(self, current_time: dt, env_description: str = "") -> Message:
+        """
+        Call this method at each time frame
+        """
+        self.current_time = current_time
+
+        self.manipulated_memory = self.memory_manipulator.manipulate_memory()
+
+        prompt = self._fill_prompt_template(env_description)
+
+        parsed_response, reaction, target = None, None, None
+        for i in range(self.max_retry):
+            try:
+                response = self.llm.agenerate_response(prompt)
+                parsed_response = self.output_parser.parse(response)
+
+                if "say(" in parsed_response.return_values["output"]:
+                    reaction, target = eval(
+                        "self._" + parsed_response.return_values["output"].strip()
+                    )
+                elif "act(" in parsed_response.return_values["output"]:
+                    reaction, target = eval(
+                        "self._" + parsed_response.return_values["output"].strip()
+                    )
+                elif "do_nothing(" in parsed_response.return_values["output"]:
+                    reaction, target = None, None
+                else:
+                    raise Exception(
+                        f"no valid parsed_response detected, "
+                        f"cur response {parsed_response.return_values['output']}"
+                    )
+                break
+
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+
+        if reaction is None:
+            reaction = "Keep doing last action ..."
+
+        message = Message(
+            content="" if reaction is None else reaction,
+            sender=self.name,
+            receiver=self.get_receiver()
+            if target is None
+            else self.get_valid_receiver(target),
+        )
+
+        self.step_cnt += 1
+
+        return message
+
+    async def astep(self, current_time: dt, env_description: str = "") -> Message:
+        """Asynchronous version of step"""
+        # use environment's time to update agent's time
+        self.current_time = current_time
+        # Before the agent step, we check current status,
+        # TODO add this func after
+        # self.check_status_passive()
+
+        self.manipulated_memory = self.memory_manipulator.manipulate_memory()
+
+        prompt = self._fill_prompt_template(env_description)
+
+        parsed_response, reaction, target = None, None, None
+        for i in range(self.max_retry):
+            try:
+                response = await self.llm.agenerate_response(prompt)
+                parsed_response = self.output_parser.parse(response)
+
+                if "say(" in parsed_response.return_values["output"]:
+                    reaction, target = eval(
+                        "self._" + parsed_response.return_values["output"].strip()
+                    )
+                elif "act(" in parsed_response.return_values["output"]:
+                    reaction, target = eval(
+                        "self._" + parsed_response.return_values["output"].strip()
+                    )
+                elif "do_nothing(" in parsed_response.return_values["output"]:
+                    reaction, target = None, None
+                else:
+                    raise Exception(
+                        f"no valid parsed_response detected, "
+                        f"cur response {parsed_response.return_values['output']}"
+                    )
+
+                break
+
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+
+        if reaction is None:
+            reaction = "Keep doing last action ..."
+
+        message = Message(
+            content="" if reaction is None else reaction,
+            sender=self.name,
+            receiver=self.get_receiver()
+            if target is None
+            else self.get_valid_receiver(target),
+        )
+
+        self.step_cnt += 1
+
+        return message
+
+    def _act(self, description=None, target=None):
+        if description is None:
+            return ""
+        if target is None:
+            reaction_content = f"{self.name} performs action: '{description}'."
+        else:
+            reaction_content = (
+                f"{self.name} performs action to {target}: '{description}'."
+            )
+        # self.environment.broadcast_observations(self, target, reaction_content)
+        return reaction_content, target
+
+    def _say(self, description, target=None):
+        if description is None:
+            return ""
+        if target is None:
+            reaction_content = f"{self.name} says: '{description}'."
+        else:
+            reaction_content = f"{self.name} says to {target}: '{description}'."
+        # self.environment.broadcast_observations(self, target, reaction_content)
+        return reaction_content, target
+
+    def get_valid_receiver(self, target: str) -> set():
+        all_agents_name = []
+        for agent in self.environment.agents:
+            all_agents_name.append(agent.name)
+
+        if not (target in all_agents_name):
+            return {"all"}
+        else:
+            return {target}
+
+    def _fill_prompt_template(self, env_description: str = "") -> str:
+        """Fill the placeholders in the prompt template
+
+        In the conversation agent, three placeholders are supported:
+        - ${agent_name}: the name of the agent
+        - ${env_description}: the description of the environment
+        - ${role_description}: the description of the role of the agent
+        - ${chat_history}: the chat history of the agent
+        """
+        input_arguments = {
+            "agent_name": self.name,
+            "role_description": self.role_description,
+            "chat_history": self.memory.to_string(add_sender_prefix=True),
+            "current_time": self.current_time,
+            "env_description": env_description,
+        }
+        return Template(self.prompt_template).safe_substitute(input_arguments)
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self, environment: BaseEnvironment) -> None:
+        """Reset the agent"""
+        self.environment = environment
+        self.memory.reset()
+        self.memory_manipulator.agent = self
+        self.memory_manipulator.memory = self.memory

agentverse/agents/simulation_agent/tool.py

@@ -0,0 +1,177 @@
+import logging
+from string import Template
+from typing import List, NamedTuple, Optional, Union
+
+from langchain.tools import BaseTool
+from pydantic import Field
+
+
+from agentverse.memory import BaseMemory, ChatHistoryMemory
+from agentverse.message import Message
+from agentverse.utils import AgentAction, AgentFinish
+
+#from . import agent_registry
+#from .base import BaseAgent
+
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+
+class ToolNotExistError(BaseException):
+    """Exception raised when parsing output from a command fails."""
+
+    def __init__(self, tool_name=""):
+        self.tool_name = tool_name
+
+    def __str__(self):
+        return f"Tool {self.tool_name} does not exist."
+
+
+@agent_registry.register("tool")
+class ToolAgent(BaseAgent):
+    tools: List[BaseTool] = Field(default=[])
+    tool_memory: BaseMemory = Field(default_factory=ChatHistoryMemory)
+    verbose: bool = Field(default=False)
+
+    def step(self, env_description: str = "") -> Message:
+        parsed_response = None
+        tool_observation = [self.tool_memory.to_string()]
+        while True:
+            prompt = self._fill_prompt_template(env_description, tool_observation)
+
+            for i in range(self.max_retry):
+                try:
+                    response = self.llm.generate_response(prompt)
+                    parsed_response = self.output_parser.parse(response)
+                    if isinstance(parsed_response, AgentAction):
+                        observation = self._call_tool(parsed_response)
+                        tool_observation.append(
+                            parsed_response.log.strip()
+                            + f"\nObservation: {observation.strip()}"
+                        )
+                    break
+                except BaseException as e:
+                    logging.error(e)
+                    logging.warning("Retrying...")
+                    continue
+            if parsed_response is None or isinstance(parsed_response, AgentFinish):
+                break
+
+        if parsed_response is None:
+            logging.error(f"{self.name} failed to generate valid response.")
+
+        self._update_tool_memory(tool_observation)
+
+        message = Message(
+            content=""
+            if parsed_response is None
+            else parsed_response.return_values["output"],
+            sender=self.name,
+            receiver=self.get_receiver(),
+        )
+        return message
+
+    async def astep(self, env_description: str = "") -> Message:
+        """Asynchronous version of step"""
+        parsed_response = None
+        # Initialize the tool_observation with tool_memory
+        tool_observation = [self.tool_memory.to_string()]
+        while True:
+            prompt = self._fill_prompt_template(env_description, tool_observation)
+
+            for i in range(self.max_retry):
+                try:
+                    response = await self.llm.agenerate_response(prompt)
+                    parsed_response = self.output_parser.parse(response)
+                    if isinstance(parsed_response, AgentAction):
+                        # If the response is an action, call the tool
+                        # and append the observation to tool_observation
+                        observation = await self._acall_tool(parsed_response)
+                        tool_observation.append(
+                            parsed_response.log.strip()
+                            + f"\nObservation: {observation.strip()}"
+                        )
+                    break
+                except BaseException as e:
+                    logging.error(e)
+                    logging.warning("Retrying...")
+                    continue
+            if parsed_response is None or isinstance(parsed_response, AgentFinish):
+                break
+
+        if parsed_response is None:
+            logging.error(f"{self.name} failed to generate valid response.")
+
+        self._update_tool_memory(tool_observation)
+
+        message = Message(
+            content=""
+            if parsed_response is None
+            else parsed_response.return_values["output"],
+            sender=self.name,
+            receiver=self.get_receiver(),
+        )
+        return message
+
+    def _call_tool(self, response: NamedTuple) -> str:
+        """Call a tool and return the output"""
+        name_to_tool = {tool.name: tool for tool in self.tools}
+        if response.tool not in name_to_tool:
+            raise ToolNotExistError(response.tool)
+        tool = name_to_tool[response.tool]
+        observation = tool.run(response.tool_input, verbose=self.verbose)
+        return observation
+
+    async def _acall_tool(self, response: NamedTuple) -> str:
+        """Call a tool and return the output"""
+        name_to_tool = {tool.name: tool for tool in self.tools}
+        if response.tool not in name_to_tool:
+            raise ToolNotExistError(response.tool)
+        tool = name_to_tool[response.tool]
+        observation = await tool.arun(response.tool_input, verbose=self.verbose)
+        return observation
+
+    def _update_tool_memory(self, tool_observation: List[str]):
+        """Update the memory of the tool"""
+        if len(tool_observation) == 1:
+            # If no tool is called this turn, do nothing
+            return
+        messages = [
+            Message(content=observation) for observation in tool_observation[1:]
+        ]
+        self.tool_memory.add_message(messages)
+
+    def _fill_prompt_template(
+        self, env_description: str = "", tool_observation: List[str] = []
+    ) -> str:
+        """Fill the placeholders in the prompt template
+
+        In the tool agent, these placeholders are supported:
+        - ${agent_name}: the name of the agent
+        - ${env_description}: the description of the environment
+        - ${role_description}: the description of the role of the agent
+        - ${chat_history}: the chat history of the agent
+        - ${tools}: the list of tools and their usage
+        - ${tool_names}: the list of tool names
+        - ${tool_observations}: the observation of the tool in this turn
+        """
+        tools = "\n".join([f"> {tool.name}: {tool.description}" for tool in self.tools])
+        tools = tools.replace("{{", "{").replace("}}", "}")
+        tool_names = ", ".join([tool.name for tool in self.tools])
+        input_arguments = {
+            "agent_name": self.name,
+            "env_description": env_description,
+            "role_description": self.role_description,
+            "chat_history": self.memory.to_string(add_sender_prefix=True),
+            "tools": tools,
+            "tool_names": tool_names,
+            "tool_observation": "\n".join(tool_observation),
+        }
+        return Template(self.prompt_template).safe_substitute(input_arguments)
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self) -> None:
+        """Reset the agent"""
+        self.memory.reset()
+        # TODO: reset receiver

agentverse/agents/tasksolving_agent/__init__.py

@@ -0,0 +1,6 @@
+from .critic import CriticAgent
+from .evaluator import EvaluatorAgent
+from .executor import ExecutorAgent
+from .manager import ManagerAgent
+from .role_assigner import RoleAssignerAgent
+from .solver import SolverAgent

agentverse/agents/tasksolving_agent/critic.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import json
+from colorama import Fore
+from agentverse.logging import get_logger
+import bdb
+from string import Template
+from typing import TYPE_CHECKING, List, Union
+
+from agentverse.message import Message
+
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+from agentverse.utils import AgentCriticism
+from agentverse.message import CriticMessage
+
+logger = get_logger()
+
+
+@agent_registry.register("critic")
+class CriticAgent(BaseAgent):
+    max_history: int = 3
+    tools: List[dict] = []
+    tool_names: List[str] = []
+    tool_descriptions: str = ""
+
+    def __init__(self, *args, **kwargs):
+        tool_config_file = kwargs.pop("tool_config", "")
+        tools = []
+        tool_names = []
+        tool_descriptions = ""
+        if tool_config_file != "":
+            try:
+                with open(tool_config_file, "r") as f:
+                    tools_dict = json.load(f)
+                tools = tools_dict["tools_json"]
+                tool_names = [t["name"] for t in tools]
+                tool_descriptions = "\n".join(
+                    [f"- {t['name']}: " + t["description"] for t in tools]
+                )
+                kwargs.update('tools', tools)
+                kwargs.update('tool_names', tool_names)
+                kwargs.update('tool_descriptions', tool_descriptions)
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Failed to load tool config file.")
+        super().__init__(
+            *args,
+            **kwargs,
+        )
+
+    def step(self, env_description: str = "") -> CriticMessage:
+        pass
+
+    async def astep(
+        self,
+        preliminary_solution: str,
+        advice: str = "No advice yet.",
+        task_description: str = "",
+        all_roles: str = "",
+        **kwargs,
+    ) -> CriticMessage:
+        """Asynchronous version of step"""
+        logger.debug("", self.name, Fore.MAGENTA)
+        prepend_prompt, append_prompt = self.get_all_prompts(
+            preliminary_solution=preliminary_solution,
+            advice=advice,
+            task_description=task_description,
+            role_description=self.role_description,
+            agent_name=self.name,
+            all_roles=all_roles,
+            # tool_names=self.tool_names,
+            tool_descriptions=self.tool_descriptions,
+        )
+        history = self.memory.to_messages(self.name, start_index=-self.max_history)
+        parsed_response: Union[AgentCriticism, None] = None
+        for i in range(self.max_retry):
+            try:
+                response = await self.llm.agenerate_response(
+                    prepend_prompt, history, append_prompt
+                )
+                parsed_response = self.output_parser.parse(response)
+                break
+            except (KeyboardInterrupt, bdb.BdbQuit):
+                raise
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+
+        message = CriticMessage(
+            content=parsed_response.criticism if parsed_response is not None else "",
+            sender=self.name,
+            sender_agent=self,
+            is_agree=parsed_response.is_agree if parsed_response is not None else False,
+        )
+        return message
+
+    def _fill_prompt_template(
+        self, preliminary_solution: str, advice: str, task_description: str
+    ) -> str:
+        """Fill the placeholders in the prompt template
+
+        In the conversation agent, three placeholders are supported:
+        - ${role_description}
+        - ${task_description}
+        - ${preliminary_solution}
+        - ${advice}
+        """
+        input_arguments = {
+            "role_description": self.role_description,
+            "task_description": task_description,
+            "preliminary_solution": preliminary_solution,
+            "advice": advice,
+        }
+        return Template(self.prompt_template).safe_substitute(input_arguments)
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self) -> None:
+        """Reset the agent"""
+        self.memory.reset()
+        # TODO: reset receiver

agentverse/agents/tasksolving_agent/evaluator.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import asyncio
+from colorama import Fore
+
+from agentverse.logging import get_logger
+import bdb
+from string import Template
+from typing import TYPE_CHECKING, List, Tuple
+
+from agentverse.message import EvaluatorMessage, Message
+
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+
+
+logger = get_logger()
+
+
+@agent_registry.register("evaluator")
+class EvaluatorAgent(BaseAgent):
+    def step(
+        self,
+        solution: str,
+        result: str,
+        task_description: str,
+        all_role_description: str,
+    ) -> EvaluatorMessage:
+        logger.debug("", self.name, Fore.MAGENTA)
+        prepend_prompt, append_prompt = self.get_all_prompts(
+            solution=solution,
+            result=result,
+            task_description=task_description,
+            all_role_description=all_role_description,
+        )
+        history = self.memory.to_messages(self.name)
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                response = self.llm.generate_response(
+                    prepend_prompt, history, append_prompt
+                )
+                parsed_response = self.output_parser.parse(response)
+                break
+            except (KeyboardInterrupt, bdb.BdbQuit):
+                raise
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+        message = EvaluatorMessage(
+            sender=self.name,
+            sender_agent=self,
+            score=parsed_response[0] if parsed_response is not None else 0,
+            advice=parsed_response[1] if parsed_response is not None else "",
+        )
+        return message
+        # return parsed_response
+
+    async def astep(self, solution: str) -> EvaluatorMessage:
+        """Asynchronous version of step"""
+        pass
+
+    def _fill_prompt_template(self, solution: str, task_description: str) -> str:
+        """Fill the placeholders in the prompt template
+
+        In the role_assigner agent, three placeholders are supported:
+        - ${task_description}
+        - ${solution}
+        """
+        input_arguments = {
+            "task_description": task_description,
+            "solution": solution,
+        }
+        return Template(self.prompt_template).safe_substitute(input_arguments)
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self) -> None:
+        """Reset the agent"""
+        self.memory.reset()
+        # TODO: reset receiver

agentverse/agents/tasksolving_agent/executor.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+from agentverse.logging import get_logger
+from colorama import Fore
+import bdb
+from string import Template
+from typing import TYPE_CHECKING, List, Any
+
+from agentverse.message import ExecutorMessage, Message, SolverMessage
+from agentverse.utils import AgentFinish, AgentAction
+
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+import requests
+
+logger = get_logger()
+
+
+@agent_registry.register("executor")
+class ExecutorAgent(BaseAgent):
+    max_history: int = 5
+
+    def step(
+        self, task_description: str, solution: str, tools: List[dict] = [], **kwargs
+    ) -> ExecutorMessage:
+        logger.debug("", self.name, Fore.MAGENTA)
+        prepend_prompt, append_prompt = self.get_all_prompts(
+            task_description=task_description,
+            solution=solution,
+            agent_name=self.name,
+            **kwargs,
+        )
+
+        history = self.memory.to_messages(self.name, start_index=-self.max_history)
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                response = self.llm.generate_response(
+                    prepend_prompt, history, append_prompt, tools
+                )
+                parsed_response = self.output_parser.parse(response)
+                break
+            except (KeyboardInterrupt, bdb.BdbQuit):
+                raise
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+        if isinstance(parsed_response, AgentFinish):
+            message = ExecutorMessage(
+                content=parsed_response.return_values["output"],
+                sender=self.name,
+                sender_agent=self,
+            )
+        elif isinstance(parsed_response, AgentAction):
+            message = ExecutorMessage(
+                content=parsed_response.log,
+                sender=self.name,
+                sender_agent=self,
+                tool_name=parsed_response.tool,
+                tool_input=parsed_response.tool_input,
+            )
+        else:
+            raise ValueError(
+                f"Error response type: {type(parsed_response)}. Only support \
+                    AgentFinish and AgentAction. Modify your output parser."
+            )
+        return message
+
+    async def astep(
+        self, task_description: str, solution: str, tools: List[dict] = [], **kwargs
+    ) -> ExecutorMessage:
+        logger.debug("", self.name, Fore.MAGENTA)
+        prepend_prompt, append_prompt = self.get_all_prompts(
+            task_description=task_description,
+            solution=solution,
+            agent_name=self.name,
+            **kwargs,
+        )
+
+        history = self.memory.to_messages(self.name, start_index=-self.max_history)
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                response = await self.llm.agenerate_response(
+                    prepend_prompt, history, append_prompt, tools
+                )
+                parsed_response = self.output_parser.parse(response)
+                break
+            except (KeyboardInterrupt, bdb.BdbQuit):
+                raise
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+            parsed_response = AgentAction(tool="", tool_input="", log="")
+        if isinstance(parsed_response, AgentFinish):
+            message = ExecutorMessage(
+                content=parsed_response.return_values["output"],
+                sender=self.name,
+                sender_agent=self,
+            )
+        elif isinstance(parsed_response, AgentAction):
+            message = ExecutorMessage(
+                content=parsed_response.log,
+                sender=self.name,
+                sender_agent=self,
+                tool_name=parsed_response.tool,
+                tool_input=parsed_response.tool_input,
+            )
+        else:
+            raise ValueError(
+                f"Error response type: {type(parsed_response)}. Only support \
+                    AgentFinish and AgentAction. Modify your output parser."
+            )
+        return message
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self) -> None:
+        """Reset the agent"""
+        self.memory.reset()
+        # TODO: reset receiver

agentverse/agents/tasksolving_agent/manager.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import asyncio
+from colorama import Fore
+
+from agentverse.logging import get_logger
+import bdb
+from string import Template
+from typing import TYPE_CHECKING, List, Tuple
+
+from agentverse.message import Message
+
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+from agentverse.utils import AgentCriticism
+
+import random
+from rapidfuzz import fuzz
+
+
+logger = get_logger()
+
+
+@agent_registry.register("manager")
+class ManagerAgent(BaseAgent):
+    prompt_template: str
+
+    def step(
+        self,
+        former_solution: str,
+        candidate_critic_opinions: List[AgentCriticism],
+        advice: str,
+        task_description: str = "",
+        previous_sentence: str = "",
+    ) -> Message:
+        logger.debug("", self.name, Fore.MAGENTA)
+
+        prompt = self._fill_prompt_template(
+            former_solution,
+            candidate_critic_opinions,
+            advice,
+            task_description,
+            previous_sentence,
+        )
+
+        logger.debug(f"Prompt:\n{prompt}", "Manager", Fore.CYAN)
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                # LLM Manager
+                # response = self.llm.generate_response(prompt)
+                # parsed_response = self.output_parser.parse(response)
+                selected_role_description = self.llm.generate_response(prompt).content
+                candidate_score_list = [
+                    fuzz.ratio(candidate.sender, selected_role_description)
+                    for candidate in candidate_critic_opinions
+                ]
+                selected_index = candidate_score_list.index(max(candidate_score_list))
+                candidate_critic_opinion = candidate_critic_opinions[selected_index]
+
+                # Random Manager
+                # parsed_response = random.choice(candidate_critic_opinions)
+                break
+            except (KeyboardInterrupt, bdb.BdbQuit):
+                raise
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+        return candidate_critic_opinion
+
+    async def astep(self, env_description: str = "") -> Message:
+        """Asynchronous version of step"""
+        pass
+
+    def _fill_prompt_template(
+        self,
+        former_solution: str,
+        candidate_critic_opinions: List[AgentCriticism],
+        advice: str,
+        task_description: str,
+        previous_sentence: str,
+    ) -> str:
+        """Fill the placeholders in the prompt template
+
+        In the role_assigner agent, three placeholders are supported:
+        - ${task_description}
+        - ${former_solution}
+        - ${critic_messages}
+        - ${advice}
+        - ${previous_sentence}
+        """
+        input_arguments = {
+            "task_description": task_description,
+            "former_solution": former_solution,
+            "previous_sentence": previous_sentence,
+            "critic_opinions": "\n".join(
+                [
+                    f"Role: {critic.sender}. {critic.sender_agent.role_description} said: {critic.content}"
+                    for critic in candidate_critic_opinions
+                ]
+            ),
+            "advice": advice,
+        }
+
+        # manger select the proper sentence
+        template = Template(self.prompt_template)
+        return template.safe_substitute(input_arguments)
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self) -> None:
+        """Reset the agent"""
+        self.memory.reset()
+        # TODO: reset receiver

agentverse/agents/tasksolving_agent/role_assigner.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import asyncio
+from colorama import Fore
+
+from agentverse.logging import get_logger
+import bdb
+from string import Template
+from typing import TYPE_CHECKING, List
+
+from agentverse.message import RoleAssignerMessage, Message
+
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+
+
+logger = get_logger()
+
+
+@agent_registry.register("role_assigner")
+class RoleAssignerAgent(BaseAgent):
+    def step(
+        self, advice: str, task_description: str, cnt_critic_agents: int
+    ) -> RoleAssignerMessage:
+        logger.debug("", self.name, Fore.MAGENTA)
+        prepend_prompt, append_prompt = self.get_all_prompts(
+            advice=advice,
+            task_description=task_description,
+            cnt_critic_agents=cnt_critic_agents,
+        )
+        history = self.memory.to_messages(self.name)
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                response = self.llm.generate_response(
+                    prepend_prompt, history, append_prompt
+                )
+                parsed_response = self.output_parser.parse(response)
+                if len(parsed_response) < cnt_critic_agents:
+                    logger.warn(
+                        f"Number of generate roles ({len(parsed_response)}) and number of group members ({cnt_critic_agents}) do not match."
+                    )
+                    logger.warn("Retrying...")
+                    continue
+                break
+            except (KeyboardInterrupt, bdb.BdbQuit):
+                raise
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+
+        message = RoleAssignerMessage(
+            content=parsed_response, sender=self.name, sender_agent=self
+        )
+        return message
+
+    async def astep(self, env_description: str = "") -> RoleAssignerMessage:
+        """Asynchronous version of step"""
+        pass
+
+    def _fill_prompt_template(
+        self, advice, task_description: str, cnt_critic_agents: int
+    ) -> str:
+        """Fill the placeholders in the prompt template
+
+        In the role_assigner agent, three placeholders are supported:
+        - ${task_description}
+        - ${cnt_critic_agnets}
+        - ${advice}
+        """
+        input_arguments = {
+            "task_description": task_description,
+            "cnt_critic_agents": cnt_critic_agents,
+            "advice": advice,
+        }
+        return Template(self.prompt_template).safe_substitute(input_arguments)
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self) -> None:
+        """Reset the agent"""
+        self.memory.reset()
+        # TODO: reset receiver

agentverse/agents/tasksolving_agent/solver.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import asyncio
+from colorama import Fore
+
+from agentverse.logging import get_logger
+import bdb
+from string import Template
+from typing import TYPE_CHECKING, List, Tuple
+
+# from agentverse.environments import PipelineEnvironment
+from agentverse.message import SolverMessage, Message, CriticMessage
+
+from agentverse.agents import agent_registry
+from agentverse.agents.base import BaseAgent
+from agentverse.utils import AgentCriticism
+
+
+logger = get_logger()
+
+
+@agent_registry.register("solver")
+class SolverAgent(BaseAgent):
+    max_history: int = 3
+
+    def step(
+        self, former_solution: str, advice: str, task_description: str = "", **kwargs
+    ) -> SolverMessage:
+        logger.debug("", self.name, Fore.MAGENTA)
+        # prompt = self._fill_prompt_template(
+        #     former_solution, critic_opinions, advice, task_description
+        # )
+        prepend_prompt, append_prompt = self.get_all_prompts(
+            former_solution=former_solution,
+            task_description=task_description,
+            advice=advice,
+            role_description=self.role_description,
+            **kwargs,
+        )
+        history = self.memory.to_messages(self.name, start_index=-self.max_history)
+        parsed_response = None
+        for i in range(self.max_retry):
+            try:
+                response = self.llm.generate_response(
+                    prepend_prompt, history, append_prompt
+                )
+                parsed_response = self.output_parser.parse(response)
+                break
+            except (KeyboardInterrupt, bdb.BdbQuit):
+                raise
+            except Exception as e:
+                logger.error(e)
+                logger.warn("Retrying...")
+                continue
+
+        if parsed_response is None:
+            logger.error(f"{self.name} failed to generate valid response.")
+
+        message = SolverMessage(
+            content=""
+            if parsed_response is None
+            else parsed_response.return_values["output"],
+            sender=self.name,
+            receiver=self.get_receiver(),
+        )
+        return message
+
+    async def astep(self, env_description: str = "") -> SolverMessage:
+        """Asynchronous version of step"""
+        pass
+
+    def _fill_prompt_template(
+        self,
+        former_solution: str,
+        critic_opinions: List[AgentCriticism],
+        advice: str,
+        task_description: str,
+    ) -> str:
+        """Fill the placeholders in the prompt template
+
+        In the role_assigner agent, three placeholders are supported:
+        - ${task_description}
+        - ${former_solution}
+        - ${critic_messages}
+        - ${advice}
+        """
+        input_arguments = {
+            "task_description": task_description,
+            "former_solution": former_solution,
+            "critic_opinions": "\n".join(
+                [
+                    f"{critic.sender_agent.role_description} said: {critic.criticism}"
+                    for critic in critic_opinions
+                ]
+            ),
+            "advice": advice,
+        }
+        # if discussion_mode:
+        #     template = Template(self.prompt_template[1])
+        # else:
+        template = Template(self.prompt_template)
+        return template.safe_substitute(input_arguments)
+
+    def add_message_to_memory(self, messages: List[Message]) -> None:
+        self.memory.add_message(messages)
+
+    def reset(self) -> None:
+        """Reset the agent"""
+        self.memory.reset()
+        # TODO: reset receiver

agentverse/agentverse.py

@@ -0,0 +1,65 @@
+import asyncio
+import logging
+from typing import List
+
+# from agentverse.agents import Agent
+from agentverse.agents.conversation_agent import BaseAgent
+from agentverse.environments import BaseEnvironment
+from agentverse.initialization import load_agent, load_environment, prepare_task_config
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.INFO,
+)
+
+openai_logger = logging.getLogger("openai")
+openai_logger.setLevel(logging.WARNING)
+
+
+class AgentVerse:
+    def __init__(self, agents: List[BaseAgent], environment: BaseEnvironment):
+        self.agents = agents
+        self.environment = environment
+
+    @classmethod
+    def from_task(cls, task: str, tasks_dir: str):
+        """Build an AgentVerse from a task name.
+        The task name should correspond to a directory in `tasks` directory.
+        Then this method will load the configuration from the yaml file in that directory.
+        """
+        # Prepare the config of the task
+        task_config = prepare_task_config(task, tasks_dir)
+
+        # Build the agents
+        agents = []
+        for agent_configs in task_config["agents"]:
+            agent = load_agent(agent_configs)
+            agents.append(agent)
+
+        # Build the environment
+        env_config = task_config["environment"]
+        env_config["agents"] = agents
+        environment = load_environment(env_config)
+
+        return cls(agents, environment)
+
+    def run(self):
+        """Run the environment from scratch until it is done."""
+        self.environment.reset()
+        while not self.environment.is_done():
+            asyncio.run(self.environment.step())
+
+    def reset(self):
+        self.environment.reset()
+        for agent in self.agents:
+            agent.reset()
+
+    def next(self, *args, **kwargs):
+        """Run the environment for one step and return the return message."""
+        return_message = asyncio.run(self.environment.step(*args, **kwargs))
+        return return_message
+
+    def update_state(self, *args, **kwargs):
+        """Run the environment for one step and return the return message."""
+        self.environment.update_state(*args, **kwargs)

agentverse/demo.py

@@ -0,0 +1,487 @@
+import base64
+import itertools
+import json
+from typing import Dict, List, Tuple
+
+import cv2
+import gradio as gr
+
+from agentverse.agentverse import AgentVerse
+from agentverse.message import Message
+
+
+def cover_img(background, img, place: Tuple[int, int]):
+    """
+    Overlays the specified image to the specified position of the background image.
+    :param background: background image
+    :param img: the specified image
+    :param place: the top-left coordinate of the target location
+    """
+    back_h, back_w, _ = background.shape
+    height, width, _ = img.shape
+    for i, j in itertools.product(range(height), range(width)):
+        if img[i, j, 3]:
+            background[place[0] + i, place[1] + j] = img[i, j, :3]
+
+
+class UI:
+    """
+    the UI of frontend
+    """
+
+    def __init__(self, task: str):
+        """
+        init a UI.
+        default number of students is 0
+        """
+        self.messages = []
+        self.task = task
+        self.backend = AgentVerse.from_task(task)
+        self.turns_remain = 0
+        self.agent_id = {
+            self.backend.agents[idx].name: idx
+            for idx in range(len(self.backend.agents))
+        }
+        self.stu_num = len(self.agent_id) - 1
+        self.autoplay = False
+        self.image_now = None
+        self.text_now = None
+        self.tot_solutions = 5
+        self.solution_status = [False] * self.tot_solutions
+
+    def get_avatar(self, idx):
+        if idx == -1:
+            img = cv2.imread("./imgs/db_diag/-1.png")
+        elif self.task == "prisoner_dilemma":
+            img = cv2.imread(f"./imgs/prison/{idx}.png")
+        elif self.task == "db_diag":
+            img = cv2.imread(f"./imgs/db_diag/{idx}.png")
+        elif "sde" in self.task:
+            img = cv2.imread(f"./imgs/sde/{idx}.png")
+        else:
+            img = cv2.imread(f"./imgs/{idx}.png")
+        base64_str = cv2.imencode(".png", img)[1].tostring()
+        return "data:image/png;base64," + base64.b64encode(base64_str).decode("utf-8")
+
+    def stop_autoplay(self):
+        self.autoplay = False
+        return (
+            gr.Button.update(interactive=False),
+            gr.Button.update(interactive=False),
+            gr.Button.update(interactive=False),
+        )
+
+    def start_autoplay(self):
+        self.autoplay = True
+        yield (
+            self.image_now,
+            self.text_now,
+            gr.Button.update(interactive=False),
+            gr.Button.update(interactive=True),
+            gr.Button.update(interactive=False),
+            *[gr.Button.update(visible=statu) for statu in self.solution_status],
+            gr.Box.update(visible=any(self.solution_status)),
+        )
+
+        while self.autoplay and self.turns_remain > 0:
+            outputs = self.gen_output()
+            self.image_now, self.text_now = outputs
+
+            yield (
+                *outputs,
+                gr.Button.update(interactive=not self.autoplay and self.turns_remain > 0),
+                gr.Button.update(interactive=self.autoplay and self.turns_remain > 0),
+                gr.Button.update(interactive=not self.autoplay and self.turns_remain > 0),
+                *[gr.Button.update(visible=statu) for statu in self.solution_status],
+                gr.Box.update(visible=any(self.solution_status))
+            )
+
+    def delay_gen_output(self):
+        yield (
+            self.image_now,
+            self.text_now,
+            gr.Button.update(interactive=False),
+            gr.Button.update(interactive=False),
+            *[gr.Button.update(visible=statu) for statu in self.solution_status],
+            gr.Box.update(visible=any(self.solution_status))
+        )
+
+        outputs = self.gen_output()
+        self.image_now, self.text_now = outputs
+
+        yield (
+            self.image_now,
+            self.text_now,
+            gr.Button.update(interactive=self.turns_remain > 0),
+            gr.Button.update(interactive=self.turns_remain > 0),
+            *[gr.Button.update(visible=statu) for statu in self.solution_status],
+            gr.Box.update(visible=any(self.solution_status))
+        )
+
+    def delay_reset(self):
+        self.autoplay = False
+        self.image_now, self.text_now = self.reset()
+        return (
+            self.image_now,
+            self.text_now,
+            gr.Button.update(interactive=True),
+            gr.Button.update(interactive=False),
+            gr.Button.update(interactive=True),
+            *[gr.Button.update(visible=statu) for statu in self.solution_status],
+            gr.Box.update(visible=any(self.solution_status))
+        )
+
+    def reset(self, stu_num=0):
+        """
+        tell backend the new number of students and generate new empty image
+        :param stu_num:
+        :return: [empty image, empty message]
+        """
+        if not 0 <= stu_num <= 30:
+            raise gr.Error("the number of students must be between 0 and 30.")
+
+        """
+        # [To-Do] Need to add a function to assign agent numbers into the backend.
+        """
+        # self.backend.reset(stu_num)
+        # self.stu_num = stu_num
+
+        """
+        # [To-Do] Pass the parameters to reset
+        """
+        self.backend.reset()
+        self.turns_remain = self.backend.environment.max_turns
+
+        if self.task == "prisoner_dilemma":
+            background = cv2.imread("./imgs/prison/case_1.png")
+        elif self.task == "db_diag":
+            background = cv2.imread("./imgs/db_diag/background.png")
+        elif "sde" in self.task:
+            background = cv2.imread("./imgs/sde/background.png")
+        else:
+            background = cv2.imread("./imgs/background.png")
+            back_h, back_w, _ = background.shape
+            stu_cnt = 0
+            for h_begin, w_begin in itertools.product(
+                    range(800, back_h, 300), range(135, back_w - 200, 200)
+            ):
+                stu_cnt += 1
+                img = cv2.imread(
+                    f"./imgs/{(stu_cnt - 1) % 11 + 1 if stu_cnt <= self.stu_num else 'empty'}.png",
+                    cv2.IMREAD_UNCHANGED,
+                )
+                cover_img(
+                    background,
+                    img,
+                    (h_begin - 30 if img.shape[0] > 190 else h_begin, w_begin),
+                )
+        self.messages = []
+        self.solution_status = [False] * self.tot_solutions
+        return [cv2.cvtColor(background, cv2.COLOR_BGR2RGB), ""]
+
+    def gen_img(self, data: List[Dict]):
+        """
+        generate new image with sender rank
+        :param data:
+        :return: the new image
+        """
+        # The following code need to be more general. This one is too task-specific.
+        # if len(data) != self.stu_num:
+        if len(data) != self.stu_num + 1:
+            raise gr.Error("data length is not equal to the total number of students.")
+        if self.task == "prisoner_dilemma":
+            img = cv2.imread("./imgs/speaking.png", cv2.IMREAD_UNCHANGED)
+            if (
+                    len(self.messages) < 2
+                    or self.messages[-1][0] == 1
+                    or self.messages[-2][0] == 2
+            ):
+                background = cv2.imread("./imgs/prison/case_1.png")
+                if data[0]["message"] != "":
+                    cover_img(background, img, (400, 480))
+            else:
+                background = cv2.imread("./imgs/prison/case_2.png")
+                if data[0]["message"] != "":
+                    cover_img(background, img, (400, 880))
+            if data[1]["message"] != "":
+                cover_img(background, img, (550, 480))
+            if data[2]["message"] != "":
+                cover_img(background, img, (550, 880))
+        elif self.task == "db_diag":
+            background = cv2.imread("./imgs/db_diag/background.png")
+            img = cv2.imread("./imgs/db_diag/speaking.png", cv2.IMREAD_UNCHANGED)
+            if data[0]["message"] != "":
+                cover_img(background, img, (750, 80))
+            if data[1]["message"] != "":
+                cover_img(background, img, (310, 220))
+            if data[2]["message"] != "":
+                cover_img(background, img, (522, 11))
+        elif "sde" in self.task:
+            background = cv2.imread("./imgs/sde/background.png")
+            img = cv2.imread("./imgs/sde/speaking.png", cv2.IMREAD_UNCHANGED)
+            if data[0]["message"] != "":
+                cover_img(background, img, (692, 330))
+            if data[1]["message"] != "":
+                cover_img(background, img, (692, 660))
+            if data[2]["message"] != "":
+                cover_img(background, img, (692, 990))
+        else:
+            background = cv2.imread("./imgs/background.png")
+            back_h, back_w, _ = background.shape
+            stu_cnt = 0
+            if data[stu_cnt]["message"] not in ["", "[RaiseHand]"]:
+                img = cv2.imread("./imgs/speaking.png", cv2.IMREAD_UNCHANGED)
+                cover_img(background, img, (370, 1250))
+            for h_begin, w_begin in itertools.product(
+                    range(800, back_h, 300), range(135, back_w - 200, 200)
+            ):
+                stu_cnt += 1
+                if stu_cnt <= self.stu_num:
+                    img = cv2.imread(
+                        f"./imgs/{(stu_cnt - 1) % 11 + 1}.png", cv2.IMREAD_UNCHANGED
+                    )
+                    cover_img(
+                        background,
+                        img,
+                        (h_begin - 30 if img.shape[0] > 190 else h_begin, w_begin),
+                    )
+                    if "[RaiseHand]" in data[stu_cnt]["message"]:
+                        # elif data[stu_cnt]["message"] == "[RaiseHand]":
+                        img = cv2.imread("./imgs/hand.png", cv2.IMREAD_UNCHANGED)
+                        cover_img(background, img, (h_begin - 90, w_begin + 10))
+                    elif data[stu_cnt]["message"] not in ["", "[RaiseHand]"]:
+                        img = cv2.imread("./imgs/speaking.png", cv2.IMREAD_UNCHANGED)
+                        cover_img(background, img, (h_begin - 90, w_begin + 10))
+
+                else:
+                    img = cv2.imread("./imgs/empty.png", cv2.IMREAD_UNCHANGED)
+                    cover_img(background, img, (h_begin, w_begin))
+        return cv2.cvtColor(background, cv2.COLOR_BGR2RGB)
+
+    def return_format(self, messages: List[Message]):
+        _format = [{"message": "", "sender": idx} for idx in range(len(self.agent_id))]
+
+        for message in messages:
+            if self.task == "db_diag":
+                content_json: dict = message.content
+                content_json["diagnose"] = f"[{message.sender}]: {content_json['diagnose']}"
+                _format[self.agent_id[message.sender]]["message"] = json.dumps(content_json)
+            elif "sde" in self.task:
+                if message.sender == "code_tester":
+                    pre_message, message_ = message.content.split("\n")
+                    message_ = "{}\n{}".format(pre_message, json.loads(message_)["feedback"])
+                    _format[self.agent_id[message.sender]]["message"] = "[{}]: {}".format(
+                        message.sender, message_
+                    )
+                else:
+                    _format[self.agent_id[message.sender]]["message"] = "[{}]: {}".format(
+                        message.sender, message.content
+                    )
+
+            else:
+                _format[self.agent_id[message.sender]]["message"] = "[{}]: {}".format(
+                    message.sender, message.content
+                )
+
+        return _format
+
+    def gen_output(self):
+        """
+        generate new image and message of next step
+        :return: [new image, new message]
+        """
+
+        # data = self.backend.next_data()
+        return_message = self.backend.next()
+        data = self.return_format(return_message)
+
+        # data.sort(key=lambda item: item["sender"])
+        """
+        # [To-Do]; Check the message from the backend: only 1 person can speak
+        """
+
+        for item in data:
+            if item["message"] not in ["", "[RaiseHand]"]:
+                self.messages.append((item["sender"], item["message"]))
+
+        message = self.gen_message()
+        self.turns_remain -= 1
+        return [self.gen_img(data), message]
+
+    def gen_message(self):
+        # If the backend cannot handle this error, use the following code.
+        message = ""
+        """
+        for item in data:
+            if item["message"] not in ["", "[RaiseHand]"]:
+                message = item["message"]
+                break
+        """
+        for sender, msg in self.messages:
+            if sender == 0:
+                avatar = self.get_avatar(0)
+            elif sender == -1:
+                avatar = self.get_avatar(-1)
+            else:
+                avatar = self.get_avatar((sender - 1) % 11 + 1)
+            if self.task == "db_diag":
+                msg_json = json.loads(msg)
+                self.solution_status = [False] * self.tot_solutions
+                msg = msg_json["diagnose"]
+                if msg_json["solution"] != "":
+                    solution: List[str] = msg_json["solution"]
+                    for solu in solution:
+                        if "query" in solu or "queries" in solu:
+                            self.solution_status[0] = True
+                            solu = solu.replace("query", '<span style="color:yellow;">query</span>')
+                            solu = solu.replace("queries", '<span style="color:yellow;">queries</span>')
+                        if "join" in solu:
+                            self.solution_status[1] = True
+                            solu = solu.replace("join", '<span style="color:yellow;">join</span>')
+                        if "index" in solu:
+                            self.solution_status[2] = True
+                            solu = solu.replace("index", '<span style="color:yellow;">index</span>')
+                        if "system configuration" in solu:
+                            self.solution_status[3] = True
+                            solu = solu.replace("system configuration",
+                                                '<span style="color:yellow;">system configuration</span>')
+                        if "monitor" in solu or "Monitor" in solu or "Investigate" in solu:
+                            self.solution_status[4] = True
+                            solu = solu.replace("monitor", '<span style="color:yellow;">monitor</span>')
+                            solu = solu.replace("Monitor", '<span style="color:yellow;">Monitor</span>')
+                            solu = solu.replace("Investigate", '<span style="color:yellow;">Investigate</span>')
+                        msg = f"{msg}<br>{solu}"
+                if msg_json["knowledge"] != "":
+                    msg = f'{msg}<hr style="margin: 5px 0"><span style="font-style: italic">{msg_json["knowledge"]}<span>'
+            else:
+                msg = msg.replace("<", "&lt;")
+                msg = msg.replace(">", "&gt;")
+            message = (
+                f'<div style="display: flex; align-items: center; margin-bottom: 10px;overflow:auto;">'
+                f'<img src="{avatar}" style="width: 5%; height: 5%; border-radius: 25px; margin-right: 10px;">'
+                f'<div style="background-color: gray; color: white; padding: 10px; border-radius: 10px;'
+                f'max-width: 70%; white-space: pre-wrap">'
+                f"{msg}"
+                f"</div></div>" + message
+            )
+        message = '<div id="divDetail" style="height:600px;overflow:auto;">' + message + "</div>"
+        return message
+
+    def submit(self, message: str):
+        """
+        submit message to backend
+        :param message: message
+        :return: [new image, new message]
+        """
+        self.backend.submit(message)
+        self.messages.append((-1, f"[User]: {message}"))
+        return self.gen_img([{"message": ""}] * len(self.agent_id)), self.gen_message()
+
+    def launch(self):
+        """
+        start a frontend
+        """
+        with gr.Blocks() as demo:
+            with gr.Row():
+                with gr.Column():
+                    image_output = gr.Image()
+                    with gr.Row():
+                        reset_btn = gr.Button("Reset")
+                        # next_btn = gr.Button("Next", variant="primary")
+                        next_btn = gr.Button("Next", interactive=False)
+                        stop_autoplay_btn = gr.Button(
+                            "Stop Autoplay", interactive=False
+                        )
+                        start_autoplay_btn = gr.Button("Start Autoplay", interactive=False)
+                    with gr.Box(visible=False) as solutions:
+                        with gr.Column():
+                            gr.HTML("Optimization Solutions:")
+                            with gr.Row():
+                                rewrite_slow_query_btn = gr.Button("Rewrite Slow Query", visible=False)
+                                add_query_hints_btn = gr.Button("Add Query Hints", visible=False)
+                                update_indexes_btn = gr.Button("Update Indexes", visible=False)
+                                tune_parameters_btn = gr.Button("Tune Parameters", visible=False)
+                                gather_more_info_btn = gr.Button("Gather More Info", visible=False)
+                # text_output = gr.Textbox()
+                text_output = gr.HTML(self.reset()[1])
+
+            # Given a botton to provide student numbers and their inf.
+            # stu_num = gr.Number(label="Student Number", precision=0)
+            # stu_num = self.stu_num
+
+            if self.task == "db_diag":
+                user_msg = gr.Textbox()
+                submit_btn = gr.Button("Submit", variant="primary")
+
+                submit_btn.click(fn=self.submit, inputs=user_msg, outputs=[image_output, text_output], show_progress=False)
+            else:
+                pass
+
+            # next_btn.click(fn=self.gen_output, inputs=None, outputs=[image_output, text_output], show_progress=False)
+            next_btn.click(
+                fn=self.delay_gen_output,
+                inputs=None,
+                outputs=[
+                    image_output,
+                    text_output,
+                    next_btn,
+                    start_autoplay_btn,
+                    rewrite_slow_query_btn,
+                    add_query_hints_btn,
+                    update_indexes_btn,
+                    tune_parameters_btn,
+                    gather_more_info_btn,
+                    solutions
+                ],
+                show_progress=False,
+            )
+
+            # [To-Do] Add botton: re-start (load different people and env)
+            # reset_btn.click(fn=self.reset, inputs=stu_num, outputs=[image_output, text_output], show_progress=False)
+            # reset_btn.click(fn=self.reset, inputs=None, outputs=[image_output, text_output], show_progress=False)
+            reset_btn.click(
+                fn=self.delay_reset,
+                inputs=None,
+                outputs=[
+                    image_output,
+                    text_output,
+                    next_btn,
+                    stop_autoplay_btn,
+                    start_autoplay_btn,
+                    rewrite_slow_query_btn,
+                    add_query_hints_btn,
+                    update_indexes_btn,
+                    tune_parameters_btn,
+                    gather_more_info_btn,
+                    solutions
+                ],
+                show_progress=False,
+            )
+
+            stop_autoplay_btn.click(
+                fn=self.stop_autoplay,
+                inputs=None,
+                outputs=[next_btn, stop_autoplay_btn, start_autoplay_btn],
+                show_progress=False,
+            )
+            start_autoplay_btn.click(
+                fn=self.start_autoplay,
+                inputs=None,
+                outputs=[
+                    image_output,
+                    text_output,
+                    next_btn,
+                    stop_autoplay_btn,
+                    start_autoplay_btn,
+                    rewrite_slow_query_btn,
+                    add_query_hints_btn,
+                    update_indexes_btn,
+                    tune_parameters_btn,
+                    gather_more_info_btn,
+                    solutions
+                ],
+                show_progress=False,
+            )
+
+        demo.queue(concurrency_count=5, max_size=20).launch()
+        # demo.launch()

agentverse/environments/__init__.py

@@ -0,0 +1,17 @@
+from typing import Dict
+from agentverse.registry import Registry
+
+
+env_registry = Registry(name="EnvironmentRegistry")
+
+
+from .base import BaseEnvironment, BaseRule
+
+# from .basic import PipelineEnvironment
+from .simulation_env.basic import BasicEnvironment
+from .simulation_env.pokemon import PokemonEnvironment
+from .simulation_env.prisoner_dilemma import PrisonerDilemmaEnvironment
+from .simulation_env.sde_team import SdeTeamEnvironment
+from .simulation_env.sde_team_given_tests import SdeTeamGivenTestsEnvironment
+
+from .tasksolving_env.basic import BasicEnvironment

agentverse/environments/base.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+from agentverse.logging import logger
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Any, Dict, List
+
+from pydantic import BaseModel
+
+# from agentverse.agents.agent import Agent
+
+if TYPE_CHECKING:
+    from agentverse.agents.base import BaseAgent
+    from agentverse.message import Message
+
+
+class BaseRule(BaseModel):
+    pass
+
+
+class BaseEnvironment(BaseModel):
+    """
+    Base class for environment.
+
+    Args:
+        agents: List of agents
+        rule: Rule for the environment
+        max_turns: Maximum number of turns
+        cnt_turn: Current turn number
+        last_messages: Messages from last turn
+        rule_params: Variables set by the rule
+    """
+
+    agents: List[BaseAgent]
+    rule: BaseRule
+    max_turns: int = 10
+    cnt_turn: int = 0
+    last_messages: List[Message] = []
+    rule_params: Dict = {}
+
+    @abstractmethod
+    async def step(self) -> List[Message]:
+        """Run one step of the environment"""
+        pass
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Reset the environment"""
+        pass
+
+    def report_metrics(self) -> None:
+        """Report useful metrics"""
+        total_spent = sum([agent.get_spend() for agent in self.agents])
+        logger.info(f"Total spent: ${total_spent}")
+        pass
+
+    def is_done(self) -> bool:
+        """Check if the environment is done"""
+        return self.cnt_turn >= self.max_turns

agentverse/environments/simulation_env/basic.py

@@ -0,0 +1,101 @@
+import asyncio
+
+# import logging
+from agentverse.logging import get_logger
+from typing import Any, Dict, List
+
+# from agentverse.agents.agent import Agent
+from agentverse.agents.simulation_agent.conversation import BaseAgent
+
+# from agentverse.environments.simulation_env.rules.base import Rule
+from agentverse.environments.simulation_env.rules.base import SimulationRule as Rule
+from agentverse.message import Message
+
+logger = get_logger()
+
+from .. import env_registry as EnvironmentRegistry
+from ..base import BaseEnvironment
+
+
+@EnvironmentRegistry.register("sim-basic")
+class BasicEnvironment(BaseEnvironment):
+    """
+    A basic environment implementing the logic of conversation.
+
+    Args:
+        agents: List of agents
+        rule: Rule for the environment
+        max_turns: Maximum number of turns
+        cnt_turn: Current turn number
+        last_messages: Messages from last turn
+        rule_params: Variables set by the rule
+    """
+
+    agents: List[BaseAgent]
+    rule: Rule
+    max_turns: int = 10
+    cnt_turn: int = 0
+    last_messages: List[Message] = []
+    rule_params: Dict = {}
+
+    def __init__(self, rule, **kwargs):
+        rule_config = rule
+        order_config = rule_config.get("order", {"type": "sequential"})
+        visibility_config = rule_config.get("visibility", {"type": "all"})
+        selector_config = rule_config.get("selector", {"type": "basic"})
+        updater_config = rule_config.get("updater", {"type": "basic"})
+        describer_config = rule_config.get("describer", {"type": "basic"})
+        rule = Rule(
+            order_config,
+            visibility_config,
+            selector_config,
+            updater_config,
+            describer_config,
+        )
+        super().__init__(rule=rule, **kwargs)
+
+    async def step(self) -> List[Message]:
+        """Run one step of the environment"""
+
+        # Get the next agent index
+        agent_ids = self.rule.get_next_agent_idx(self)
+
+        # Generate current environment description
+        env_descriptions = self.rule.get_env_description(self)
+
+        # Generate the next message
+        messages = await asyncio.gather(
+            *[self.agents[i].astep(env_descriptions[i]) for i in agent_ids]
+        )
+
+        # Some rules will select certain messages from all the messages
+        selected_messages = self.rule.select_message(self, messages)
+        self.last_messages = selected_messages
+        self.print_messages(selected_messages)
+
+        # Update the memory of the agents
+        self.rule.update_memory(self)
+
+        # Update the set of visible agents for each agent
+        self.rule.update_visible_agents(self)
+
+        self.cnt_turn += 1
+
+        return selected_messages
+
+    def print_messages(self, messages: List[Message]) -> None:
+        for message in messages:
+            if message is not None:
+                # logging.info(f"{message.sender}: {message.content}")
+                logger.info(f"{message.sender}: {message.content}")
+
+    def reset(self) -> None:
+        """Reset the environment"""
+        self.cnt_turn = 0
+        self.rule.reset()
+        for agent in self.agents:
+            agent.reset()
+
+    def is_done(self) -> bool:
+        """Check if the environment is done"""
+        return self.cnt_turn >= self.max_turns

agentverse/environments/simulation_env/pokemon.py

@@ -0,0 +1,222 @@
+import asyncio
+import datetime
+import logging
+from typing import Any, Dict, List, Optional, Set
+
+# from agentverse.agents.agent import Agent
+from agentverse.agents.simulation_agent.conversation import BaseAgent
+
+# from agentverse.environments.simulation_env.rules.base import Rule
+from agentverse.environments.simulation_env.rules.base import SimulationRule as Rule
+from agentverse.message import Message
+
+from .. import env_registry as EnvironmentRegistry
+from ..base import BaseEnvironment
+
+
+@EnvironmentRegistry.register("pokemon")
+class PokemonEnvironment(BaseEnvironment):
+    """
+    An environment for Pokémon demo.
+
+    Args:
+        agents: List of agents
+        locations: A dict of locations to agents within them
+        rule: Rule for the environment
+        max_turns: Maximum number of turns
+        cnt_turn: Current turn number
+        last_messages: Messages from last turn
+        rule_params: Variables set by the rule
+    """
+
+    agents: List[BaseAgent]
+    locations_to_agents: Dict[str, Set[str]]
+    # locations_descriptions: Dict[str, str]
+    time: datetime.datetime = datetime.datetime(2021, 1, 1, 8, 0, 0)
+    rule: Rule
+    max_turns: int = 10
+    cnt_turn: int = 0
+    last_messages: List[Message] = []
+    rule_params: Dict = {}
+
+    def __init__(self, rule, locations, **kwargs):
+        rule_config = rule
+        order_config = rule_config.get("order", {"type": "sequential"})
+        visibility_config = rule_config.get("visibility", {"type": "all"})
+        selector_config = rule_config.get("selector", {"type": "basic"})
+        updater_config = rule_config.get("updater", {"type": "basic"})
+        describer_config = rule_config.get("describer", {"type": "basic"})
+        rule = Rule(
+            order_config,
+            visibility_config,
+            selector_config,
+            updater_config,
+            describer_config,
+        )
+        locations_to_agents = {}
+        # locations_descriptions = {}
+        locations_config = locations
+        for loc in locations_config:
+            locations_to_agents[loc["name"]] = set(loc["init_agents"])
+            # locations_descriptions[loc["name"]] = loc["description"]
+        super().__init__(
+            rule=rule,
+            locations_to_agents=locations_to_agents,
+            # locations_descriptions=locations_descriptions,
+            **kwargs,
+        )
+
+    async def step(
+        self,
+        is_player: bool = False,
+        player_content: str = None,
+        receiver: str = None,
+        receiver_id: Optional[int] = None,
+        agent_ids: Optional[List[int]] = None,
+    ) -> List[Message]:
+        """Run one step of the environment"""
+
+        # Get the next agent index
+        # time.sleep(8)
+        # return [Message(content="Test", sender="May", receiver=["May"])]
+        if is_player:
+            return await self._respond_to_player(player_content, receiver, receiver_id)
+        else:
+            return await self._routine_step(agent_ids)
+
+    async def _routine_step(self, agent_ids) -> List[Message]:
+        self.rule.update_visible_agents(self)
+
+        # agent_ids = self.rule.get_next_agent_idx(self)
+
+        # Generate current environment description
+        env_descriptions = self.rule.get_env_description(self)
+
+        # Generate the next message
+        messages = await asyncio.gather(
+            *[self.agents[i].astep(env_descriptions[i]) for i in agent_ids]
+        )
+        # messages = self.get_test_messages()
+
+        # Some rules will select certain messages from all the messages
+        selected_messages = self.rule.select_message(self, messages)
+
+        # Update the memory of the agents
+        self.last_messages = selected_messages
+        self.rule.update_memory(self)
+        self.print_messages(selected_messages)
+
+        self.cnt_turn += 1
+        self.time += datetime.timedelta(minutes=5)
+
+        return selected_messages
+
+    async def _respond_to_player(
+        self,
+        player_content: str = None,
+        receiver: str = None,
+        receiver_id: Optional[int] = None,
+    ) -> List[Message]:
+        if receiver_id is None:
+            for agent in self.agents:
+                if agent.name == receiver:
+                    receiver_id = agent.agent_id
+                    break
+        agent_ids = [receiver_id]
+        agent_name = receiver
+        player_message = Message(
+            sender="Brenden", content=player_content, receiver=[agent_name]
+        )
+
+        # Update the set of visible agents for each agent
+        self.rule.update_visible_agents(self)
+
+        # Generate current environment description
+        env_descriptions = self.rule.get_env_description(self, player_content)
+
+        # Generate the next message
+        messages = await asyncio.gather(
+            *[self.agents[i].astep(env_descriptions[i]) for i in agent_ids]
+        )
+
+        # Some rules will select certain messages from all the messages
+        # selected_messages = self.rule.select_message(self, messages)
+
+        # Update the memory of the agents
+        self.last_messages = [player_message, *messages]
+        self.rule.update_memory(self)
+        self.print_messages(messages)
+
+        self.cnt_turn += 1
+
+        return messages
+
+    def update_state(self, agent_location: Dict[str, str]):
+        for agent_name, location in agent_location.items():
+            # original_location = self.get_agent_to_location()[agent_name]
+            # self.locations_to_agents[original_location].remove(agent_name)
+            self.locations_to_agents[location].add(agent_name)
+
+    def get_agent_to_location(self) -> Dict[str, str]:
+        ret = {}
+        for location, agent_names in self.locations_to_agents.items():
+            for agent in agent_names:
+                ret[agent] = location
+        return ret
+
+    def print_messages(self, messages: List[Message]) -> None:
+        for message in messages:
+            if message is not None:
+                logging.info(f"{message.sender}: {message.content}")
+
+    def reset(self) -> None:
+        """Reset the environment"""
+        self.cnt_turn = 0
+        self.rule.reset()
+        for agent in self.agents:
+            agent.reset()
+
+    def is_done(self) -> bool:
+        """Check if the environment is done"""
+        return self.cnt_turn >= self.max_turns
+
+    def get_test_messages(self) -> List[Message]:
+        messages = [
+            Message(
+                content='{"to": "Birch", "action": "Speak", "text": "Hi!!!"}',
+                sender="May",
+                receiver={"May", "Birch"},
+                tool_response=[],
+            ),
+            Message(
+                content='{"to": "May", "text": "Good morning, May! How is your research going?", "action": "Speak"}',
+                sender="Birch",
+                receiver={"May", "Birch"},
+                tool_response=[],
+            ),
+            Message(
+                content='{"to": "Pokémon Center", "action": "MoveTo"}',
+                sender="Steven",
+                receiver={"Steven"},
+                tool_response=[],
+            ),
+            Message(
+                content='{"to": "Shop", "last_time": "10 minutes", "action": "MoveTo"}',
+                sender="Maxie",
+                receiver={"Maxie"},
+                tool_response=[],
+            ),
+            Message(
+                content='{"to": "Pok\\u00e9mon Center", "action": "MoveTo"}',
+                sender="Archie",
+                receiver={"Archie"},
+                tool_response=[],
+            ),
+            Message(
+                content='{"to": "Shop", "action": "MoveTo"}',
+                sender="Joseph",
+                receiver={"Joseph"},
+                tool_response=[],
+            ),
+        ]
+        return messages

agentverse/environments/simulation_env/prisoner_dilemma.py

@@ -0,0 +1,49 @@
+import asyncio
+import logging
+from typing import Any, Dict, List
+
+# from agentverse.agents.agent import Agent
+from agentverse.agents.simulation_agent.conversation import BaseAgent
+
+# from agentverse.environments.simulation_env.rules.base import Rule
+from agentverse.environments.simulation_env.rules.base import SimulationRule as Rule
+from agentverse.message import Message
+
+from .. import env_registry as EnvironmentRegistry
+from .basic import BasicEnvironment
+
+
+@EnvironmentRegistry.register("prisoner_dilemma")
+class PrisonerDilemmaEnvironment(BasicEnvironment):
+    """
+    An environment for prisoner dilemma.
+    """
+
+    async def step(self) -> List[Message]:
+        """Run one step of the environment"""
+
+        # Get the next agent index
+        agent_ids = self.rule.get_next_agent_idx(self)
+
+        # Generate current environment description
+        env_descriptions = self.rule.get_env_description(self)
+
+        # Generate the next message
+        messages = await asyncio.gather(
+            *[self.agents[i].astep(self, env_descriptions[i]) for i in agent_ids]
+        )
+
+        # Some rules will select certain messages from all the messages
+        selected_messages = self.rule.select_message(self, messages)
+        self.last_messages = selected_messages
+        self.print_messages(selected_messages)
+
+        # Update the memory of the agents
+        self.rule.update_memory(self)
+
+        # Update the set of visible agents for each agent
+        self.rule.update_visible_agents(self)
+
+        self.cnt_turn += 1
+
+        return selected_messages

agentverse/environments/simulation_env/reflection.py

@@ -0,0 +1,128 @@
+import asyncio
+import logging
+from typing import Any, Dict, List
+
+from datetime import datetime as dt
+import datetime
+
+from pydantic import Field
+
+from agentverse.agents.simulation_agent.conversation import BaseAgent
+
+# from agentverse.environments.simulation_env.rules.base import Rule
+from agentverse.environments.simulation_env.rules.base import SimulationRule as Rule
+from agentverse.message import Message
+
+from . import env_registry as EnvironmentRegistry
+from ..base import BaseEnvironment
+
+from pydantic import validator
+
+
+@EnvironmentRegistry.register("reflection")
+class ReflectionEnvironment(BaseEnvironment):
+    """
+    Environment used in Observation-Planning-Reflection agent architecture.
+
+    Args:
+        agents: List of agents
+        rule: Rule for the environment
+        max_turns: Maximum number of turns
+        cnt_turn: Current turn number
+        last_messages: Messages from last turn
+        rule_params: Variables set by the rule
+        current_time
+        time_delta: time difference between steps
+    """
+
+    agents: List[BaseAgent]
+    rule: Rule
+    max_turns: int = 10
+    cnt_turn: int = 0
+    last_messages: List[Message] = []
+    rule_params: Dict = {}
+    current_time: dt = dt.now()
+    time_delta: int = 120
+    #
+
+    # @validator("time_delta")
+    # def convert_str_to_timedelta(cls, string):
+    #
+    #     return datetime.timedelta(seconds=int(string))
+
+    def __init__(self, rule, **kwargs):
+        rule_config = rule
+        order_config = rule_config.get("order", {"type": "sequential"})
+        visibility_config = rule_config.get("visibility", {"type": "all"})
+        selector_config = rule_config.get("selector", {"type": "basic"})
+        updater_config = rule_config.get("updater", {"type": "basic"})
+        describer_config = rule_config.get("describer", {"type": "basic"})
+        rule = Rule(
+            order_config,
+            visibility_config,
+            selector_config,
+            updater_config,
+            describer_config,
+        )
+
+        super().__init__(rule=rule, **kwargs)
+
+    async def step(self) -> List[Message]:
+        """Run one step of the environment"""
+
+        logging.log(logging.INFO, f"Tick tock. Current time: {self.current_time}")
+
+        # Get the next agent index
+        agent_ids = self.rule.get_next_agent_idx(self)
+
+        # Generate current environment description
+        env_descriptions = self.rule.get_env_description(self)
+
+        # Generate the next message
+        messages = await asyncio.gather(
+            *[
+                self.agents[i].astep(self.current_time, env_descriptions[i])
+                for i in agent_ids
+            ]
+        )
+
+        # Some rules will select certain messages from all the messages
+        selected_messages = self.rule.select_message(self, messages)
+        self.last_messages = selected_messages
+        self.print_messages(selected_messages)
+
+        # Update the memory of the agents
+        self.rule.update_memory(self)
+
+        # Update the set of visible agents for each agent
+        self.rule.update_visible_agents(self)
+
+        self.cnt_turn += 1
+
+        # update current_time
+        self.tick_tock()
+
+        return selected_messages
+
+    def print_messages(self, messages: List[Message]) -> None:
+        for message in messages:
+            if message is not None:
+                logging.info(f"{message.sender}: {message.content}")
+
+    def reset(self) -> None:
+        """Reset the environment"""
+        self.cnt_turn = 0
+        self.rule.reset()
+        BaseAgent.update_forward_refs()
+        for agent in self.agents:
+            agent.reset(environment=self)
+
+    def is_done(self) -> bool:
+        """Check if the environment is done"""
+        return self.cnt_turn >= self.max_turns
+
+    def tick_tock(self) -> None:
+        """Increment the time"""
+        self.current_time = self.current_time + datetime.timedelta(
+            seconds=self.time_delta
+        )

agentverse/environments/simulation_env/rules/__init__.py

@@ -0,0 +1 @@
+from .base import SimulationRule

agentverse/environments/simulation_env/rules/base.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from agentverse.environments.simulation_env.rules.describer import (
+    BaseDescriber,
+    describer_registry,
+)
+from agentverse.environments.simulation_env.rules.order import BaseOrder, order_registry
+from agentverse.environments.simulation_env.rules.selector import (
+    BaseSelector,
+    selector_registry,
+)
+from agentverse.environments.simulation_env.rules.updater import (
+    BaseUpdater,
+    updater_registry,
+)
+from agentverse.environments.simulation_env.rules.visibility import (
+    BaseVisibility,
+    visibility_registry,
+)
+from agentverse.environments import BaseRule
+
+if TYPE_CHECKING:
+    from agentverse.environments.base import BaseEnvironment
+
+from agentverse.message import Message
+
+
+# class Rule(BaseModel):
+class SimulationRule(BaseRule):
+    """
+    Rule for the environment. It controls the speaking order of the agents
+    and maintain the set of visible agents for each agent.
+    """
+
+    order: BaseOrder
+    visibility: BaseVisibility
+    selector: BaseSelector
+    updater: BaseUpdater
+    describer: BaseDescriber
+
+    def __init__(
+        self,
+        order_config,
+        visibility_config,
+        selector_config,
+        updater_config,
+        describer_config,
+    ):
+        order = order_registry.build(**order_config)
+        visibility = visibility_registry.build(**visibility_config)
+        selector = selector_registry.build(**selector_config)
+        updater = updater_registry.build(**updater_config)
+        describer = describer_registry.build(**describer_config)
+        super().__init__(
+            order=order,
+            visibility=visibility,
+            selector=selector,
+            updater=updater,
+            describer=describer,
+        )
+
+    def get_next_agent_idx(
+        self, environment: BaseEnvironment, *args, **kwargs
+    ) -> List[int]:
+        """Return the index of the next agent to speak"""
+        return self.order.get_next_agent_idx(environment, *args, **kwargs)
+
+    def update_visible_agents(
+        self, environment: BaseEnvironment, *args, **kwargs
+    ) -> None:
+        """Update the set of visible agents for the agent"""
+        self.visibility.update_visible_agents(environment, *args, **kwargs)
+
+    def select_message(
+        self, environment: BaseEnvironment, messages: List[Message], *args, **kwargs
+    ) -> List[Message]:
+        """Select a set of valid messages from all the generated messages"""
+        return self.selector.select_message(environment, messages, *args, **kwargs)
+
+    def update_memory(self, environment: BaseEnvironment, *args, **kwargs) -> None:
+        """For each message, add it to the memory of the agent who is able to see that message"""
+        self.updater.update_memory(environment, *args, **kwargs)
+
+    def get_env_description(
+        self, environment: BaseEnvironment, *args, **kwargs
+    ) -> List[str]:
+        """Return the description of the environment for each agent"""
+        return self.describer.get_env_description(environment, *args, **kwargs)
+
+    def reset(self) -> None:
+        self.order.reset()
+        self.visibility.reset()
+        self.selector.reset()
+        self.updater.reset()
+        self.describer.reset()

agentverse/environments/simulation_env/rules/describer/__init__.py

@@ -0,0 +1,9 @@
+from agentverse.registry import Registry
+
+describer_registry = Registry(name="DescriberRegistry")
+
+from .base import BaseDescriber
+from .basic import BasicDescriber
+from .classroom import ClassroomDescriber
+from .pokemon import PokemonDescriber
+from .prisoner import PrisonerDescriber

agentverse/environments/simulation_env/rules/describer/base.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List
+
+from pydantic import BaseModel
+
+from . import describer_registry as DescriberRegistry
+from abc import abstractmethod
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+class BaseDescriber(BaseModel):
+    @abstractmethod
+    def get_env_description(
+        self, environment: BaseEnvironment, *args, **kwargs
+    ) -> List[str]:
+        """Return the environment description for each agent"""
+        pass
+
+    def reset(self) -> None:
+        pass

agentverse/environments/simulation_env/rules/describer/basic.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List
+
+from . import describer_registry as DescriberRegistry
+from .base import BaseDescriber
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@DescriberRegistry.register("basic")
+class BasicDescriber(BaseDescriber):
+    def get_env_description(self, environment: BaseEnvironment) -> List[str]:
+        """Return the environment description for each agent"""
+        return ["" for _ in range(len(environment.agents))]

agentverse/environments/simulation_env/rules/describer/classroom.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List
+from string import Template
+
+from . import describer_registry as DescriberRegistry
+from .basic import BasicDescriber
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@DescriberRegistry.register("classroom")
+class ClassroomDescriber(BasicDescriber):
+    start_prompt: str
+    end_prompt: str
+
+    def get_env_description(self, environment: BaseEnvironment) -> List[str]:
+        if not environment.rule_params.get("is_grouped", False):
+            if environment.rule_params.get("is_grouped_ended", False):
+                # If the group discussion is just ended
+                environment.rule_params["is_grouped_ended"] = False
+                return [self.end_prompt for _ in range(len(environment.agents))]
+            else:
+                return super().get_env_description(environment)
+        description = []
+        for i, agent in enumerate(environment.agents):
+            if i == 0:
+                # Professor will not participate in group discussion
+                description.append("")
+            else:
+                description.append(
+                    Template(self.start_prompt).safe_substitute(
+                        {"receiver_name": ", ".join(agent.receiver)}
+                    )
+                )
+        return description
+
+    def reset(self) -> None:
+        pass

agentverse/environments/simulation_env/rules/describer/pokemon.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List, Optional, Dict
+from copy import deepcopy
+
+from . import describer_registry as DescriberRegistry
+from .base import BaseDescriber
+
+if TYPE_CHECKING:
+    from agentverse.environments.pokemon import PokemonEnvironment
+
+
+@DescriberRegistry.register("pokemon")
+class PokemonDescriber(BaseDescriber):
+    def get_env_description(
+        self,
+        environment: PokemonEnvironment,
+        player_content: str = "",
+    ) -> List[str]:
+        time = environment.time
+        if player_content == "":
+            agent_to_location = environment.get_agent_to_location()
+            descriptions = []
+            for agent in environment.agents:
+                description = ""
+                if agent.name not in agent_to_location:
+                    # Agent is on the way to a location
+                    descriptions.append("")
+                    continue
+                location = agent_to_location[agent.name]
+                agents_in_same_loc = deepcopy(environment.locations_to_agents[location])
+                agents_in_same_loc.remove(agent.name)
+                agents_in_same_loc = list(agents_in_same_loc)
+                description += f"It is now {time}. You are at {location}."
+                if len(agents_in_same_loc) == 0:
+                    description += " There is no one else here."
+                elif len(agents_in_same_loc) == 1:
+                    description += f" {agents_in_same_loc[0]} is also here."
+                else:
+                    other_agents = ", ".join(agents_in_same_loc)
+                    description += f" {other_agents} are also here."
+                # description += " The locations you can go to include: \n"
+                # for loc, dsec in environment.locations_descriptions.items():
+                #     description += f"{loc}: {dsec}\n"
+                descriptions.append(description)
+            return descriptions
+        else:
+            description = ""
+            description += f"It is now {time}. Brendan is talking to you.\n"
+            description += f"[Brendan]: {player_content}\n"
+            return [description for _ in range(len(environment.agents))]

agentverse/environments/simulation_env/rules/describer/prisoner.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List
+
+from . import describer_registry as DescriberRegistry
+from .base import BaseDescriber
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@DescriberRegistry.register("prisoner")
+class PrisonerDescriber(BaseDescriber):
+    switch_func = {
+        "Both Suspects": "Suspect2",
+        "Suspect1": "Suspect2",
+        "Suspect2": "Suspect1",
+    }
+    receiver: str = "Both Suspects"
+
+    def get_env_description(self, environment: BaseEnvironment) -> List[str]:
+        if environment.cnt_turn == 0:
+            environment.agents[0].set_receiver({"all"})
+            environment.agents[1].set_receiver({"Police", "Suspect1"})
+            environment.agents[2].set_receiver({"Police", "Suspect2"})
+
+        # only police have to choose to talk to suspect1 or suspect
+        description = []
+        for i, agent in enumerate(environment.agents):
+            if i == 0:
+                # police -> suspect1 -> police -> suspect2
+                if environment.cnt_turn % 2 == 1:
+                    description.append("")
+                    continue
+
+                # Police will have to choose talk to which suspect
+                description.append(f"You are now talking to {self.receiver}")
+
+                receiver = "all" if self.receiver == "Both Suspects" else self.receiver
+                self.receiver = self.switch_func[self.receiver]
+                agent.set_receiver({receiver})
+
+            else:
+                description.append("")
+
+        return description
+
+    def reset(self) -> None:
+        pass

agentverse/environments/simulation_env/rules/order/__init__.py

@@ -0,0 +1,11 @@
+from agentverse.registry import Registry
+order_registry = Registry(name="OrderRegistry")
+
+from .base import BaseOrder
+from .sequential import SequentialOrder
+from .random import RandomOrder
+from .concurrent import ConcurrentOrder
+from .classroom import ClassroomOrder
+from .prisoner import PrisonerOrder
+from .sde_team import SdeTeamOrder
+from .sde_team_given_tests import SdeTeamGivenTestsOrder

agentverse/environments/simulation_env/rules/order/base.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Any, List
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+class BaseOrder(BaseModel):
+    @abstractmethod
+    def get_next_agent_idx(self, environment: BaseEnvironment) -> List[int]:
+        """Return the index of the next agent to speak"""
+
+    def reset(self) -> None:
+        pass

agentverse/environments/simulation_env/rules/order/classroom.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import logging
+import re
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from . import order_registry as OrderRegistry
+from .base import BaseOrder
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@OrderRegistry.register("classroom")
+class ClassroomOrder(BaseOrder):
+    """The order for a classroom discussion
+    The agents speak in the following order:
+    1. The professor speaks first
+    2. Then the professor can continue to speak, and the students can raise hands
+    3. The professor can call on a student, then the student can speak or ask a question
+    4. In the group discussion, the students in the group can speak in turn
+    """
+
+    def get_next_agent_idx(self, environment: BaseEnvironment) -> List[int]:
+        # `is_grouped_ended`: whether the group discussion just ended
+        # `is_grouped`: whether it is currently in a group discussion
+        if environment.rule_params.get("is_grouped_ended", False):
+            return [0]
+        if environment.rule_params.get("is_grouped", False):
+            return self.get_next_agent_idx_grouped(environment)
+        else:
+            return self.get_next_agent_idx_ungrouped(environment)
+
+    def get_next_agent_idx_ungrouped(self, environment: BaseEnvironment) -> List[int]:
+        if len(environment.last_messages) == 0:
+            # If the class just begins or no one speaks in the last turn, we let only the professor speak
+            return [0]
+        elif len(environment.last_messages) == 1:
+            message = environment.last_messages[0]
+            sender = message.sender
+            content = message.content
+            if sender.startswith("Professor"):
+                if content.startswith("[CallOn]"):
+                    # 1. professor calls on someone, then the student should speak
+                    result = re.search(r"\[CallOn\] Yes, ([sS]tudent )?(\w+)", content)
+                    if result is not None:
+                        name_to_id = {
+                            agent.name[len("Student ") :]: i
+                            for i, agent in enumerate(environment.agents)
+                        }
+                        return [name_to_id[result.group(2)]]
+                else:
+                    # 2. professor normally speaks, then anyone can act
+                    return list(range(len(environment.agents)))
+            elif sender.startswith("Student"):
+                # 3. student ask question after being called on, or
+                # 4. only one student raises hand, and the professor happens to listen
+                # 5. the group discussion is just over, and there happens to be only a student speaking in the last turn
+                return [0]
+        else:
+            # If len(last_messages) > 1, then
+            # 1. there must be at least one student raises hand or speaks.
+            # 2. the group discussion is just over.
+            return [0]
+        assert (
+            False
+        ), f"Should not reach here, last_messages: {environment.last_messages}"
+
+    def get_next_agent_idx_grouped(self, environment: BaseEnvironment) -> List[int]:
+        # Get the grouping information
+        # groups: A list of list of agent ids, the i-th list contains
+        #   the agent ids in the i-th group
+        # group_speaker_mapping: A mapping from group id to the id of
+        #   the speaker in the group
+        # `groups` should be set in the corresponding `visibility`,
+        # and `group_speaker_mapping` should be maintained here.
+        if "groups" not in environment.rule_params:
+            logging.warning(
+                "The environment is grouped, but the grouping information is not provided."
+            )
+        groups = environment.rule_params.get(
+            "groups", [list(range(len(environment.agents)))]
+        )
+        group_speaker_mapping = environment.rule_params.get(
+            "group_speaker_mapping", {i: 0 for i in range(len(groups))}
+        )
+
+        # For grouped environment, we let the students speak in turn within each group
+        next_agent_idx = []
+        for group_id in range(len(groups)):
+            speaker_index = group_speaker_mapping[group_id]
+            speaker = groups[group_id][speaker_index]
+            next_agent_idx.append(speaker)
+
+        # Maintain the `group_speaker_mapping`
+        for k, v in group_speaker_mapping.items():
+            group_speaker_mapping[k] = (v + 1) % len(groups[k])
+        environment.rule_params["group_speaker_mapping"] = group_speaker_mapping
+
+        return next_agent_idx

agentverse/environments/simulation_env/rules/order/concurrent.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+from . import order_registry as OrderRegistry
+from .base import BaseOrder
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@OrderRegistry.register("concurrent")
+class ConcurrentOrder(BaseOrder):
+    """
+    The agents speak concurrently
+    """
+
+    def get_next_agent_idx(self, environment: BaseEnvironment) -> List[int]:
+        return list(range(len(environment.agents)))

agentverse/environments/simulation_env/rules/order/prisoner.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import logging
+import re
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from . import order_registry as OrderRegistry
+from .base import BaseOrder
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@OrderRegistry.register("prisoner")
+class PrisonerOrder(BaseOrder):
+    """The order for a classroom discussion
+    The agents speak in the following order:
+    1. The professor speaks first
+    2. Then the professor can continue to speak, and the students can raise hands
+    3. The professor can call on a student, then the student can speak or ask a question
+    4. In the group discussion, the students in the group can speak in turn
+    """
+
+    # try police, prisoner1 prisoner2 first
+
+    last_prisoner_index: int = 1
+    switch_func: dict = {1: 2, 2: 1}
+
+    def get_next_agent_idx(self, environment: BaseEnvironment) -> List[int]:
+        if len(environment.last_messages) == 0:
+            # If the game just begins or , we let only the police speak
+            return [0]
+        elif len(environment.last_messages) == 1:
+            message = environment.last_messages[0]
+            sender = message.sender
+            content = message.content
+            if sender.startswith("Police"):
+                next_prisoner = self.last_prisoner_index
+                self.last_prisoner_index = self.switch_func[self.last_prisoner_index]
+                return [next_prisoner]
+            elif sender.startswith("Suspect"):
+                # 3. when one prisoner made his action, let the police tell another prisoner
+                return [0]
+        else:
+            # If len(last_messages) > 1, then
+            # 1. there must be at least one student raises hand or speaks.
+            # 2. the group discussion is just over.
+            return [0]

agentverse/environments/simulation_env/rules/order/random.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import random
+from typing import TYPE_CHECKING, List
+
+from . import order_registry as OrderRegistry
+from .base import BaseOrder
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@OrderRegistry.register("random")
+class RandomOrder(BaseOrder):
+    """
+    Order for random conversation
+    The agents speak in a random order
+    """
+
+    def get_next_agent_idx(self, environment: BaseEnvironment) -> List[int]:
+        return [random.randint(0, len(environment.agents) - 1)]

agentverse/environments/simulation_env/rules/order/sde_team.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import logging
+import re
+import random
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from . import order_registry as OrderRegistry
+from .base import BaseOrder
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@OrderRegistry.register("sde_team")
+class SdeTeamOrder(BaseOrder):
+    """The order for a code problem solving
+    """
+    next_agent_idx: int = 2
+    
+    def get_next_agent_idx(self, environment: BaseEnvironment) -> List[int]:
+        if self.next_agent_idx == 2:
+            self.next_agent_idx = 0
+            return [2] * 5 # TODO set the number in yaml
+        elif self.next_agent_idx == 0:
+            self.next_agent_idx = 1
+            return [0]
+        elif self.next_agent_idx == 1:
+            self.next_agent_idx = 0
+            return [1]

agentverse/environments/simulation_env/rules/order/sde_team_given_tests.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import logging
+import re
+import random
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from . import order_registry as OrderRegistry
+from .base import BaseOrder
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@OrderRegistry.register("sde_team_given_tests")
+class SdeTeamGivenTestsOrder(BaseOrder):
+    """The order for a code problem solving given unit tests
+    0 - code writer
+    1 - code tester
+    2 - code reviewer
+    """
+    next_agent_idx: int = 0
+    
+    def get_next_agent_idx(self, environment: BaseEnvironment) -> List[int]:
+        if self.next_agent_idx == 0:
+            self.next_agent_idx = 1
+            return [0]
+        elif self.next_agent_idx == 1:
+            self.next_agent_idx = 2
+            return [1]
+        elif self.next_agent_idx == 2:
+            self.next_agent_idx = 0
+            return [2]
+        else:
+            raise ValueError("Invalid next_agent_idx: {}".format(self.next_agent_idx))

agentverse/environments/simulation_env/rules/order/sequential.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+from . import order_registry as OrderRegistry
+from .base import BaseOrder
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@OrderRegistry.register("sequential")
+class SequentialOrder(BaseOrder):
+    """
+    Order for sequential conversation
+    The agents speak in a round-robin fashion
+    """
+
+    next_agent_idx: int = 0
+
+    def get_next_agent_idx(self, environment: BaseEnvironment) -> List[int]:
+        """Return the index of the next agent to speak"""
+        ret = self.next_agent_idx
+        self.next_agent_idx = (self.next_agent_idx + 1) % len(environment.agents)
+        return [ret]
+
+    def reset(self) -> None:
+        self.next_agent_idx = 0

agentverse/environments/simulation_env/rules/selector/__init__.py

@@ -0,0 +1,10 @@
+from agentverse.registry import Registry
+
+selector_registry = Registry(name="SelectorRegistry")
+
+from .base import BaseSelector
+from .basic import BasicSelector
+from .classroom import ClassroomSelector
+from .sde_team import SdeTeamSelector
+from .sde_team_given_tests import SdeTeamGivenTestsSelector
+from .pokemon import PokemonSelector

agentverse/environments/simulation_env/rules/selector/base.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+from pydantic import BaseModel
+
+from agentverse.message import Message
+
+from . import selector_registry as SelectorRegistry
+from abc import abstractmethod
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@SelectorRegistry.register("base")
+class BaseSelector(BaseModel):
+    """
+    Base class for all selecters
+    """
+
+    @abstractmethod
+    def select_message(
+        self, environment: BaseEnvironment, messages: List[Message]
+    ) -> List[Message]:
+        """Selects a set of valid messages from all messages"""
+        pass
+
+    def reset(self) -> None:
+        pass

agentverse/environments/simulation_env/rules/selector/basic.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+from agentverse.message import Message
+
+from . import selector_registry as SelectorRegistry
+from .base import BaseSelector
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@SelectorRegistry.register("basic")
+class BasicSelector(BaseSelector):
+    """
+    Base class for all selecters
+    """
+
+    def select_message(
+        self, environment: BaseEnvironment, messages: List[Message]
+    ) -> List[Message]:
+        """Selects a set of valid messages from all messages"""
+        return messages
+
+    def reset(self) -> None:
+        pass

agentverse/environments/simulation_env/rules/selector/classroom.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+from agentverse.message import Message
+
+from . import selector_registry as SelectorRegistry
+from .base import BaseSelector
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@SelectorRegistry.register("classroom")
+class ClassroomSelector(BaseSelector):
+    def select_message(
+        self, environment: BaseEnvironment, messages: List[Message]
+    ) -> List[Message]:
+        selected = []
+        for message in messages:
+            if message.sender.startswith("Student"):
+                if message.content.startswith("[RaiseHand]"):
+                    message.content = "[RaiseHand]"
+                    selected.append(message)
+                elif message.content != "" or len(message.tool_response) > 0:
+                    selected.append(message)
+            elif message.sender.startswith("Professor"):
+                # If the professor launch a group discussion, then we
+                # brutely discard the student's message in this turn
+                if message.content.startswith("[GroupDiscuss]"):
+                    return [message]
+                selected.append(message)
+
+        # If some student speak while the professor is speaking, then
+        # we brutely discard the student's message in this turn
+        if (
+            len(selected) > 1
+            and selected[0].sender.startswith("Professor")
+            and selected[0].content != ""
+        ):
+            filtered_selected = []
+            filtered_selected.append(selected[0])
+            for message in selected[1:]:
+                if message.content.startswith("[RaiseHand]"):
+                    filtered_selected.append(message)
+            selected = filtered_selected
+        return selected

agentverse/environments/simulation_env/rules/selector/code_api.py

@@ -0,0 +1,97 @@
+import io
+import sys
+import ast
+import json
+import astunparse
+import concurrent.futures
+import traceback
+
+
+def get_call_str(assert_statement: str) -> str:
+    call_str = ast.parse(assert_statement).body[0].test.left # type: ignore
+    return astunparse.unparse(call_str).strip()
+
+def get_output(func: str, assert_statement: str) -> str:
+    try:
+        func_call = get_call_str(assert_statement)
+        try:
+            exec(func, globals())
+            output = eval(func_call)
+            return output
+        except Exception as e:
+            return str(e)
+    except:
+        return "get_call_str error"
+
+def worker(code, globals=None, locals=None):
+    old_stdout = sys.stdout
+    redirected_output = sys.stdout = io.StringIO()
+    if locals is None:
+        locals = {}
+    try:
+        # TODO: exec(code, globals, locals) could be buggy
+        # In cases where both import statement and function exits in the code, if the locals are given, 
+        # the code will not find the imported package. 
+        # For example,
+        # code = "import math\ndef f(x):\n\treturn math.pow(x, 2)\nassert f(2) == 4"
+        # It will return "NameError: name 'math' is not defined"
+        exec(code, locals, locals)
+        stdout = redirected_output.getvalue()
+        return stdout, globals, locals
+    except Exception as e:
+        trace_str = traceback.format_exc()
+        return f"Error: {trace_str}", globals, locals
+    finally:
+        sys.stdout = old_stdout  # restore the original stdout
+        
+def execute_code(code: str) -> str:
+    """Execute a snippet of python code and return the output or the error message.
+    """
+    timeout = 5
+    try:
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            future = executor.submit(worker, code)
+            result, _, _ = future.result(timeout)
+            return result
+    except concurrent.futures.TimeoutError:
+        return "Timeout"
+
+def execute_unit_tests(func_impl: str, tests: str) -> str:
+    """Run a python function on a bunch of unit tests tests and return detailed feedback.
+    """
+    # tests = eval(tests)
+    # assert type(tests) == list
+
+    # Combine function code and assert statement
+    func_test_list = [f'{func_impl}\n{test}' for test in tests]
+
+    # Run the tests and collect the results
+    success_tests = []
+    failed_tests = []
+    is_passing = True
+    num_tests = len(func_test_list)
+    for i in range(num_tests):
+        output = execute_code(func_test_list[i])
+        if output == "Timeout":
+            failed_tests += [f"{tests[i]} # output: Timeout"]
+            is_passing = False
+        elif output.startswith("Error: "):
+            # print(output)
+            func_output = get_output(func_impl, tests[i])
+            if func_output == "get_call_str error":
+                func_output = output
+            failed_tests += [f"{tests[i]} # output: {func_output}"]
+            is_passing = False
+        else:
+            success_tests += [tests[i]]
+                
+    feedback = "Tested passed:\n\n"
+    for test in success_tests:
+        feedback += f"{test}\n\n"
+    feedback += "Tests failed:\n\n"
+    for test in failed_tests:
+        feedback += f"{test}\n\n"
+        
+    return json.dumps({"is_passing": is_passing, 
+            "feedback": feedback})
+

agentverse/environments/simulation_env/rules/selector/pokemon.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+import numpy as np
+import json
+
+from agentverse.message import Message
+
+from . import selector_registry as SelectorRegistry
+from .base import BaseSelector
+
+if TYPE_CHECKING:
+    from agentverse.environments import PokemonEnvironment
+
+
+@SelectorRegistry.register("pokemon")
+class PokemonSelector(BaseSelector):
+    """
+    Selector for Pokemon environment
+    """
+
+    def select_message(
+        self, environment: PokemonEnvironment, messages: List[Message]
+    ) -> List[Message]:
+        valid = []
+        talk_matrix = np.zeros((len(environment.agents), len(environment.agents)))
+        agent_to_idx = {agent.name: i for i, agent in enumerate(environment.agents)}
+        for i, message in enumerate(messages):
+            try:
+                content = json.loads(message.content)
+            except json.decoder.JSONDecodeError:
+                valid.append(0)
+                continue
+            if content["action"] == "Speak":
+                try:
+                    if "to" not in content:
+                        # If the model does not generate receiver, then we discard the message
+                        valid.append(0)
+                    elif content["to"] in agent_to_idx:
+                        # TODO: allow talk to a list of agents
+                        valid.append(1)
+                        # talk_matrix[i][j] = 1 ==> i talk to j
+                        talk_matrix[agent_to_idx[message.sender]][
+                            agent_to_idx[content["to"]]
+                        ] = 1
+                    else:
+                        # If the receiver is not in the environment, then we discard the message
+                        valid.append(0)
+                except:
+                    valid.append(0)
+                    continue
+            elif content["action"] == "MoveTo":
+                # If the agent move to a location that does not exist, then we discard the message
+                valid.append(
+                    "to" in content and content["to"] in environment.locations_to_agents
+                )
+            else:
+                valid.append(1)
+        selected_messages = []
+        for i, message in enumerate(messages):
+            content = json.loads(message.content)
+            sender_idx = agent_to_idx[message.sender]
+            if valid[i] == 0:
+                selected_messages.append(Message())
+                continue
+            if content["action"] == "MoveTo":
+                if np.sum(talk_matrix[:, sender_idx]) > 0:
+                    # If someone talk to this agent, then we discard the move action
+                    selected_messages.append(Message())
+                else:
+                    selected_messages.append(message)
+            elif content["action"] == "Speak":
+                receiver_idx = agent_to_idx[content["to"]]
+                if talk_matrix[sender_idx][receiver_idx] == 0:
+                    # If this agent talk to someone who also talk to this agent, and we
+                    # select the message from this agent, then we discard the message
+                    selected_messages.append(Message())
+                    continue
+                if np.sum(talk_matrix[receiver_idx, :]) > 0:
+                    if talk_matrix[receiver_idx][sender_idx] == 1:
+                        # If the receiver talk to this agent, then we randomly select one message
+                        if sender_idx < receiver_idx:
+                            if np.random.random() < 0.5:
+                                selected_messages.append(message)
+                                talk_matrix[receiver_idx][sender_idx] = 0
+                            else:
+                                selected_messages.append(Message())
+                                talk_matrix[sender_idx][receiver_idx] = 0
+                        else:
+                            print("Shouldn't happen")
+                    else:
+                        # If the receiver talk to other agent, we still talk to the receiver (?)
+                        selected_messages.append(message)
+                else:
+                    selected_messages.append(message)
+            else:
+                selected_messages.append(message)
+        return selected_messages

agentverse/environments/simulation_env/rules/selector/sde_team.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+from agentverse.message import Message
+
+from . import selector_registry as SelectorRegistry
+from .base import BaseSelector
+
+import json
+import re
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+    
+def extract(content: str, keyword: str):
+    result = ""
+    flag = False
+    for line in content.split('\n'):
+        if line.strip().startswith(keyword):
+            flag = True
+            continue
+        if flag:
+            result += line
+            result += "\n"
+    return result
+        
+    
+@SelectorRegistry.register("sde_team")
+class SdeTeamSelector(BaseSelector):
+    def select_message(self, environment: BaseEnvironment, messages: List[Message]) -> List[Message]:
+        last_sender = environment.last_messages[0].sender
+        selected = messages
+
+        if last_sender == "unit_test_generator":
+            unit_tests = set()
+            for message in selected:
+                unit_test = extract(message.content, "<unit test>:")
+                if unit_test not in unit_tests:
+                    unit_tests.add(extract(message.content, "<unit test>:"))
+            unit_tests = list(unit_tests)
+            environment.rule_params["unit_tests"] = str(unit_tests)
+            new_message = Message(
+                            content="",
+                            sender="unit_test_generator",
+                            receiver=[],
+                        )   # TODO: set the content of the message
+            selected = [new_message]
+
+        elif last_sender == "code_writer":
+            cur_code = extract(selected[0].content, "<code>:")
+            environment.rule_params["code"] = cur_code
+            
+            from .code_api import execute_unit_tests
+            feedback = execute_unit_tests(environment.rule_params["code"], eval(environment.rule_params["unit_tests"]))
+            
+            environment.rule_params["feedback"] = feedback
+            selected[0].content = f"<current code>:\n\n{cur_code}\n\n<unit test feedback>:\n{feedback}"
+            f_dict = json.loads(feedback)
+            if f_dict["is_passing"]:
+                environment.rule_params["end_flag"] = True
+        
+        elif last_sender == "code_reviewer":
+            code_review = selected[0].content
+            cur_code = environment.rule_params["code"]
+            selected[0].content = f"<current code>:\n\n{cur_code}\n\n{code_review}"
+            feedback = environment.rule_params["feedback"]
+            f_dict = json.loads(feedback)
+            if f_dict["is_passing"]:
+                environment.rule_params["end_flag"] = True
+
+        return selected

agentverse/environments/simulation_env/rules/selector/sde_team_given_tests.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+from agentverse.message import Message
+
+from . import selector_registry as SelectorRegistry
+from .base import BaseSelector
+
+import json
+import re
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+    
+def extract(content: str, keyword: str):
+    result = ""
+    flag = False
+    for line in content.split('\n'):
+        if line.strip().startswith(keyword):
+            flag = True
+            continue
+        if flag:
+            result += line
+            result += "\n"
+    return result
+        
+    
+@SelectorRegistry.register("sde_team_given_tests")
+class SdeTeamGivenTestsSelector(BaseSelector):
+    def select_message(self, environment: BaseEnvironment, messages: List[Message]) -> List[Message]:
+        last_sender = environment.last_messages[0].sender
+        selected = messages
+
+        if last_sender == "code_writer":
+            cur_code = extract(selected[0].content, "<code>:")
+            environment.rule_params["code"] = cur_code
+            selected[0].content = f"<current code>:\n{cur_code}"
+        
+        elif last_sender == "code_tester":
+            
+            from .code_api import execute_unit_tests
+            feedback = execute_unit_tests(environment.rule_params["code"], eval(environment.unit_tests))
+            environment.rule_params["feedback"] = feedback
+            selected[0].content = f"<unit test feedback>:\n{feedback}"
+            
+            f_dict = json.loads(feedback)
+            if f_dict["is_passing"]:
+                environment.rule_params["end_flag"] = True
+            
+        elif last_sender == "code_reviewer":
+            code_review = selected[0].content
+            cur_code = environment.rule_params["code"]
+            selected[0].content = f"{code_review}"
+
+        return selected

agentverse/environments/simulation_env/rules/updater/__init__.py

@@ -0,0 +1,9 @@
+from agentverse.registry import Registry
+
+updater_registry = Registry(name="UpdaterRegistry")
+
+from .base import BaseUpdater
+from .basic import BasicUpdater
+from .classroom import ClassroomUpdater
+from .sde_team import SdeTeamUpdater
+from .pokemon import PokemonUpdater

agentverse/environments/simulation_env/rules/updater/base.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Tuple
+
+from pydantic import BaseModel
+
+# from agentverse.agents import Agent
+from abc import abstractmethod
+
+from . import updater_registry as UpdaterRegistry
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@UpdaterRegistry.register("base")
+class BaseUpdater(BaseModel):
+    """
+    The base class of updater class.
+    """
+
+    @abstractmethod
+    def update_memory(self, environment: BaseEnvironment):
+        pass
+
+    def reset(self):
+        pass

agentverse/environments/simulation_env/rules/updater/basic.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Tuple
+
+from . import updater_registry as UpdaterRegistry
+from .base import BaseUpdater
+from agentverse.message import Message
+from agentverse.logging import get_logger
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+    from agentverse.agents import BaseAgent
+
+logger = get_logger()
+
+
+@UpdaterRegistry.register("basic")
+class BasicUpdater(BaseUpdater):
+    """
+    The basic version of updater.
+    The messages will be seen by all the receiver specified in the message.
+    """
+
+    def update_memory(self, environment: BaseEnvironment):
+        added = False
+        for message in environment.last_messages:
+            if len(message.tool_response) > 0:
+                self.add_tool_response(
+                    message.sender, environment.agents, message.tool_response
+                )
+            if message.content == "":
+                continue
+            added |= self.add_message_to_all_agents(environment.agents, message)
+        # If no one speaks in this turn. Add an empty message to all agents
+        if not added:
+            for agent in environment.agents:
+                agent.add_message_to_memory([Message(content="[Silence]")])
+
+    def add_tool_response(
+        self,
+        name: str,
+        agents: List[BaseAgent],
+        tool_response: List[str],
+    ):
+        for agent in agents:
+            if agent.name != name:
+                continue
+            if agent.tool_memory is not None:
+                agent.tool_memory.add_message(tool_response)
+            break
+
+    def add_message_to_all_agents(
+        self, agents: List[BaseAgent], message: Message
+    ) -> bool:
+        if "all" in message.receiver:
+            # If receiver is all, then add the message to all agents
+            for agent in agents:
+                agent.add_message_to_memory([message])
+            return True
+        else:
+            # If receiver is not all, then add the message to the specified agents
+            receiver_set = message.receiver
+            for agent in agents:
+                if agent.name in receiver_set:
+                    agent.add_message_to_memory([message])
+                    receiver_set.remove(agent.name)
+            if len(receiver_set) > 0:
+                missing_receiver = ", ".join(list(receiver_set))
+                # raise ValueError(
+                #    "Receiver {} not found. Message discarded".format(missing_receiver)
+                # )
+                logger.warn(
+                    "Receiver {} not found. Message discarded".format(missing_receiver)
+                )
+            return True