Update app.py

app.py

@@ -112,7 +112,7 @@ def get_container_logs_action(repo_id, profile, token):
     if not (repo_id and profile and token):
         return "⚠️ Please log in and create a Space first."
     # Add a short delay before fetching run logs, build might just finish
-    time.sleep(5)
+    time.sleep(5) # Added delay for robustness
     return _fetch_space_logs_level(repo_id, "run", token.token)
 
 
@@ -159,18 +159,32 @@ STATE_CHECKING_LOGS_BUILD = "checking_logs_build"
 STATE_CHECKING_LOGS_RUN = "checking_logs_run"
 STATE_DEBUGGING_CODE = "debugging_code"
 STATE_UPLOADING_FIXED_APP_PY = "uploading_fixed_app_py"
-STATE_COMPLETE = "complete"
+STATE_COMPLETE = "complete" # Added a final state
 
 MAX_DEBUG_ATTEMPTS = 3
 
-def update_chat(history, bot_message):
-    """Helper to add a bot message and yield state."""
-    # Ensure last user message is in history
-    if history and history[-1][1] is None:
-         history[-1][1] = "" # Add empty bot response to last user message if none exists
+def update_chat(history: list[list[str | None]], bot_message: str) -> list[list[str | None]]:
+    """Helper to set the bot's response for the last user message."""
+    # Assume the last entry was just added with history.append([message, None])
+    # If history is empty, this is an error in logic flow, but add safety.
+    if history:
+        # Ensure the last message is indeed a user message awaiting response
+        if history[-1][1] is None:
+            history[-1][1] = bot_message
+        else:
+             # This case means we're trying to add a bot response but the last message
+             # already has one. This might happen if a step is re-triggered.
+             # Append as a new bot-only message as a fallback.
+             # In a strict state machine, this might indicate a flow error.
+             # But for robustness, let's add it.
+             history.append([None, bot_message])
+    else:
+        # This shouldn't happen - update_chat should always be called after a user message is added.
+        print("Warning: update_chat called with empty history.")
+        history.append([None, bot_message]) # As a fallback, add a bot-only message
+
+    return history # Return the modified history list
 
