Update app.py

app.py

@@ -229,54 +229,35 @@ 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."}]
 
 # Helper function to update send button interactivity based on prereqs
-# MODIFIED: Signature takes exactly 4 inputs, removed *args, **kwargs
-# FIXED: Return type hint corrected to gr.update
-# MODIFIED: Internal logic uses bool() check for simplicity
-# REMOVED: problematic debug prints involving *args, **kwargs
+# Now takes explicit arguments, intended to be called with all necessary states.
 def check_send_button_ready(hf_profile: gr.OAuthProfile | None, hf_token: gr.OAuthToken | None, gemini_key: str | None, gemini_model: str | None) -> gr.update:
     """Checks if HF login and Gemini configuration are complete and returns update for button interactivity."""
     # --- START ENHANCED DEBUGGING LOGS ---
     print("\n--- check_send_button_ready START ---")
-    print(f"  Received hf_profile: Type={type(hf_profile)}, Value={'present' if hf_profile else 'None'}")
-    print(f"  Received hf_token: Type={type(hf_token)}, Value={'present' if hf_token else 'None'}")
+    print(f"  Received hf_profile: {hf_profile is not None}")
+    print(f"  Received hf_token: {hf_token is not None}")
     # For api_key, print part of the key if not None/empty for verification
     api_key_display = gemini_key[:5] if isinstance(gemini_key, str) and gemini_key else ('Empty String' if isinstance(gemini_key, str) and gemini_key == "" else 'None')
-    print(f"  Received gemini_key: Type={type(gemini_key)}, Value={api_key_display}")
-    print(f"  Received gemini_model: Type={type(gemini_model)}, Value={gemini_model}")
+    print(f"  Received gemini_key: Value starts with '{api_key_display}'")
+    print(f"  Received gemini_model: {gemini_model}")
     # --- END ENHANCED DEBUGGING LOGS ---
 
     is_logged_in    = hf_profile is not None and hf_token is not None
     # Use bool() check for simplicity - handles None and "" correctly
     is_gemini_ready = bool(gemini_key) and bool(gemini_model)
 
-    # --- CONTINUED DEBUGGING LOGS ---
-    print(f"  HF check: ({hf_profile is not None} and {hf_token is not None}) = {is_logged_in}")
-    print(f"  Gemini check: ({bool(gemini_key)} and {bool(gemini_model)}) = {is_gemini_ready}")
-    print("--- END CONTINUED DEBUGGING LOGS ---")
-
     is_ready = is_logged_in and is_gemini_ready
-    print(f"check_send_button_ready - HF Ready: {is_logged_in}, Gemini Ready: {is_gemini_ready}, Button Ready (boolean): {is_ready}")
+    print(f"check_send_button_ready - HF Ready: {is_logged_in}, Gemini Ready: {is_gemini_ready}, Button Ready: {is_ready}")
     print("--- check_send_button_ready END ---\n")
 
-    # FIXED: Call gr.update instead of gr.Button.update
     return gr.update(interactive=is_ready)
 
-# --- New wrapper function to handle the state updates and call the checker ---
-# This function takes the required state variables as inputs and calls the logic.
-# This function will be called by the .change() events of the state variables.
-def update_send_button_state(profile: gr.OAuthProfile | None, token: gr.OAuthToken | None, gemini_key: str | None, gemini_model: str | None) -> gr.update:
-    """Wrapper function to read states and update button interactivity."""
-    # Pass the received values directly to the core check function
-    return check_send_button_ready(profile, token, gemini_key, gemini_model)
-# --- End of new wrapper function ---
-
 
 # This is the main generator function for the workflow, triggered by the 'Send' button
 # NOTE: This function MUST accept ALL state variables as inputs that it might need to modify or pass through.
 # It MUST also yield/return ALL state variables in the same order they appear in the `outputs` list of the `.click()` event.
