Made entire workflow async with proper async/await patterns

app.py

@@ -755,152 +755,108 @@ class EnhancedAIFlow(AIFlow):
         self.file_manager = FileOperationManager(self.context_manager)
         self.command_manager = CommandManager(self.context_manager)
         self.rule_system = RuleSystem()
+        self.flow_graph = nx.DiGraph()
+        self.contexts: Dict[AgentRole, AgentContext] = {}
+        self.global_context = {}
     
-    async def _execute_agent_task(self, role: AgentRole, prompt: str) -> str:
-        """Execute a specific agent's task with the given prompt."""
-        try:
-            if role == AgentRole.ARCHITECT:
-                agent = get_agent("architect")
-            elif role == AgentRole.UI_DESIGNER:
-                agent = get_agent("ui_designer")
-            elif role == AgentRole.BACKEND_DEVELOPER:
-                agent = get_agent("backend_developer")
-            elif role == AgentRole.DATABASE_ENGINEER:
-                agent = get_agent("database_engineer")
-            elif role == AgentRole.SECURITY_EXPERT:
-                agent = get_agent("security_expert")
-            elif role == AgentRole.QA_ENGINEER:
-                agent = get_agent("qa_engineer")
-            elif role == AgentRole.DEVOPS_ENGINEER:
-                agent = get_agent("devops_engineer")
-            elif role == AgentRole.DOCUMENTATION_WRITER:
-                agent = get_agent("documentation_writer")
-            else:
-                raise ValueError(f"Unknown agent role: {role}")
-            
-            # Execute the agent's task
-            result = agent.run(prompt)
-            
-            # Log the execution
-            logger.info(f"Agent {role.value} completed task")
-            
-            return result
-        except Exception as e:
-            logger.error(f"Agent {role.value} failed: {str(e)}")
-            raise
-    
-    async def execute_flow(self, requirements: str):
+    async def execute_flow(self, requirements: str) -> str:
         """Execute the AI Flow and build the project."""
         try:
-            # Initialize flow with context
+            # Initialize flow with requirements
             self.initialize_flow()
-            self.global_context["requirements"] = requirements
-            self.context_manager.global_context.update(self.global_context)
-            
-            # Add memory of requirements
-            self.context_manager.add_memory(
-                "requirements",
-                requirements,
-                {"timestamp": datetime.now()}
-            )
-            
-            # Get all paths through the flow graph
-            paths = list(nx.all_simple_paths(
-                self.flow_graph,
-                AgentRole.ARCHITECT,
-                AgentRole.DOCUMENTATION_WRITER
-            ))
-            
-            # Execute paths with enhanced context
-            await self._execute_paths(paths)
             
             # Extract app name
             app_name = requirements.split()[0].lower().replace(" ", "_")
             
-            # Create basic project structure
-            structure = {
-                "frontend": {
-                    "src": {
-                        "components": {},
-                        "pages": {},
-                        "styles": {},
-                    },
-                    "package.json": "{\n  \"name\": \"frontend\",\n  \"version\": \"1.0.0\"\n}",
-                    "README.md": "# Frontend\n"
-                },
-                "backend": {
-                    "src": {
-                        "routes": {},
-                        "models": {},
-                        "controllers": {},
-                    },
-                    "requirements.txt": "fastapi\nuvicorn\n",
-                    "README.md": "# Backend\n"
-                },
-                "README.md": f"# {app_name}\nGenerated by AI Flow\n"
-            }
+            # 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)
             
-            # Build the project
-            await self.project_builder.create_project(app_name, structure)
+            # 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 download
+            # Create downloadable output
             download_path = self.output_manager.create_download(app_name)
             
-            # Compile results
-            results = self._compile_results()
-            
-            # Add download information to results
-            results += f"""
-## Download
-Your project has been created and is ready for download:
-- File: {Path(download_path).name}
-- Size: {self.output_manager.downloads[Path(download_path).name]['size'] / 1024:.1f} KB
-- Created: {self.output_manager.downloads[Path(download_path).name]['created_at'].strftime('%Y-%m-%d %H:%M:%S')}
+            return f"""
+# {app_name.title()} - Generated Application
+
+## Project Structure
+```
+{project_structure}
+```
+
+## Documentation
+{documentation}
 
+## Download
 To download your project, use this path: {download_path}
 
-## Project Structure
-The following files have been created:
+## Next Steps
+1. Review the generated architecture and components
+2. Set up the development environment
+3. Implement the components following the provided structure
+4. Run the test suite
+5. Deploy using the provided configurations
+
+## Support
+For any issues or questions, please refer to the documentation or create an issue in the repository.
 """
-            for file_path in self.project_builder.file_manifest:
-                rel_path = file_path.relative_to(self.project_builder.current_build)
-                results += f"- {rel_path}\n"
-            
-            return results
-            
         except Exception as e:
-            logger.error(f"Enhanced flow execution failed: {str(e)}")
+            logger.error(f"Failed to execute flow: {str(e)}")
             raise
-        
+        finally:
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+    
+    async def _execute_paths(self, paths: List[List[AgentRole]]):
+        """Execute all paths in the flow graph."""
+        try:
+            for path in paths:
+                for role in path:
+                    if self.contexts[role].state != FlowState.COMPLETED:
+                        await self._execute_agent(role)
+        except Exception as e:
+            logger.error(f"Failed to execute paths: {str(e)}")
+            raise
+    
     async def _execute_agent(self, role: AgentRole):
         """Execute a single agent's tasks with enhanced context."""
