Update app.py

app.py

@@ -7,6 +7,7 @@ import requests
 
 import gradio as gr
 import google.generativeai as genai
+from google.generativeai import types # Import types for configuration and tools
 
 from huggingface_hub import create_repo, list_models, upload_file, constants
 from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status
@@ -53,7 +54,7 @@ def create_space_action(repo_name: str, sdk: str, profile: gr.OAuthProfile, toke
         iframe = f'<iframe src="{url}" width="100%" height="500px"></iframe>'
         return repo_id, iframe
     except Exception as e:
-        raise RuntimeError(f"Failed to create Space {repo_id}: {e}")
+        raise RuntimeError(f"Failed to create Space `{repo_id}`: {e}")
 
 def upload_file_to_space_action(
     file_obj: io.StringIO, # Specify type hint for clarity
@@ -74,7 +75,7 @@ def upload_file_to_space_action(
             repo_type="space"
         )
     except Exception as e:
-        raise RuntimeError(f"Failed to upload `{path_in_repo}` to {repo_id}: {e}")
+        raise RuntimeError(f"Failed to upload `{path_in_repo}` to `{repo_id}`: {e}")
 
 def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
     """Fetches build or run logs for a Space."""
@@ -82,9 +83,9 @@ def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
          return f"Cannot fetch {level} logs: repo_id or token missing."
     jwt_url  = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
     try:
-        r        = get_session().get(jwt_url, headers=build_hf_headers(token=token))
+        r = get_session().get(jwt_url, headers=build_hf_headers(token=token))
         hf_raise_for_status(r) # Raise HTTPError for bad responses (4xx or 5xx)
-        jwt      = r.json()["token"]
+        jwt = r.json()["token"]
         logs_url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
         lines, count = [], 0
         # Using stream=True is good for potentially large logs
@@ -99,8 +100,8 @@ def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
                 payload = raw[len(b"data: "):]
                 try:
                     event = json.loads(payload.decode())
-                    ts    = event.get("timestamp", "")
-                    txt   = event.get("data", "").strip()
+                    ts = event.get("timestamp", "")
+                    txt = event.get("data", "").strip()
                     if txt:
                         lines.append(f"[{ts}] {txt}")
                         count += 1
@@ -110,7 +111,7 @@ def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
         return "\n".join(lines) if lines else f"No {level} logs found."
     except Exception as e:
         # Catching generic exception is acceptable for helper functions
-        return f"Error fetching {level} logs: {e}"
+        return f"Error fetching {level} logs for `{repo_id}`: {e}"
 
 
 def get_build_logs_action(repo_id, profile, token):
@@ -130,33 +131,52 @@ def get_container_logs_action(repo_id, profile, token):
     return _fetch_space_logs_level(repo_id, "run", token.token)
 
 
-# --- Google Gemini integration with model selection ---
+# --- Google Gemini integration with model selection and grounding ---
 
 def configure_gemini(api_key: str | None, model_name: str | None) -> str:
     """Configures the Gemini API and checks if the model is accessible."""
     if not api_key:
-        return "Gemini API key is not set."
+        return "⚠️ Gemini API key is not set."
     if not model_name:
-        return "Please select a Gemini model."
+        return "⚠️ Please select a Gemini model."
     try:
         genai.configure(api_key=api_key)
         # Attempt a simple call to verify credentials and model availability
         # This will raise an exception if the key is invalid or model not found
         genai.GenerativeModel(model_name).generate_content("ping", stream=False)
-        return f"Gemini configured successfully with **{model_name}**."
+        return f"✅ Gemini configured successfully with **{model_name}**."
     except Exception as e:
-        return f"Error configuring Gemini: {e}"
+        return f"❌ Error configuring Gemini: {e}"
 
-def call_gemini(prompt: str, api_key: str, model_name: str) -> str:
-    """Calls the Gemini API with a given prompt."""
+def call_gemini(prompt: str, api_key: str, model_name: str, use_grounding: bool = False) -> str:
+    """Calls the Gemini API with a given prompt, optionally using grounding."""
     if not api_key or not model_name:
         raise ValueError("Gemini API key or model not set.")
     try:
         genai.configure(api_key=api_key)
         model = genai.GenerativeModel(model_name)
+
+        # Define tools for grounding if requested.
+        # Using genai.types.GoogleSearch() is recommended for Gemini 2.0+
+        # and is backwards compatible with 1.5 for retrieval.
+        tools_config = [types.Tool(google_search=types.GoogleSearch())] if use_grounding else None
+
         # Using generate_content and stream=False for simplicity here
-        response = model.generate_content(prompt, stream=False)
+        response = model.generate_content(
+            prompt,
+            stream=False,
+            tools=tools_config # Pass the tools configuration
+        )
+        # Check if response is blocked
+        if response.prompt_feedback and response.prompt_feedback.block_reason:
+             raise RuntimeError(f"Gemini API call blocked: {response.prompt_feedback.block_reason}")
+        if not response.candidates:
+             raise RuntimeError("Gemini API call returned no candidates.")
+
+        # If response.candidates is not empty, get the text
+        # Using response.text is a convenient way to get text from the first candidate part
         return response.text or "" # Return empty string if no text
+
     except Exception as e:
         # Re-raising as RuntimeError for the workflow to catch and manage
         raise RuntimeError(f"Gemini API call failed: {e}")
@@ -205,6 +225,7 @@ def ai_workflow_chat(
     app_description_state: str | None,
     repo_name_state: str | None,
     generated_code_state: str | None,
+    use_grounding_state: bool, # New input for grounding state
     # Absorb potential extra args passed by Gradio event listeners (e.g. old value, event data)
     *args,
     **kwargs
@@ -219,6 +240,7 @@ def ai_workflow_chat(
     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)
 ]:
     """
     Generator function to handle the AI workflow state machine.
@@ -231,6 +253,7 @@ def ai_workflow_chat(
     app_desc = app_description_state
     repo_name = repo_name_state
     generated_code = generated_code_state
+    use_grounding = use_grounding_state # Unpack grounding state
 
     # Keep copies of potentially updated UI elements passed as inputs to update them later
     updated_preview = preview_html
@@ -246,8 +269,9 @@ 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
     yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-           attempts, app_desc, repo_name, generated_code)
+           attempts, app_desc, repo_name, generated_code, use_grounding)
 
     try:
         # --- State Machine Logic based on the current 'state' variable ---
@@ -258,14 +282,14 @@ def ai_workflow_chat(
                 history = add_bot_message(history, "Workflow paused: Please log in to Hugging Face first.")
                 # Yield updated history and current state, then exit for this click
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code)
+                       attempts, app_desc, repo_name, generated_code, use_grounding)
                 return # Exit the generator for this click
 
             if not (gemini_api_key and gemini_model):
                  history = add_bot_message(history, "Workflow paused: Please enter your API key and select a Gemini model.")
                  # Yield updated history and current state, then exit for this click
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code)
+                        attempts, app_desc, repo_name, generated_code, use_grounding)
                  return # Exit the generator for this click
 
             # Look for specific commands in the user's message
@@ -280,7 +304,7 @@ def ai_workflow_chat(
                 history = add_bot_message(history, "Workflow reset.")
                 # Yield updated history and reset state variables to their initial values
                 yield (history, None, STATE_IDLE, "<p>No Space created yet.</p>", "", "", 0,
-                       None, None, None)
+                       None, None, None, use_grounding) # Include use_grounding
                 # No return needed after yield in this generator pattern; execution for this click ends here.
 
             elif generate_match:
@@ -294,7 +318,7 @@ def ai_workflow_chat(
                 app_desc = new_app_desc
                 # Yield updated history and state variables
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code)
+                       attempts, app_desc, repo_name, generated_code, use_grounding)
                 # No return needed
 
             elif create_match:
@@ -306,7 +330,7 @@ def ai_workflow_chat(
                  repo_name = new_repo_name
                  # Yield updated history and state variables
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code)
+                        attempts, app_desc, repo_name, generated_code, use_grounding)
                  # No return needed
 
             elif "create" in message.lower() and not repo_id:
@@ -315,7 +339,7 @@ def ai_workflow_chat(
                 state = STATE_AWAITING_REPO_NAME # Transition to the state where we wait for the name
                 # Yield updated history and state
                 yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code)
+                       attempts, app_desc, repo_name, generated_code, use_grounding)
                 # No return needed
 
             else:
@@ -323,7 +347,7 @@ def ai_workflow_chat(
                 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 (history, repo_id, state, updated_preview, updated_run, updated_build,
-                       attempts, app_desc, repo_name, generated_code)
+                       attempts, app_desc, repo_name, generated_code, use_grounding)
                 # No return needed
 
 
@@ -331,11 +355,12 @@ def ai_workflow_chat(
              # User's message is expected to be the repo name
              new_repo_name = message.strip()
              # Basic validation for Hugging Face repo name format
-             if not new_repo_name or re.search(r'[^a-zA-Z0-9_-]', new_repo_name):
-                 history = add_bot_message(history, "Invalid name. Please provide a single word/slug for the Space name (letters, numbers, underscores, hyphens only).")
+             # 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
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code)
+                        attempts, app_desc, repo_name, generated_code, use_grounding)
                  # No return needed
 
              else:
@@ -345,7 +370,7 @@ def ai_workflow_chat(
                  # Yield updated history, state, and repo name.
                  # 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)
+                        attempts, app_desc, repo_name, generated_code, use_grounding)
                  # No return needed
 
         # Note: Each 'elif' block below represents a distinct step in the workflow triggered
@@ -356,7 +381,7 @@ def ai_workflow_chat(
              if not repo_name:
                  history = add_bot_message(history, "Internal error: Repo name missing for creation. Resetting.")
                  yield (history, None, STATE_IDLE, "<p>Error creating space.</p>", "", "", 0,
-                       None, None, None)
+                       None, None, None, use_grounding) # Include use_grounding
                  # No return needed
 
              else:
@@ -369,38 +394,41 @@ def ai_workflow_chat(
                      state = STATE_GENERATING_CODE # Transition to the next state
                      # Yield updated state variables and history
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code)
+                            attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                      # 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)
+                           None, None, None, use_grounding) # Include use_grounding
                      # No return needed
 
 
         elif state == STATE_GENERATING_CODE:
              # Define the prompt for Gemini based on the app description or a default
-             prompt_desc = app_desc if app_desc else 'a Gradio image-blur test app with upload and slider controls'
+             prompt_desc = app_desc if app_desc else f'a simple {space_sdk} app'
              prompt = f"""
 You are an AI assistant specializing in Hugging Face Spaces using the {space_sdk} SDK.
 Generate a full, single-file Python app based on:
 '{prompt_desc}'
-Return **only** the python code block for app.py. Do not include any extra text, explanations, or markdown outside the code block.
+Ensure the code is runnable as `app.py` in a Hugging Face Space using the `{space_sdk}` SDK. Include necessary imports and setup.
+Return **only** the python code block for `app.py`. Do not include any extra text, explanations, or markdown outside the code block.
 """
              try:
-                 history = add_bot_message(history, "🧠 Generating `app.py` code with Gemini...")
+                 history = add_bot_message(history, f"🧠 Generating `{prompt_desc}` `{space_sdk}` app (`app.py`) code with Gemini...")
                  # 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)
+                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
 
-                 # Perform the Gemini API call to generate code
-                 code = call_gemini(prompt, gemini_api_key, gemini_model)
+                 # Perform the Gemini API call to generate code, optionally using grounding
+                 code = call_gemini(prompt, gemini_api_key, 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()
 
@@ -412,14 +440,14 @@ Return **only** the python code block for app.py. Do not include any extra text,
                  generated_code = code # Store the generated code in the state variable for the next step
                  # Yield updated state variables and history
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code)
+                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include 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)
+                       None, None, None, use_grounding) # Include use_grounding
                  # No return needed
 
 
@@ -429,14 +457,14 @@ Return **only** the python code block for app.py. Do not include any extra text,
              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)
+                        None, None, None, use_grounding) # Include use_grounding
                   # No return needed
 
              else:
                  history = add_bot_message(history, "☁️ Uploading `app.py`...")
                  # Yield to show message before the upload action
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code)
+                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
 
                  try:
                      # Perform the file upload action
@@ -446,14 +474,14 @@ Return **only** the python code block for app.py. Do not include any extra text,
                      generated_code = None # Clear the stored code after use to free memory/state space
                      # Yield updated state variables and history
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code)
+                            attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                      # No return needed
 
                  except Exception as e:
-                     history = add_bot_message(history, f"❌ Error uploading app.py: {e}. Click 'reset'.")
+                     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)
+                           None, None, None, use_grounding) # Include use_grounding
                      # No return needed
 
 
@@ -461,27 +489,18 @@ Return **only** the python code block for app.py. Do not include any extra text,
              history = add_bot_message(history, "📄 Generating `requirements.txt`...")
              # Yield to show message before generating requirements
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code)
+                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include 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"]
-             if app_desc: # Check app_desc for keywords only if it's not None
+             # Add essential libraries regardless of description keywords or grounding
+             reqs_list.extend(["google-generativeai", "huggingface_hub", "Pillow", "numpy", "pandas", "requests"]) # Assume common ones might be needed
+             # Add scikit-image and opencv if description implies image processing
+             if app_desc:
                  app_desc_lower = app_desc.lower()
-                 if "google.generativeai" in app_desc_lower or "gemini" in app_desc_lower or gemini_api_key:
-                      reqs_list.append("google-generativeai")
-                 if "requests" in app_desc_lower:
-                      reqs_list.append("requests")
-                 # Add common libraries if description suggests they might be needed
-                 if "image" in app_desc_lower or "upload" in app_desc_lower or "blur" in app_desc_lower or "vision" in app_desc_lower:
-                     reqs_list.append("Pillow") # Pillow is a common image processing library
-                 if "numpy" in app_desc_lower: reqs_list.append("numpy")
-                 if "pandas" in app_desc_lower: reqs_list.append("pandas")
-                 # Add scikit-image and opencv if image processing is heavily implied
                  if any(lib in app_desc_lower for lib in ["scikit-image", "skimage", "cv2", "opencv"]):
                       reqs_list.extend(["scikit-image", "opencv-python"]) # Note: opencv-python for pip
 
-             # Add essential libraries regardless of description keywords
-             reqs_list.append("huggingface_hub") # Needed for interaction helpers if used in app
 
              # Use dict.fromkeys to get unique items while preserving insertion order (Python 3.7+)
              reqs_list = list(dict.fromkeys(reqs_list))
@@ -493,7 +512,7 @@ Return **only** the python code block for app.py. Do not include any extra text,
              generated_code = reqs_content # Store requirements content
              # Yield updated state variables and history
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code)
+                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
              # No return needed
 
 
@@ -503,14 +522,14 @@ Return **only** the python code block for app.py. Do not include any extra text,
             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)
+                        None, None, None, use_grounding) # Include use_grounding
                  # No return needed
 
             else:
                  history = add_bot_message(history, "☁️ Uploading `requirements.txt`...")
                  # Yield to show message before upload
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code)
+                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
 
                  try:
                      # Perform requirements file upload
@@ -520,21 +539,21 @@ Return **only** the python code block for app.py. Do not include any extra text,
                      generated_code = None # Clear content after use
                      # Yield updated state variables and history
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code)
+                            attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                      # No return needed
 
                  except Exception as e:
-                     history = add_bot_message(history, f"❌ Error uploading requirements.txt: {e}. Click 'reset'.")
+                     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)
+                           None, None, None, use_grounding) # Include use_grounding
                      # No return needed
 
         elif state == STATE_GENERATING_README:
             history = add_bot_message(history, "📝 Generating `README.md`...")
             # Yield message before generating README
             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code)
+                   attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
 
             # Generate simple README content with Space metadata header
             readme_title = repo_name if repo_name else "My Awesome Space"
@@ -562,7 +581,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
             generated_code = readme_content # Store README content
             # Yield updated state variables and history
             yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                   attempts, app_desc, repo_name, generated_code)
+                   attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
             # No return needed
 
 
@@ -572,14 +591,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)
+                        None, None, None, use_grounding) # Include use_grounding
                  # No return needed
 
             else:
                  history = add_bot_message(history, "☁️ Uploading `README.md`...")
                  # Yield message before upload
                  yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                        attempts, app_desc, repo_name, generated_code)
+                        attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
 
                  try:
                      # Perform README file upload
@@ -589,21 +608,21 @@ This Space was automatically generated by an AI workflow using Google Gemini and
                      generated_code = None # Clear content after use
                      # Yield updated state variables and history
                      yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                            attempts, app_desc, repo_name, generated_code)
+                            attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                      # No return needed
 
                  except Exception as e:
-                     history = add_bot_message(history, f"❌ Error uploading README.md: {e}. Click 'reset'.")
+                     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)
+                           None, None, None, use_grounding) # Include 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 (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code)
+                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
 
              # Fetch build logs from HF Space
              build_logs_text = get_build_logs_action(repo_id, hf_profile, hf_token)
@@ -615,7 +634,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)
+                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                   # No return needed
 
              else:
@@ -623,7 +642,7 @@ 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)
+                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                   # No return needed
 
 
@@ -631,7 +650,7 @@ This Space was automatically generated by an AI workflow using Google Gemini and
              history = add_bot_message(history, "🔍 Fetching container logs...")
              # Yield message before fetching logs (includes a delay)
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code)
+                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
 
              # Fetch container logs from HF Space
              container_logs_text = get_container_logs_action(repo_id, hf_profile, hf_token)
@@ -644,7 +663,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)
+                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                   # No return needed
 
              elif ("error" in updated_run.lower() or "exception" in updated_run.lower()) and attempts >= MAX_DEBUG_ATTEMPTS:
@@ -653,7 +672,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)
+                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                   # No return needed
 
              else:
@@ -662,7 +681,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)
+                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                   # No return needed
 
 
@@ -670,26 +689,28 @@ 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...")
              # Yield message before Gemini API call
              yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                    attempts, app_desc, repo_name, generated_code)
+                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
 
              # Construct prompt for Gemini including the container logs
              debug_prompt = f"""
 You are debugging a {space_sdk} Space. The goal is to fix the code in `app.py` based on the container logs provided.
 
 Here are the container logs:
-Use code with caution.
-Python
 {updated_run}
 Generate the *complete, fixed* content for `app.py` based on these logs.
 Return **only** the python code block for app.py. Do not include any extra text, explanations, or markdown outside the code block.
 """
              try:
-                  # Call Gemini to generate the corrected code
-                  fix_code = call_gemini(debug_prompt, gemini_api_key, gemini_model)
+                  # Call Gemini to generate the corrected code, optionally using grounding
+                  # Note: Grounding might be less effective for debugging based *only* on logs,
+                  # but we include the option as requested.
+                  fix_code = call_gemini(debug_prompt, gemini_api_key, gemini_model, use_grounding=use_grounding)
                   fix_code = fix_code.strip()
                   # Clean up potential markdown formatting
                   if fix_code.startswith("```python"):
                       fix_code = fix_code[len("```python"):].strip()
+                  if fix_code.startswith("```"):
+                       fix_code = fix_code[len("```"):].strip()
                   if fix_code.endswith("```"):
                       fix_code = fix_code[:-len("```")].strip()
 
@@ -701,14 +722,14 @@ Return **only** the python code block for app.py. Do not include any extra text,
                   generated_code = fix_code # Store the generated fix code
                   # Yield updated state, code, and variables
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code)
+                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include 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)
+                         None, None, None, use_grounding) # Include use_grounding
                   # No return needed
 
         elif state == STATE_UPLOADING_FIXED_APP_PY:
@@ -717,14 +738,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)
+                        None, None, None, use_grounding) # Include use_grounding
                   # No return needed
 
              else:
                   history = add_bot_message(history, "☁️ Uploading fixed `app.py`...")
                   # Yield message before upload
                   yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                         attempts, app_desc, repo_name, generated_code)
+                         attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
 
                   try:
                       # Perform the upload of the fixed app.py
@@ -734,21 +755,21 @@ Return **only** the python code block for app.py. Do not include any extra text,
                       generated_code = None # Clear code after use
                       # Yield updated state, code, and variables
                       yield (history, repo_id, state, updated_preview, updated_run, updated_build,
-                             attempts, app_desc, repo_name, generated_code)
+                             attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
                       # No return needed
 
                   except Exception as e:
-                      history = add_bot_message(history, f"❌ Error uploading fixed app.py: {e}. Click 'reset'.")
+                      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)
+                            None, None, None, use_grounding) # Include 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)
+                    attempts, app_desc, repo_name, generated_code, use_grounding) # Include use_grounding
              # No return needed
 
 
@@ -759,7 +780,7 @@ Return **only** the python code block for app.py. Do not include any extra text,
         print(f"Critical Error in state {state}: {e}") # Log the error for debugging purposes
         # Yield an error state and reset essential workflow variables on critical failure
         yield (history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0,
-               None, None, None)
+               None, None, None, use_grounding) # Include use_grounding
         # No return needed after yield
 
 
@@ -769,7 +790,8 @@ 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)
     hf_profile   = gr.State(None)
     hf_token     = gr.State(None)
-    gemini_key   = gr.State(None)
+    # FIX: Initialize gemini_key state from env var on load
+    gemini_key   = gr.State(os.environ.get("GOOGLE_API_KEY"))
     gemini_model = gr.State("gemini-1.5-flash") # Default selected model
     repo_id      = gr.State(None) # Stores the ID of the created Space
     workflow     = gr.State(STATE_IDLE) # Stores the current state of the AI workflow
@@ -778,6 +800,8 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     app_description = gr.State(None) # Stores the user's initial description of the desired app
     repo_name_state = gr.State(None) # Stores the chosen repository name for the Space
     generated_code_state = gr.State(None) # Temporary storage for generated file content (app.py, reqs, README)
+    # New State variable for grounding checkbox
+    use_grounding_state = gr.State(False)
 
     with gr.Row():
         # Sidebar column for inputs and status displays
@@ -788,21 +812,22 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
             login_btn    = gr.LoginButton(variant="huggingface")
 
             # Initial load event to check login status (if cached)
+            # This also ensures the gemini_key state is potentially read from env var
             ai_builder_tab.load(show_profile, outputs=login_status)
             # Update status display when login button reports success
             login_btn.click(show_profile, outputs=login_status)
 
-            gr.Markdown("## Google AI Studio API Key")
+            gr.Markdown("## Google AI Studio / Gemini")
             # Textbox for Gemini API key. Read from environment variable if available.
             gemini_input  = gr.Textbox(
                 label="API Key",
                 type="password", # Hides input for security
                 interactive=True,
-                value=os.environ.get("GOOGLE_API_KEY") # Pre-fill if GOOGLE_API_KEY env var is set
+                value=os.environ.get("GOOGLE_API_KEY"), # Pre-fill if GOOGLE_API_KEY env var is set
+                info="Get your key from Google AI Studio"
             )
             gemini_status = gr.Markdown("") # Display Gemini configuration status
 
-            gr.Markdown("## Gemini Model")
             # Radio buttons to select the Gemini model
             model_selector = gr.Radio(
                 choices=[
@@ -816,12 +841,26 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
             )
 
             # Configure Gemini status on initial load (if API key env var is set)
+            # This uses the *initial* value of gemini_key state
             ai_builder_tab.load(
                 configure_gemini,
                 inputs=[gemini_key, gemini_model],
                 outputs=[gemini_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
+            use_grounding_checkbox.change(
+                lambda v: v, inputs=use_grounding_checkbox, outputs=use_grounding_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)
@@ -876,10 +915,10 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                     status_str = "✅ Ready to send commands."
                 else:
                      status_str = " ".join(status_parts)
-                     if not status_str: # Fallback, should not be needed if not is_ready
+                     # Fallback, should not be needed if not is_ready, but good practice
+                     if not status_str:
                           status_str = "Checking prerequisites..."
 
-
                 # gr.update is used to dynamically change a component's properties
                 return gr.update(interactive=is_ready), status_str
 
@@ -958,14 +997,16 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
                     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
-                    debug_attempts, app_description, repo_name_state, generated_code_state # Other State variables
+                    debug_attempts, app_description, repo_name_state, generated_code_state, # Other State variables
+                    use_grounding_state # Add the new grounding state input
                 ],
                 # Outputs are updated by the values yielded from the generator
                 outputs=[
                     chatbot, # Update Chatbot with new messages
                     repo_id, workflow, # Update workflow State variables
                     iframe, run_txt, build_txt, # Update UI outputs
-                    debug_attempts, app_description, repo_name_state, generated_code_state # Update other State variables
+                    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)
                 ]
             ).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
@@ -1006,4 +1047,4 @@ if __name__ == "__main__":
     os.makedirs(os.environ["GRADIO_TEMP_DIR"], exist_ok=True) # Ensure the directory exists
 
     # Launch the Gradio UI
-    ai_builder_tab.launch()
+    ai_builder_tab.launch()