Merge pull request #95 from yiranwu0/gdswarm

feat: Refine Game Design Agent App with AG2's Swarm

ai_agent_tutorials/ai_game_design_agent_team/README.md

@@ -1,6 +1,6 @@
 # AI Game Design Agent Team 🎮
 
-The AI Game Design Agent Team is a collaborative game design system powered by Autogen's AI Agent framework. This app generates comprehensive game concepts through the coordination of multiple specialized AI agents, each focusing on different aspects of game design based on user inputs such as game type, target audience, art style, and technical requirements. This is built on Autogen's GroupChat methods and AssistantAgent features, run through initiate_chat() method.
+The AI Game Design Agent Team is a collaborative game design system powered by [AG2](https://github.com/ag2ai/ag2?tab=readme-ov-file)(formerly AutoGen)'s AI Agent framework. This app generates comprehensive game concepts through the coordination of multiple specialized AI agents, each focusing on different aspects of game design based on user inputs such as game type, target audience, art style, and technical requirements. This is built on AG2's new swarm feature run through initiate_chat() method.
 
 ## Features
 
@@ -17,6 +17,7 @@ The AI Game Design Agent Team is a collaborative game design system powered by A
   - Visual and audio direction
   - Technical specifications and requirements
   - Development timeline and budget considerations
+  - Coherent game design from the team.
 
 - **Customizable Input Parameters**:
   - Game type and target audience
@@ -25,7 +26,9 @@ The AI Game Design Agent Team is a collaborative game design system powered by A
   - Development constraints (time, budget)
   - Core mechanics and gameplay features
 
-- **Interactive Results**: Results are presented in expandable sections for easy navigation and reference
+- **Interactive Results**: 
+   - Quick show of game design ideas from each agent 
+   - Detailed results are presented in expandable sections for easy navigation and reference
 
 ## How to Run
 

ai_agent_tutorials/ai_game_design_agent_team/game_design_agent_team.py

@@ -1,10 +1,13 @@
 import asyncio
 import streamlit as st
-from autogen_agentchat.agents import AssistantAgent
+from autogen import (
-from autogen_agentchat.teams import RoundRobinGroupChat
+    SwarmAgent,
-from autogen_agentchat.conditions import TextMentionTermination
+    SwarmResult,
-from autogen_agentchat.ui import Console
+    initiate_swarm_chat,
-from autogen_ext.models.openai import OpenAIChatCompletionClient
+    OpenAIWrapper,
+    AFTER_WORK,
+    UPDATE_SYSTEM_MESSAGE
+)
 
 # Initialize session state
 if 'output' not in st.session_state:
@@ -111,15 +114,43 @@ if st.button("Generate Game Concept"):
             - Detail Level: {depth}
             """
 
-            # Configure OpenAI model client with the API key
+            llm_config = {"config_list": [{"model": "gpt-4o-mini","api_key": api_key}]}
 
-            model_client = OpenAIChatCompletionClient(model="gpt-4o-mini", api_key=api_key, temperature=0.0)
+            # initialize context variables
+            context_variables = {
+                "story": None,
+                "gameplay": None,
+                "visuals": None,
+                "tech": None,
+            }
 
-            # Define agents with detailed system prompts
+            # define functions to be called by the agents
-            story_agent = AssistantAgent(
+            def update_story_overview(story_summary:str, context_variables:dict) -> SwarmResult:
-                "story_agent",
+                """Keep the summary as short as possible."""
-                model_client=model_client,
+                context_variables["story"] = story_summary
-                system_message="""
+                st.sidebar.success('Story overview: ' + story_summary)
+                return SwarmResult(agent="gameplay_agent", context_variables=context_variables)
+                
+            def update_gameplay_overview(gameplay_summary:str, context_variables:dict) -> SwarmResult:
+                """Keep the summary as short as possible."""
+                context_variables["gameplay"] = gameplay_summary
+                st.sidebar.success('Gameplay overview: ' + gameplay_summary)
+                return SwarmResult(agent="visuals_agent", context_variables=context_variables)
+
+            def update_visuals_overview(visuals_summary:str, context_variables:dict) -> SwarmResult:
+                """Keep the summary as short as possible."""
+                context_variables["visuals"] = visuals_summary
+                st.sidebar.success('Visuals overview: ' + visuals_summary)
+                return SwarmResult(agent="tech_agent", context_variables=context_variables)
+
+            def update_tech_overview(tech_summary:str, context_variables:dict) -> SwarmResult:
+                """Keep the summary as short as possible."""
+                context_variables["tech"] = tech_summary
+                st.sidebar.success('Tech overview: ' + tech_summary)
+                return SwarmResult(agent="story_agent", context_variables=context_variables)
+
+            system_messages = {
+                "story_agent": """
             You are an experienced game story designer specializing in narrative design and world-building. Your task is to:
             1. Create a compelling narrative that aligns with the specified game type and target audience.
             2. Design memorable characters with clear motivations and character arcs.
@@ -127,14 +158,8 @@ if st.button("Generate Game Concept"):
             4. Plan story progression and major plot points.
             5. Integrate the narrative with the specified mood/atmosphere.
             6. Consider how the story supports the core gameplay mechanics.
-                Provide your response in a detailed, well-structured report format. Do not use XML tags.
+                """,
-                """
+                "gameplay_agent": """
-            )
-
-            gameplay_agent = AssistantAgent(
-                "gameplay_agent",
-                model_client=model_client,
-                system_message="""
             You are a senior game mechanics designer with expertise in player engagement and systems design. Your task is to:
             1. Design core gameplay loops that match the specified game type and mechanics.
             2. Create progression systems (character development, skills, abilities).
@@ -143,14 +168,8 @@ if st.button("Generate Game Concept"):
             5. Design multiplayer interactions if applicable.
             6. Specify game modes and difficulty settings.
             7. Consider the budget and development time constraints.
-                Provide your response in a detailed, well-structured report format. Do not use XML tags.
+                """,
-                """
+                "visuals_agent": """
-            )
-
-            visuals_agent = AssistantAgent(
-                "visuals_agent",
-                model_client=model_client,
-                system_message="""
             You are a creative art director with expertise in game visual and audio design. Your task is to:
             1. Define the visual style guide matching the specified art style.
             2. Design character and environment aesthetics.
@@ -159,14 +178,8 @@ if st.button("Generate Game Concept"):
             5. Consider technical constraints of chosen platforms.
             6. Align visual elements with the game's mood/atmosphere.
             7. Work within the specified budget constraints.
-                Provide your response in a detailed, well-structured report format. Do not use XML tags.
+                """,
-                """
+                "tech_agent": """
-            )
-
-            tech_agent = AssistantAgent(
-                name="tech_agent",
-                model_client=model_client,
-                system_message="""
             You are a technical director with extensive game development experience. Your task is to:
             1. Recommend appropriate game engine and development tools.
             2. Define technical requirements for all target platforms.
@@ -175,30 +188,89 @@ if st.button("Generate Game Concept"):
             5. Estimate resource requirements within the budget.
             6. Consider scalability and performance optimization.
             7. Plan for multiplayer infrastructure if applicable.
-                Provide your response in a detailed, well-structured report format. Do not use XML tags.
                 """
-            )
-
-            # Function to run agents sequentially
-            def run_agents_sequentially(task):
-                # Task agent provides the task to each agent one by one
-                story_response = asyncio.run(story_agent.run(task=task)).messages[-1].content
-                gameplay_response = asyncio.run(gameplay_agent.run(task=task)).messages[-1].content
-                visuals_response = asyncio.run(visuals_agent.run(task=task)).messages[-1].content
-                tech_response = asyncio.run(tech_agent.run(task=task)).messages[-1].content
-
-                return {
-                    "story": story_response,
-                    "gameplay": gameplay_response,
-                    "visuals": visuals_response,
-                    "tech": tech_response,
             }
 
-            # Run the agents sequentially and capture their responses
+            def update_system_message_func(agent: SwarmAgent, messages) -> str:
-            individual_responses = run_agents_sequentially(task)
+                """"""
+                system_prompt = system_messages[agent.name]
+
+                current_gen = agent.name.split("_")[0]
+                if agent._context_variables.get(current_gen) is None:
+                    system_prompt += f"Call the update function provided to first provide a 2-3 sentence summary of your ideas on {current_gen.upper()} based on the context provided."
+                    agent.llm_config['tool_choice'] = {"type": "function", "function": {"name": f"update_{current_gen}_overview"}}
+                    agent.client = OpenAIWrapper(**agent.llm_config)
+                else:
+                    # remove the tools to avoid the agent from using it and reduce cost
+                    agent.llm_config["tools"] = None
+                    agent.llm_config['tool_choice'] = None
+                    agent.client = OpenAIWrapper(**agent.llm_config)
+                    # the agent has given a summary, now it should generate a detailed response
+                    system_prompt += f"\n\nYour task\nYou task is write the {current_gen} part of the report. Do not include any other parts. Do not use XML tags.\nStart your reponse with: '## {current_gen.capitalize()} Design'."    
+                    
+                    # Remove all messages except the first one with less cost
+                    k = list(agent._oai_messages.keys())[-1]
+                    agent._oai_messages[k] = agent._oai_messages[k][:1]
+
+                system_prompt += f"\n\n\nBelow are some context for you to refer to:"
+                # Add context variables to the prompt
+                for k, v in agent._context_variables.items():
+                    if v is not None:
+                        system_prompt += f"\n{k.capitalize()} Summary:\n{v}"
+
+                return system_prompt
+            
+            state_update = UPDATE_SYSTEM_MESSAGE(update_system_message_func)
+
+            # Define agents
+            story_agent = SwarmAgent(
+                "story_agent", 
+                llm_config=llm_config,
+                functions=update_story_overview,
+                update_agent_state_before_reply=[state_update]
+            )
+
+            gameplay_agent = SwarmAgent(
+                "gameplay_agent",
+                llm_config= llm_config,
+                functions=update_gameplay_overview,
+                update_agent_state_before_reply=[state_update]
+            )
+
+            visuals_agent = SwarmAgent(
+                "visuals_agent",
+                llm_config=llm_config,
+                functions=update_visuals_overview,
+                update_agent_state_before_reply=[state_update]
+            )
+
+            tech_agent = SwarmAgent(
+                name="tech_agent",
+                llm_config=llm_config,
+                functions=update_tech_overview,
+                update_agent_state_before_reply=[state_update]
+            )
+
+            story_agent.register_hand_off(AFTER_WORK(gameplay_agent))
+            gameplay_agent.register_hand_off(AFTER_WORK(visuals_agent))
+            visuals_agent.register_hand_off(AFTER_WORK(tech_agent))
+            tech_agent.register_hand_off(AFTER_WORK(story_agent))
+
+            result, _, _ = initiate_swarm_chat(
+                initial_agent=story_agent,
+                agents=[story_agent, gameplay_agent, visuals_agent, tech_agent],
+                user_agent=None,
+                messages=task,
+                max_rounds=13,
+            )
 
             # Update session state with the individual responses
-            st.session_state.output = individual_responses
+            st.session_state.output = {
+                'story': result.chat_history[-4]['content'],
+                'gameplay': result.chat_history[-3]['content'],
+                'visuals': result.chat_history[-2]['content'],
+                'tech': result.chat_history[-1]['content']
+            }
 
         # Display success message after completion
         st.success('✨ Game concept generated successfully!')
@@ -216,25 +288,3 @@ if st.button("Generate Game Concept"):
         with st.expander("Technical Recommendations"):
             st.markdown(st.session_state.output['tech'])
 
-        termination = TextMentionTermination("TERMINATE")
-
-        groupchat = RoundRobinGroupChat(
-            [story_agent, gameplay_agent, visuals_agent, tech_agent],
-            termination_condition=termination,
-            max_turns=4, # Each agent gets 1 turn
-        )
-
-        # Function to run the agent collaboration
-        def run_agents(task):
-            result = asyncio.run(Console(groupchat.run_stream(task=task)))
-            responses = {}
-            for message in result.messages:
-                if message.source == "story_agent":
-                    responses["story"] = message.content
-                elif message.source == "gameplay_agent":
-                    responses["gameplay"] = message.content
-                elif message.source == "visuals_agent":
-                    responses["visuals"] = message.content
-                elif message.source == "tech_agent":
-                    responses["tech"] = message.content
-            return responses

ai_agent_tutorials/ai_game_design_agent_team/requirements.txt

@@ -1,3 +1,2 @@
 streamlit==1.41.1
-autogen-agentchat>=0.4.2,<0.5
+autogen
-autogen-ext[openai]>=0.4.2,<0.5