Update app.py

app.py

@@ -1,91 +1,294 @@
-import re
-import json
-import time
-import requests
-import importlib.metadata
-import gradio as gr
-from huggingface_hub import (
-    create_repo, upload_file, list_models, constants
-)
-from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status
-from google import genai
-from google.genai.types import Tool, GenerateContentConfig, GoogleSearch
-
-# — USER INFO & MODEL LISTING —
-
-def show_profile(profile: gr.OAuthProfile | None) -> str:
-    return f"✅ Logged in as **{profile.username}**" if profile else "*Not logged in.*"
-
-def list_private_models(
-    profile: gr.OAuthProfile | None,
-    oauth_token: gr.OAuthToken | None
-) -> str:
-    if not profile or not oauth_token:
-        return "Please log in to see your models."
-    models = [
-        f"{m.id} ({'private' if m.private else 'public'})"
-        for m in list_models(author=profile.username, token=oauth_token.token)
-    ]
-    return "No models found." if not models else "Models:\n\n" + "\n - ".join(models)
+# ... (existing imports and utility functions like show_profile, list_private_models,
+#      get_sdk_version, extract_code, classify_errors, _get_space_jwt,
+#      fetch_logs, check_iframe remain the same)
+
+# --- AGENT PROMPTS & FUNCTIONS ---
+
+# Updated/New System Prompts
+SYSTEM_ORCHESTRATOR = {
+    "role": "system",
+    "content": (
+        "You are **Orchestrator Agent**, the project manager. "
+        "Your role is to guide the development process from user request to a deployed HF Space application. "
+        "You will delegate tasks to other agents (Architect, Code-Gen, Debug) based on the project state. "
+        "Your inputs will include the current project state (requirements, plan, files, logs, feedback, status). "
+        "Your output should be a decision on the next step/task, potentially summarizing information for the next agent."
+        "Possible next steps could be: 'PLANNING', 'CODING - {task_description}', 'DEBUGGING', 'COMPLETE', 'FAILED'."
+        "Analyze the feedback and logs carefully to decide the next coding task if debugging is needed."
+    )
+}
+
+SYSTEM_ARCHITECT = {
+    "role": "system",
+    "content": (
+        "You are **Architect Agent**, the lead planner. "
+        "Given the user requirements, your task is to devise a high-level plan for the application. "
+        "Outline the main features, suggest a file structure (e.g., `app.py`, `utils.py`), and identify key components needed. "
+        "Output the plan clearly, perhaps using bullet points or a simple structure. Do NOT write code."
+    )
+}
+
+SYSTEM_CODEGEN = {
+    "role": "system",
+    "content": (
+        "You are **Code‑Gen Agent**, a proactive AI developer. "
+        "Your sole responsibility is to author and correct code files based on the plan and assigned task. "
+        "You will receive the full project state, including the plan, existing files, and debug feedback. "
+        "Based on the current task assigned by the Orchestrator, write or modify the necessary code *only* in the specified file(s). "
+        "Output the *full content* of the updated file(s) in markdown code blocks, clearly indicating the filename(s)."
+        "Example output format: \n\n`app.py`\n```python\n# app code here\n```\n\n`utils.py`\n```python\n# utils code here\n```"
+        "Focus *only* on outputting the code blocks for the files you are modifying/creating. Do not add extra commentary outside the code blocks."
+    )
+}
+
+SYSTEM_DEBUG = {
+    "role": "system",
+    "content": (
+        "You are **Debug Agent**, a meticulous code reviewer and tester. "
+        "You have access to the full project state: requirements, plan, code files, build logs, and run logs. "
+        "Your task is to analyze the logs and code in the context of the plan and requirements. "
+        "Identify errors, potential issues, missing features based on the plan, and suggest concrete improvements or fixes. "
+        "If the application appears to be working and meets the plan/requirements (and no significant errors in logs), state 'All clear. Project appears complete.' "
+        "Otherwise, provide actionable feedback for the Code-Gen agent, referencing file names and line numbers where possible. Format feedback clearly."
+        "Example feedback: 'Error in `app.py`: ModuleNotFoundError for 'missing_library'. Add 'missing_library' to `requirements.txt`.\nIssue: The 'plot_data' function in `utils.py` doesn't handle empty input data per the plan.' "
+        "Do NOT write or suggest code directly in your feedback unless it's a very small, inline fix suggestion. Focus on *what* needs fixing/adding and *why*."
+    )
+}
 
-# — UTILITIES —
+# --- AGENT FUNCTIONS ---
 
