Update app.py

app.py

@@ -219,7 +219,7 @@ def add_bot_message(history: list[dict], bot_message: str) -> list[dict]:
 
 # Helper function to control send button interactivity and prerequisite status text
 # This function is triggered by changes in login status and Gemini configuration
-# Moved OUTSIDE the gr.Blocks context
+# Defined OUTSIDE the gr.Blocks context
 def update_send_button_state(
     profile: gr.OAuthProfile | None,
     token: gr.OAuthToken | None,
@@ -257,7 +257,7 @@ def update_send_button_state(
     return gr.update(interactive=is_ready), status_str
 
 # Add an initial welcome message to the chatbot (defined outside Blocks to be called by load chain)
-# Moved OUTSIDE the gr.Blocks context
+# Defined OUTSIDE the gr.Blocks context
 def greet():
     return [{"role": "assistant", "content": "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."}]
 
@@ -275,29 +275,30 @@ def ai_workflow_chat(
     repo_id_state:  str | None,
     workflow_state: str,
     space_sdk:      str,
-    preview_html:   str, # Passed in to maintain its value in State
-    container_logs: str, # Passed in to maintain its value in State
-    build_logs:     str, # Passed in to maintain its value in State
+    # NOTE: UI component values are passed *by value* to the generator
+    preview_html:   str, # Value from iframe HTML
+    container_logs: str, # Value from run_txt Textbox
+    build_logs:     str, # Value from build_txt Textbox
     debug_attempts_state: int,
     app_description_state: str | None,
     repo_name_state: str | None,
     generated_code_state: str | None,
-    use_grounding_state: bool, # New input for grounding state
+    use_grounding_state: bool, # Value from use_grounding_checkbox
     # Absorb potential extra args passed by Gradio event listeners (e.g. old value, event data)
     *args,
     **kwargs
 ) -> tuple[
-    list[dict],       # 0: Updated chat history
-    str | None,       # 1: Updated repo_id
-    str,              # 2: Updated workflow state
-    str,              # 3: Updated iframe HTML
-    str,              # 4: Updated container logs
-    str,              # 5: Updated build logs
-    int,              # 6: Updated debug attempts count
-    str | None,       # 7: Updated app description
-    str | None,       # 8: Updated repo name
-    str | None,       # 9: Updated generated code (for temporary storage)
-    bool,             # 10: Updated use_grounding_state (must return all state vars passed in)
+    list[dict],       # 0: Updated chat history (for chatbot)
+    str | None,       # 1: Updated repo_id (for repo_id state)
+    str,              # 2: Updated workflow state (for workflow state)
+    str,              # 3: Updated iframe HTML (for iframe UI component)
+    str,              # 4: Updated container logs (for run_txt UI component)
+    str,              # 5: Updated build logs (for build_txt UI component)
+    int,              # 6: Updated debug attempts count (for debug_attempts state)
+    str | None,       # 7: Updated app description (for app_description state)
+    str | None,       # 8: Updated repo name (for repo_name_state state)
+    str | None,       # 9: Updated generated code (for generated_code_state state)
+    bool,             # 10: Updated use_grounding_state (for use_grounding_state state)
 ]:
     """
     Generator function to handle the AI workflow state machine.
@@ -313,6 +314,7 @@ def ai_workflow_chat(
     use_grounding = use_grounding_state # Unpack grounding state
 
     # Keep copies of potentially updated UI elements passed as inputs to update them later
+    # These are the *current values* of the UI components as of the button click
     updated_preview = preview_html
     updated_build = build_logs
     updated_run = container_logs
@@ -326,7 +328,7 @@ def ai_workflow_chat(
 
     # Yield immediately to update the chat UI with the user's message
     # This provides immediate feedback to the user while the AI processes
-    # Ensure all state variables are yielded back in the correct order
+    # Ensure all state variables and UI outputs are yielded back in the correct order
     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
            attempts, app_desc, repo_name, generated_code, use_grounding)
 
@@ -361,6 +363,7 @@ def ai_workflow_chat(
                 # Reset the workflow state and associated variables
                 history = add_bot_message(history, "Workflow reset.")
                 # Yield updated history and reset state variables to their initial values
+                # Also reset UI outputs to their initial state
                 yield (history, None, STATE_IDLE, "<p>No Space created yet.</p>", "", "", 0,
                        None, None, None, False) # Reset use_grounding to default False, and other states to None/default
                 # No return needed after yield in this generator pattern; execution for this click ends here.
@@ -374,7 +377,7 @@ def ai_workflow_chat(
                 state = STATE_CREATING_SPACE
                 repo_name = new_repo_name
                 app_desc = new_app_desc
-                # Yield updated history and state variables
+                # Yield updated history and state variables (pass UI outputs through)
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                        attempts, app_desc, repo_name, generated_code, use_grounding)
                 # No return needed
@@ -386,7 +389,7 @@ def ai_workflow_chat(
                  # Update state variables for the next step (creation)
                  state = STATE_CREATING_SPACE
                  repo_name = new_repo_name
-                 # Yield updated history and state variables
+                 # Yield updated history and state variables (pass UI outputs through)
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                         attempts, app_desc, repo_name, generated_code, use_grounding)
                  # No return needed
@@ -395,7 +398,7 @@ def ai_workflow_chat(
                 # User wants to create but didn't specify a name yet
                 history = add_bot_message(history, "Okay, what should the Space be called? (e.g., `my-awesome-app`)")
                 state = STATE_AWAITING_REPO_NAME # Transition to the state where we wait for the name
-                # Yield updated history and state
+                # Yield updated history and state (pass UI outputs through)
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                        attempts, app_desc, repo_name, generated_code, use_grounding)
                 # No return needed
@@ -403,7 +406,7 @@ def ai_workflow_chat(
             else:
                 # Command not recognized in IDLE state
                 history = add_bot_message(history, "Command not recognized. Try 'generate me a gradio app called myapp', or 'reset'.")
-                # Yield updated history and current state
+                # Yield updated history and current state (pass UI outputs through)
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                        attempts, app_desc, repo_name, generated_code, use_grounding)
                 # No return needed
@@ -416,7 +419,7 @@ def ai_workflow_chat(
              # Allow letters, numbers, hyphens, underscores, max 100 chars (HF limit check)
              if not new_repo_name or re.search(r'[^a-zA-Z0-9_-]', new_repo_name) or len(new_repo_name) > 100:
                  history = add_bot_message(history, "Invalid name. Please provide a single word/slug for the Space name (letters, numbers, underscores, hyphens only, max 100 chars).")
-                 # Stay in AWAITING_REPO_NAME state and yield message
+                 # Stay in AWAITING_REPO_NAME state and yield message (pass UI outputs through)
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                         attempts, app_desc, repo_name, generated_code, use_grounding)
                  # No return needed
@@ -425,7 +428,7 @@ def ai_workflow_chat(
                  history = add_bot_message(history, f"Using Space name `{new_repo_name}`. Creating Space `{hf_profile.username}/{new_repo_name}`...")
                  state = STATE_CREATING_SPACE # Transition state to creation
                  repo_name = new_repo_name # Store the validated repo name
-                 # Yield updated history, state, and repo name.
+                 # Yield updated history, state, and repo name. UI outputs remain unchanged for now.
                  # The next click will proceed from the STATE_CREATING_SPACE block.
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                         attempts, app_desc, repo_name, generated_code, use_grounding)
@@ -438,8 +441,9 @@ def ai_workflow_chat(
              # Ensure repo_name is available (it should have been set in the previous step)
              if not repo_name:
                  history = add_bot_message(history, "Internal error: Repo name missing for creation. Resetting.")
+                 # Reset relevant states and UI outputs on critical error
                  yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
-                       None, None, None, use_grounding) # Include use_grounding
+                       None, None, None, use_grounding) # Pass grounding state through
                  # No return needed
 
              else:
@@ -450,16 +454,16 @@ def ai_workflow_chat(
                      repo_id = new_repo_id # Store the official repo_id
                      history = add_bot_message(history, f"✅ Space `{repo_id}` created. Click 'Send' to generate and upload code.")
                      state = STATE_GENERATING_CODE # Transition to the next state
-                     # Yield updated state variables and history
+                     # Yield updated state variables and history, and the new iframe HTML
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                            attempts, app_desc, repo_name, generated_code, use_grounding) # Pass logs and grounding through
                      # No return needed
 
                  except Exception as e:
                      history = add_bot_message(history, f"❌ Error creating space: {e}. Click 'reset'.")
                      # Yield error message and reset state on failure
                      yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
-                           None, None, None, use_grounding) # Include use_grounding
+                           None, None, None, use_grounding) # Pass logs and grounding through
                      # No return needed
 
 
@@ -479,7 +483,7 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                      history = add_bot_message(history, "(Using Grounding with Google Search)")
                  # Yield to show message before the potentially time-consuming API call
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                        attempts, app_desc, repo_name, generated_code, use_grounding)
 
                  # Perform the Gemini API call to generate code, optionally using grounding
                  code = call_gemini(prompt, gemini_api_key, gemini_model, use_grounding=use_grounding)
@@ -498,16 +502,16 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                  history = add_bot_message(history, "✅ `app.py` code generated. Click 'Send' to upload.")
                  state = STATE_UPLOADING_APP_PY # Transition to the upload state
                  generated_code = code # Store the generated code in the state variable for the next step
-                 # Yield updated state variables and history
+                 # Yield updated state variables and history (pass UI outputs and other states through)
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                        attempts, app_desc, repo_name, generated_code, use_grounding)
                  # No return needed
 
              except Exception as e:
                  history = add_bot_message(history, f"❌ Error generating code: {e}. Click 'reset'.")
                  # Yield error message and reset state on failure
                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                       None, None, None, use_grounding) # Include use_grounding
+                       None, None, None, use_grounding)
                  # No return needed
 
 
@@ -517,14 +521,14 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
              if not code_to_upload:
                   history = add_bot_message(history, "Internal error: No code to upload. Resetting.")
                   yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding) # Include use_grounding
+                        None, None, None, use_grounding)
                   # No return needed
 
              else:
                  history = add_bot_message(history, "☁️ Uploading `app.py`...")
-                 # Yield to show message before the upload action
+                 # Yield to show message before the upload action (pass UI outputs and states through)
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                        attempts, app_desc, repo_name, generated_code, use_grounding)
 
                  try:
                      # Perform the file upload action
@@ -532,24 +536,24 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                      history = add_bot_message(history, "✅ Uploaded `app.py`. Click 'Send' to generate requirements.")
                      state = STATE_GENERATING_REQUIREMENTS # Transition state
                      generated_code = None # Clear the stored code after use to free memory/state space
-                     # Yield updated state variables and history
+                     # Yield updated state variables and history (pass UI outputs and other states through)
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                            attempts, app_desc, repo_name, generated_code, use_grounding)
                      # No return needed
 
                  except Exception as e:
                      history = add_bot_message(history, f"❌ Error uploading `app.py`: {e}. Click 'reset'.")
                      # Yield error message and reset state on failure
                      yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding) # Include use_grounding
+                           None, None, None, use_grounding)
                      # No return needed
 
 
         elif state == STATE_GENERATING_REQUIREMENTS:
              history = add_bot_message(history, "📄 Generating `requirements.txt`...")
-             # Yield to show message before generating requirements
+             # Yield to show message before generating requirements (pass UI outputs and states through)
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                    attempts, app_desc, repo_name, generated_code, use_grounding)
 
              # Logic to determine required packages based on SDK and keywords in the app description
              reqs_list = ["gradio"] if space_sdk == "gradio" else ["streamlit"]
@@ -583,15 +587,14 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
              # Sort alphabetically for cleaner requirements.txt
              reqs_list.sort()
 
-
              reqs_content = "\n".join(reqs_list) + "\n"
 
              history = add_bot_message(history, "✅ `requirements.txt` generated. Click 'Send' to upload.")
              state = STATE_UPLOADING_REQUIREMENTS # Transition state
              generated_code = reqs_content # Store requirements content
-             # Yield updated state variables and history
+             # Yield updated state variables and history (pass UI outputs and other states through)
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                    attempts, app_desc, repo_name, generated_code, use_grounding)
                      # No return needed
 
 
@@ -601,14 +604,14 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
             if not reqs_content_to_upload:
                  history = add_bot_message(history, "Internal error: No requirements content to upload. Resetting.")
                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding) # Include use_grounding
+                        None, None, None, use_grounding)
                  # No return needed
 
             else:
                  history = add_bot_message(history, "☁️ Uploading `requirements.txt`...")
-                 # Yield to show message before upload
+                 # Yield message before upload (pass UI outputs and states through)
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                        attempts, app_desc, repo_name, generated_code, use_grounding)
 
                  try:
                      # Perform requirements file upload
@@ -616,23 +619,23 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                      history = add_bot_message(history, "✅ Uploaded `requirements.txt`. Click 'Send' to generate README.")
                      state = STATE_GENERATING_README # Transition state
                      generated_code = None # Clear content after use
-                     # Yield updated state variables and history
+                     # Yield updated state variables and history (pass UI outputs and other states through)
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                            attempts, app_desc, repo_name, generated_code, use_grounding)
                      # No return needed
 
                  except Exception as e:
                      history = add_bot_message(history, f"❌ Error uploading `requirements.txt`: {e}. Click 'reset'.")
                      # Yield error message and reset state on failure
                      yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding) # Include use_grounding
+                           None, None, None, use_grounding)
                      # No return needed
 
         elif state == STATE_GENERATING_README:
             history = add_bot_message(history, "📝 Generating `README.md`...")
-            # Yield message before generating README
+            # Yield message before generating README (pass UI outputs and states through)
             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                   attempts, app_desc, repo_name, generated_code, use_grounding)
 
             # Generate simple README content with Space metadata header
             readme_title = repo_name if repo_name else "My Awesome Space"
@@ -658,9 +661,9 @@ This Space was automatically generated by an AI workflow using Google Gemini and
             history = add_bot_message(history, "✅ `README.md` generated. Click 'Send' to upload.")
             state = STATE_UPLOADING_README # Transition state
             generated_code = readme_content # Store README content
-            # Yield updated state variables and history
+            # Yield updated state variables and history (pass UI outputs and other states through)
             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                   attempts, app_desc, repo_name, generated_code, use_grounding)
             # No return needed
 
 
@@ -670,14 +673,14 @@ This Space was automatically generated by an AI workflow using Google Gemini and
             if not readme_content_to_upload:
                  history = add_bot_message(history, "Internal error: No README content to upload. Resetting.")
                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding) # Include use_grounding
+                        None, None, None, use_grounding)
                  # No return needed
 
             else:
                  history = add_bot_message(history, "☁️ Uploading `README.md`...")
-                 # Yield message before upload
+                 # Yield message before upload (pass UI outputs and states through)
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                        attempts, app_desc, repo_name, generated_code, use_grounding)
 
                  try:
                      # Perform README file upload
@@ -685,23 +688,23 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                      history = add_bot_message(history, "✅ Uploaded `README.md`. All files uploaded. Space is now building. Click 'Send' to check build logs.")
                      state = STATE_CHECKING_LOGS_BUILD # Transition to checking build logs
                      generated_code = None # Clear content after use
-                     # Yield updated state variables and history
+                     # Yield updated state variables and history (pass UI outputs and other states through)
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                            attempts, app_desc, repo_name, generated_code, use_grounding)
                      # No return needed
 
                  except Exception as e:
                      history = add_bot_message(history, f"❌ Error uploading `README.md`: {e}. Click 'reset'.")
                      # Yield error message and reset state on failure
                      yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding) # Include use_grounding
+                           None, None, None, use_grounding)
                      # No return needed
 
         elif state == STATE_CHECKING_LOGS_BUILD:
              history = add_bot_message(history, "🔍 Fetching build logs...")
-             # Yield message before fetching logs (which includes a delay)
+             # Yield message before fetching logs (which includes a delay) (pass UI outputs and states through)
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                    attempts, app_desc, repo_name, generated_code, use_grounding)
 
              # Fetch build logs from HF Space
              build_logs_text = get_build_logs_action(repo_id, hf_profile, hf_token)
@@ -713,7 +716,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   state = STATE_CHECKING_LOGS_RUN # Transition even on build error, to see if container starts
                   # Yield updated state, logs, and variables
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                         attempts, app_desc, repo_name, generated_code, use_grounding)
                   # No return needed
 
              else:
@@ -721,15 +724,15 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   state = STATE_CHECKING_LOGS_RUN # Transition to next log check
                   # Yield updated state, logs, and variables
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                         attempts, app_desc, repo_name, generated_code, use_grounding)
                   # No return needed
 
 
         elif state == STATE_CHECKING_LOGS_RUN:
              history = add_bot_message(history, "🔍 Fetching container logs...")
-             # Yield message before fetching logs (includes a delay)
+             # Yield message before fetching logs (includes a delay) (pass UI outputs and states through)
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                    attempts, app_desc, repo_name, generated_code, use_grounding)
 
              # Fetch container logs from HF Space
              container_logs_text = get_container_logs_action(repo_id, hf_profile, hf_token)
@@ -742,7 +745,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   state = STATE_DEBUGGING_CODE # Transition to the debugging state
                   # Yield updated state, logs, attempts, and variables
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                         attempts, app_desc, repo_name, generated_code, use_grounding)
                   # No return needed
 
              elif ("error" in updated_run.lower() or "exception" in updated_run.lower()) and attempts >= MAX_DEBUG_ATTEMPTS:
@@ -751,7 +754,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   state = STATE_COMPLETE # Workflow ends on failure after attempts
                   # Yield updated state, logs, attempts, and variables
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                         attempts, app_desc, repo_name, generated_code, use_grounding)
                   # No return needed
 
              else:
@@ -760,7 +763,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   state = STATE_COMPLETE # Workflow ends on success
                   # Yield updated state, logs, attempts, and variables
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                         attempts, app_desc, repo_name, generated_code, use_grounding)
                   # No return needed
 
 
@@ -768,9 +771,9 @@ This Space was automatically generated by an AI workflow using Google Gemini and
              history = add_bot_message(history, f"🧠 Calling Gemini to generate fix based on logs...")
              if use_grounding:
                   history = add_bot_message(history, "(Using Grounding with Google Search)")
-             # Yield message before Gemini API call
+             # Yield message before Gemini API call (pass UI outputs and states through)
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                    attempts, app_desc, repo_name, generated_code, use_grounding)
 
              # Construct prompt for Gemini including the container logs
              debug_prompt = f"""
@@ -803,16 +806,16 @@ Return **only** the python code block for app.py. Do not include any extra text,
                   history = add_bot_message(history, "✅ Fix code generated. Click 'Send' to upload.")
                   state = STATE_UPLOADING_FIXED_APP_PY # Transition to the upload state for the fix
                   generated_code = fix_code # Store the generated fix code
-                  # Yield updated state, code, and variables
+                  # Yield updated state, code, and variables (pass UI outputs and states through)
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                         attempts, app_desc, repo_name, generated_code, use_grounding)
                   # No return needed
 
              except Exception as e:
                   history = add_bot_message(history, f"❌ Error generating debug code: {e}. Click 'reset'.")
                   # Yield error message and reset state on failure
                   yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                         None, None, None, use_grounding) # Include use_grounding
