# app.py ā€” FINAL VERSION (with Download and Files Tab) import gradio as gr import time import os from backend import ( create_user, get_user_by_username, verify_password, create_project, get_user_projects, get_project, queue_job, ) import psutil import torch # --- UI HELPER FUNCTIONS --- def get_system_stats(): """Returns a markdown string with current RAM and VRAM usage.""" ram_used = round(psutil.virtual_memory().used / (1024**3), 2) ram_total = round(psutil.virtual_memory().total / (1024**3), 2) stats_str = f"šŸ“Š System RAM: {ram_used}GB / {ram_total}GB" if torch.cuda.is_available(): vram_reserved = round(torch.cuda.memory_reserved() / (1024**3), 2) vram_total = round(torch.cuda.get_device_properties(0).total_memory / (1024**3), 2) stats_str += f" | VRAM: {vram_reserved}GB / {vram_total}GB" return stats_str def format_projects_list(projects): """Formats a list of project records into a readable markdown string.""" if not projects: return "No projects yet. Start one on the left!" project_lines = [] for p in projects: status_icon = {"queued": "šŸ•’", "running": "āš™ļø", "completed": "āœ…", "failed": "āŒ"}.get(p['status'], "ā“") project_lines.append(f"{status_icon} **[{p['id']}] {p['title']}** ({p['status']})") return "\n".join(project_lines) 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) return [(f"[{p['id']}] {p['title']}", p['id']) for p in projects] if projects else [] def format_files_list(projects): """Formats completed projects into downloadable file links.""" completed_projects = [p for p in projects if p['status'] == 'completed'] if not completed_projects: return "No completed projects yet. Your finished work will appear here!" file_lines = [] for p in completed_projects: zip_path = p.get('zip_path') if zip_path and os.path.exists(zip_path): file_lines.append(f"šŸ“„ **[{p['id']}] {p['title']}** - [Download ZIP]({zip_path})") else: file_lines.append(f"āš ļø **[{p['id']}] {p['title']}** - ZIP file missing") return "\n".join(file_lines) # --- CORE LOGIC HANDLERS FOR UI --- 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, welcome_message = None, "" if user and verify_password(password, user['password_hash']): user_id = user['id'] welcome_message = f"Welcome back, {username}! šŸ‘Øā€šŸ’»" else: new_user_id = create_user(username, password) if new_user_id: user_id = new_user_id welcome_message = f"Account created! Welcome, {username}!" else: 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) return ( user_state, gr.update(visible=False), gr.update(visible=True), gr.update(value=welcome_message), gr.update(value=format_projects_list(projects)), gr.update(choices=get_project_choices(user_id), value=None), gr.update(value=format_files_list(projects)) # For files tab ) return user_state, gr.update(), gr.update() 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, gr.update(), gr.update() project_id = create_project(user_id, "New AI Project", prompt) if not project_id: return "Error: Could not create project.", "", "", None, gr.update(), gr.update() queue_job(project_id, user_id, prompt) projects = get_user_projects(user_id) return ( 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"), gr.update(value=format_files_list(projects)) # Update files 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(), gr.update() projects = get_user_projects(user_id) return ( gr.update(value=format_projects_list(projects)), gr.update(choices=get_project_choices(user_id)), gr.update(value=format_files_list(projects)) ) # --- CUSTOM GENERATOR LOOPS FOR LIVE UPDATES --- def system_stats_loop(): """A generator that yields system stats every 5 seconds, forever.""" while True: yield get_system_stats() time.sleep(5) def stream_logs(project_id): """A generator that streams logs and provides a download link when complete.""" if not project_id: yield "ā¬…ļø Select a project from the dropdown to see its live logs.", gr.update(visible=False) return # On first run for a new selection, ensure the download button is hidden yield "Fetching logs...", gr.update(visible=False) while True: project = get_project(project_id) if project: logs = project['logs'] if project['logs'] else f"šŸ•’ Project #{project_id} is queued or starting... Logs will appear here shortly." # Check for terminal states if project['status'] == 'completed': final_log = logs + "\n\nāœ… **Project Completed!** Your file is ready for download." zip_path = project.get('zip_path') if zip_path and os.path.exists(zip_path): yield final_log, gr.update(value=zip_path, visible=True) else: yield logs + "\n\nāš ļø **Warning:** Project completed, but the ZIP file was not found.", gr.update(visible=False) break # Stop the loop elif project['status'] == 'failed': yield logs, gr.update(visible=False) break # Stop the loop # Otherwise, just stream the current logs yield logs, gr.update(visible=False) else: yield f"Error: Project #{project_id} not found.", gr.update(visible=False) break time.sleep(3) # --- GRADIO UI DEFINITION --- with gr.Blocks(title="Code Agents Pro", theme=gr.themes.Soft()) as demo: user_state = gr.State({"user_id": None, "username": ""}) gr.Markdown("# šŸ¤– Code Agents Pro ā€” Your Automated AI Software Team") with gr.Group(visible=True) as login_ui: 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() with gr.Group(visible=False) as main_ui_container: welcome_msg = gr.Markdown() with gr.Tabs() as tabs: 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...'") 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.") 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) with gr.Group(elem_classes=["logs-container"]): logs_display = gr.Markdown("Logs will appear here in real-time.", elem_classes=["monospace"]) # File download component, initially hidden download_zip_ui = gr.File(label="Download Project ZIP", visible=False, interactive=False) # NEW FILES TAB with gr.Tab("šŸ“ Files", id="files_tab"): gr.Markdown("### šŸ“¦ Download Completed Projects") refresh_files_btn = gr.Button("šŸ”„ Refresh Files") files_display = gr.Markdown("Login to see your completed projects.") ram_monitor = gr.Markdown() # ------------------------------ EVENT WIRING ------------------------------ 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, files_display # New files tab content ] ) start_btn.click( fn=start_project_flow, inputs=[project_prompt, user_state], outputs=[ status_msg, project_prompt, projects_display, project_selector, tabs, files_display # Update files tab when new project starts ] ) refresh_all_btn.click( fn=refresh_all_projects, inputs=[user_state], outputs=[ projects_display, project_selector, files_display # Also refresh files tab ] ) # NEW: Refresh files tab button refresh_files_btn.click( fn=lambda user_state: gr.update(value=format_files_list(get_user_projects(user_state.get('user_id', 0)))) if user_state.get('user_id') else gr.update(), inputs=[user_state], outputs=[files_display] ) project_selector.change( fn=stream_logs, inputs=project_selector, outputs=[logs_display, download_zip_ui] ) demo.load( fn=system_stats_loop, outputs=ram_monitor ) gr.HTML(""" """) demo.queue().launch(server_name="0.0.0.0", server_port=7860)