Fixed app generation workflow with proper async execution and error handling

app.py

@@ -777,7 +777,7 @@ class EnhancedAIFlow(AIFlow):
         self.contexts: Dict[AgentRole, AgentContext] = {}
         self.global_context = {}
         self.requirements = ""
-    
+        
     def initialize_flow(self):
         """Initialize the AI Flow with agent relationships and dependencies."""
         # Create nodes for each agent role
@@ -884,7 +884,7 @@ Focus: setup, API docs, guides"""
         except Exception as e:
             logger.error(f"Failed to generate prompt for {role}: {str(e)}")
             raise
-    
+            
     async def execute_flow(self, requirements: str) -> str:
         """Execute the AI Flow and build the project."""
         try:
@@ -892,20 +892,17 @@ Focus: setup, API docs, guides"""
             self.requirements = requirements
             self.initialize_flow()
             
-            # Extract app name
+            # Extract app name from requirements
             app_name = requirements.split()[0].lower().replace(" ", "_")
             
             # Execute agents in parallel where possible
             paths = list(nx.all_simple_paths(self.flow_graph, AgentRole.ARCHITECT, AgentRole.DOCUMENTATION_WRITER))
-            await self._execute_paths(paths)
+            results = await self._execute_paths(paths)
             
             # Generate project structure and documentation
             project_structure = generate_project_structure(app_name, self.contexts[AgentRole.ARCHITECT].artifacts)
             documentation = generate_documentation(app_name, requirements, self.contexts[AgentRole.DOCUMENTATION_WRITER].artifacts)
             
-            # Create downloadable output
-            download_path = self.output_manager.create_download(app_name)
-            
             return f"""
 # {app_name.title()} - Generated Application
 
@@ -917,9 +914,6 @@ Focus: setup, API docs, guides"""
 ## Documentation
 {documentation}
 
-## Download
-To download your project, use this path: {download_path}
-
 ## Next Steps
 1. Review the generated architecture and components
 2. Set up the development environment
@@ -936,90 +930,102 @@ For any issues or questions, please refer to the documentation or create an issu
         finally:
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
-    
-    async def _execute_paths(self, paths: List[List[AgentRole]]):
+
+    async def _execute_paths(self, paths: List[List[AgentRole]]) -> List[str]:
         """Execute all paths in the flow graph."""
         try:
+            # Execute paths in parallel
+            tasks = []
             for path in paths:
                 for role in path:
-                    if self.contexts[role].state != FlowState.COMPLETED:
-                        await self._execute_agent(role)
+                    if self.contexts[role].state == FlowState.PENDING:
+                        tasks.append(self._execute_agent(role))
+                        self.contexts[role].state = FlowState.RUNNING
+            
+            # Wait for all tasks to complete
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            # Process results
+            for result in results:
+                if isinstance(result, Exception):
+                    raise result
+            
+            return results
+            
         except Exception as e:
             logger.error(f"Failed to execute paths: {str(e)}")
             raise
-    
-    async def _execute_agent(self, role: AgentRole):
+
+    async def _execute_agent(self, role: AgentRole) -> str:
         """Execute a single agent's tasks with enhanced context."""
         try:
-            context = self.contexts[role]
-            if context.state == FlowState.COMPLETED:
-                return
-            
-            context.state = FlowState.RUNNING
+            # Generate prompt
             prompt = self._generate_prompt(role)
             
-            # Execute agent task
+            # Execute agent's task
             result = await self._execute_agent_task(role, prompt)
             
-            # Update context with results
-            context.artifacts.update(result)
-            context.state = FlowState.COMPLETED
+            # Update context
+            self.contexts[role].state = FlowState.COMPLETED
+            self.contexts[role].artifacts["output"] = result
+            
+            return result
             
         except Exception as e:
-            context.state = FlowState.FAILED
             logger.error(f"Failed to execute agent {role}: {str(e)}")
+            self.contexts[role].state = FlowState.FAILED
             raise
-    
-    async def _execute_agent_task(self, role: AgentRole, prompt: str) -> Dict[str, str]:
+
+    async def _execute_agent_task(self, role: AgentRole, prompt: str) -> str:
         """Execute a specific agent's task with the given prompt."""
         try:
-            # Get the appropriate agent for the role
+            # Get agent
             agent = get_agent(role)
             
             # Execute the agent's task
             result = await asyncio.to_thread(agent.run, prompt)
             
             # Process and return the result
-            return {"output": result}
+            return result
             
         except Exception as e:
-            logger.error(f"Failed to execute agent task for {role}: {str(e)}")
+            logger.error(f"Agent task execution failed for {role}: {str(e)}")
             raise
 
 # Update the multi_agent_workflow function to use AI Flows
 async def multi_agent_workflow(requirements: str) -> str:
-    """
-    Execute a multi-agent workflow using AI Flows to generate a complex app.
-    Args:
-        requirements (str): App requirements.
-    Returns:
-        str: Generated app code and API details.
-    """
+    """Execute a multi-agent workflow using AI Flows to generate a complex app."""
     try:
-        # Initialize AI Flow
+        # Create AI Flow instance
         ai_flow = EnhancedAIFlow()
         
-        # Execute the flow with requirements
+        # Generate the app
         result = await ai_flow.execute_flow(requirements)
         
-        # Return the compiled results
         return result
     except Exception as e:
-        logger.error(f"Failed to execute multi-agent workflow: {str(e)}")
+        logger.error(f"Multi-agent workflow failed: {str(e)}")
         raise
 
 # Update the app_generator function to handle async execution
 async def app_generator(requirements: str) -> Dict[str, str]:
     """Generate an app based on the provided requirements using AI Flows."""
     try:
-        # Execute the multi-agent workflow
-        result = await multi_agent_workflow(requirements)
+        # Create AI Flow instance
+        ai_flow = EnhancedAIFlow()
+        
+        # Generate the app
+        result = await ai_flow.execute_flow(requirements)
+        
+        # Create downloadable output
+        download_path = ai_flow.output_manager.create_download("generated_app")
+        
         return {
             "output": result,
-            "download_path": None
+            "download_path": str(download_path) if download_path else None
         }
     except Exception as e:
-        logger.error(f"Failed to generate app: {str(e)}")
+        logger.error(f"App generation failed: {str(e)}")
         raise
 
 async def stream_output(requirements, progress=gr.Progress()):
@@ -1049,41 +1055,29 @@ async def stream_output(requirements, progress=gr.Progress()):
         
         try:
             # Run the app generator with a timeout
-            result = await asyncio.wait_for(
-                app_generator(requirements),
-                timeout=60  # 60 second timeout
-            )
-            
-            if result:
-                # Create download file
-                try:
-                    download_path = result.get("download_path")
-                    if download_path:
-                        stream_handler.update(" Generation completed successfully!", "Success")
-                        yield result.get("output", ""), download_path, "\n".join(stream_handler.output), "Success"
-                    else:
-                        error_msg = " Failed to create download file"
-                        stream_handler.update(error_msg, "Error")
-                        yield result.get("output", ""), None, "\n".join(stream_handler.output), "Error"
-                except Exception as e:
-                    error_msg = f" Error creating download: {str(e)}"
-                    stream_handler.update(error_msg, "Error")
-                    yield result.get("output", ""), None, "\n".join(stream_handler.output), "Error"
-            else:
-                error_msg = " Failed to generate app"
-                stream_handler.update(error_msg, "Error")
-                yield "", None, "\n".join(stream_handler.output), "Error"
+            async with asyncio.timeout(60):  # 60 second timeout
+                result = await app_generator(requirements)
                 
+                # Update output with result
+                if result["output"]:
+                    stream_handler.update("\n" + result["output"], "Completed")
+                    yield result["output"], result["download_path"], "\n".join(stream_handler.output), "Completed"
+                else:
+                    raise Exception("No output generated")
+                    
         except asyncio.TimeoutError:
-            error_msg = " Generation timed out after 60 seconds"
-            stream_handler.update(error_msg, "Error")
-            yield "", None, "\n".join(stream_handler.output), "Error"
+            stream_handler.update("\nApp generation timed out after 60 seconds", "Failed")
+            yield None, None, "\n".join(stream_handler.output), "Failed"
+            raise
             
     except Exception as e:
-        error_msg = f" Error: {str(e)}"
-        logger.error(error_msg)
-        stream_handler.update(error_msg, "Error")
-        yield error_msg, None, "\n".join(stream_handler.output), "Error"
+        error_msg = f"\nError: {str(e)}"
+        stream_handler.update(error_msg, "Failed")
+        yield None, None, "\n".join(stream_handler.output), "Failed"
+        raise
+    finally:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
 class StreamHandler:
     """Handles streaming output for the Gradio interface."""
@@ -1193,41 +1187,29 @@ with gr.Blocks(theme=gr.themes.Soft()) as ui:
             
             try:
                 # Run the app generator with a timeout
-                result = await asyncio.wait_for(
-                    app_generator(requirements),
-                    timeout=60  # 60 second timeout
-                )
-                
-                if result:
-                    # Create download file
-                    try:
-                        download_path = result.get("download_path")
-                        if download_path:
-                            stream_handler.update(" Generation completed successfully!", "Success")
-                            yield result.get("output", ""), download_path, "\n".join(stream_handler.output), "Success"
-                        else:
-                            error_msg = " Failed to create download file"
-                            stream_handler.update(error_msg, "Error")
-                            yield result.get("output", ""), None, "\n".join(stream_handler.output), "Error"
-                    except Exception as e:
-                        error_msg = f" Error creating download: {str(e)}"
-                        stream_handler.update(error_msg, "Error")
-                        yield result.get("output", ""), None, "\n".join(stream_handler.output), "Error"
-                else:
-                    error_msg = " Failed to generate app"
-                    stream_handler.update(error_msg, "Error")
-                    yield "", None, "\n".join(stream_handler.output), "Error"
+                async with asyncio.timeout(60):  # 60 second timeout
+                    result = await app_generator(requirements)
                     
+                    # Update output with result
+                    if result["output"]:
+                        stream_handler.update("\n" + result["output"], "Completed")
+                        yield result["output"], result["download_path"], "\n".join(stream_handler.output), "Completed"
+                    else:
+                        raise Exception("No output generated")
+                        
             except asyncio.TimeoutError:
-                error_msg = " Generation timed out after 60 seconds"
-                stream_handler.update(error_msg, "Error")
-                yield "", None, "\n".join(stream_handler.output), "Error"
+                stream_handler.update("\nApp generation timed out after 60 seconds", "Failed")
+                yield None, None, "\n".join(stream_handler.output), "Failed"
+                raise
                 
         except Exception as e:
-            error_msg = f" Error: {str(e)}"
-            logger.error(error_msg)
-            stream_handler.update(error_msg, "Error")
-            yield error_msg, None, "\n".join(stream_handler.output), "Error"
+            error_msg = f"\nError: {str(e)}"
+            stream_handler.update(error_msg, "Failed")
+            yield None, None, "\n".join(stream_handler.output), "Failed"
+            raise
+        finally:
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
 
     def cancel_generation():
         """Cancel the current generation process."""