Fixed stream output and app generation

app.py

@@ -1009,22 +1009,82 @@ async def multi_agent_workflow(requirements: str) -> str:
         raise
 
 # Update the app_generator function to handle async execution
-async def app_generator(requirements: str) -> str:
-    """
-    Generate an app based on the provided requirements using AI Flows.
-    Args:
-        requirements (str): App requirements.
-    Returns:
-        str: Generated app code and API details.
-    """
+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)
-        return result
+        return {
+            "output": result,
+            "download_path": None
+        }
     except Exception as e:
         logger.error(f"Failed to generate app: {str(e)}")
         raise
 
+async def stream_output(requirements, progress=gr.Progress()):
+    """Stream the output during app generation."""
+    try:
+        # Initialize
+        stream_handler.update(" Starting app generation...", "Initializing")
+        yield "Starting...", None, " Starting app generation...", "Initializing"
+        
+        # Update progress
+        phases = [
+            (" Analyzing requirements...", "Analyzing"),
+            (" Generating architecture...", "Designing"),
+            (" Creating project structure...", "Creating"),
+            (" Implementing features...", "Implementing"),
+            (" Finalizing...", "Finalizing")
+        ]
+        
+        for msg, status in progress.tqdm(phases):
+            stream_handler.update(msg, status)
+            yield None, None, "\n".join(stream_handler.output), status
+            await asyncio.sleep(1)  # Non-blocking sleep
+        
+        # Generate the app
+        stream_handler.update(" Running AI Flow system...", "Processing")
+        yield None, None, "\n".join(stream_handler.output), "Processing"
+        
+        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"
+                
+        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"
+            
+    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"
+
 class StreamHandler:
     """Handles streaming output for the Gradio interface."""
     
@@ -1124,34 +1184,44 @@ with gr.Blocks(theme=gr.themes.Soft()) as ui:
             
             for msg, status in progress.tqdm(phases):
                 stream_handler.update(msg, status)
-                yield None, None, stream_handler.output[-1], status
+                yield None, None, "\n".join(stream_handler.output), status
                 await asyncio.sleep(1)  # Non-blocking sleep
             
             # Generate the app
             stream_handler.update(" Running AI Flow system...", "Processing")
-            result = await app_generator(requirements)
-            
-            # Extract download path and logs
-            download_path = None
-            logs = []
+            yield None, None, "\n".join(stream_handler.output), "Processing"
             
-            for line in result.split('\n'):
-                if line.startswith("To download your project, use this path:"):
-                    download_path = line.split(": ")[1].strip()
-                    stream_handler.update(f" Generated download file: {Path(download_path).name}", "Processing")
-                if line.startswith("2025-") or line.startswith("INFO") or line.startswith("ERROR"):
-                    logs.append(line)
-                    stream_handler.update(line.split(" - ")[-1], "Processing")
-                    yield None, None, "\n".join(stream_handler.output), "Processing"
-                await asyncio.sleep(0.1)  # Small delay to prevent UI freezing
-            
-            if download_path and Path(download_path).exists():
-                stream_handler.update(" App generated successfully! Download is ready.", "Complete")
-                yield result, download_path, "\n".join(stream_handler.output), "Complete"
-            else:
-                error_msg = " Failed to generate download file"
+            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"
+                    
+            except asyncio.TimeoutError:
+                error_msg = " Generation timed out after 60 seconds"
                 stream_handler.update(error_msg, "Error")
-                yield result, None, "\n".join(stream_handler.output), "Error"
+                yield "", None, "\n".join(stream_handler.output), "Error"
                 
         except Exception as e:
             error_msg = f" Error: {str(e)}"