-# MODIFIED: Input parameter names align with renamed states for clarity
-# MODIFIED: Output tuple includes renamed gemini state variables
-# REMOVED: *args, **kwargs from signature
+# Added back *args, **kwargs to the generator function signature and yield for robustness,
+# as the previous attempt suggested this might be necessary for state consistency within the generator's lifecycle.
 def ai_workflow_chat(
     message: str,
     history: list[dict],
@@ -297,7 +278,7 @@ def ai_workflow_chat(
     repo_name_state: str | None,
     generated_code_state: str | None,
     use_grounding_state: bool, # Value from use_grounding_checkbox
-    # Absorb potential extra args passed by Gradio event listeners (e.g. old value, event data) - RE-ADDED for safety in generator context
+    # Absorb potential extra args passed by Gradio event listeners
     *args,
     **kwargs
 ) -> tuple[
@@ -331,6 +312,27 @@ def ai_workflow_chat(
     current_gemini_key = gemini_api_key_state
     current_gemini_model = gemini_model_state
 
+    # --- START DEBUGGING ai_workflow_chat inputs ---
+    print("\n--- ai_workflow_chat START (Inputs received) ---")
+    print(f"  message: {message}")
+    print(f"  history len: {len(history)}")
+    print(f"  hf_profile: {hf_profile is not None}")
+    print(f"  hf_token: {hf_token is not None}")
+    api_key_display = current_gemini_key[:5] if isinstance(current_gemini_key, str) and current_gemini_key else ('Empty String' if isinstance(current_gemini_key, str) and current_gemini_key == "" else 'None')
+    print(f"  gemini_api_key_state: Value starts with '{api_key_display}'")
+    print(f"  gemini_model_state: {current_gemini_model}")
+    print(f"  repo_id_state: {repo_id_state}") # Check value here
+    print(f"  workflow_state: {workflow_state}")
+    print(f"  space_sdk: {space_sdk}")
+    print(f"  use_grounding_state: {use_grounding_state}")
+    print(f"  debug_attempts_state: {debug_attempts_state}")
+    print(f"  app_description_state: {app_description_state}")
+    print(f"  repo_name_state: {repo_name_state}")
+    print(f"  generated_code_state: {'Present' if generated_code_state is not None else 'None'}")
+    # print(f"  *args: {args}") # Optional: print extra args if needed
+    # print(f"  **kwargs: {kwargs}") # Optional: print extra kwargs if needed
+    print("--- END DEBUGGING ai_workflow_chat inputs ---\n")
+
 
     # 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
@@ -349,6 +351,7 @@ def ai_workflow_chat(
     # This provides immediate feedback to the user while the AI processes
     # Ensure all state variables and UI outputs are yielded back in the correct order
     # Include gemini_api_key_state and gemini_model_state in the yield tuple
+    # The yielded tuple must match the send_btn.click outputs list exactly.
     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
            attempts, app_desc, repo_name, generated_code, use_grounding,
            current_gemini_key, current_gemini_model) # Explicitly pass back current state values
@@ -358,20 +361,21 @@ def ai_workflow_chat(
 
         # Although button interactivity prevents reaching here without key/model,
         # the checks remain as a safeguard for the workflow logic itself.
+        # Use the local variables derived from state inputs (current_gemini_key, current_gemini_model)
         if not (hf_profile and hf_token):
             history = add_bot_message(history, "Workflow paused: Please log in to Hugging Face first.")
+            # Re-yield state to update chat and keep current state values
             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                    attempts, app_desc, repo_name, generated_code, use_grounding,
-                   current_gemini_key, current_gemini_model) # Explicitly pass back current state values
+                   current_gemini_key, current_gemini_model)
             return # Stop workflow execution for this click
 
-        # Check if API key is non-empty string and model is set
-        # Use the local variables derived from state inputs
         if not (isinstance(current_gemini_key, str) and current_gemini_key != "" and current_gemini_model):
              history = add_bot_message(history, "Workflow cannot start: Please ensure your Gemini API key is entered and a model is selected.")
+             # Re-yield state to update chat and keep current state values
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                     attempts, app_desc, repo_name, generated_code, use_grounding,
-                    current_gemini_key, current_gemini_model) # Explicitly pass back current state values
+                    current_gemini_key, current_gemini_model) # Correctly yielding back the received (likely empty/None) Gemini states
              return # Stop workflow execution for this click
 
 