-def get_sdk_version(sdk_choice: str) -> str:
-    pkg = "gradio" if sdk_choice == "gradio" else "streamlit"
-    try:
-        return importlib.metadata.version(pkg)
-    except importlib.metadata.PackageNotFoundError:
-        return "UNKNOWN"
-
-def extract_code(text: str) -> str:
-    blocks = re.findall(r"```(?:\w*\n)?([\s\S]*?)```", text)
-    return blocks[-1].strip() if blocks else text.strip()
-
-def classify_errors(logs: str) -> str:
-    errs = set()
-    for line in logs.splitlines():
-        if "SyntaxError" in line:
-            errs.add("syntax")
-        elif "ImportError" in line or "ModuleNotFoundError" in line:
-            errs.add("import")
-        elif "Traceback" in line or "Exception" in line:
-            errs.add("runtime")
-    return ", ".join(errs) or "unknown"
-
-# — HF SPACE LOGGING —
-
-def _get_space_jwt(repo_id: str) -> str:
-    url = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
-    r = get_session().get(url, headers=build_hf_headers())
-    hf_raise_for_status(r)
-    return r.json()["token"]
-
-def fetch_logs(repo_id: str, level: str) -> str:
-    jwt = _get_space_jwt(repo_id)
-    url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
-    lines = []
-    with get_session().get(url, headers=build_hf_headers(token=jwt), stream=True) as resp:
-        hf_raise_for_status(resp)
-        for raw in resp.iter_lines():
-            if raw.startswith(b"data: "):
-                try:
-                    ev = json.loads(raw[len(b"data: "):].decode())
-                    ts, txt = ev.get("timestamp",""), ev.get("data","")
-                    lines.append(f"[{ts}] {txt}")
-                except:
-                    continue
-    return "\n".join(lines)
-
-def check_iframe(url: str, timeout: int = 5) -> bool:
-    try:
-        return requests.get(url, timeout=timeout).status_code == 200
-    except:
-        return False
+def run_agent(client, model_name, system_prompt, user_input, config, chat_history=None):
+    """Helper to run a single agent interaction."""
+    if chat_history is None:
+        chat_history = []
+
+    messages = chat_history + [{"role": "user", "content": user_input}]
+
+    # Gemini expects alternating user/model roles starting with user,
+    # but system prompts are outside this. We'll prepend system instructions
+    # to the *first* user message if no history exists, or manage roles manually.
+    # A simpler way for few-turn interactions is to put the system prompt
+    # at the start of the *first* user message, or use a dedicated system role
+    # if the model supports it (Gemini sometimes treats the first entry specially).
+    # Let's stick to prepending system info to the first message for robustness.
+    if not chat_history:
+         user_input = system_prompt["content"] + "\n\n" + user_input
+         messages = [{"role": "user", "content": user_input}] # Start fresh with combined prompt
 
