import os import asyncio from pydantic import BaseModel, Field from typing import List, Dict, Any from swarms import Agent from swarm_models import OpenAIChat from dotenv import load_dotenv from swarms.utils.formatter import formatter # Load environment variables load_dotenv() # Get OpenAI API key api_key = os.getenv("OPENAI_API_KEY") # Define Pydantic schema for agent outputs class AgentOutput(BaseModel): """Schema for capturing the output of each agent.""" agent_name: str = Field(..., description="The name of the agent") message: str = Field( ..., description="The agent's response or contribution to the group chat", ) metadata: Dict[str, Any] = Field( default_factory=dict, description="Additional metadata about the agent's response", ) class GroupChat: """ GroupChat class to enable multiple agents to communicate in an asynchronous group chat. Each agent is aware of all other agents, every message exchanged, and the social context. """ def __init__( self, name: str, description: str, agents: List[Agent], max_loops: int = 1, ): """ Initialize the GroupChat. Args: name (str): Name of the group chat. description (str): Description of the purpose of the group chat. agents (List[Agent]): A list of agents participating in the chat. max_loops (int): Maximum number of loops to run through all agents. """ self.name = name self.description = description self.agents = agents self.max_loops = max_loops self.chat_history = ( [] ) # Stores all messages exchanged in the chat formatter.print_panel( f"Initialized GroupChat '{self.name}' with {len(self.agents)} agents. Max loops: {self.max_loops}", title="Groupchat Swarm", ) async def _agent_conversation( self, agent: Agent, input_message: str ) -> AgentOutput: """ Facilitate a single agent's response to the chat. Args: agent (Agent): The agent responding. input_message (str): The message triggering the response. Returns: AgentOutput: The agent's response captured in a structured format. """ formatter.print_panel( f"Agent '{agent.agent_name}' is responding to the message: {input_message}", title="Groupchat Swarm", ) response = await asyncio.to_thread(agent.run, input_message) output = AgentOutput( agent_name=agent.agent_name, message=response, metadata={"context_length": agent.context_length}, ) # logger.debug(f"Agent '{agent.agent_name}' response: {response}") return output async def _run(self, initial_message: str) -> List[AgentOutput]: """ Execute the group chat asynchronously, looping through all agents up to max_loops. Args: initial_message (str): The initial message to start the chat. Returns: List[AgentOutput]: The responses of all agents across all loops. """ formatter.print_panel( f"Starting group chat '{self.name}' with initial message: {initial_message}", title="Groupchat Swarm", ) self.chat_history.append( {"sender": "System", "message": initial_message} ) outputs = [] for loop in range(self.max_loops): formatter.print_panel( f"Group chat loop {loop + 1}/{self.max_loops}", title="Groupchat Swarm", ) for agent in self.agents: # Create a custom input message for each agent, sharing the chat history and social context input_message = ( f"Chat History:\n{self._format_chat_history()}\n\n" f"Participants:\n" + "\n".join( [ f"- {a.agent_name}: {a.system_prompt}" for a in self.agents ] ) + f"\n\nNew Message: {initial_message}\n\n" f"You are '{agent.agent_name}'. Remember to keep track of the social context, who is speaking, " f"and respond accordingly based on your role: {agent.system_prompt}." ) # Collect agent's response output = await self._agent_conversation( agent, input_message ) outputs.append(output) # Update chat history with the agent's response self.chat_history.append( { "sender": agent.agent_name, "message": output.message, } ) formatter.print_panel( "Group chat completed. All agent responses captured.", title="Groupchat Swarm", ) return outputs def run(self, task: str, *args, **kwargs): return asyncio.run(self.run(task, *args, **kwargs)) def _format_chat_history(self) -> str: """ Format the chat history for agents to understand the context. Returns: str: The formatted chat history as a string. """ return "\n".join( [ f"{entry['sender']}: {entry['message']}" for entry in self.chat_history ] ) def __str__(self) -> str: """String representation of the group chat's outputs.""" return self._format_chat_history() def to_json(self) -> str: """JSON representation of the group chat's outputs.""" return [ {"sender": entry["sender"], "message": entry["message"]} for entry in self.chat_history ] # Example Usage if __name__ == "__main__": load_dotenv() # Get the OpenAI API key from the environment variable api_key = os.getenv("OPENAI_API_KEY") # Create an instance of the OpenAIChat class model = OpenAIChat( openai_api_key=api_key, model_name="gpt-4o-mini", temperature=0.1, ) # Example agents agent1 = Agent( agent_name="Financial-Analysis-Agent", system_prompt="You are a financial analyst specializing in investment strategies.", llm=model, max_loops=1, autosave=False, dashboard=False, verbose=True, dynamic_temperature_enabled=True, user_name="swarms_corp", retry_attempts=1, context_length=200000, output_type="string", streaming_on=False, ) agent2 = Agent( agent_name="Tax-Adviser-Agent", system_prompt="You are a tax adviser who provides clear and concise guidance on tax-related queries.", llm=model, max_loops=1, autosave=False, dashboard=False, verbose=True, dynamic_temperature_enabled=True, user_name="swarms_corp", retry_attempts=1, context_length=200000, output_type="string", streaming_on=False, ) # Create group chat group_chat = GroupChat( name="Financial Discussion", description="A group chat for financial analysis and tax advice.", agents=[agent1, agent2], ) # Run the group chat asyncio.run( group_chat.run( "How can I establish a ROTH IRA to buy stocks and get a tax break? What are the criteria? What do you guys think?" ) )