+                         None, None, None, use_grounding)
                   # No return needed
 
         elif state == STATE_UPLOADING_FIXED_APP_PY:
@@ -821,14 +824,14 @@ Return **only** the python code block for app.py. Do not include any extra text,
              if not fixed_code_to_upload:
                   history = add_bot_message(history, "Internal error: No fixed code available to upload. Resetting.")
                   yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                        None, None, None, use_grounding) # Include use_grounding
+                        None, None, None, use_grounding)
                   # No return needed
 
              else:
                   history = add_bot_message(history, "☁️ Uploading fixed `app.py`...")
-                  # Yield message before upload
+                  # Yield message before upload (pass UI outputs and states through)
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                         attempts, app_desc, repo_name, generated_code, use_grounding)
 
                   try:
                       # Perform the upload of the fixed app.py
@@ -836,23 +839,23 @@ Return **only** the python code block for app.py. Do not include any extra text,
                       history = add_bot_message(history, "✅ Fixed `app.py` uploaded. Space will rebuild. Click 'Send' to check logs again.")
                       state = STATE_CHECKING_LOGS_RUN # Go back to checking run logs after uploading the fix
                       generated_code = None # Clear code after use
-                      # Yield updated state, code, and variables
+                      # Yield updated state, code, and variables (pass UI outputs and states through)
                       yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                             attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                             attempts, app_desc, repo_name, generated_code, use_grounding)
                       # No return needed
 
                   except Exception as e:
                       history = add_bot_message(history, f"❌ Error uploading fixed `app.py`: {e}. Click 'reset'.")
                       # Yield error message and reset state on failure
                       yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                            None, None, None, use_grounding) # Include use_grounding