-# — CORE LOOP WITH TWO AGENTS —
+    # For subsequent turns, just append the new user input
+    if chat_history:
+         messages = chat_history + [{"role": "user", "content": user_input}]
+
+
+    # Ensure messages strictly alternate user/model roles for Gemini
+    # This is a common pain point. Let's just send the last few relevant messages + system context
+    # A better approach is to pass the *full* history but carefully structure it.
+    # For simplicity here, let's just pass the system prompt and the immediate user input,
+    # letting the agent rely on the *structured* project_state for context, not chat history.
+    # This shifts complexity from prompt history management to state management.
+
+    # Revised approach: Pass relevant project_state info *as* the user input
+    # The agent's core prompt (SYSTEM_*) defines its role and what to do with the input state.
+    full_prompt = f"{system_prompt['content']}\n\nProject State:\n{user_input}"
+
+    try:
+        resp = client.models.generate_content(
+            model=model_name,
+            contents=[{"role": "user", "content": full_prompt}], # Only send the current state input
+            config=config
+        )
+        resp.raise_for_status() # Raise an exception for bad status codes
+        return resp.text.strip(), [{"role": "user", "content": full_prompt}, {"role": "model", "content": resp.text.strip()}] # Return response and the turn
+    except Exception as e:
+        print(f"Agent failed: {e}")
+        return f"ERROR: Agent failed - {e}", None # Indicate failure
+
+def run_planner(client, project_state, config):
+    print("Running Planner Agent...")
+    input_state_str = f"Requirements: {project_state['requirements']}"
+    response_text, _ = run_agent(
+        client=client,
+        model_name="gemini-2.5-flash-preview-04-17", # Or a stronger model if needed for planning
+        system_prompt=SYSTEM_ARCHITECT,
+        user_input=input_state_str,
+        config=config,
+        # chat_history=project_state.get('planner_chat', []) # Optional: keep chat history per agent
+    )
+    # project_state['planner_chat'] = planner_chat # Update chat history
+
+    if response_text.startswith("ERROR:"):
+        return response_text # Propagate error
+
+    project_state['plan'] = response_text
+    print("Planner Output:", response_state['plan'])
+    return "Plan generated."
+
+def run_codegen(client, project_state, config):
+    print(f"Running Code-Gen Agent for task: {project_state['current_task']}...")
+    # Code-Gen needs requirements, plan, existing files, and debug feedback
+    input_state_str = (
+        f"Current Task: {project_state['current_task']}\n\n"
+        f"Requirements:\n{project_state['requirements']}\n\n"
+        f"Plan:\n{project_state['plan']}\n\n"
+        f"Existing Files:\n{json.dumps(project_state['files'], indent=2)}\n\n"
+        f"Debug Feedback:\n{project_state['feedback'] or 'None'}\n"
+    )
+    response_text, _ = run_agent(
+        client=client,
+        model_name="gemini-2.5-flash-preview-04-17", # Or a stronger model
+        system_prompt=SYSTEM_CODEGEN,
+        user_input=input_state_str,
+        config=config,
+        # chat_history=project_state.get('codegen_chat', []) # Optional
+    )
+    # project_state['codegen_chat'] = codegen_chat # Update chat history
+
+    if response_text.startswith("ERROR:"):
+        return response_text # Propagate error
+
+    # Parse the response text to extract code blocks for potentially multiple files
+    files_updated = {}
+    # Simple regex to find blocks like `filename`\n```...\n```
+    # This needs to handle multiple blocks
+    blocks = re.findall(r"(`[^`]*`)\s*```(?:\w*\n)?([\s\S]*?)```", response_text)
+
+    if not blocks:
+         # If no code blocks are found, maybe the agent just gave commentary or failed
+         # Treat this as an error or a failure to produce code
+         print("Code-Gen Agent did not output any code blocks.")
+         return "ERROR: Code-Gen Agent failed to output code blocks."
+
+    for filename_match, code_content in blocks:
+        filename = filename_match.strip('`').strip()
+        if filename: # Ensure filename is not empty
+            files_updated[filename] = code_content.strip()
+
+    if not files_updated:
+         print("Code-Gen Agent outputted blocks but couldn't parse filenames.")
+         return "ERROR: Code-Gen Agent outputted blocks but couldn't parse filenames."
+
+
+    project_state['files'].update(files_updated) # Update existing files or add new ones
+    print(f"Code-Gen Agent updated files: {list(files_updated.keys())}")
+
+    # Quick syntax check on updated files
+    for fn, code in files_updated.items():
+        if fn.endswith('.py'):
+            try:
+                compile(code, fn, "exec")
+                print(f"Syntax check passed for {fn}")
+            except SyntaxError as e:
+                print(f"Syntax Error in {fn}: {e}")
+                return f"ERROR: SyntaxError in {fn}: {e}"
+
+    return "Code generated/updated."
+
+def run_debugger(client, project_state, config):
+    print("Running Debug Agent...")
+    # Debugger needs requirements, plan, files, logs
+    input_state_str = (
+        f"Requirements:\n{project_state['requirements']}\n\n"
+        f"Plan:\n{project_state['plan']}\n\n"
+        f"Current Files:\n{json.dumps(project_state['files'], indent=2)}\n\n"
+        f"Build Logs:\n{project_state['logs'].get('build', 'No build logs.')}\n\n"
+        f"Run Logs:\n{project_state['logs'].get('run', 'No run logs.')}\n"
+        f"Iframe Check: {'Responded OK' if project_state.get('iframe_ok') else 'Not responding or checked yet.'}"
+    )
+    response_text, _ = run_agent(
+        client=client,
+        model_name="gemini-2.5-flash-preview-04-17", # Flash is good for analysis
+        system_prompt=SYSTEM_DEBUG,
+        user_input=input_state_str,
+        config=config,
+        # chat_history=project_state.get('debug_chat', []) # Optional
+    )
+    # project_state['debug_chat'] = debug_chat # Update chat history
+
+    if response_text.startswith("ERROR:"):
+        return response_text # Propagate error
+
+    project_state['feedback'] = response_text
+    print("Debug Agent Feedback:", project_state['feedback'])
+    return "Feedback generated."
+
+def run_orchestrator(client, project_state, config):
+    print("Running Orchestrator Agent...")
+    # The Orchestrator's prompt is implicit in the state it receives and the function logic
+    # We *could* make the Orchestrator an LLM call too, but a state machine is often more reliable
+    # for controlling flow. Let's use Python code as the Orchestrator for now.
+
+    status_message = "Orchestrator deciding next step..."
+
+    # Based on current status and feedback, decide next task
+    if project_state['status'] == 'START':
+        project_state['current_task'] = 'PLANNING'
+        project_state['status'] = 'In Progress'
+        status_message = "Starting project: Planning phase."
+
+    elif project_state['current_task'] == 'PLANNING':
+        # Assume run_planner was just executed
+        if project_state['plan']:
+             project_state['current_task'] = 'CODING - Initial Implementation'
+             status_message = "Planning complete. Moving to initial coding."
+        else:
+             # Planner failed to produce a plan
+             project_state['status'] = 'Failed'
+             status_message = "Planning failed. Could not generate project plan."
+
+    elif project_state['current_task'].startswith('CODING'):
+        # Assume run_codegen was just executed
+        # Next step is always pushing after coding
+        project_state['current_task'] = 'PUSHING'
+        status_message = "Coding phase complete. Preparing to push changes."
+
+    elif project_state['current_task'] == 'PUSHING':
+         # Assume files were just pushed in the main loop
+         # Next step is fetching logs
+         project_state['current_task'] = 'LOGGING'
+         status_message = "Push complete. Fetching logs."
+
+    elif project_state['current_task'] == 'LOGGING':
+        # Assume logs were fetched in the main loop
+        # Next step is debugging
+        project_state['current_task'] = 'DEBUGGING'
+        status_message = "Logs fetched. Analyzing with Debug Agent."
+
+    elif project_state['current_task'] == 'DEBUGGING':
+        # Assume run_debugger was just executed
+        feedback = project_state['feedback']
+        build_logs = project_state['logs'].get('build', '')
+        run_logs = project_state['logs'].get('run', '')
+        iframe_ok = project_state.get('iframe_ok', False)
+
+        # Analyze feedback and logs to decide next step
+        if "All clear. Project appears complete." in feedback or \
+           ("ERROR" not in build_logs.upper() and "ERROR" not in run_logs.upper() and iframe_ok):
+            project_state['status'] = 'Complete'
+            project_state['current_task'] = 'FINISHED'
+            status_message = "Debug Agent reports clear. Project appears complete."
+        elif project_state['attempt_count'] >= 6: # Max attempts reached
+            project_state['status'] = 'Failed'
+            project_state['current_task'] = 'FINISHED'
+            status_message = f"Max attempts ({project_state['attempt_count']}/6) reached. Project failed."
+        else:
+            # Errors or issues found, need more coding/debugging
+            project_state['current_task'] = 'CODING - Addressing Feedback'
+            status_message = "Debug Agent found issues. Returning to Coding phase to address feedback."
+
+    else:
+         # Should not happen, maybe a cleanup or failure state
+         project_state['status'] = 'Failed'
+         status_message = f"Orchestrator in unknown state: {project_state['current_task']}"
+
+
+    project_state['status_message'] = status_message
+    print(f"Orchestrator Status: {project_state['status']} - Task: {project_state['current_task']}")
+
+
+# --- MAIN HANDLER REVISION ---
 
 def handle_user_message(
-    history,
+    history, # Keep history for chatbot display, but agents use project_state
     sdk_choice: str,
     gemini_api_key: str,
     grounding_enabled: bool,
@@ -95,148 +298,227 @@ def handle_user_message(
     oauth_token: gr.OAuthToken | None
 ):
     if not profile or not oauth_token:
-        return history + [{"role":"assistant","content":"⚠️ Please log in first."}], "", "", "<p>No Space yet.</p>"
+        return history + [{"role":"assistant","content":"⚠️ Please log in first."}], "", "", "<p>No Space yet.</p>", "Waiting for prompt..."
 
-    client     = genai.Client(api_key=gemini_api_key)
-    code_fn    = "app.py" if sdk_choice=="gradio" else "streamlit_app.py"
-    repo_id    = f"{profile.username}/{profile.username}-auto-space"
+    client = genai.Client(api_key=gemini_api_key)
+    repo_id = f"{profile.username}/{profile.username}-auto-space"
     iframe_url = f"https://huggingface.co/spaces/{repo_id}"
-
-    # — SYSTEM PROMPTS —  
-    system_code = {
-        "role":"system",
-        "content":(
-            "You are **Code‑Gen Agent**, a proactive AI developer. Your sole responsibility is to author "
-            f"and correct the entire `{code_fn}` file in a single markdown code block—no extra commentary. "
-            "You have permission to edit files, push updates to the HF Space, and optimize code. "
-            "After each push, await build & run logs before making further changes."
-        )
-    }
-    system_debug = {
-        "role":"system",
-        "content":(
-            "You are **Debug Agent**, a meticulous code reviewer. You can read all files, logs, and the app "
-            "preview, but you **cannot** modify or push code. Your task is to analyze the latest code + logs "
-            "and return concise, actionable feedback or “All clear.”"
-        )
+    sdk_version = get_sdk_version(sdk_choice)
+    code_fn = "app.py" if sdk_choice == "gradio" else "streamlit_app.py" # Still need a primary file name
+
+    # Initialize or load project state
+    # For this implementation, we'll initialize fresh for each new user message
+    # A more complex app might persist state or load from HF Space
+    project_state = {
+        'requirements': history[-1]['content'] if history else "", # Get the latest user message
+        'plan': '',
+        'files': {}, # Use a dict to store multiple file contents {filename: code}
+        'logs': {'build': '', 'run': ''},
+        'feedback': '',
+        'current_task': 'START',
+        'status': 'In Progress',
+        'status_message': 'Initializing...',
+        'attempt_count': 0,
+        'sdk_choice': sdk_choice,
+        'sdk_version': sdk_version,
+        'repo_id': repo_id,
+        'iframe_url': iframe_url,
+        'main_app_file': code_fn, # Keep track of the main app file name convention
+        # Optional: Per-agent chat histories for more complex interactions
+        # 'planner_chat': [],
+        # 'codegen_chat': [],
+        # 'debug_chat': [],
     }
 
-    # initialize each agent’s conversation
-    code_chat  = [system_code] + history[:]
-    debug_chat = [system_debug] + history[:]
-
-    build_logs = run_logs = ""
-    backoff = 1
-
-    for attempt in range(1, 7):
-        tools = [Tool(google_search=GoogleSearch())] if grounding_enabled else []
-        cfg = GenerateContentConfig(
-            tools=tools,
-            response_modalities=["TEXT"],
-            temperature=temperature,
-            max_output_tokens=max_output_tokens,
-        )
+    # Add initial user message to history for chatbot display
+    if history and history[-1].get("role") == "user":
+        # This assumes the Gradio chatbot already added the user message
+        pass # History should already have the user input
+    else:
+        # Handle cases where history might be empty or last turn wasn't user
+        # (e.g. first message, or if previous turn failed)
+        # This depends on how Gradio handles the initial history input
+        # For standard Gradio chatbot, history will be like [{"role": "user", "content": "..."}]
+        pass
+
+
+    cfg = GenerateContentConfig(
+        tools=[Tool(google_search=GoogleSearch())] if grounding_enabled else [],
+        response_modalities=["TEXT"],
+        temperature=temperature,
+        max_output_tokens=max_output_tokens, # Ensure this is high enough for code
+    )
 
-        # --- 1) Code‑Gen generates or updates code ---
-        resp_code = client.models.generate_content(
-            model="gemini-2.5-flash-preview-04-17",
-            contents=[m["content"] for m in code_chat],
-            config=cfg
-        )
-        code = extract_code(resp_code.text)
-        code_chat.append({"role":"assistant","content":code})
-        debug_chat.append({"role":"assistant","content":code})
-
-        # quick syntax check
-        try:
-            compile(code, code_fn, "exec")
-        except SyntaxError as e:
-            code_chat.append({
-                "role":"user",
-                "content": f"SyntaxError caught: {e}. Please correct `{code_fn}` only."
-            })
-            time.sleep(backoff); backoff = min(backoff*2, 30)
-            continue
-
-        # write & push to HF Space
-        sdk_version = get_sdk_version(sdk_choice)
-        files = {
-            code_fn: code,
-            "README.md": f"""---
-title: Wuhp Auto Space
+    # Main Orchestration Loop
+    while project_state['status'] == 'In Progress' and project_state['attempt_count'] < 7:
+        print(f"\n--- Attempt {project_state['attempt_count'] + 1} ---")
+        print(f"Current Task: {project_state['current_task']}")
+        current_task = project_state['current_task']
+        feedback_message = "" # Message to add to chatbot history about step outcome
+
+        if current_task == 'START':
+             run_orchestrator(client, project_state, cfg) # Move to PLANNING
+
+        elif current_task == 'PLANNING':
+             result = run_planner(client, project_state, cfg)
+             feedback_message = f"**Plan:**\n\n{project_state['plan']}"
+             if result.startswith("ERROR:"):
+                 project_state['status'] = 'Failed'
+                 project_state['status_message'] = result
+             run_orchestrator(client, project_state, cfg) # Decide next based on plan result
+
+        elif current_task.startswith('CODING'):
+             # If this is the very first coding task, ensure the main app file exists
+             # even if the planner didn't explicitly list it.
+             if project_state['main_app_file'] not in project_state['files'] and project_state['attempt_count'] == 0:
+                 project_state['files'][project_state['main_app_file']] = f"# Initial {sdk_choice} app file\n" # Start with something
+
+             result = run_codegen(client, project_state, cfg)
+             if result.startswith("ERROR:"):
+                 project_state['status'] = 'Failed'
+                 project_state['status_message'] = result
+             else:
+                 # Show the code that was just generated/updated in the chatbot history
+                 code_output = "\n".join([f"`{fn}`\n```python\n{code}\n```" for fn, code in project_state['files'].items()])
+                 feedback_message = f"**Generated/Updated Code:**\n\n{code_output}"
+
+             run_orchestrator(client, project_state, cfg) # Move to PUSHING or Failed
+
+        elif current_task == 'PUSHING':
+            # Save and push all files in project_state['files']
+            try:
+                # Create/update repo first
+                create_repo(repo_id=repo_id, token=oauth_token.token,
+                            exist_ok=True, repo_type="space", space_sdk=sdk_choice)
+
+                # Write and upload all files
+                for fn, content in project_state['files'].items():
+                     # Write to a temporary file or in the current directory
+                     with open(fn, "w") as f:
+                         f.write(content)
+                     upload_file(
+                         path_or_fileobj=fn, path_in_repo=fn,
+                         repo_id=repo_id, token=oauth_token.token,
+                         repo_type="space"
+                     )
+
+                # Ensure requirements.txt and README are also pushed, or handled by CodeGen
+                # For simplicity, let's add boilerplate here if not generated by CodeGen
+                if 'requirements.txt' not in project_state['files']:
+                     req_content = "pandas\n" + ("streamlit\n" if sdk_choice=="streamlit" else "gradio\n")
+                     with open("requirements.txt", "w") as f:
+                         f.write(req_content)
+                     upload_file(path_or_fileobj="requirements.txt", path_in_repo="requirements.txt",
+                                 repo_id=repo_id, token=oauth_token.token, repo_type="space")
+
+                if 'README.md' not in project_state['files']:
+                    readme_content = f"""---
+title: {profile.username}'s Auto Space
 emoji: 🐢
 sdk: {sdk_choice}
 sdk_version: {sdk_version}
-app_file: {code_fn}
+app_file: {project_state['main_app_file']}
 pinned: false
 ---
-""",
-            "requirements.txt": "pandas\n" + ("streamlit\n" if sdk_choice=="streamlit" else "gradio\n")
-        }
-        for fn, content in files.items():
-            with open(fn, "w") as f:
-                f.write(content)
-
-        create_repo(repo_id=repo_id, token=oauth_token.token,
-                    exist_ok=True, repo_type="space", space_sdk=sdk_choice)
-        for fn in files:
-            upload_file(
-                path_or_fileobj=fn, path_in_repo=fn,
-                repo_id=repo_id, token=oauth_token.token,
-                repo_type="space"
-            )
-
-        # fetch logs
-        build_logs = fetch_logs(repo_id, "build")
-        run_logs   = fetch_logs(repo_id, "run")
-        err_types  = classify_errors(build_logs + "\n" + run_logs)
-
-        # --- 2) Debug‑Agent reviews code & logs ---
-        debug_input = (
-            f"🏷 **Attempt {attempt}**\n"
-            f"Error types: {err_types}\n\n"
-            f"**Build logs:**\n{build_logs}\n\n"
-            f"**Run logs:**\n{run_logs}\n\n"
-            "If there are no errors, reply “All clear.” Otherwise, list your recommended changes."
-        )
-        debug_chat.append({"role":"user","content":debug_input})
-        resp_debug = client.models.generate_content(
-            model="gemini-2.5-flash-preview-04-17",
-            contents=[m["content"] for m in debug_chat],
-            config=cfg
-        )
-        feedback = resp_debug.text.strip()
-        debug_chat.append({"role":"assistant","content":feedback})
-
-        # check for success
-        if "ERROR" not in build_logs.upper() and \
-           "ERROR" not in run_logs.upper() and \
-           check_iframe(iframe_url):
-            break
-
-        # feed debug feedback back to Code‑Gen
-        code_chat.append({
-            "role":"user",
-            "content": f"🔧 Debug feedback:\n{feedback}\nPlease output the full corrected `{code_fn}` code block only."
-        })
-        time.sleep(backoff); backoff = min(backoff*2, 30)
-
-    # prepare UI outputs
-    messages = [
-        {"role": m["role"], "content": m["content"]}
-        for m in code_chat if m["role"] != "system"
-    ]
-    iframe_html = (
+"""
+                    with open("README.md", "w") as f:
+                        f.write(readme_content)
+                    upload_file(path_or_fileobj="README.md", path_in_repo="README.md",
+                                 repo_id=repo_id, token=oauth_token.token, repo_type="space")
+
+
+                print(f"Pushed {len(project_state['files'])} files to {repo_id}")
+                feedback_message = f"Pushed code to HF Space **{repo_id}**."
+                run_orchestrator(client, project_state, cfg) # Move to LOGGING
+
+            except Exception as e:
+                project_state['status'] = 'Failed'
+                project_state['status_message'] = f"ERROR: Failed to push to HF Space: {e}"
+                feedback_message = project_state['status_message']
+                print(feedback_message)
+
+
+        elif current_task == 'LOGGING':
+            # Fetch logs and check iframe status
+            try:
+                project_state['logs']['build'] = fetch_logs(repo_id, "build")
+                project_state['logs']['run'] = fetch_logs(repo_id, "run")
+                project_state['iframe_ok'] = check_iframe(iframe_url)
+                feedback_message = "**Logs Fetched.** Analyzing..."
+                print("Logs fetched.")
+                run_orchestrator(client, project_state, cfg) # Move to DEBUGGING
+            except Exception as e:
+                 project_state['status'] = 'Failed'
+                 project_state['status_message'] = f"ERROR: Failed to fetch logs or check iframe: {e}"
+                 feedback_message = project_state['status_message']
+                 print(feedback_message)
+
+
+        elif current_task == 'DEBUGGING':
+            result = run_debugger(client, project_state, cfg)
+            feedback_message = f"**Debug Feedback:**\n\n{project_state['feedback']}"
+            if result.startswith("ERROR:"):
+                 project_state['status'] = 'Failed'
+                 project_state['status_message'] = result
+            run_orchestrator(client, project_state, cfg) # Move to Decision/CODING/COMPLETE/FAILED
+
+        elif current_task == 'FINISHED':
+            # Exit loop state
+            pass
+
+        else:
+            # Unknown task
+            project_state['status'] = 'Failed'
+            project_state['status_message'] = f"ERROR: Orchestrator entered an unknown task state: {current_task}"
+            feedback_message = project_state['status_message']
+            print(feedback_message)
+
+
+        # Add the outcome of the last step to the chatbot history
+        if feedback_message:
+             history.append({"role": "assistant", "content": feedback_message})
+             # Provide status updates as part of the assistant messages or via the status_message output
+             history.append({"role": "assistant", "content": f"Project Status: {project_state['status_message']}"})
+
+
+        # Increment attempt count after a cycle of Coding/Pushing/Logging/Debugging
+        # Only increment if we went through the core loop steps, not planning or initial start
+        if current_task in ['DEBUGGING', 'PLANNING']: # End of a major cycle
+             project_state['attempt_count'] += 1
+             # Add a small backoff before the next attempt loop starts
+             if project_state['status'] == 'In Progress':
+                 backoff = min(project_state['attempt_count'] * 5, 30) # Linear or exponential backoff
+                 print(f"Waiting for {backoff} seconds before next attempt...")
+                 time.sleep(backoff)
+
+
+    # --- End of Orchestration Loop ---
+
+    # Prepare final UI outputs
+    final_messages = history[:] # Use the accumulated history
+    final_build_logs = project_state['logs'].get('build', '')
+    final_run_logs = project_state['logs'].get('run', '')
+
+    # Final iframe check and HTML
+    final_iframe_html = (
         f'<iframe src="{iframe_url}" width="100%" height="500px"></iframe>'
-        + ("" if check_iframe(iframe_url)
-           else "<p style='color:red;'>⚠️ iframe not responding.</p>")
+        + ("" if project_state.get('iframe_ok') else "<p style='color:red;'>⚠️ iframe not responding or checked.</p>")
     )
-    return messages, build_logs, run_logs, iframe_html
 
-# — SIMPLE UI WITH HIGHER MAX TOKENS —
+    # Add a final status message to the chatbot history
+    final_messages.append({"role": "assistant", "content": f"**Project Outcome:** {project_state['status']} - {project_state['status_message']}"})
+    if project_state['status'] == 'Complete':
+        final_messages.append({"role": "assistant", "content": "✅ Application deployed successfully (likely)! Check the preview above."})
+    else:
+         final_messages.append({"role": "assistant", "content": "❌ Project failed to complete. Review logs and feedback for details."})
+
+    return final_messages, final_build_logs, final_run_logs, final_iframe_html, project_state['status_message']
+
+
+# --- SIMPLE UI WITH HIGHER MAX TOKENS & STATUS DISPLAY —
 
-with gr.Blocks(title="HF Space Auto‑Builder") as demo:
-    gr.Markdown("## 🐢 HF Space Auto‑Builder\n1) Sign in  2) Enter prompt  3) Watch code, logs & preview")
+with gr.Blocks(title="HF Space Auto‑Builder (Team AI)") as demo:
+    gr.Markdown("## 🐢 HF Space Auto‑Builder (Team AI)\n1) Sign in  2) Enter prompt  3) Watch progress, logs & preview")
 
     login_btn = gr.LoginButton(variant="huggingface", size="lg")
     status_md = gr.Markdown("*Not logged in.*")
@@ -246,37 +528,52 @@ with gr.Blocks(title="HF Space Auto‑Builder") as demo:
     login_btn.click(show_profile,     inputs=None, outputs=status_md)
     login_btn.click(list_private_models, inputs=None, outputs=models_md)
 
-    sdk_choice = gr.Radio(["gradio","streamlit"], value="gradio", label="SDK")
-    api_key    = gr.Textbox(label="Gemini API Key", type="password")
-    grounding  = gr.Checkbox(label="Enable grounding", value=False)
-    temp       = gr.Slider(0,1,value=0.2, label="Temperature")
-    max_tokens = gr.Number(value=2048, label="Max output tokens (set up to Gemini’s limit)")
+    with gr.Row():
+        with gr.Column(scale=1):
+            sdk_choice = gr.Radio(["gradio","streamlit"], value="gradio", label="SDK")
+            api_key    = gr.Textbox(label="Gemini API Key", type="password")
+            grounding  = gr.Checkbox(label="Enable grounding", value=False)
+            temp       = gr.Slider(0,1,value=0.2, label="Temperature")
+            max_tokens = gr.Number(value=4096, label="Max output tokens (recommend >= 4096)", minimum=512) # Increased recommendation
 
-    chatbot   = gr.Chatbot(type="messages")
-    user_in   = gr.Textbox(placeholder="Your prompt…", label="Prompt", lines=1)
-    send_btn  = gr.Button("Send")
-    build_box = gr.Textbox(label="Build logs", lines=5, interactive=False)
-    run_box   = gr.Textbox(label="Run logs",   lines=5, interactive=False)
-    preview   = gr.HTML("<p>No Space yet.</p>")
+        with gr.Column(scale=2):
+            project_status_md = gr.Markdown("Waiting for prompt...")
+            chatbot   = gr.Chatbot(type="messages", label="Team Communication & Status")
+            user_in   = gr.Textbox(placeholder="Your prompt for the application...", label="Application Requirements", lines=2)
+            send_btn  = gr.Button("Start Development Team")
 
+    # Separate accordions for logs and preview
+    with gr.Accordion("Logs", open=False):
+        build_box = gr.Textbox(label="Build logs", lines=5, interactive=False)
+        run_box   = gr.Textbox(label="Run logs",   lines=5, interactive=False)
+        refresh_btn = gr.Button("Refresh Logs Only") # Add refresh button for logs
+
+    with gr.Accordion("App Preview", open=True):
+        preview   = gr.HTML("<p>No Space yet.</p>")
+
+
+    # Update the button click handler
+    # It will now return the updated chatbot history, logs, preview, and the project status
     send_btn.click(
         fn=handle_user_message,
-        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens],
-        outputs=[chatbot, build_box, run_box, preview]
+        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, login_btn, login_btn],
+        outputs=[chatbot, build_box, run_box, preview, project_status_md]
     )
     user_in.submit(
-        fn=handle_user_message,
-        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens],
-        outputs=[chatbot, build_box, run_box, preview]
+         fn=handle_user_message,
+        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, login_btn, login_btn],
+        outputs=[chatbot, build_box, run_box, preview, project_status_md]
     )
 
-    refresh_btn = gr.Button("Refresh Logs")
+    # Handler for refreshing logs manually
     refresh_btn.click(
         fn=lambda profile, token: (
             fetch_logs(f"{profile.username}/{profile.username}-auto-space", "build"),
             fetch_logs(f"{profile.username}/{profile.username}-auto-space", "run")
         ),
+        inputs=[login_btn, login_btn], # Pass profile and token from login button state
         outputs=[build_box, run_box]
     )
 
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+
+    demo.launch(server_name="0.0.0.0", server_port=7860)