Update app.py

app.py

@@ -79,8 +79,7 @@ def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
         jwt      = r.json()["token"]
         logs_url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
         lines, count = [], 0
-        # Increase timeout for potentially slow log fetching
-        with get_session().get(logs_url, headers=build_hf_headers(token=jwt), stream=True, timeout=30) as resp: # Increased timeout
+        with get_session().get(logs_url, headers=build_hf_headers(token=jwt), stream=True, timeout=30) as resp:
             hf_raise_for_status(resp)
             for raw in resp.iter_lines():
                 if count >= 200:
@@ -106,16 +105,13 @@ def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
 def get_build_logs_action(repo_id, profile, token):
     if not (repo_id and profile and token):
         return "⚠️ Cannot fetch build logs: log in and create a Space first."
-    # Add a short delay before fetching build logs after upload
-    # The space needs a moment to detect file changes and start building
-    time.sleep(5) # Added delay
+    time.sleep(5)
     return _fetch_space_logs_level(repo_id, "build", token.token)
 
 def get_container_logs_action(repo_id, profile, token):
     if not (repo_id and profile and token):
         return "⚠️ Cannot fetch container logs: log in and create a Space first."
-    # Add a delay before fetching run logs, build might just finish
-    time.sleep(10) # Increased delay, build can take a while
+    time.sleep(10)
     return _fetch_space_logs_level(repo_id, "run", token.token)
 
 
@@ -165,10 +161,8 @@ STATE_COMPLETE = "complete"
 
 MAX_DEBUG_ATTEMPTS = 3
 
-# Modified to add a new assistant message
 def add_bot_message(history: list[dict], bot_message: str) -> list[dict]:
     """Helper to add a new assistant message to the history."""
-    # Chatbot(type='messages') expects [{'role': ..., 'content': ...}]
     history.append({"role": "assistant", "content": bot_message})
     return history
 