+                            None, None, None, use_grounding)
                       # No return needed
 
         elif state == STATE_COMPLETE:
              # If in the complete state, the workflow is finished for this project.
              # Subsequent clicks just add user messages; we simply yield the current state.
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
+                    attempts, app_desc, repo_name, generated_code, use_grounding)
              # No return needed
 
 
@@ -871,6 +874,7 @@ Return **only** the python code block for app.py. Do not include any extra text,
 
 with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     # Gradio State variables - these persist their values across user interactions (clicks)
+    # Define these first as they might be used in default values for components
     hf_profile   = gr.State(None)
     hf_token     = gr.State(None)
     # FIX: Initialize gemini_key state from env var on load
@@ -890,27 +894,13 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
         # Sidebar column for inputs and status displays
         with gr.Column(scale=1, min_width=300):
             gr.Markdown("## Hugging Face Login")
+            # Define login_status before it's used in login_btn.click outputs
             login_status = gr.Markdown("*Not logged in.*")
             # Hugging Face Login Button
             login_btn    = gr.LoginButton(variant="huggingface")
 
-            # Update status display when login button reports success
-            # Chain to also update the Send button/prereq status
-            # This uses the hf_profile and hf_token *state* variables as inputs
-            login_btn.click(
-                 # The LoginButton outputs a tuple (OAuthProfile, OAuthToken) on success
-                 lambda x: (x[0], x[1]),
-                 inputs=[login_btn],
-                 outputs=[hf_profile, hf_token] # Update these State variables
-            ).then( # Chain the next action after state is updated
-                 # Pass all relevant state variables to re-evaluate prerequisites
-                 update_send_button_state,
-                 inputs=[hf_profile, hf_token, gemini_key, gemini_model],
-                 outputs=[send_btn, prereq_status] # Update button interactivity and status text
-            )
-
             gr.Markdown("## Google AI Studio / Gemini")
