Spaces:
Running
Running
| # 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(""" | |
| <style> | |
| .logs-container { height: 550px; 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> | |
| """) | |
| demo.queue().launch(server_name="0.0.0.0", server_port=7860) |