@@ -384,15 +388,10 @@ def ai_workflow_chat(
             create_match = re.search(r'create (?:a|an)? space called (\w+)', message, re.I)
 
             if reset_match:
-                # 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
-                # gemini_key and gemini_model are passed back with their *current* (likely valid) state
+                # Reset relevant states and UI outputs, passing through current Gemini state
                 yield (history, None, STATE_IDLE, "<p>No Space created yet.</p>", "", "", 0,
-                       None, None, None, False, # Reset use_grounding to default False, other states to None/default
-                       current_gemini_key, current_gemini_model) # Explicitly pass back current state values
-                # No return needed after yield in this generator pattern; execution for this click ends here.
+                       None, None, None, False, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
             elif generate_match:
                 # User requested generation with description and name
@@ -403,11 +402,10 @@ def ai_workflow_chat(
                 state = STATE_CREATING_SPACE
                 repo_name = new_repo_name
                 app_desc = new_app_desc
-                # Yield updated history and state variables (pass UI outputs through)
+                # Yield updated state variables, passing others through
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                       current_gemini_key, current_gemini_model)
-                # No return needed
+                       current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
             elif create_match:
                  # User requested simple space creation with a name
@@ -415,30 +413,27 @@ def ai_workflow_chat(
                  history = add_bot_message(history, f"Acknowledged: '{message}'. Starting workflow to create 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 and state variables (pass UI outputs through)
+                 # Yield updated state variables, passing others through
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model)
-                 # No return needed
+                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
             elif "create" in message.lower() and not repo_id:
                 # 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 (pass UI outputs through)
+                # Yield updated state, passing others through
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                       current_gemini_key, current_gemini_model)
-                # No return needed
+                       current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
             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 (pass UI outputs through)
+                # Yield current state, passing others through
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                        attempts, app_desc, repo_name, generated_code, use_grounding,
-                       current_gemini_key, current_gemini_model)
-                # No return needed
+                       current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
         elif state == STATE_AWAITING_REPO_NAME:
@@ -451,54 +446,41 @@ def ai_workflow_chat(
                  # 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,
-                        current_gemini_key, current_gemini_model)
-                 # No return needed
+                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              else:
                  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. UI outputs remain unchanged for now.
-                 # The next click will proceed from the STATE_CREATING_SPACE block.
+                 # Yield updated state variables, passing others through
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model)
-                 # No return needed
+                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
-        # Note: Each 'elif' block below represents a distinct step in the workflow triggered
-        # when the 'state' variable matches its condition on a button click.
 
         elif state == STATE_CREATING_SPACE:
              # 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
+                 # Reset state on error, passing through current Gemini state
                  yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
-                       None, None, None, use_grounding,
-                       current_gemini_key, current_gemini_model) # Pass grounding state through
-                 # No return needed
-
+                       None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
              else:
                  try:
-                     # Perform the action to create the Space on Hugging Face
                      new_repo_id, iframe_html = create_space_action(repo_name, space_sdk, hf_profile, hf_token)
-                     updated_preview = iframe_html # Update the iframe content to show the new space
-                     repo_id = new_repo_id # Store the official repo_id
+                     updated_preview = iframe_html
+                     repo_id = new_repo_id # Update repo_id state variable
                      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, and the new iframe HTML
+                     state = STATE_GENERATING_CODE
+                     # Yield updated state variables, passing others through
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                             attempts, app_desc, repo_name, generated_code, use_grounding,
-                            current_gemini_key, current_gemini_model) # Pass logs and grounding through
-                     # No return needed
-
+                            current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
                  except Exception as e:
                      history = add_bot_message(history, f"❌ Error creating space: {e}. Click 'reset'.")