-    history.append([None, bot_message])
-    return history
 
 def ai_workflow_chat(
     message: str,
@@ -188,7 +202,7 @@ def ai_workflow_chat(
     debug_attempts_state: int,
     app_description_state: str | None, # Persist initial request
     repo_name_state: str | None,       # Persist chosen name
-    generated_code_state: str | None,  # Persist generated code between steps
+    generated_code_state: str | None,  # Temporarily stores generated code or file content
 ) -> tuple[
     list[list[str | None]], # history
     str | None,             # repo_id
@@ -207,15 +221,16 @@ def ai_workflow_chat(
     attempts = debug_attempts_state
     app_desc = app_description_state
     repo_name = repo_name_state
-    generated_code = generated_code_state
+    generated_code = generated_code_state # This slot is reused for different file contents
+
 
     updated_preview = preview_html
     updated_build = build_logs
     updated_run = container_logs
 
-    # Add user message to history for context
+    # Add user message to history with a placeholder for the bot's response
+    # We yield immediately after this to show the user's message
     history.append([message, None])
-    # Yield immediately to show user message
     yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
 
@@ -242,43 +257,50 @@ def ai_workflow_chat(
                 history = update_chat(history, "Workflow reset.")
                 # Reset all state variables
                 yield history, None, STATE_IDLE, "<p>No Space created yet.</p>", "", "", 0, None, None, None
-                return # End workflow
+                return # End workflow for this trigger
 
             elif generate_match:
                 new_repo_name = generate_match.group(1)
                 new_app_desc = message # Store the full request
                 history = update_chat(history, f"Acknowledged: '{message}'. Starting workflow to create Space `{hf_profile.username}/{new_repo_name}`.")
                 # Transition to creating space state, passing name and description
-                yield history, repo_id, STATE_CREATING_SPACE, updated_preview, updated_run, updated_build, attempts, new_app_desc, new_repo_name, generated_code
+                state = STATE_CREATING_SPACE
+                repo_name = new_repo_name
+                app_desc = new_app_desc
+                # Yield state change and bot message
+                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+
 
             elif create_match:
                  new_repo_name = create_match.group(1)
                  history = update_chat(history, f"Acknowledged: '{message}'. Starting workflow to create Space `{hf_profile.username}/{new_repo_name}`.")
                  # Transition to creating space state, just passing the name (desc will be default)
-                 yield history, repo_id, STATE_CREATING_SPACE, updated_preview, updated_run, updated_build, attempts, app_desc, new_repo_name, generated_code # Use existing app_desc or None
+                 state = STATE_CREATING_SPACE
+                 repo_name = new_repo_name
+                 # app_desc remains None or existing
+                 # Yield state change and bot message
+                 yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
             elif "create" in message.lower() and not repo_id: # Generic create trigger
                 history = update_chat(history, "Okay, what should the Space be called? (e.g., `my-awesome-app`)")
                 # Transition to awaiting name state
-                yield history, repo_id, STATE_AWAITING_REPO_NAME, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Stay in this state
+                state = STATE_AWAITING_REPO_NAME
+                # Yield state change and bot message
+                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
             else:
                 # Handle other chat messages if needed, or just respond unknown
                 history = update_chat(history, "Command not recognized. Try 'generate me a gradio app called myapp', or 'reset'.")
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
                 # Stay in IDLE state
+                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                return # End workflow for this trigger
 
-        elif state == STATE_AWAITING_REPO_NAME:
-            new_repo_name = message.strip()
-            if not new_repo_name or re.search(r'\s', new_repo_name): # Basic validation for repo name
-                 history = update_chat(history, "Invalid name. Please provide a single word/slug for the Space name.")
-                 yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Stay in this state
-            else:
-                history = update_chat(history, f"Using Space name `{new_repo_name}`. Creating Space `{hf_profile.username}/{new_repo_name}`...")
-                # Transition to creating space state, pass the received name
-                yield history, repo_id, STATE_CREATING_SPACE, updated_preview, updated_run, updated_build, attempts, app_desc, new_repo_name, generated_code
 
-        elif state == STATE_CREATING_SPACE:
+        # --- Workflow steps triggered by state ---
+        # Note: These steps assume the state was transitioned into in the previous yield.
+        # A new click triggers the function again, and the state machine picks up here.
+
+        if state == STATE_CREATING_SPACE:
             # This state is triggered when we *already have* the repo_name in state
             if not repo_name: # Safety check
                 history = update_chat(history, "Internal error: Repo name missing for creation. Resetting.")
@@ -288,15 +310,19 @@ def ai_workflow_chat(
             try:
                 new_repo_id, iframe_html = create_space_action(repo_name, space_sdk, hf_profile, hf_token)
                 updated_preview = iframe_html
-                history = update_chat(history, f"✅ Space `{new_repo_id}` created.")
-                # Transition to generating code state, update repo_id and preview
-                yield history, new_repo_id, STATE_GENERATING_CODE, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                repo_id = new_repo_id # Update repo_id state variable
+                history = update_chat(history, f"✅ Space `{repo_id}` created. Click 'Send' to generate and upload code.")
+                # Transition to generating code state
+                state = STATE_GENERATING_CODE
+                # Yield state change and bot message
+                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
             except Exception as e:
-                history = update_chat(history, f"❌ Error creating space: {e}")
+                history = update_chat(history, f"❌ Error creating space: {e}. Click 'reset'.")
                 # Reset state on failure
                 yield history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0, None, None, None
 
+
         elif state == STATE_GENERATING_CODE:
             # Use the stored app description or a default
             prompt_desc = app_desc if app_desc else 'a Gradio image-blur test app with upload and slider controls'
@@ -308,7 +334,7 @@ Return **only** the python code block for app.py. Do not include any extra text,
 """
             try:
                 history = update_chat(history, "🧠 Generating `app.py` code with Gemini...")
-                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Update UI to show "Generating..."
+                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Yield to show "Generating..." message
 
                 code = call_gemini(prompt, gemini_api_key, gemini_model)
                 # Clean markdown and whitespace
@@ -323,13 +349,17 @@ Return **only** the python code block for app.py. Do not include any extra text,
 
                 history = update_chat(history, "✅ `app.py` code generated. Click 'Send' to upload.")
                 # Transition to uploading state, store the generated code
-                yield history, repo_id, STATE_UPLOADING_APP_PY, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, code # Pass code in state
+                state = STATE_UPLOADING_APP_PY
+                generated_code = code # Store code in state
+                # Yield state change and bot message
+                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
             except Exception as e:
                 history = update_chat(history, f"❌ Error generating code: {e}. Click 'reset'.")
                 # Reset state on failure
                 yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
 
+
         elif state == STATE_UPLOADING_APP_PY:
             # Use the generated_code stored in state
             if not generated_code:
@@ -338,18 +368,22 @@ Return **only** the python code block for app.py. Do not include any extra text,
                  return
 
             history = update_chat(history, "☁️ Uploading `app.py`...")
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, None # Yield to show message, clear generated_code state
+            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Yield to show message
 
             try:
                 upload_file_to_space_action(io.StringIO(generated_code), "app.py", repo_id, hf_profile, hf_token)
                 history = update_chat(history, "✅ Uploaded `app.py`. Click 'Send' to generate requirements.")
-                # Transition to generating requirements
-                yield history, repo_id, STATE_GENERATING_REQUIREMENTS, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, None # Keep generated_code as None
+                # Transition to generating requirements, clear the temporary code storage
+                state = STATE_GENERATING_REQUIREMENTS
+                generated_code = None # Clear temporary storage
+                # Yield state change and bot message
+                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
             except Exception as e:
                 history = update_chat(history, f"❌ Error uploading app.py: {e}. Click 'reset'.")
                 yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
 
+
         elif state == STATE_GENERATING_REQUIREMENTS:
             history = update_chat(history, "📄 Generating `requirements.txt`...")
             yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Yield to show message
@@ -357,21 +391,28 @@ Return **only** the python code block for app.py. Do not include any extra text,
             # Simple heuristic for requirements based on SDK and common needs
             reqs_list = ["gradio"] if space_sdk == "gradio" else ["streamlit"]
             # Add common deps if likely used (could parse code, but simpler heuristic for demo)
-            if "google.generativeai" in (generated_code or "") or gemini_api_key: # Check if Gemini was used for code or if key is set
+            if "google.generativeai" in str(generated_code_state) or gemini_api_key: # Check if Gemini was used for code generation OR if key is set
                  reqs_list.append("google-generativeai")
-            if "requests" in (generated_code or ""):
+            if "requests" in str(generated_code_state):
                  reqs_list.append("requests")
-            reqs_list.append("huggingface_hub") # Needed for log fetching etc if done inside the space itself (though not in this current app's space code)
+            reqs_list.append("huggingface_hub") # Needed for log fetching etc if done inside the space itself (though not currently)
+            # Add Pillow for image processing if it's an image app (common requirement)
+            if "image" in str(app_desc).lower() or "upload" in str(app_desc).lower():
+                reqs_list.append("Pillow")
+
 
             reqs_content = "\n".join(reqs_list) + "\n"
 
             history = update_chat(history, "✅ `requirements.txt` generated. Click 'Send' to upload.")
             # Transition to uploading requirements, store content temporarily
-            yield history, repo_id, STATE_UPLOADING_REQUIREMENTS, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, reqs_content # Pass content in state (abusing generated_code slot)
+            state = STATE_UPLOADING_REQUIREMENTS
+            generated_code = reqs_content # Pass content in state (reusing generated_code slot)
+            # Yield state change and bot message
+            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
 
         elif state == STATE_UPLOADING_REQUIREMENTS:
-             # Use content stored in state (abusing generated_code slot)
+             # Use content stored in state (reusing generated_code slot)
             reqs_content_to_upload = generated_code # Get content from state
             if not reqs_content_to_upload:
                  history = update_chat(history, "Internal error: No requirements content to upload. Resetting.")
@@ -379,13 +420,17 @@ Return **only** the python code block for app.py. Do not include any extra text,
                  return
 
             history = update_chat(history, "☁️ Uploading `requirements.txt`...")
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, None # Yield, clear temp state
+            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Yield, keep temp state
 
             try:
                 upload_file_to_space_action(io.StringIO(reqs_content_to_upload), "requirements.txt", repo_id, hf_profile, hf_token)
                 history = update_chat(history, "✅ Uploaded `requirements.txt`. Click 'Send' to generate README.")
-                # Transition to generating README
-                yield history, repo_id, STATE_GENERATING_README, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, None
+                # Transition to generating README, clear the temporary storage
+                state = STATE_GENERATING_README
+                generated_code = None # Clear temporary storage
+                # Yield state change and bot message
+                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+
 
             except Exception as e:
                 history = update_chat(history, f"❌ Error uploading requirements.txt: {e}. Click 'reset'.")
@@ -397,18 +442,23 @@ Return **only** the python code block for app.py. Do not include any extra text,
 
             # Generate a simple README based on app_desc or repo_name
             readme_title = repo_name if repo_name else "My Awesome Space"
-            readme_description = app_desc if app_desc else "This Hugging Face Space hosts an AI-generated application."
+            # Use app_desc if available, otherwise a generic description
+            readme_description = app_desc if app_desc else f"This Hugging Face Space hosts an AI-generated {space_sdk} application."
 
             readme_content = f"# {readme_title}\n\n{readme_description}\n\n" \
-                             "This Space was automatically generated by an AI workflow.\n"
+                             "This Space was automatically generated by an AI workflow.\n\n" \
+                             f"Built with the {space_sdk} SDK.\n"
 
             history = update_chat(history, "✅ `README.md` generated. Click 'Send' to upload.")
             # Transition to uploading README, store content temporarily
-            yield history, repo_id, STATE_UPLOADING_README, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, readme_content # Pass content in state
+            state = STATE_UPLOADING_README
+            generated_code = readme_content # Pass content in state
+            # Yield state change and bot message
+            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
 
         elif state == STATE_UPLOADING_README:
-            # Use content stored in state (abusing generated_code slot)
+            # Use content stored in state (reusing generated_code slot)
             readme_content_to_upload = generated_code # Get content from state
             if not readme_content_to_upload:
                  history = update_chat(history, "Internal error: No README content to upload. Resetting.")
@@ -416,13 +466,17 @@ Return **only** the python code block for app.py. Do not include any extra text,
                  return
 
             history = update_chat(history, "☁️ Uploading `README.md`...")
-            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, None # Yield, clear temp state
+            yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Yield, keep temp state
 
             try:
                 upload_file_to_space_action(io.StringIO(readme_content_to_upload), "README.md", repo_id, hf_profile, hf_token)
-                history = update_chat(history, "✅ Uploaded `README.md`. Files uploaded. Space is now building. Click 'Send' to check build logs.")
-                # Transition to checking build logs
-                yield history, repo_id, STATE_CHECKING_LOGS_BUILD, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, None
+                history = update_chat(history, "✅ Uploaded `README.md`. All files uploaded. Space is now building. Click 'Send' to check build logs.")
+                # Transition to checking build logs, clear the temporary storage
+                state = STATE_CHECKING_LOGS_BUILD
+                generated_code = None # Clear temporary storage
+                # Yield state change and bot message
+                yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+
 
             except Exception as e:
                 history = update_chat(history, f"❌ Error uploading README.md: {e}. Click 'reset'.")
@@ -440,13 +494,20 @@ Return **only** the python code block for app.py. Do not include any extra text,
              # Simple check: if build logs contain "Error" or "Exception", might indicate build issue.
              # More robust would involve checking build status via API, but logs are simpler for demo.
              # Assuming successful build leads to container logs check.
+             # Check updated_build content for errors
              if "Error" in updated_build or "Exception" in updated_build:
-                  history = update_chat(history, "⚠️ Build logs may contain errors. Please inspect. Click 'Send' to check container logs (app might still start).")
-                  yield history, repo_id, STATE_CHECKING_LOGS_RUN, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Transition to run logs check
+                  history = update_chat(history, "⚠️ Build logs may contain errors. Please inspect above. Click 'Send' to check container logs (app might still start).")
+                  # Transition to run logs check
+                  state = STATE_CHECKING_LOGS_RUN
+                  # Yield state change and message, include updated build logs
+                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
              else:
                   history = update_chat(history, "✅ Build logs fetched. Click 'Send' to check container logs.")
-                  yield history, repo_id, STATE_CHECKING_LOGS_RUN, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Transition to run logs check
+                  # Transition to run logs check
+                  state = STATE_CHECKING_LOGS_RUN
+                  # Yield state change and message, include updated build logs
+                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
 
         elif state == STATE_CHECKING_LOGS_RUN:
@@ -461,17 +522,24 @@ Return **only** the python code block for app.py. Do not include any extra text,
                   attempts += 1
                   history = update_chat(history, f"❌ Errors detected in container logs. Attempting debug fix #{attempts}/{MAX_DEBUG_ATTEMPTS}. Click 'Send' to proceed.")
                   # Transition to debugging state, increment attempts
-                  yield history, repo_id, STATE_DEBUGGING_CODE, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Keep other states same
+                  state = STATE_DEBUGGING_CODE
+                  # Yield state change and message, include updated run logs
+                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
              elif ("Error" in updated_run or "Exception" in updated_run) and attempts >= MAX_DEBUG_ATTEMPTS:
                   history = update_chat(history, f"❌ Errors detected in container logs. Max debug attempts ({MAX_DEBUG_ATTEMPTS}) reached. Please inspect logs manually or click 'reset'.")
-                  # Transition to complete/idle after failed attempts
-                  yield history, repo_id, STATE_COMPLETE, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                  # Transition to complete state after failed attempts
+                  state = STATE_COMPLETE # Indicate workflow finished with errors
+                  # Yield state change and message
+                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
              else:
                   history = update_chat(history, "✅ App appears to be running successfully! Check the iframe above. Click 'reset' to start a new project.")
-                  # Transition to complete/idle on success
-                  yield history, repo_id, STATE_COMPLETE, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+                  # Transition to complete state on success
+                  state = STATE_COMPLETE # Indicate workflow finished successfully
+                  # Yield state change and message
+                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+
 
         elif state == STATE_DEBUGGING_CODE:
              history = update_chat(history, f"🧠 Calling Gemini to generate fix based on logs...")
@@ -501,7 +569,10 @@ Return **only** the python code block for app.py. Do not include any extra text,
 
                   history = update_chat(history, "✅ Fix code generated. Click 'Send' to upload.")
                   # Transition to uploading fixed code, pass the fixed code
-                  yield history, repo_id, STATE_UPLOADING_FIXED_APP_PY, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, fix_code # Pass fix_code in state
+                  state = STATE_UPLOADING_FIXED_APP_PY
+                  generated_code = fix_code # Store fix_code in state
+                  # Yield state change and message
+                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
 
              except Exception as e:
@@ -510,7 +581,7 @@ Return **only** the python code block for app.py. Do not include any extra text,
                   yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
 
         elif state == STATE_UPLOADING_FIXED_APP_PY:
-             # Use the fixed code stored in state (abusing generated_code slot)
+             # Use the fixed code stored in state (reusing generated_code slot)
              fixed_code_to_upload = generated_code # Get code from state
              if not fixed_code_to_upload:
                   history = update_chat(history, "Internal error: No fixed code available to upload. Resetting.")
@@ -518,13 +589,17 @@ Return **only** the python code block for app.py. Do not include any extra text,
                   return
 
              history = update_chat(history, "☁️ Uploading fixed `app.py`...")
-             yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, None # Yield, clear temp state
+             yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code # Yield, keep temp state
 
              try:
                   upload_file_to_space_action(io.StringIO(fixed_code_to_upload), "app.py", repo_id, hf_profile, hf_token)
                   history = update_chat(history, "✅ Fixed `app.py` uploaded. Space will rebuild. Click 'Send' to check logs again.")
-                  # Transition back to checking run logs (after rebuild)
-                  yield history, repo_id, STATE_CHECKING_LOGS_RUN, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, None # Keep generated_code as None
+                  # Transition back to checking run logs (after rebuild), clear temporary storage
+                  state = STATE_CHECKING_LOGS_RUN
+                  generated_code = None # Clear temporary storage
+                  # Yield state change and message
+                  yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+
 
              except Exception as e:
                   history = update_chat(history, f"❌ Error uploading fixed app.py: {e}. Click 'reset'.")
@@ -533,22 +608,33 @@ Return **only** the python code block for app.py. Do not include any extra text,
         elif state == STATE_COMPLETE:
              # App is successfully deployed or failed after attempts.
              # User should click reset or start a new command.
-             # Just yield the current state.
-             yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+             # Just yield the current state to update UI if needed.
+             # The message for STATE_COMPLETE is set in the state it transitions from.
+             pass # No state change needed, just yield current state at the end of try block
+
 
-        else: # Should not happen
-            history = update_chat(history, f"Internal error: Unknown state '{state}'. Resetting.")
-            yield history, None, STATE_IDLE, "<p>Unknown state error.</p>", "", "", 0, None, None, None
+        # If we reached here and the state wasn't handled (e.g., from a previous error state)
+        # or if a yield didn't happen in the previous block (logic error)
+        # The initial yield after adding the user message ensures something is always sent back.
 
 
     except Exception as e:
         # Catch-all for unexpected exceptions in any state
-        history = update_chat(history, f"Workflow step failed unexpectedly: {e}. Click 'Send' to re-attempt this step or 'reset'.")
-        # Stay in the current state to allow re-attempt, or maybe go to idle?
-        # Let's go to idle on unexpected errors to prevent getting stuck.
+        # This might mean the state machine logic itself failed or a core function raised unexpectedly.
+        error_message = f"Workflow step failed unexpectedly ({state}): {e}. Click 'Send' to re-attempt this step or 'reset'."
+        history = update_chat(history, error_message)
+        print(f"Critical Error in state {state}: {e}") # Print to server logs
+        # Transition to idle state on unexpected errors to prevent getting stuck.
         yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
 
 
+    # Final yield after the try block ensures the final state is returned
+    # after a successful step completes or after an error is caught.
+    # This is particularly important for states like STATE_COMPLETE.
+    yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+
+
+
 # --- Build the Gradio UI ---
 
 with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
@@ -563,7 +649,7 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     debug_attempts = gr.State(0)       # Counter for debug attempts
     app_description = gr.State(None)  # Stores the user's original request string
     repo_name_state = gr.State(None)  # Stores the parsed repo name
-    generated_code_state = gr.State(None) # Temporarily stores generated code or file content
+    generated_code_state = gr.State(None) # Temporarily stores generated code or file content (reused)
 
     with gr.Row():
         # Sidebar
@@ -574,10 +660,9 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
 
             # Initial load to check login status
             ai_builder_tab.load(show_profile, outputs=login_status)
-            # Update status on login click
+            # Update status and state on login click
             login_btn.click(show_profile, outputs=login_status)
-            # Store profile and token in state on login click
-            login_btn.click(lambda p, t: (p, t), outputs=[hf_profile, hf_token])
+            login_btn.click(lambda p, t: (p, t), inputs=[login_btn], outputs=[hf_profile, hf_token]) # Use login_btn output directly
 
             gr.Markdown("## Google AI Studio API Key")
             gemini_input  = gr.Textbox(label="API Key", type="password", interactive=True) # Ensure interactive
@@ -591,9 +676,9 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                 choices=[
                     ("Gemini 2.5 Flash Preview 04-17", "gemini-2.5-flash-preview-04-17"),
                     ("Gemini 2.5 Pro Preview 03-25",   "gemini-2.5-pro-preview-03-25"),
-                    ("Gemini 2.0 Flash",               "gemini-2.0-flash"), # Add more recent models if available
-                    ("Gemini 2.0 Flash‑Lite",          "gemini-2.0-flash-lite"),
-                    ("Gemini 1.5 Flash",               "gemini-1.5-flash"),
+                    ("Gemini 1.5 Flash",               "gemini-1.5-flash"), # Keep relevant models
+                    ("Gemini 1.5 Pro",                 "gemini-1.5-pro"),
+                    ("Gemini 1.0 Pro",                 "gemini-1.0-pro"),
                 ],
                 value="gemini-2.5-flash-preview-04-17",
                 label="Select model",
@@ -603,6 +688,7 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
             model_selector.change(lambda m: m, inputs=model_selector, outputs=gemini_model)
 
             # Configure Gemini status on load and when key/model changes
+            # Note: These handlers *update the status text*, they don't block the workflow.
             ai_builder_tab.load(
                 configure_gemini,
                 inputs=[gemini_key, gemini_model],
@@ -630,19 +716,20 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
 
         # Main content
         with gr.Column(scale=3):
-            chatbot    = gr.Chatbot()
+            # Corrected Chatbot initialization
+            chatbot    = gr.Chatbot(type='messages') # Added type='messages'
             user_input = gr.Textbox(placeholder="Type your message…", interactive=True) # Ensure interactive
             send_btn   = gr.Button("Send", interactive=False)
 
             # Logic to enable send button only when logged in and API key is set
-            # This logic is handled reactively by the load and change events below
-            def update_send_button_state(profile, token, key, model):
+            # This function determines the interactive state of the button
+            def update_send_button_state(profile: gr.OAuthProfile | None, token: gr.OAuthToken | None, key: str | None, model: str | None):
                 is_logged_in = profile is not None and token is not None
-                is_gemini_ready = key is not None and model is not None
-                # Simple check - a more robust check would involve calling configure_gemini
-                # but for UI responsiveness, this is okay.
+                is_gemini_ready = key is not None and model is not None # Basic check
+                # Could add a check if configure_gemini returned success last time
                 return gr.update(interactive=is_logged_in and is_gemini_ready)
 
+            # Update button state on load and whenever relevant inputs change
             ai_builder_tab.load(
                 update_send_button_state,
                 inputs=[hf_profile, hf_token, gemini_key, gemini_model],
@@ -664,12 +751,12 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                 outputs=[send_btn]
             )
 
-
             iframe    = gr.HTML("<p>No Space created yet.</p>")
             build_txt = gr.Textbox(label="Build Logs", lines=10, interactive=False, value="") # Initialize empty
             run_txt   = gr.Textbox(label="Container Logs", lines=10, interactive=False, value="") # Initialize empty
 
             # The main event handler for the Send button
+            # It maps inputs/outputs to the ai_workflow_chat generator function
             send_btn.click(
                 ai_workflow_chat,
                 inputs=[
@@ -692,20 +779,30 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                  outputs=user_input
             )
 
-            # Link state variables to UI status displays
+            # Link state variables to UI status displays (reactive updates)
+            # These update the UI components whenever the State variables they listen to change.
             workflow.change(lambda s: s, inputs=workflow, outputs=status_text)
             repo_id.change(lambda r: r if r else "None", inputs=repo_id, outputs=repo_id_text)
-            # Update logs and iframe when their state variables change (e.g., after fetch/create)
-            iframe.change(lambda h: h, inputs=iframe, outputs=iframe) # Not strictly needed, already linked via click outputs
-            build_txt.change(lambda t: t, inputs=build_txt, outputs=build_txt) # Not strictly needed
-            run_txt.change(lambda t: t, inputs=run_txt, outputs=run_txt) # Not strictly needed
-
-
-# Add an initial message to the chatbot on load
-def greet():
-    return [[None, "Welcome! Please log in to Hugging Face and provide your Google AI Studio API key to start building Spaces. Once ready, type 'generate me a gradio app called myapp' or 'create' to begin."]]
-
-ai_builder_tab.load(greet, outputs=chatbot)
+            # The logs and iframe are updated directly by the `send_btn.click` output,
+            # but adding reactive updates from the state variables can sometimes help
+            # ensure consistency if state changes are yielded before the UI components are
+            # explicitly updated in the same yield tuple.
+            # iframe.change(lambda h: h, inputs=iframe, outputs=iframe) # Already linked via click outputs
+            # build_txt.change(lambda t: t, inputs=build_txt, outputs=build_txt) # Already linked
+            # run_txt.change(lambda t: t, inputs=run_txt, outputs=run_txt) # Already linked
+
+
+    # Add an initial message to the chatbot on load
+    # THIS CALL MUST BE INSIDE the with gr.Blocks() block
+    def greet():
+        # Returning a list of lists in the format [user_msg, bot_msg] for Chatbot
+        # An initial message is often just a bot message, format can be [None, bot_msg]
+        # but Chatbot(type='messages') expects {'role': ..., 'content': ...}
+        # For the initial message, the tuple format might still work, or we structure it.
+        # Let's try the tuple format first as it was used before.
+        return [[None, "Welcome! Please log in to Hugging Face and provide your Google AI Studio API key to start building Spaces. Once ready, type 'generate me a gradio app called myapp' or 'create' to begin."]]
+
+    ai_builder_tab.load(greet, outputs=chatbot)
 
 
 if __name__ == "__main__":