Zqm dev

agentverse/agents/__init__.py

@@ -5,4 +5,6 @@
 
 from .base import BaseAgent
 from .conversation_agent import ConversationAgent
+from .prisoner_agent import PrisonerAgent
+from .police_agent import PoliceAgent
 from .tool_agent import ToolAgent

agentverse/agents/conversation_agent.py

@@ -1,23 +1,29 @@
 import logging
 from string import Template
-from typing import List
+
+
+from typing import List, TYPE_CHECKING
+
 
 from agentverse.message import Message
 
 from . import agent_registry
 from .base import BaseAgent
 
 
+if TYPE_CHECKING:
+    from agentverse.environments.base import BaseEnvironment
+
 @agent_registry.register("conversation")
 class ConversationAgent(BaseAgent):
-    def step(self, env_description: str = "") -> Message:
+    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(response)
+                parsed_response = self.output_parser.parse(self, environment, response)
                 break
             except Exception as e:
                 logging.error(e)
@@ -36,15 +42,15 @@ def step(self, env_description: str = "") -> Message:
         )
         return message
 
-    async def astep(self, env_description: str = "") -> 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(response)
+                parsed_response = self.output_parser.parse(self, environment, response)
                 break
             except Exception as e:
                 logging.error(e)

agentverse/agents/police_agent.py

@@ -0,0 +1,45 @@
+import logging
+from string import Template
+
+
+from typing import List, TYPE_CHECKING
+
+
+from agentverse.message import Message
+
+from . import agent_registry
+from .base import BaseAgent
+from .conversation_agent import ConversationAgent
+
+
+if TYPE_CHECKING:
+    from agentverse.environments.base import BaseEnvironment
+
+@agent_registry.register("police")
+class PoliceAgent(ConversationAgent):
+    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)
+

agentverse/agents/prisoner_agent.py

@@ -0,0 +1,47 @@
+import logging
+from string import Template
+
+
+from typing import List, TYPE_CHECKING
+
+
+from agentverse.message import Message
+
+from . import agent_registry
+from .base import BaseAgent
+from .conversation_agent import ConversationAgent
+
+
+if TYPE_CHECKING:
+    from agentverse.environments.base import BaseEnvironment
+
+@agent_registry.register("prisoner")
+class PrisonerAgent(ConversationAgent):
+    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/environments/basic.py

@@ -59,7 +59,7 @@ async def step(self) -> List[Message]:
 
         # Generate the next message
         messages = await asyncio.gather(
-            *[self.agents[i].astep(env_descriptions[i]) for i in agent_ids]
+            *[self.agents[i].astep(self, env_descriptions[i]) for i in agent_ids]
         )
 
         # Some rules will select certain messages from all the messages

agentverse/environments/rules/describer/__init__.py

@@ -5,3 +5,4 @@
 from .base import BaseDescriber
 from .basic import BasicDescriber
 from .classroom import ClassroomDescriber
+from .prisoner import PrisonerDescriber

agentverse/environments/rules/describer/prisoner.py

@@ -0,0 +1,51 @@
+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/rules/order/__init__.py

@@ -6,3 +6,4 @@
 from .random import RandomOrder
 from .concurrent import ConcurrentOrder
 from .classroom import ClassroomOrder
+from .prisoner import PrisonerOrder

agentverse/environments/rules/order/prisoner.py

@@ -0,0 +1,50 @@
+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/rules/visibility/__init__.py

@@ -6,3 +6,4 @@
 from .base import BaseVisibility
 from .all import AllVisibility
 from .classroom import ClassroomVisibility
+from .prisoner import PrisonerVisibility

agentverse/environments/rules/visibility/prisoner.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import random
+from typing import TYPE_CHECKING, Any, List, Union
+
+from . import visibility_registry as VisibilityRegistry
+from .base import BaseVisibility
+
+if TYPE_CHECKING:
+    from agentverse.environments import BaseEnvironment
+
+
+@VisibilityRegistry.register("prisoner")
+class PrisonerVisibility(BaseVisibility):
+    """
+    Visibility function for classroom, supports group discussion.
+
+    Args:
+        student_per_group:
+            The number of students per group.
+        num_discussion_turn:
+            The number of turns for group discussion.
+        grouping:
+            The grouping information. If it is a string, then it should be a
+            grouping method, options are ["random", "sequential"]. If it is a
+            list of list of int, then it should be the grouping information.
+    """
+
+    current_turn: int = 0
+
+    def update_visible_agents(self, environment: BaseEnvironment):
+
+        self.update_receiver(environment, reset=False)
+
+    def update_receiver(self, environment: BaseEnvironment, reset=False):
+        if reset:
+            for agent in environment.agents:
+                agent.set_receiver(["all"])
+        else:
+           # 0:police 1: prisoner1 2: prisoner2
+           #  environment.agents[0].set_receiver({"Police", "Suspect1", "Suspect2"})
+           #  environment.agents[1].set_receiver({"Police", "Suspect1"})
+           #  environment.agents[2].set_receiver({"Police", "Suspect2"})
+
+            # we update receiver in environment
+            pass
+
+    def reset(self):
+        self.current_turn = 0

agentverse/initialization.py

@@ -26,9 +26,11 @@
 
 def load_llm(llm_config: Dict):
     llm_type = llm_config.pop("llm_type", "text-davinci-003")
+
     return llm_registry.build(llm_type, **llm_config)
 
 
+
 def load_memory(memory_config: Dict):
     memory_type = memory_config.pop("memory_type", "chat_history")
     return memory_registry.build(memory_type, **memory_config)