-                     # Yield error message and reset state on failure
+                     # Reset state on error, passing through current Gemini state
                      yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
-                           None, None, None, use_grounding,
-                           current_gemini_key, current_gemini_model) # Pass logs and grounding through
-                     # No return needed
+                           None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
         elif state == STATE_GENERATING_CODE:
@@ -515,42 +497,33 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                  history = add_bot_message(history, f"🧠 Generating `{prompt_desc}` `{space_sdk}` app (`app.py`) code with Gemini...")
                  if use_grounding:
                      history = add_bot_message(history, "(Using Grounding with Google Search)")
-                 # Yield to show message before the potentially time-consuming API call
+                 # Yield message before API call
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model)
+                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
-                 # Perform the Gemini API call to generate code, optionally using grounding
-                 # Use the current_gemini_key and current_gemini_model derived from state inputs
                  code = call_gemini(prompt, current_gemini_key, current_gemini_model, use_grounding=use_grounding)
                  code = code.strip()
-                 # Clean up common markdown code block formatting if present
-                 if code.startswith("```python"):
-                     code = code[len("```python"):].strip()
-                 if code.startswith("```"): # Handle generic code blocks too
-                      code = code[len("```"):].strip()
-                 if code.endswith("```"):
-                     code = code[:-len("```")].strip()
+                 # Clean up markdown
+                 if code.startswith("```python"): code = code[len("```python"):].strip()
+                 if code.startswith("```"): code = code[len("```"):].strip()
+                 if code.endswith("```"): code = code[:-len("```")].strip()
 
-                 if not code:
-                      raise ValueError("Gemini returned empty code.")
+                 if not code: raise ValueError("Gemini returned empty code.")
 
                  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 (pass UI outputs and other states through)
+                 state = STATE_UPLOADING_APP_PY
+                 generated_code = code
+                 # Yield updated state variables
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model)
-                 # No return needed
+                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              except Exception as e:
                  history = add_bot_message(history, f"❌ Error generating code: {e}. Click 'reset'.")
-                 # Yield error message and reset state on failure
+                 # Reset state on error, passing through current Gemini state
                  yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                       None, None, None, use_grounding,
-                       current_gemini_key, current_gemini_model)
-                 # No return needed
+                       None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
         elif state == STATE_UPLOADING_APP_PY:
@@ -559,44 +532,35 @@ 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,
-                        current_gemini_key, current_gemini_model)
-                  # No return needed
-
+                        None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
              else:
                  history = add_bot_message(history, "☁️ Uploading `app.py`...")
-                 # Yield to show message before the upload action (pass UI outputs and states through)
+                 # Yield message before upload
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                         attempts, app_desc, repo_name, generated_code, use_grounding,
-                        current_gemini_key, current_gemini_model)
-
+                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
                  try:
-                     # Perform the file upload action
                      upload_file_to_space_action(io.StringIO(code_to_upload), "app.py", repo_id, hf_profile, hf_token)
                      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 (pass UI outputs and other states through)
+                     state = STATE_GENERATING_REQUIREMENTS
+                     generated_code = None
+                     # Yield updated state variables
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                             attempts, app_desc, repo_name, generated_code, use_grounding,
-                            current_gemini_key, current_gemini_model)
-                     # No return needed
-
+                            current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
                  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
+                     # Reset state on error, passing through current Gemini state
                      yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-                           None, None, None, use_grounding,
-                           current_gemini_key, current_gemini_model)
-                     # No return needed
+                           None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
         elif state == STATE_GENERATING_REQUIREMENTS:
              history = add_bot_message(history, "📄 Generating `requirements.txt`...")
-             # Yield to show message before generating requirements (pass UI outputs and states through)
+             # Yield message before generating requirements
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
                     attempts, app_desc, repo_name, generated_code, use_grounding,
-                    current_gemini_key, current_gemini_model)
+                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              # Logic to determine required packages based on SDK and keywords in the app description
              reqs_list = ["gradio"] if space_sdk == "gradio" else ["streamlit"]