-        context = self.contexts[role]
-        context.state = FlowState.RUNNING
-        
         try:
-            # Get agent-specific prompt with context
-            prompt = self._get_agent_prompt(role)
+            context = self.contexts[role]
+            if context.state == FlowState.COMPLETED:
+                return
             
-            # Add current context to prompt
-            prompt += f"\n\nContext:\n{json.dumps(self.context_manager.global_context, indent=2)}"
+            context.state = FlowState.RUNNING
+            prompt = self._generate_prompt(role)
             
-            # Execute agent's task
+            # Execute agent task
             result = await self._execute_agent_task(role, prompt)
             
-            # Store results with context
-            context.artifacts["output"] = result
+            # Update context with results
+            context.artifacts.update(result)
             context.state = FlowState.COMPLETED
             
-            # Update memory
-            self.context_manager.add_memory(
-                f"agent_result_{role.value}",
-                result,
-                {"role": role.value, "timestamp": datetime.now()}
-            )
-            
         except Exception as e:
             context.state = FlowState.FAILED
-            context.feedback.append(str(e))
+            logger.error(f"Failed to execute agent {role}: {str(e)}")
+            raise
+    
+    async def _execute_agent_task(self, role: AgentRole, prompt: str) -> Dict[str, str]:
+        """Execute a specific agent's task with the given prompt."""
+        try:
+            # Get the appropriate agent for the role
+            agent = get_agent(role)
+            
+            # Execute the agent's task
+            result = await agent.arun(prompt)
+            
+            # Process and return the result
+            return {"output": result}
+            
+        except Exception as e:
+            logger.error(f"Failed to execute agent task for {role}: {str(e)}")
             raise
 
 # Update the multi_agent_workflow function to use AI Flows
@@ -913,50 +869,20 @@ async def multi_agent_workflow(requirements: str) -> str:
         str: Generated app code and API details.
     """
     try:
-        # Initialize and execute AI Flow
-        flow = EnhancedAIFlow()
-        results = await flow.execute_flow(requirements)
-        
-        # Extract app name
-        app_name = requirements.split()[0].lower().replace(" ", "_")
+        # Initialize AI Flow
+        ai_flow = EnhancedAIFlow()
         
-        # Generate project structure and documentation
-        project_structure = generate_project_structure(app_name, flow.contexts[AgentRole.ARCHITECT].artifacts)
-        documentation = generate_documentation(app_name, requirements, flow.contexts[AgentRole.DOCUMENTATION_WRITER].artifacts)
+        # Execute the flow with requirements
+        result = await ai_flow.execute_flow(requirements)
         
-        return f"""
-# {app_name.title()} - Generated Application
-
-## Project Structure
-```
-{project_structure}
-```
-
-{results}
-
-## Documentation
-{documentation}
-
-## Next Steps
-1. Review the generated architecture and components
-2. Set up the development environment
-3. Implement the components following the provided structure
-4. Run the test suite
-5. Deploy using the provided configurations
-
-## Support
-For any issues or questions, please refer to the documentation or create an issue in the repository.
-"""
+        # Return the compiled results
+        return result
     except Exception as e:
-        error_msg = f"Workflow failed: {str(e)}"
-        logger.error(error_msg)
-        return error_msg
-    finally:
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        logger.error(f"Failed to execute multi-agent workflow: {str(e)}")
+        raise
 
 # Update the app_generator function to handle async execution
-def app_generator(requirements: str):
+async def app_generator(requirements: str) -> str:
     """
     Generate an app based on the provided requirements using AI Flows.
     Args:
@@ -964,20 +890,13 @@ def app_generator(requirements: str):
     Returns:
         str: Generated app code and API details.
     """
-    if not requirements or len(requirements.strip()) == 0:
-        return "Please provide app requirements."
-    
     try:
-        # Run the async workflow in a new event loop
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        return loop.run_until_complete(multi_agent_workflow(requirements))
+        # Execute the multi-agent workflow
+        result = await multi_agent_workflow(requirements)
+        return result
     except Exception as e:
-        error_msg = f"App generation failed: {str(e)}"
-        logger.error(error_msg)
-        return error_msg
-    finally:
-        loop.close()
+        logger.error(f"Failed to generate app: {str(e)}")
+        raise
 
 class StreamHandler:
     """Handles streaming output for the Gradio interface."""