Update app.py

app.py

@@ -29,9 +29,9 @@ class SpaceSystemCrew:
         self.mechanical_engineer = None
         self.aerospace_manager = None
         self.current_agent = None
-        self.final_design = None
+        self.final_refined_text = None
 
-    def initialize_agents(self, topic: str):
+    def initialize_agents(self, space_system_description: str):
         if not self.api_key:
             raise ValueError("OpenAI API key is required")
 
@@ -40,8 +40,8 @@ class SpaceSystemCrew:
 
         self.system_architect = Agent(
             role="System Architect",
-            goal=f"Design the architecture of the {topic} space system",
-            backstory="Expert in space system design, integrating mechanical and propulsion components into a cohesive system",
+            goal=f"Analyze and refine the space system concept based on the provided description",
+            backstory="Expert in space system design, ensuring logical consistency and feasibility",
             allow_delegation=False,
             verbose=True,
             llm=llm
@@ -49,7 +49,7 @@ class SpaceSystemCrew:
 
         self.mechanical_engineer = Agent(
             role="Mechanical Engineer",
-            goal=f"Develop structural and propulsion mechanisms for {topic}",
+            goal=f"Refine structural and propulsion mechanisms based on the provided description",
             backstory="Specialist in high-velocity launch systems and orbital stabilization",
             allow_delegation=False,
             verbose=True,
@@ -58,60 +58,57 @@ class SpaceSystemCrew:
 
         self.aerospace_manager = Agent(
             role="Aerospace Mission Manager",
-            goal=f"Ensure mission feasibility, safety, and execution for {topic}",
+            goal=f"Ensure mission feasibility and refine execution details based on the provided description",
             backstory="Oversees project feasibility, mission planning, and ensures integration with existing aerospace infrastructure",
             allow_delegation=False,
             verbose=True,
             llm=llm
         )
 
-    def create_tasks(self, topic: str) -> List[Task]:
+    def create_tasks(self, space_system_description: str) -> List[Task]:
         architect_task = Task(
-            description=f"""Develop the complete architecture for the {topic} space system, including:
-            1. High-level design of the five-stage system
-            2. Identification of key components and subsystems
-            3. Integration with existing space infrastructure
-            4. Energy efficiency and sustainability considerations""",
-            expected_output="A comprehensive space system architecture plan",
+            description=f"""Analyze and refine the provided space system concept, ensuring:
+            1. Logical consistency and feasibility
+            2. Integration with existing space infrastructure
+            3. Improved clarity and technical accuracy""",
+            expected_output="A refined version of the space system concept",
             agent=self.system_architect
         )
 
         mechanical_task = Task(
-            description="""Design the structural and propulsion aspects of the space system, including:
-            1. Engineering specifications for the spin-launcher mechanism
-            2. Payload stabilization and orbital insertion strategies
-            3. Material selection and thermal considerations
-            4. Fuel efficiency and mechanical reliability""",
-            expected_output="Detailed propulsion and structural design specifications",
+            description="""Enhance the structural and propulsion aspects described in the provided description, including:
+            1. Refining spin-launcher specifications
+            2. Optimizing payload stabilization and orbital insertion methods
+            3. Ensuring energy efficiency and mechanical reliability""",
+            expected_output="Refined propulsion and structural details",
             agent=self.mechanical_engineer
         )
 
         aerospace_task = Task(
-            description="""Assess the feasibility and mission execution strategies, including:
-            1. Launch site selection and environmental impact assessment
-            2. Safety protocols and risk mitigation strategies
-            3. Crew and payload rendezvous procedures
-            4. Operational procedures for space missions post-docking""",
-            expected_output="A validated mission feasibility and execution plan",
+            description="""Evaluate and refine mission execution strategies based on the provided description, ensuring:
+            1. Practical feasibility of launch and docking operations
+            2. Safety enhancements and risk mitigation measures
+            3. Seamless operational procedures for post-docking activities""",
+            expected_output="A refined mission feasibility and execution plan",
             agent=self.aerospace_manager
         )
 
         return [architect_task, mechanical_task, aerospace_task]
 
-    async def process_design(self, topic: str) -> Generator[List[Dict], None, None]:
+    async def process_design(self, space_system_description: str) -> Generator[List[Dict], None, None]:
         try:
-            self.initialize_agents(topic)
+            self.initialize_agents(space_system_description)
             self.current_agent = "System Architect"
 
             yield [{
                 "role": "assistant",
-                "content": "Starting work on your space system design...",
+                "content": "Starting analysis and refinement of the space system...",
                 "metadata": {"title": "🚀 Process Started"}
             }]
 
             crew = Crew(
                 agents=[self.system_architect, self.mechanical_engineer, self.aerospace_manager],
-                tasks=self.create_tasks(topic),
+                tasks=self.create_tasks(space_system_description),
                 verbose=True
             )
 
@@ -133,6 +130,11 @@ class SpaceSystemCrew:
                 if messages:
                     yield messages
                 await asyncio.sleep(0.1)
+
+            self.final_refined_text = "\n\n".join([msg["content"] for msg in self.message_queue.get_messages() if msg["role"] == "assistant"])
+            print("Final Refined Space System Design:")
+            print(self.final_refined_text)
+
         except Exception as e:
             yield [{
                 "role": "assistant",
@@ -145,7 +147,7 @@ def create_demo():
     space_crew = None
 
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🚀 AI Space System Design Crew")
+        gr.Markdown("# 🚀 AI Space System Refinement Crew")
 
         openai_api_key = gr.Textbox(
             label='OpenAI API Key',
@@ -154,58 +156,27 @@ def create_demo():
             interactive=True
         )
 
-        chatbot = gr.Chatbot(
-            label="Design Process",
-            height=700,
-            type="messages",
-            show_label=True,
-            visible=False,
-            avatar_images=(None, "https://avatars.githubusercontent.com/u/170677839?v=4"),
-            render_markdown=True
+        space_system_description = gr.Textbox(
+            label="Space System Description",
+            placeholder="Enter space system details...",
+            lines=10,
+            interactive=True
         )
 
-        with gr.Row(equal_height=True):
-            topic = gr.Textbox(
-                label="Space System Name",
-                placeholder="Enter system name...",
-                scale=4,
-                visible=False
-            )
-            btn = gr.Button("Design System", variant="primary", scale=1, visible=False)
+        btn = gr.Button("Refine System", variant="primary")
 
-        async def process_input(topic, history, api_key):
+        async def process_input(space_system_description, history, api_key):
             nonlocal space_crew
             if not api_key:
-                history = history or []
-                history.append({
-                    "role": "assistant",
-                    "content": "Please provide an OpenAI API key.",
-                    "metadata": {"title": "❌ Error"}
-                })
-                yield history
-                return
+                return [{"role": "assistant", "content": "Please provide an OpenAI API key."}]
 
             if space_crew is None:
                 space_crew = SpaceSystemCrew(api_key=api_key)
 
-            history = history or []
-            history.append({"role": "user", "content": f"Design a space system for: {topic}"})
-            yield history
-
-            try:
-                async for messages in space_crew.process_design(topic):
-                    history.extend(messages)
-                    yield history
-            except Exception as e:
-                history.append({
-                    "role": "assistant",
-                    "content": f"An error occurred: {str(e)}",
-                    "metadata": {"title": "❌ Error"}
-                })
-                yield history
-
-        openai_api_key.submit(lambda: {openai_api_key: gr.Textbox(visible=False), chatbot: gr.Chatbot(visible=True), topic: gr.Textbox(visible=True), btn: gr.Button(visible=True)}, None, [openai_api_key, chatbot, topic, btn])
-        btn.click(process_input, [topic, chatbot, openai_api_key], [chatbot])
+            async for messages in space_crew.process_design(space_system_description):
+                yield messages
+
+        btn.click(process_input, [space_system_description, openai_api_key], [])
 
     return demo