@@ -639,8 +603,7 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
              # 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,
-                    current_gemini_key, current_gemini_model)
-                     # No return needed
+                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
         elif state == STATE_UPLOADING_REQUIREMENTS:
@@ -649,17 +612,13 @@ 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,
-                        current_gemini_key, current_gemini_model)
-                 # No return needed
-
+                        None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
             else:
                  history = add_bot_message(history, "☁️ Uploading `requirements.txt`...")
                  # 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,
-                        current_gemini_key, current_gemini_model)
-
+                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
                  try:
                      # Perform requirements file upload
                      upload_file_to_space_action(io.StringIO(reqs_content_to_upload), "requirements.txt", repo_id, hf_profile, hf_token)
@@ -669,23 +628,19 @@ Return **only** the python code block for `app.py`. Do not include any extra tex
                      # 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,
-                            current_gemini_key, current_gemini_model)
-                     # No return needed
-
+                            current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
                  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,
-                           current_gemini_key, current_gemini_model)
-                     # No return needed
+                           None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
         elif state == STATE_GENERATING_README:
             history = add_bot_message(history, "📝 Generating `README.md`...")
             # 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,
-                   current_gemini_key, current_gemini_model)
+                   current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
             # Generate simple README content with Space metadata header
             readme_title = repo_name if repo_name else "My Awesome Space"
@@ -714,8 +669,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
             # 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,
-                   current_gemini_key, current_gemini_model)
-            # No return needed
+                   current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
         elif state == STATE_UPLOADING_README:
@@ -724,17 +678,13 @@ 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,
-                        current_gemini_key, current_gemini_model)
-                 # No return needed
-
+                        None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
             else:
                  history = add_bot_message(history, "☁️ Uploading `README.md`...")
                  # 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,
-                        current_gemini_key, current_gemini_model)
-
+                        current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
                  try:
                      # Perform README file upload
                      upload_file_to_space_action(io.StringIO(readme_content_to_upload), "README.md", repo_id, hf_profile, hf_token)
@@ -744,23 +694,19 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                      # 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,
-                            current_gemini_key, current_gemini_model)
-                     # No return needed
-
+                            current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
                  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,
-                           current_gemini_key, current_gemini_model)
-                     # No return needed
+                           None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
         elif state == STATE_CHECKING_LOGS_BUILD:
              history = add_bot_message(history, "🔍 Fetching build logs...")
              # 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,
-                    current_gemini_key, current_gemini_model)
+                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              # Fetch build logs from HF Space
              build_logs_text = get_build_logs_action(repo_id, hf_profile, hf_token)
@@ -773,8 +719,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   # 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,
-                         current_gemini_key, current_gemini_model)
-                  # No return needed
+                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              else:
                   history = add_bot_message(history, "✅ Build logs fetched. Click 'Send' to check container logs.")
@@ -782,8 +727,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   # 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,
-                         current_gemini_key, current_gemini_model)
-                  # No return needed
+                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
         elif state == STATE_CHECKING_LOGS_RUN:
@@ -791,7 +735,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
              # 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,
-                    current_gemini_key, current_gemini_model)
+                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              # Fetch container logs from HF Space
              container_logs_text = get_container_logs_action(repo_id, hf_profile, hf_token)
@@ -805,8 +749,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   # 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,
-                         current_gemini_key, current_gemini_model)
-                  # No return needed
+                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              elif ("error" in updated_run.lower() or "exception" in updated_run.lower()) and attempts >= MAX_DEBUG_ATTEMPTS:
                   # Max debug attempts reached
@@ -815,8 +758,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   # 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,
-                         current_gemini_key, current_gemini_model)
-                  # No return needed
+                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              else:
                   # No significant errors found in logs, assume success
@@ -825,8 +767,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                   # 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,
-                         current_gemini_key, current_gemini_model)
-                  # No return needed
+                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
         elif state == STATE_DEBUGGING_CODE:
