agentverse/environments/simulation_env/pokemon.py

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