Update app.py

app.py

@@ -1,5 +1,6 @@
-# app.py — REFACTORED & BACKWARD-COMPATIBLE VERSION
+# app.py — FINAL VERSION WITH CUSTOM GENERATOR LOOPS (BACKWARD-COMPATIBLE)
 import gradio as gr
+import time # Import the time module for sleeping
 from backend import (
     # User Management
     create_user,
@@ -16,9 +17,6 @@ from backend import (
 import psutil
 import torch
 
-# Gradio's session state will manage user data for each session
-# We will use gr.State to hold session-specific data
-
 # --- UI HELPER FUNCTIONS ---
 
 def get_system_stats():
@@ -45,9 +43,8 @@ def format_projects_list(projects):
 
 def get_project_choices(user_id):
     """Fetches user projects and formats them for a dropdown."""
-    if not user_id:
-        return []
-    projects = get_user_projects(user_id, limit=20) # Increase limit for dropdown
+    if not user_id: return []
+    projects = get_user_projects(user_id, limit=20)
     return [(f"[{p['id']}] {p['title']}", p['id']) for p in projects] if projects else []
 
 # --- CORE LOGIC HANDLERS FOR UI ---
@@ -55,8 +52,7 @@ def get_project_choices(user_id):
 def login_and_setup(username, password, user_state):
     """Handles login/signup and prepares the main UI if successful."""
     user = get_user_by_username(username)
-    user_id = None
-    welcome_message = ""
+    user_id, welcome_message = None, ""
 
     if user and verify_password(password, user['password_hash']):
         user_id = user['id']
@@ -67,27 +63,21 @@ def login_and_setup(username, password, user_state):
             user_id = new_user_id
             welcome_message = f"Account created! Welcome, {username}!"
         else:
-            # Login failed and username is taken
             return user_state, gr.update(), gr.update(value="Login failed or username is taken.")
 
     if user_id:
         user_state['user_id'] = user_id
         user_state['username'] = username
-        
         projects = get_user_projects(user_id)
-        project_list_md = format_projects_list(projects)
-        dropdown_choices = get_project_choices(user_id)
         
-        # Hide login UI, show main UI, and populate it with data
         return (
             user_state,
-            gr.update(visible=False), # Hide login UI
-            gr.update(visible=True),  # Show main UI
+            gr.update(visible=False),
+            gr.update(visible=True),
             gr.update(value=welcome_message),
-            gr.update(value=project_list_md),
-            gr.update(choices=dropdown_choices, value=None)
+            gr.update(value=format_projects_list(projects)),
+            gr.update(choices=get_project_choices(user_id), value=None)
         )
-    # This path is for the failed login/signup
     return user_state, gr.update(), gr.update()
 
 
@@ -95,167 +85,144 @@ def start_project_flow(prompt, user_state):
     """The complete, seamless flow for starting a new project."""
     user_id = user_state.get('user_id')
     if not user_id:
-        return "Error: Not logged in.", "", "", None, None, gr.update()
+        return "Error: Not logged in.", "", "", None, gr.update()
 
-    # 1. Create project in DB
     project_id = create_project(user_id, "New AI Project", prompt)
     if not project_id:
-        return "Error: Could not create project.", "", "", None, None, gr.update()
+        return "Error: Could not create project.", "", "", None, gr.update()
         
-    # 2. Queue the backend job
     queue_job(project_id, user_id, prompt)
     
-    # 3. Refresh UI elements with the new project data
     projects = get_user_projects(user_id)
-    project_list_md = format_projects_list(projects)
-    dropdown_choices = get_project_choices(user_id)
-
-    # 4. Return updates for ALL relevant components
     return (
-        f"✅ Project #{project_id} started! Switching to logs...", # Status message
-        gr.update(value=""), # Clear the prompt box
-        gr.update(value=project_list_md), # Update dashboard project list
-        gr.update(choices=dropdown_choices, value=project_id), # Update dropdown and AUTO-SELECT the new project
-        gr.Tabs(selected="logs_tab"), # Switch to the Logs tab
+        f"✅ Project #{project_id} started! Switching to logs...",
+        gr.update(value=""),
+        gr.update(value=format_projects_list(projects)),
+        gr.update(choices=get_project_choices(user_id), value=project_id),
+        gr.Tabs(selected="logs_tab"),
     )
 
 def refresh_all_projects(user_state):
     """Manually refresh project lists on both tabs."""
     user_id = user_state.get('user_id')
-    if not user_id:
-        return gr.update(), gr.update()
-        
-    projects = get_user_projects(user_id)
-    project_list_md = format_projects_list(projects)
-    dropdown_choices = get_project_choices(user_id)
+    if not user_id: return gr.update(), gr.update()
     
-    return gr.update(value=project_list_md), gr.update(choices=dropdown_choices)
+    projects = get_user_projects(user_id)
+    return gr.update(value=format_projects_list(projects)), gr.update(choices=get_project_choices(user_id))
+
+# --- CUSTOM GENERATOR LOOPS (THE FIX) ---
+
+def system_stats_loop():
+    """A generator that yields system stats every 5 seconds, forever."""
+    while True:
+        yield get_system_stats()
+        time.sleep(5)
 
-def get_live_logs(project_id):
-    """Fetches logs for the selected project. Designed for live polling."""
+def stream_logs(project_id):
+    """A generator that streams logs for a selected project every 3 seconds."""
     if not project_id:
-        return "⬅️ Select a project from the dropdown to see its live logs."
-    
+        yield "⬅️ Select a project from the dropdown to see its live logs."
+        return # Stop the generator if no project is selected
+
+    # Yield an initial message immediately
     project = get_project(project_id)
     if not project or not project['logs']:
-        return f"🕒 Project #{project_id} is queued or starting... Logs will appear here shortly."
-        
-    return project['logs']
+        yield f"🕒 Project #{project_id} is queued or starting... Logs will appear here shortly."
+    
+    # Loop to continuously poll for new logs
+    while True:
+        project = get_project(project_id)
+        if project:
+            yield project['logs']
+            # If project is finished, stop polling to save resources
+            if project['status'] in ['completed', 'failed']:
+                break
+        else:
+            yield f"Error: Project #{project_id} not found."
+            break
+        time.sleep(3)
+
 
 # --- GRADIO UI DEFINITION ---
 
 with gr.Blocks(title="Code Agents Pro", theme=gr.themes.Soft()) as demo:
-    # Session state to hold user_id and username
     user_state = gr.State({"user_id": None, "username": ""})
     
     gr.Markdown("# 🤖 Code Agents Pro — Your Automated AI Software Team")
 
-    # --- PAGE 1: LOGIN UI ---
     with gr.Group(visible=True) as login_ui:
-        with gr.Row():
-            gr.Markdown("### 🔐 Login or Sign Up")
+        gr.Markdown("### 🔐 Login or Sign Up")
         username_input = gr.Textbox(label="Username", placeholder="Choose a unique username")
         password_input = gr.Textbox(label="Password", type="password", placeholder="Enter your password")
         login_btn = gr.Button("Login / Sign Up", variant="primary")
         login_msg = gr.Markdown()
 
-    # --- PAGE 2: MAIN APPLICATION UI (Initially Hidden) ---
     with gr.Group(visible=False) as main_ui_container:
         welcome_msg = gr.Markdown()
         
         with gr.Tabs() as tabs:
-            # --- DASHBOARD TAB ---
             with gr.Tab("🚀 Dashboard", id="dashboard_tab"):
                 with gr.Row():
                     with gr.Column(scale=2):
                         gr.Markdown("### ✨ Start a New Project")
-                        project_prompt = gr.Textbox(
-                            label="Describe your project goal in detail",
-                            lines=8,
-                            placeholder="e.g., 'Build a Python FastAPI server with a PostgreSQL database using SQLAlchemy. It needs a /users endpoint for creating and listing users.'"
-                        )
+                        project_prompt = gr.Textbox(label="Describe your project goal in detail", lines=8, placeholder="e.g., 'Build a Python FastAPI server...'")
                         start_btn = gr.Button("🚀 Start AI Team Work", variant="primary")
                         status_msg = gr.Markdown()
-
                     with gr.Column(scale=3):
                         gr.Markdown("### 📁 Your Projects")
                         refresh_all_btn = gr.Button("🔄 Refresh Project List")
                         projects_display = gr.Markdown("Login to see your projects.")
                 
-            # --- AGENT LOGS TAB ---
             with gr.Tab("🧠 Agent Logs", id="logs_tab"):
                 gr.Markdown("### 📜 View Agent Conversation Logs")
-                project_selector = gr.Dropdown(
-                    label="Select Project to View",
-                    choices=[],
-                    interactive=True,
-                )
-                # --- THIS IS THE FIX ---
-                # Replaced gr.Box with the backward-compatible gr.Group
+                project_selector = gr.Dropdown(label="Select Project to View", choices=[], interactive=True)
                 with gr.Group(elem_classes=["logs-container"]):
                     logs_display = gr.Markdown("Logs will appear here in real-time.", elem_classes=["monospace"])
         
-        # System stats monitor at the bottom
         ram_monitor = gr.Markdown()
 
     # ------------------------------ EVENT WIRING ------------------------------
-
-    # 1. Login Action
+    
     login_btn.click(
         fn=login_and_setup,
         inputs=[username_input, password_input, user_state],
         outputs=[user_state, login_ui, main_ui_container, welcome_msg, projects_display, project_selector]
     )
     
-    # 2. Start Project Action (THE SEAMLESS FLOW)
     start_btn.click(
         fn=start_project_flow,
         inputs=[project_prompt, user_state],
         outputs=[status_msg, project_prompt, projects_display, project_selector, tabs]
     )
 
-    # 3. Manual Refresh Action
     refresh_all_btn.click(
         fn=refresh_all_projects,
         inputs=[user_state],
         outputs=[projects_display, project_selector]
     )
     
-    # 4. Action when dropdown value changes MANUALLY
+    # --- THE FIX IN ACTION ---
+    # When the dropdown changes (either by user or by the start_project_flow),
+    # it triggers the stream_logs generator, starting our custom loop.
     project_selector.change(
-        fn=get_live_logs,
+        fn=stream_logs,
         inputs=project_selector,
         outputs=logs_display
     )
 
-    # 5. LIVE POLLING for system stats and logs
-    demo.load(fn=get_system_stats, outputs=ram_monitor, every=5)
+    # Use a generator for the system stats loop, triggered on app load.
+    # This is the backward-compatible way to do a recurring event.
     demo.load(
-        fn=get_live_logs,
-        inputs=[project_selector], # It will use the current value of the dropdown
-        outputs=[logs_display],
-        every=3 # Refresh logs every 3 seconds
+        fn=system_stats_loop,
+        outputs=ram_monitor
     )
 
     # --- CSS STYLING ---
     gr.HTML("""
     <style>
-        .logs-container {
-            height: 600px;
-            overflow-y: auto;
-            border: 1px solid #E5E7EB;
-            border-radius: 8px;
-            padding: 8px;
-            background-color: #F9FAFB;
-        }
-        .monospace {
-            font-family: 'SF Mono', 'Consolas', 'Courier New', monospace;
-            font-size: 14px;
-            white-space: pre-wrap;
-            word-wrap: break-word;
-        }
+        .logs-container { height: 600px; overflow-y: auto; border: 1px solid #E5E7EB; border-radius: 8px; padding: 8px; background-color: #F9FAFB; }
+        .monospace { font-family: 'SF Mono', 'Consolas', 'Courier New', monospace; font-size: 14px; white-space: pre-wrap; word-wrap: break-word; }
     </style>
     """)
 
-# Use queue() for per-user state management, essential for multi-user apps
 demo.queue().launch(server_name="0.0.0.0", server_port=7860)