@@ -836,7 +777,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
              # 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,
-                    current_gemini_key, current_gemini_model)
+                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              # Construct prompt for Gemini including the container logs
              debug_prompt = f"""
@@ -873,16 +814,14 @@ Return **only** the python code block for app.py. Do not include any extra text,
                   # 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,
-                         current_gemini_key, current_gemini_model)
-                  # No return needed
+                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
              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,
-                         current_gemini_key, current_gemini_model)
-                  # No return needed
+                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
         elif state == STATE_UPLOADING_FIXED_APP_PY:
              # Retrieve the fixed code from the state variable
@@ -890,16 +829,13 @@ 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,
-                        current_gemini_key, current_gemini_model)
-                  # No return needed
-
+                        None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
              else:
                   history = add_bot_message(history, "☁️ Uploading fixed `app.py`...")
                   # 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,
-                         current_gemini_key, current_gemini_model)
+                         current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
                   try:
                       # Perform the upload of the fixed app.py
@@ -910,24 +846,20 @@ Return **only** the python code block for app.py. Do not include any extra text,
                       # 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,
-                             current_gemini_key, current_gemini_model)
-                      # No return needed
+                             current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
                   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,
-                            current_gemini_key, current_gemini_model)
-                      # No return needed
+                            None, None, None, use_grounding, current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
         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,
-                    current_gemini_key, current_gemini_model)
-             # No return needed
+                    current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
     except Exception as e:
@@ -935,10 +867,10 @@ Return **only** the python code block for app.py. Do not include any extra text,
         error_message = f"Workflow step failed unexpectedly ({state}): {e}. Click 'Send' to re-attempt this step or 'reset'."
         history = add_bot_message(history, error_message)
         print(f"Critical Error in state {state}: {e}") # Log the error for debugging purposes
+        # On unexpected error, reset to IDLE, but pass through the current Gemini state
         yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
                None, None, None, use_grounding,
-               current_gemini_key, current_gemini_model) # Include use_grounding and Gemini states
-        # No return needed after yield
+               current_gemini_key, current_gemini_model) # Correctly yielding back the received Gemini states
 
 
 # --- Build the Gradio UI ---
@@ -948,9 +880,9 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     # Define these first as they might be used in default values for components
     hf_profile   = gr.State(None)
     hf_token     = gr.State(None)
-    # RENAMED: Initialize gemini_api_key_state to empty string
+    # Initialize gemini_api_key_state to empty string
     gemini_api_key_state   = gr.State("")  # start with no key
-    # RENAMED: Initialize gemini_model_state
+    # Initialize gemini_model_state
     gemini_model_state     = gr.State("gemini-1.5-flash") # Default selected model
 
     repo_id      = gr.State(None) # Stores the ID of the created Space
@@ -975,7 +907,7 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
 
             gr.Markdown("## Google AI Studio / Gemini")
             # Define gemini_input and gemini_status before they are used in change handlers
-            # CHANGED: Blank out textbox on load and update info text
+            # Blank out textbox on load and update info text
             gemini_input  = gr.Textbox(
                 label="Your Google AI Studio API Key", # Changed label
                 type="password", # Hides input for security
@@ -1021,7 +953,7 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
             chatbot    = gr.Chatbot(type='messages', label="AI Workflow Chat")
             user_input = gr.Textbox(placeholder="Type your message…", interactive=True)
             # Define send_btn here, BEFORE it's used in send_button_update_output
-            send_btn   = gr.Button("Send") # Starts disabled by default (interactive=False)
+            send_btn   = gr.Button("Send", interactive=False) # Start disabled by default
 
 
             # Define iframe, build_txt, run_txt after send_btn
@@ -1033,75 +965,73 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
 
     # --- Define Event Handlers and Chains AFTER all components and required lists are defined ---
 
-    # Define the inputs used for checking prerequisites using the RENAMED states
-    # This list should be defined AFTER the state variables
+    # Define the inputs used for checking prerequisites
+    # This list provides the values to the check_send_button_ready function
     send_button_interactive_binding_inputs = [
         hf_profile,
         hf_token,
-        gemini_api_key_state, # Use new state name
-        gemini_model_state    # Use new state name
+        gemini_api_key_state,
+        gemini_model_state
     ]
     # Define the output for updating the send button interactivity
-    # FIXED: Moved definition to *after* send_btn is defined
     send_button_update_output = [send_btn]
 
 
-    # Trigger update_send_button_state whenever any prerequisite state changes
-    # This uses the new wrapper function which expects exactly 4 inputs.
-    # The inputs list ensures those 4 state values are passed.
+    # Trigger check_send_button_ready whenever any prerequisite state changes
+    # Bind the check function directly. The inputs list passes the required state values.
     hf_profile.change(
-        update_send_button_state, # Call the wrapper function
+        check_send_button_ready, # Call the core check function directly
         inputs=send_button_interactive_binding_inputs, # Pass all 4 prerequisite states
         outputs=send_button_update_output, # Update only the send button
     )
     hf_token.change(
-        update_send_button_state, # Call the wrapper function
+        check_send_button_ready, # Call the core check function directly
         inputs=send_button_interactive_binding_inputs, # Pass all 4 prerequisite states
         outputs=send_button_update_output, # Update only the send button
     )
     # gemini_input.change updates gemini_api_key_state, which triggers gemini_api_key_state.change
     gemini_api_key_state.change(
-        update_send_button_state, # Call the wrapper function
+        check_send_button_ready, # Call the core check function directly
         inputs=send_button_interactive_binding_inputs, # Pass all 4 prerequisite states
         outputs=send_button_update_output, # Update only the send button
     )
     # model_selector.change updates gemini_model_state, which triggers gemini_model_state.change
     gemini_model_state.change(
-        update_send_button_state, # Call the wrapper function
+        check_send_button_ready, # Call the core check function directly
         inputs=send_button_interactive_binding_inputs, # Pass all 4 prerequisite states
         outputs=send_button_update_output, # Update only the send button
     )
 
 
-    # Handle login button click: Update profile/token state -> Their .change handlers trigger update_send_button_state
+    # Handle login button click: Update profile/token state -> Their .change handlers trigger check_send_button_ready
     login_btn.click(
-         lambda x: (x[0], x[1]),
+         lambda x: (x[0], x[1]), # Lambda takes the LoginButton output (profile, token tuple) and returns it
          inputs=[login_btn],
-         outputs=[hf_profile, hf_token] # Updating these states will trigger their .change handlers
+         outputs=[hf_profile, hf_token] # Update state variables
     )
 
-    # Handle Gemini Key Input change: Update key state -> Configure Gemini status -> State .change handler triggers update_send_button_state
+    # Handle Gemini Key Input change: Update key state -> Configure Gemini status -> State .change handler triggers check_send_button_ready
     gemini_input.change(
-        lambda k: k,
+        lambda k: k, # Lambda takes the textbox value and returns it
         inputs=[gemini_input],
-        outputs=[gemini_api_key_state] # Updating this state will trigger its .change handler
+        outputs=[gemini_api_key_state] # Update state variable
     ).then(
-        # Use RENAMED state variables as inputs to configure_gemini
+        # Configure Gemini using the updated state variables
         configure_gemini,
         inputs=[gemini_api_key_state, gemini_model_state],
-        outputs=[gemini_status] # Update Gemini status based on new key
+        outputs=[gemini_status] # Update Gemini status display
     )
 
-    # Handle Gemini Model Selector change: Update model state -> Configure Gemini status -> State .change handler triggers update_send_button_state
+    # Handle Gemini Model Selector change: Update model state -> Configure Gemini status -> State .change handler triggers check_send_button_ready
     model_selector.change(
-        lambda m: m,
+        lambda m: m, # Lambda takes the radio value and returns it
         inputs=[model_selector],
-        outputs=[gemini_model_state] # Updating this state will trigger its .change handler
+        outputs=[gemini_model_state] # Update state variable
     ).then(
-        # Use RENAMED state variables as inputs to configure_gemini
+        # Configure Gemini using the updated state variables
         configure_gemini,
         inputs=[gemini_api_key_state, gemini_model_state],
-        outputs=[gemini_status] # Update Gemini status based on new model
+        outputs=[gemini_status] # Update Gemini status display
     )
 
 
@@ -1119,6 +1049,8 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     workflow.change(lambda s: s, inputs=workflow, outputs=status_text)
 
     # Link Repo ID State variable change to UI status display
+    # This binding seems correct and should update the Textbox with the value of repo_id state.
+    # The API key appearing here means the API key value is somehow ending up *in* the repo_id state.
     repo_id.change(lambda r: r if r else "None", inputs=repo_id, outputs=repo_id_text)
 
 
@@ -1126,31 +1058,25 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     # 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
-    # MODIFIED: Inputs use renamed gemini state variables
-    # MODIFIED: Outputs include renamed gemini state variables
-    # Added back *args, **kwargs to the generator function signature and yield for robustness,
-    # as the previous attempt suggested this might be necessary for state consistency within the generator's lifecycle.
+    # Ensure inputs and outputs match the ai_workflow_chat signature and yield tuple EXACTLY.
     send_btn.click(
         ai_workflow_chat, # The generator function to run
         inputs=[
             user_input, chatbot, # UI component inputs (message, current chat history)
             hf_profile, hf_token, # HF State variables
-            # Pass RENAMED gemini state variables as inputs
-            gemini_api_key_state, gemini_model_state,
+            gemini_api_key_state, gemini_model_state, # Gemini State variables
             repo_id, workflow, sdk_state, # Workflow State variables
-            # UI component inputs whose *current values* are needed by the generator
             iframe, run_txt, build_txt, # UI component inputs (current values)
             debug_attempts, app_description, repo_name_state, generated_code_state, # Other State variables
             use_grounding_state # Grounding state input
         ],
         outputs=[
-            chatbot, # Update Chatbot with new messages
-            repo_id, workflow, # Update workflow State variables
-            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, # Update the grounding state output
-            # Include RENAMED gemini state variables in outputs to ensure consistency
-            gemini_api_key_state, gemini_model_state
+            chatbot, # Updates Chatbot
+            repo_id, workflow, # Updates State variables (repo_id, workflow)
+            iframe, run_txt, build_txt, # Updates UI components (iframe, logs)
+            debug_attempts, app_description, repo_name_state, generated_code_state, # Updates other State variables
+            use_grounding_state, # Updates Grounding state
+            gemini_api_key_state, gemini_model_state # Updates Gemini State variables
         ]
     ).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
@@ -1164,20 +1090,18 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     # This chain runs once when the app loads
     ai_builder_tab.load(
         # Action 1: Show profile (loads cached login if available)
-        # This calls show_profile, which updates the login_status UI.
-        # The LoginButton component itself implicitly handles updating hf_profile/hf_token states on load if cached creds exist.
         show_profile,
         inputs=None,
         outputs=login_status # Updates UI
     ).then(
         # Action 2: Configure Gemini using initial state (gemini_api_key_state is "", model_state is default)
         configure_gemini,
-        inputs=[gemini_api_key_state, gemini_model_state], # Use RENAMED state variables as inputs
+        inputs=[gemini_api_key_state, gemini_model_state],
         outputs=[gemini_status] # Update Gemini status display
     ).then(
         # Action 3: After initial load checks, update the button state based on initial states
-        # Use the new wrapper function here as well.
-        update_send_button_state,
+        # Call check_send_button_ready directly on load with initial state values
+        check_send_button_ready,
         inputs=send_button_interactive_binding_inputs, # Pass all 4 prerequisite states
         outputs=send_button_update_output, # Update the send button
     ).then(