From 03efb87cd7187046cfc0c8d3a4348a9edef469dc Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Sun, 9 Jun 2024 18:06:51 -0700 Subject: [PATCH] Chat room example for event-based agents (#61) --- examples/chat_room.py | 161 ++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 1 + 2 files changed, 162 insertions(+) create mode 100644 examples/chat_room.py diff --git a/examples/chat_room.py b/examples/chat_room.py new file mode 100644 index 000000000..c737cffec --- /dev/null +++ b/examples/chat_room.py @@ -0,0 +1,161 @@ +"""This is an example of a chat room with AI agents. It demonstrates how to use the +`TypeRoutedAgent` class to create custom agents that can use custom message types, +and interact with other using event-based messaging.""" + +import argparse +import asyncio +import json +from dataclasses import dataclass + +from agnext.application import SingleThreadedAgentRuntime +from agnext.chat.memory import BufferedChatMemory, ChatMemory +from agnext.chat.types import TextMessage +from agnext.chat.utils import convert_messages_to_llm_messages +from agnext.components import TypeRoutedAgent, message_handler +from agnext.components.models import ChatCompletionClient, OpenAI, SystemMessage +from agnext.core import AgentRuntime, CancellationToken +from colorama import Fore, Style, init + + +# Define a custom message type for chat room messages. +@dataclass +class ChatRoomMessage(TextMessage): # type: ignore + pass + + +sep = "-" * 50 + +init(autoreset=True) + + +# Define a custom agent that can handle chat room messages. +class ChatRoomAgent(TypeRoutedAgent): # type: ignore + def __init__( # type: ignore + self, + name: str, + description: str, + runtime: AgentRuntime, # type: ignore + background_story: str, + memory: ChatMemory, # type: ignore + model_client: ChatCompletionClient, # type: ignore + color: str = Style.RESET_ALL, + ) -> None: # type: ignore + super().__init__(name, description, runtime) + system_prompt = f"""Your name is {name}. +Your background story is: +{background_story} + +Now you are in a chat room with other users. +You can send messages to the chat room by typing your message below. +You do not need to respond to every message. +Use the following JSON format to provide your thought on the latest message and choose whether to respond: +{{ + "thought": "Your thought on the message", + "respond": , + "response": "Your response to the message or None if you choose not to respond." +}} +""" + self._system_messages = [SystemMessage(system_prompt)] + self._memory = memory + self._client = model_client + self._color = color + + @message_handler() # type: ignore + async def on_chat_room_message(self, message: ChatRoomMessage, cancellation_token: CancellationToken) -> None: # type: ignore + # Save the message to memory as structured JSON. + from_message = TextMessage( + content=json.dumps({"sender": message.source, "content": message.content}), source=message.source + ) + self._memory.add_message(from_message) + + # Get a response from the model. + raw_response = await self._client.create( + self._system_messages + convert_messages_to_llm_messages(self._memory.get_messages(), self_name=self.name), + json_output=True, + ) + assert isinstance(raw_response.content, str) + + # Save the response to memory. + self._memory.add_message(ChatRoomMessage(source=self.name, content=raw_response.content)) + + # Parse the response. + data = json.loads(raw_response.content) + respond = data.get("respond") + response = data.get("response") + + # Publish the response if needed. + if respond is True or str(respond).lower().strip() == "true": + self._publish_message(ChatRoomMessage(source=self.name, content=str(response))) + print(f"{sep}\n{self._color}{self.name}:{Style.RESET_ALL}\n{response}") + + +# Define a chat room with participants -- the runtime is the chat room. +def chat_room(runtime: AgentRuntime) -> None: # type: ignore + _ = ChatRoomAgent( + name="Alice", + description="Alice in the chat room.", + runtime=runtime, + background_story="Alice is a software engineer who loves to code.", + memory=BufferedChatMemory(buffer_size=10), + model_client=OpenAI(model="gpt-4-turbo"), # type: ignore + color=Fore.CYAN, + ) + + _ = ChatRoomAgent( + name="Bob", + description="Bob in the chat room.", + runtime=runtime, + background_story="Bob is a data scientist who loves to analyze data.", + memory=BufferedChatMemory(buffer_size=10), + model_client=OpenAI(model="gpt-4-turbo"), # type: ignore + color=Fore.GREEN, + ) + + _ = ChatRoomAgent( + name="Charlie", + description="Charlie in the chat room.", + runtime=runtime, + background_story="Charlie is a designer who loves to create art.", + memory=BufferedChatMemory(buffer_size=10), + model_client=OpenAI(model="gpt-4-turbo"), # type: ignore + color=Fore.MAGENTA, + ) + + +async def get_user_input(prompt: str) -> str: + loop = asyncio.get_event_loop() + return await loop.run_in_executor(None, input, prompt) + + +async def main(user_name: str, wait_seconds: int) -> None: + runtime = SingleThreadedAgentRuntime() + chat_room(runtime) + while True: + # TODO: allow user to input at any time while runtime is running. + # Get user input and send messages to the chat room. + user_input = await get_user_input(f"{sep}\nYou:\n") + if user_input.strip(): + # Publish user message if it is not empty. + runtime.publish_message(ChatRoomMessage(source=user_name, content=user_input)) + # Wait for agents to respond. + while runtime.unprocessed_messages: + await runtime.process_next() + await asyncio.sleep(wait_seconds) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run a chat room simulation.") + parser.add_argument( + "--user-name", + type=str, + default="Host", + help="The name of the user who is participating in the chat room.", + ) + parser.add_argument( + "--wait-seconds", + type=int, + default=5, + help="The number of seconds to wait between processing messages.", + ) + args = parser.parse_args() + asyncio.run(main(args.user_name, args.wait_seconds)) diff --git a/pyproject.toml b/pyproject.toml index fd4f5665b..86013fa52 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dev = [ "tavily-python", "aiofiles", "types-aiofiles", + "colorama", ] docs = ["sphinx", "furo", "sphinxcontrib-apidoc", "myst-parser"]