-            # Textbox for Gemini API key. Read from environment variable if available.
+            # Define gemini_input and gemini_status before they are used in change handlers
             gemini_input  = gr.Textbox(
                 label="API Key",
                 type="password", # Hides input for security
@@ -920,7 +910,7 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
             )
             gemini_status = gr.Markdown("") # Display Gemini configuration status
 
-            # Radio buttons to select the Gemini model
+            # Define model_selector before it's used in its change handler
             model_selector = gr.Radio(
                 choices=[
                     ("Gemini 1.5 Flash", "gemini-1.5-flash"),
@@ -932,8 +922,58 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                 interactive=True
             )
 
-            # Gemini Key Input change: update key state -> configure Gemini -> update send button
-            # This uses gemini_input value as input to update gemini_key state
+            # Define grounding checkbox before its change handler
+            use_grounding_checkbox = gr.Checkbox(
+                label="Enable Grounding with Google Search",
+                value=False, # Default to off
+                interactive=True,
+                info="Use Google Search results to inform Gemini's response (may improve factuality)."
+            )
+
+            gr.Markdown("## Space SDK")
+            # Define sdk_selector before its change handler
+            sdk_selector = gr.Radio(choices=["gradio","streamlit"], value="gradio", label="Template SDK", interactive=True)
+
+            gr.Markdown("## Workflow Status")
+            # Define status_text and repo_id_text before they are used in change handlers
+            status_text = gr.Textbox(label="Current State", value=STATE_IDLE, interactive=False)
+            repo_id_text = gr.Textbox(label="Current Space ID", value="None", interactive=False)
+
+            # Define prereq_status before it's used in handlers that update it
+            prereq_status = gr.Markdown("Checking prerequisites...")
+
+
+        # Main content area column
+        with gr.Column(scale=3):
+            # Define chatbot, user_input, send_btn before send_btn.click
+            chatbot    = gr.Chatbot(type='messages', label="AI Workflow Chat")
+            user_input = gr.Textbox(placeholder="Type your message…", interactive=True)
+            # Define send_btn before its click handler
+            send_btn   = gr.Button("Send", interactive=False) # Starts disabled until prereqs are met
+
+            # Define iframe, build_txt, run_txt before they are used in send_btn.click inputs/outputs
+            # These are UI components, NOT State variables
+            iframe    = gr.HTML("<p>No Space created yet.</p>") # HTML element for the Space iframe
+            build_txt = gr.Textbox(label="Build Logs", lines=10, interactive=False, value="", max_lines=20) # Set max_lines for scrollability
+            run_txt   = gr.Textbox(label="Container Logs", lines=10, interactive=False, value="", max_lines=20) # Set max_lines for scrollability
+
+
+            # --- Define Event Handlers and Chains AFTER all components are defined ---
+
+            # Handle login button click
+            login_btn.click(
+                 # The LoginButton outputs a tuple (OAuthProfile, OAuthToken) on success
+                 lambda x: (x[0], x[1]),
+                 inputs=[login_btn],
+                 outputs=[hf_profile, hf_token] # Update these State variables
+            ).then( # Chain the next action after state is updated
+                 # Pass all relevant state variables to re-evaluate prerequisites
+                 update_send_button_state,
+                 inputs=[hf_profile, hf_token, gemini_key, gemini_model],
+                 outputs=[send_btn, prereq_status] # Update button interactivity and status text
+            )
+
+            # Handle Gemini Key Input change: update key state -> configure Gemini -> update send button
             gemini_input.change(
                 lambda k: k, inputs=[gemini_input], outputs=[gemini_key] # Update gemini_key state
             ).then(
@@ -943,8 +983,7 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                 update_send_button_state, inputs=[hf_profile, hf_token, gemini_key, gemini_model], outputs=[send_btn, prereq_status] # Update button/prereq status
             )
 
-            # Gemini Model Selector change: update model state -> configure Gemini -> update send button
-            # This uses model_selector value as input to update gemini_model state
+            # Handle Gemini Model Selector change: update model state -> configure Gemini -> update send button
             model_selector.change(
                 lambda m: m, inputs=[model_selector], outputs=[gemini_model] # Update gemini_model state
             ).then(
@@ -954,66 +993,47 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                 update_send_button_state, inputs=[hf_profile, hf_token, gemini_key, gemini_model], outputs=[send_btn, prereq_status] # Update button/prereq status
             )
 
-            # New checkbox for optional grounding
-            use_grounding_checkbox = gr.Checkbox(
-                label="Enable Grounding with Google Search",
-                value=False, # Default to off
-                interactive=True,
-                info="Use Google Search results to inform Gemini's response (may improve factuality)."
-            )
-            # Link checkbox change to update the state variable
-            # No need to update send button status here, as grounding is not a core prerequisite
+            # Handle Grounding checkbox change: update grounding state
             use_grounding_checkbox.change(
                 lambda v: v, inputs=use_grounding_checkbox, outputs=use_grounding_state
             )
 
+            # Handle SDK selector change: update sdk state
+            sdk_selector.change(
+                lambda s: s, inputs=sdk_selector, outputs=sdk_state
+            )
 
-            gr.Markdown("## Space SDK")
-            # Radio buttons to select the Space SDK (Gradio or Streamlit)
-            sdk_selector = gr.Radio(choices=["gradio","streamlit"], value="gradio", label="Template SDK", interactive=True)
-            # Update the sdk_state state variable when the selection changes
-            sdk_selector.change(lambda s: s, inputs=sdk_selector, outputs=sdk_state)
-
-            gr.Markdown("## Workflow Status")
-            # Textboxes to display the current workflow state and Space ID
-            status_text = gr.Textbox(label="Current State", value=STATE_IDLE, interactive=False)
-            repo_id_text = gr.Textbox(label="Current Space ID", value="None", interactive=False)
-
-            # --- Prerequisite Status Indicator ---
-            # Markdown to show if prerequisites (HF login, Gemini key) are met
-            prereq_status = gr.Markdown("Checking prerequisites...")
+            # Link Workflow State variable change to UI status display
+            workflow.change(lambda s: s, inputs=workflow, outputs=status_text)
 
+            # Link Repo ID State variable change to UI status display
+            repo_id.change(lambda r: r if r else "None", inputs=repo_id, outputs=repo_id_text)
 
-        # Main content area column
-        with gr.Column(scale=3):
-            # Chatbot to display the conversation and workflow messages
-            chatbot    = gr.Chatbot(type='messages', label="AI Workflow Chat")
-            # Textbox for user input messages
-            user_input = gr.Textbox(placeholder="Type your message…", interactive=True)
-            # Button to send the user message and trigger the workflow step
-            send_btn   = gr.Button("Send", interactive=False) # Starts disabled until prereqs are met
 
             # The main event handler for the Send button
             # This .click() event triggers the ai_workflow_chat generator function
-            # Inputs are read from UI components and State variables
-            # Outputs are updated by the values yielded from the generator
             send_btn.click(
                 ai_workflow_chat, # The generator function to run
+                # Inputs are read from UI components AND State variables
                 inputs=[
-                    user_input, chatbot, # UI inputs (message, current chat history)
+                    user_input, chatbot, # UI component inputs (message, current chat history)
                     hf_profile, hf_token, # HF State variables
                     gemini_key, gemini_model, # Gemini State variables
                     repo_id, workflow, sdk_state, # Workflow State variables
-                    iframe, run_txt, build_txt, # UI outputs whose current values are needed by the generator
+                    # UI component inputs whose *current values* are needed by the generator
+                    # These are NOT State variables with the same names
+                    iframe, run_txt, build_txt,
                     debug_attempts, app_description, repo_name_state, generated_code_state, # Other State variables
                     use_grounding_state # Add the new grounding state input
                 ],
+                # Outputs update UI components AND State variables.
+                # The order MUST match the tuple yielded by the generator function.
                 outputs=[
                     chatbot, # Update Chatbot with new messages
                     repo_id, workflow, # Update workflow State variables
-                    iframe, run_txt, build_txt, # Update UI outputs
+                    iframe, run_txt, build_txt, # Update UI component outputs
                     debug_attempts, app_description, repo_name_state, generated_code_state, # Update other State variables
-                    use_grounding_state # Add the new grounding state output (generators must yield/return all state they modify/pass through)
+                    use_grounding_state # Update the grounding state output (generators must yield/return all state they modify/pass through)
                 ]
             ).success( # Chain a .success() event to run *after* the .click() handler completes without error
                  # Clear the user input textbox after the message is sent and processed
@@ -1022,18 +1042,13 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                  outputs=user_input # Update the user input textbox
             )
 
-            # Link State variables' changes to UI status displays (reactive updates)
-            # When the 'workflow' state variable changes, update the text in status_text
-            workflow.change(lambda s: s, inputs=workflow, outputs=status_text)
-            # When the 'repo_id' state variable changes, update the text in repo_id_text
-            repo_id.change(lambda r: r if r else "None", inputs=repo_id, outputs=repo_id_text)
 
-    # --- Initial Load Event Chain (Defined INSIDE gr.Blocks) ---
+    # --- Initial Load Event Chain (Defined INSIDE gr.Blocks, AFTER components) ---
     # This chain runs once when the app loads
     ai_builder_tab.load(
-        # Action 1: Show profile (loads cached login if available)
+        # Action 1: Show profile (loads cached login if available), does NOT need inputs
         show_profile,
-        inputs=None, # This trigger doesn't need inputs from components
+        inputs=None,
         outputs=login_status # Update login status markdown
     ).then(
         # Action 2: Configure Gemini using the initial state values (from env var if set)