@@ -209,19 +203,19 @@ def ai_workflow_chat(
     repo_name = repo_name_state
     generated_code = generated_code_state
 
+
     updated_preview = preview_html
     updated_build = build_logs
     updated_run = container_logs
 
     # Add user message to history in the new format
     user_message_entry = {"role": "user", "content": message}
-    # Optional: add user name if logged in
     if hf_profile and hf_profile.username:
          user_message_entry["name"] = hf_profile.username
     history.append(user_message_entry)
 
-    # Yield immediately to show user message and indicate thinking (optional - can skip this first yield)
-    # yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
+    # Yield immediately to show user message - this is the first yield
+    yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
 
     try:
@@ -245,7 +239,6 @@ def ai_workflow_chat(
 
             if reset_match:
                 history = add_bot_message(history, "Workflow reset.")
-                # Reset all state variables
                 yield history, None, STATE_IDLE, "<p>No Space created yet.</p>", "", "", 0, None, None, None
                 return
 
@@ -366,21 +359,20 @@ Return **only** the python code block for app.py. Do not include any extra text,
             yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
 
             reqs_list = ["gradio"] if space_sdk == "gradio" else ["streamlit"]
-            if "google.generativeai" in str(app_desc).lower() or "gemini" in str(app_desc).lower() or gemini_api_key: # Check desc too
+            if "google.generativeai" in str(app_desc).lower() or "gemini" in str(app_desc).lower() or gemini_api_key:
                  reqs_list.append("google-generativeai")
-            if "requests" in str(app_desc).lower(): # Check desc too
+            if "requests" in str(app_desc).lower():
                  reqs_list.append("requests")
             reqs_list.append("huggingface_hub")
-            if "image" in str(app_desc).lower() or "upload" in str(app_desc).lower() or "blur" in str(app_desc).lower():
+            if "image" in str(app_desc).lower() or "upload" in str(app_desc).lower() or "blur" in str(app_desc).lower() or "vision" in str(app_desc).lower(): # Added vision
                 reqs_list.append("Pillow")
-
-            # Add any other common libraries the LLM might generate code for
             if "numpy" in str(app_desc).lower(): reqs_list.append("numpy")
             if "pandas" in str(app_desc).lower(): reqs_list.append("pandas")
-            # Add more heuristics if needed
+            if "scikit-image" in str(app_desc).lower() or "skimage" in str(app_desc).lower() or "cv2" in str(app_desc).lower() or "opencv-python" in str(app_desc).lower(): # Common image processing deps
+                 reqs_list.append("scikit-image") # skimage
+                 reqs_list.append("opencv-python") # cv2
 
-            # Remove duplicates
-            reqs_list = list(dict.fromkeys(reqs_list))
+            reqs_list = list(dict.fromkeys(reqs_list)) # Remove duplicates
 
             reqs_content = "\n".join(reqs_list) + "\n"
 
@@ -459,7 +451,7 @@ Return **only** the python code block for app.py. Do not include any extra text,
              build_logs_text = get_build_logs_action(repo_id, hf_profile, hf_token)
              updated_build = build_logs_text
 
-             if "Error" in updated_build or "Exception" in updated_build or "Build failed" in updated_build: # Added "Build failed" check
+             if "Error" in updated_build or "Exception" in updated_build or "Build failed" in updated_build:
                   history = add_bot_message(history, "⚠️ Build logs indicate potential issues. Please inspect above. Click 'Send' to check container logs (app might still start).")
                   state = STATE_CHECKING_LOGS_RUN
                   yield history, repo_id, state, updated_preview, updated_run, updated_build, attempts, app_desc, repo_name, generated_code
@@ -551,27 +543,26 @@ Return **only** the python code block for app.py. Do not include any extra text,
                   yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
 
         elif state == STATE_COMPLETE:
-             # Workflow finished successfully or after max attempts.
-             # Stay in this state until reset.
-             pass
+             # Workflow finished. Stay in this state until reset.
+             # The message is set in the state that transitioned to COMPLETE.
+             pass # No further action needed in this state
 
 
     except Exception as e:
+        # Catch-all for unexpected exceptions in any state
         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}")
+        # Transition to idle state on unexpected errors
         yield history, None, STATE_IDLE, updated_preview, updated_run, updated_build, 0, None, None, None
-        return # End generator on critical error
+        return # End generator
 
 
-    # Final yield after the try block for states that didn't explicitly yield at the end
-    # This catches states like STATE_COMPLETE or steps that finish successfully without needing another click
-    # However, in this state machine design, each step *requires* a click to proceed, so
-    # every state block should ideally end with a yield after setting the *next* state,
-    # or yield the current state if waiting for user action (like checking logs).
-    # Let's ensure each path explicitly yields the updated state.
-    # Removed the final implicit yield here, as each state handler should yield.
-    pass
+    # Note: Each state block should end with a yield to update the UI/state.
+    # If a state block finishes without yielding, the generator will stop until the next call.
+    # By yielding the state variables *after* adding the user message, and then yielding
+    # again after state transitions/actions *within* the state logic, we ensure updates.
+    # The final yield after the try/except is removed as yields should happen within the state logic.
 
 
 # --- Build the Gradio UI ---
@@ -581,14 +572,14 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     hf_profile   = gr.State(None)
     hf_token     = gr.State(None)
     gemini_key   = gr.State(None)
-    gemini_model = gr.State("gemini-1.5-flash") # Switched to a more recent stable model
+    gemini_model = gr.State("gemini-1.5-flash")
     repo_id      = gr.State(None)
     workflow     = gr.State(STATE_IDLE)
     sdk_state    = gr.State("gradio")
     debug_attempts = gr.State(0)
     app_description = gr.State(None)
     repo_name_state = gr.State(None)
-    generated_code_state = gr.State(None) # Reused for different file contents
+    generated_code_state = gr.State(None) # Reused
 
     with gr.Row():
         # Sidebar
@@ -598,11 +589,11 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
             login_btn    = gr.LoginButton(variant="huggingface")
 
             # Initial load to check login status
+            # Removed _preprocess=False
             ai_builder_tab.load(show_profile, outputs=login_status)
             # Update status on login click
             login_btn.click(show_profile, outputs=login_status)
             # Store profile and token in state on login click
-            # This lambda takes the output of the LoginButton (profile, token) as 'x' and returns it
             login_btn.click(lambda x: x, inputs=[login_btn], outputs=[hf_profile, hf_token])
 
             gr.Markdown("## Google AI Studio API Key")
@@ -614,12 +605,9 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
             gr.Markdown("## Gemini Model")
             model_selector = gr.Radio(
                 choices=[
-                    ("Gemini 1.5 Flash", "gemini-1.5-flash"), # Prioritize 1.5 Flash
+                    ("Gemini 1.5 Flash", "gemini-1.5-flash"),
                     ("Gemini 1.5 Pro", "gemini-1.5-pro"),
                     ("Gemini 1.0 Pro", "gemini-1.0-pro"),
-                    # Remove preview models as they can be unstable
-                    # ("Gemini 2.5 Flash Preview 04-17", "gemini-2.5-flash-preview-04-17"),
-                    # ("Gemini 2.5 Pro Preview 03-25",   "gemini-2.5-pro-preview-03-25"),
                 ],
                 value="gemini-1.5-flash",
                 label="Select model",
@@ -628,13 +616,13 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
             model_selector.change(lambda m: m, inputs=model_selector, outputs=gemini_model)
 
             # Configure Gemini status on load and when key/model changes
-            # Added _preprocess=False to avoid implicit arguments causing TypeErrors
+            # Removed _preprocess=False from load
             ai_builder_tab.load(
                 configure_gemini,
                 inputs=[gemini_key, gemini_model],
-                outputs=[gemini_status],
-                _preprocess=False # Tell Gradio not to pass implicit event output
+                outputs=[gemini_status]
             )
+            # Keep _preprocess=False on change events as they have implicit outputs (the new value)
             gemini_input.change(
                 configure_gemini,
                 inputs=[gemini_key, gemini_model],
@@ -659,29 +647,28 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
 
         # Main content
         with gr.Column(scale=3):
-            # Corrected Chatbot initialization
-            chatbot    = gr.Chatbot(type='messages', label="AI Workflow Chat") # Added type='messages' and label
+            chatbot    = gr.Chatbot(type='messages', label="AI Workflow Chat")
             user_input = gr.Textbox(placeholder="Type your message…", interactive=True)
             send_btn   = gr.Button("Send", interactive=False)
 
             # Logic to enable send button only when logged in and API key is set
-            # Added _preprocess=False to avoid implicit arguments causing TypeErrors
             def update_send_button_state(profile: gr.OAuthProfile | None, token: gr.OAuthToken | None, key: str | None, model: str | None):
                 is_logged_in = profile is not None and token is not None
                 is_gemini_ready = key is not None and model is not None
                 return gr.update(interactive=is_logged_in and is_gemini_ready)
 
+            # Removed _preprocess=False from load
             ai_builder_tab.load(
                 update_send_button_state,
                 inputs=[hf_profile, hf_token, gemini_key, gemini_model],
-                outputs=[send_btn],
-                _preprocess=False # Tell Gradio not to pass implicit event output
+                outputs=[send_btn]
             )
+            # Keep _preprocess=False on click/change events with explicit inputs
             login_btn.click(
                 update_send_button_state,
                 inputs=[hf_profile, hf_token, gemini_key, gemini_model],
                 outputs=[send_btn],
-                _preprocess=False # Tell Gradio not to pass implicit event output
+                _preprocess=False
             )
             gemini_input.change(
                 update_send_button_state,
@@ -729,11 +716,10 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
 
 
     # Add an initial message to the chatbot on load
-    # THIS CALL IS NOW INSIDE the with gr.Blocks() block
     def greet():
-         # Return a list containing a single assistant message dictionary
         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."}]
 
+    # THIS CALL IS INSIDE the with gr.Blocks() block
     ai_builder_tab.load(greet, outputs=chatbot)
 
 
@@ -741,7 +727,7 @@ if __name__ == "__main__":
     # Optional: Configure retries for huggingface_hub requests
     # from requests.adapters import HTTPAdapter
     # from urllib3.util.retry import Retry
-    # retry_strategy = Retry(total=5, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504]) # Increased retries
+    # retry_strategy = Retry(total=5, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504])
     # adapter = HTTPAdapter(max_retries=retry_strategy)
     # get_session().mount("http://", adapter)
     # get_session().mount("https://", adapter)