Update app.py

app.py

@@ -22,7 +22,7 @@ GEMINI_MODELS = {
     "gemini-1.5-flash-8b": ("Gemini 1.5 Flash 8B", "High volume and lower intelligence tasks."),
     "gemini-2.0-flash": ("Gemini 2.0 Flash", "Next generation features, speed, thinking, realtime streaming, and multimodal generation."),
     "gemini-2.0-flash-lite": ("Gemini 2.0 Flash-Lite", "Cost efficiency and low latency."),
-    # Note: Preview models might have shorter lifespans or different capabilities
+    # Note: Preview models might have shorter lifespans or different capabilities. Uncomment if you want to include them.
     # "gemini-2.5-flash-preview-04-17": ("Gemini 2.5 Flash Preview (04-17)", "Adaptive thinking, cost efficiency."),
     # "gemini-2.5-pro-preview-03-25": ("Gemini 2.5 Pro Preview (03-25)", "Enhanced thinking and reasoning, multimodal understanding, advanced coding, and more."),
 }
@@ -159,22 +159,23 @@ def configure_gemini(api_key: str | None, model_name: str | None) -> str:
     # Check for empty string "" as well as None
     if not api_key:
         return "⚠️ Gemini API key is not set."
-    if not model_name:
-        return "⚠️ Please select a Gemini model."
+    # Check if model_name is None or not a valid key in GEMINI_MODELS
+    if not model_name or model_name not in GEMINI_MODELS:
+        return "⚠️ Please select a valid 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)
         # This message indicates the API call *for configuration check* was successful
-        return f"✅ Gemini configured successfully with **{GEMINI_MODELS.get(model_name, ('Unknown Model', ''))[0]}**."
+        return f"✅ Gemini configured successfully with **{GEMINI_MODELS[model_name][0]}**."
     except Exception as e:
         # This message indicates the API call *for configuration check* failed
         return f"❌ Error configuring Gemini: {e}"
 
 def get_model_description(model_name: str | None) -> str:
     """Retrieves the description for a given model name."""
-    if model_name is None:
+    if model_name is None or model_name not in GEMINI_MODELS:
         return "Select a model to see its description."
     # Use .get with a default value to handle cases where the key might not be found
     return GEMINI_MODELS.get(model_name, (model_name, "No description available."))[1]
@@ -255,9 +256,8 @@ 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
-# Now takes explicit arguments, intended to be called with all necessary states.
-# Modified to accept extra positional and keyword arguments passed by Gradio event chains
-def check_send_button_ready(hf_profile: gr.OAuthProfile | None, hf_token: gr.OAuthToken | None, gemini_key: str | None, gemini_model: str | None, *args, **kwargs) -> gr.update:
+# This function has the clean signature it expects. Wrappers handle Gradio's argument passing.
+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 ---")
@@ -267,8 +267,6 @@ def check_send_button_ready(hf_profile: gr.OAuthProfile | None, hf_token: gr.OAu
     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: Value starts with '{api_key_display}'")
     print(f"  Received gemini_model: {gemini_model}")
-    print(f"  Received extra args: {args}") # Added debug for extra args
-    print(f"  Received extra kwargs: {kwargs}") # Added debug for extra kwargs
     # --- END ENHANCED DEBUGGING LOGS ---
 
     is_logged_in    = hf_profile is not None and hf_token is not None
@@ -281,27 +279,87 @@ def check_send_button_ready(hf_profile: gr.OAuthProfile | None, hf_token: gr.OAu
 
     return gr.update(interactive=is_ready)
 
-# Wrapper function for configure_gemini when called in a .then() chain
-# It accepts the output of the previous function as the first arg and discards it.
-def wrapper_configure_gemini(prev_output, api_key: str | None, model_name: str | None) -> str:
-     return configure_gemini(api_key, model_name)
-
-# Wrapper function for greet when called in a .then() chain
-# It accepts the output of the previous function as the first arg and discards it.
-def wrapper_greet(prev_output):
-     return greet()
-
-# Wrapper function for get_model_description when called in a .then() chain
-# It accepts the output of the previous function as the first arg and discards it.
-def wrapper_get_model_description(prev_output, model_name: str | None) -> str:
-    return get_model_description(model_name)
+# --- Wrappers to handle Gradio's argument passing in event chains ---
+# These wrappers accept whatever Gradio passes (*args, **kwargs) and call the target function
+# with the specific arguments it expects, extracted from *args based on the expected call signature.
+
+# Wrapper for functions called in .then() chains with specific inputs list: expects (prev_output, *input_values)
+# e.g., .then(wrapper, inputs=[s1, s2]) -> wrapper receives (prev_out, s1_val, s2_val)
+def wrapper_from_then_inputs(func, num_inputs):
+    def wrapped(*args, **kwargs):
+        # We expect num_inputs values at the end of *args, after prev_output (index 0)
+        if len(args) > num_inputs:
+            required_args = args[-num_inputs:]
+            try:
+                return func(*required_args)
+            except Exception as e:
+                print(f"Error calling wrapped function {func.__name__} with args {required_args}: {e}")
+                # Provide a fallback or re-raise depending on context
+                if func == configure_gemini: return f"❌ Error configuring Gemini: {e}"
+                if func == get_model_description: return f"Error getting description: {e}"
+                if func == check_send_button_ready: return gr.update(interactive=False)
+                raise # Re-raise if no specific fallback
+        else:
+            print(f"Warning: wrapper_from_then_inputs for {func.__name__} received unexpected args (expecting at least {num_inputs+1}): {args}")
+            # Provide a fallback or re-raise
+            if func == configure_gemini: return "❌ Error configuring Gemini: unexpected arguments received."
+            if func == get_model_description: return "No description available (unexpected arguments received)."
+            if func == check_send_button_ready: return gr.update(interactive=False)
+            raise ValueError(f"Unexpected args received for {func.__name__}: {args}")
+    return wrapped
+
+# Wrapper for functions called by .change() trigger with specific inputs list: expects (changed_value, *input_values)
+# e.g., component.change(wrapper, inputs=[s1, s2]) -> wrapper receives (changed_val, s1_val, s2_val)
+def wrapper_from_change_inputs(func, num_inputs):
+    def wrapped(*args, **kwargs):
+         # We expect num_inputs values at the end of *args, after the changed_value (index 0)
+         if len(args) > num_inputs:
+              required_args = args[-num_inputs:]
+              try:
+                   return func(*required_args)
+              except Exception as e:
+                   print(f"Error calling wrapped function {func.__name__} with args {required_args}: {e}")
+                   if func == check_send_button_ready: return gr.update(interactive=False)
+                   raise # Re-raise if no specific fallback
+         else:
+              print(f"Warning: wrapper_from_change_inputs for {func.__name__} received unexpected args (expecting at least {num_inputs+1}): {args}")
+              if func == check_send_button_ready: return gr.update(interactive=False)
+              raise ValueError(f"Unexpected args received for {func.__name__}: {args}")
+    return wrapped
+
+
+# Wrapper for functions called in .then() chains with inputs=None: expects (prev_output,)
+# e.g., .then(wrapper, inputs=None) -> wrapper receives (prev_out,)
+def wrapper_from_prev_output(func):
+    def wrapped(*args, **kwargs):
+        # We expect only prev_output, or sometimes nothing if the chain starts
+        if len(args) >= 0: # Just accept anything here
+            try:
+                 # The target function expects 0 args, so call it with no args
+                 return func()
+            except Exception as e:
+                 print(f"Error calling wrapped function {func.__name__} with args {args}: {e}")
+                 # Provide a fallback or re-raise
+                 if func == greet: return [{"role": "assistant", "content": f"❌ Error loading initial message: {e}"}]
+                 raise # Re-raise if no specific fallback
+        else:
+             print(f"Warning: wrapper_from_prev_output for {func.__name__} received unexpected args: {args}")
+             if func == greet: return [{"role": "assistant", "content": "❌ Error loading initial message: unexpected arguments received."}]
+             raise ValueError(f"Unexpected args received for {func.__name__}: {args}")
+    return wrapped
+
+# Instantiate specific wrappers using the generic ones
+wrapper_check_button_change = wrapper_from_change_inputs(check_send_button_ready, 4) # Expects (changed, s1, s2, s3, s4)
+wrapper_check_button_then = wrapper_from_then_inputs(check_send_button_ready, 4)   # Expects (prev_out, s1, s2, s3, s4)
+
+wrapper_configure_gemini_then = wrapper_from_then_inputs(configure_gemini, 2) # Expects (prev_out, s1, s2) -> api_key, model_name
+wrapper_get_model_description_then = wrapper_from_then_inputs(get_model_description, 1) # Expects (prev_out, s1) -> model_name
+
+wrapper_greet_then = wrapper_from_prev_output(greet) # Expects (prev_out,), needs 0 args
 
 
 # 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.
-# 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.
+# Inputs and Outputs list must match exactly. The generator receives values from the inputs list.
 def ai_workflow_chat(
     message: str,
     history: list[dict],
@@ -323,8 +381,8 @@ def ai_workflow_chat(
     generated_code_state: str | None,
     use_grounding_state: bool, # Value from use_grounding_checkbox
     # Absorb potential extra args passed by Gradio event listeners (important for generators)
-    *args,
-    **kwargs
+    *args, # Generators might receive extra args, need to accept them but don't need to yield unless they are state
+    **kwargs # Generators might receive extra kwargs
 ) -> tuple[
     list[dict],       # 0: Updated chat history (for chatbot)
     str | None,       # 1: Updated repo_id (for repo_id state)
@@ -373,8 +431,8 @@ def ai_workflow_chat(
     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(f"  *args (unexpected by generator): {args}") # Added debug for unexpected args
+    print(f"  **kwargs (unexpected by generator): {kwargs}") # Added debug for unexpected kwargs
     print("--- END DEBUGGING ai_workflow_chat inputs ---\n")
 
 
@@ -1011,8 +1069,7 @@ 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
-    # This list provides the values to the check_send_button_ready function
+    # Define the inputs used for checking prerequisites (These are State components)
     send_button_interactive_binding_inputs = [
         hf_profile,
         hf_token,
@@ -1024,61 +1081,73 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
 
 
     # Trigger check_send_button_ready whenever any prerequisite state changes
-    # Bind the check function directly. The inputs list passes the required state values.
-    # The check_send_button_ready function signature now handles the implicit first argument.
+    # Use the specific change wrapper which expects (changed_value, *input_values)
     hf_profile.change(
-        check_send_button_ready,
+        wrapper_check_button_change,
+        inputs=send_button_interactive_binding_inputs, # Pass all 4 prerequisite states
+        outputs=send_button_update_output, # Update only the send button
+    )
+    hf_token.change(
+        wrapper_check_button_change,
         inputs=send_button_interactive_binding_inputs,
         outputs=send_button_update_output,
     )
-    hf_token.change(
-        check_send_button_ready,
+    # When gemini_api_key_state changes (updated by gemini_input.change), check button readiness
+    gemini_api_key_state.change(
+        wrapper_check_button_change,
         inputs=send_button_interactive_binding_inputs,
         outputs=send_button_update_output,
     )
-    # gemini_input.change updates gemini_api_key_state (in the first lambda).
-    # The .then() chain attached to gemini_input.change will then trigger.
-    # The function in the first .then() will receive the *output* of the previous lambda + its inputs.
-    # This requires the wrapper for configure_gemini.
+    # When gemini_model_state changes (updated by model_selector.change), check button readiness
+    gemini_model_state.change(
+        wrapper_check_button_change,
+        inputs=send_button_interactive_binding_inputs,
+        outputs=send_button_update_output,
+    )
+
+
+    # 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 takes the LoginButton output (profile, token tuple) and returns it
+         inputs=[login_btn], # Pass the LoginButton itself to get its output
+         outputs=[hf_profile, hf_token] # Update state variables
+    )
+
+    # Handle Gemini Key Input change: Update key state -> Configure Gemini status -> Update send button state
     gemini_input.change(
-        lambda k: k, # Updates gemini_api_key_state
-        inputs=[gemini_input],
-        outputs=[gemini_api_key_state] # This output becomes the implicit first arg for the .then() chain
+        # Lambda receives the new value of gemini_input (1 arg) because inputs=None (implied)
+        lambda new_key_value: new_key_value,
+        inputs=None, # Only need the new value of the changed component
+        outputs=[gemini_api_key_state] # This output becomes the implicit first arg for the next .then() in this chain
     ).then(
         # Configure Gemini using the updated state variables
-        wrapper_configure_gemini, # Use the wrapper
+        # Use the then_inputs wrapper which expects (prev_output, api_key_val, model_name_val)
+        wrapper_configure_gemini_then,
         inputs=[gemini_api_key_state, gemini_model_state], # Explicitly pass the required states
-        outputs=[gemini_status] # Update Gemini status display
-    ).then(
-        # Trigger check_send_button_ready after status update
-        # check_send_button_ready signature now handles the implicit previous output + explicit inputs
-        check_send_button_ready,
-        inputs=send_button_interactive_binding_inputs,
-        outputs=send_button_update_output
+        outputs=[gemini_status] # Update Gemini status display. This output becomes the implicit first arg for the next .then()
     )
+    # The .then chain continues from the outputs of the configure_gemini call, handled by gemini_api_key_state.change handler above
 
     # Handle Gemini Model Selector change: Update model state -> Update description -> Configure Gemini status -> Update send button state
     model_selector.change(
-        lambda m: m, # Updates gemini_model_state
-        inputs=[model_selector],
-        outputs=[gemini_model_state] # This output becomes the implicit first arg for the next .then() in this chain
+        # Lambda receives the new value of model_selector (1 arg)
+        lambda new_model_name: new_model_name,
+        inputs=None, # Only need the new value of the changed component
+        outputs=[gemini_model_state] # This output becomes the implicit first arg for the next .then()
     ).then(
         # Update the model description display
-        wrapper_get_model_description, # Use the wrapper
+        # Use the then_inputs wrapper which expects (prev_output, model_name_val)
+        wrapper_get_model_description_then,
         inputs=[gemini_model_state], # Get the new state value
-        outputs=[model_description_text] # Update description UI
+        outputs=[model_description_text] # Update description UI. This output becomes implicit first arg for next .then()
     ).then(
         # Configure Gemini using the updated state variables
-        wrapper_configure_gemini, # Use the wrapper
+        # Use the then_inputs wrapper which expects (prev_output, api_key_val, model_name_val)
+        wrapper_configure_gemini_then,
         inputs=[gemini_api_key_state, gemini_model_state], # Explicitly pass the required states
-        outputs=[gemini_status] # Update Gemini status display
-    ).then(
-        # Trigger check_send_button_ready after status update
-        # check_send_button_ready signature now handles the implicit previous output + explicit inputs
-        check_send_button_ready,
-        inputs=send_button_interactive_binding_inputs,
-        outputs=send_button_update_output
+        outputs=[gemini_status] # Update Gemini status display. This output becomes the implicit first arg for the next .then()
     )
+    # The .then chain continues from the outputs of the configure_gemini call, handled by gemini_model_state.change handler above
 
 
     # Handle Grounding checkbox change: update grounding state
@@ -1092,17 +1161,24 @@ with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
     )
 
     # Link Workflow State variable change to UI status display
-    workflow.change(lambda s: s, inputs=workflow, outputs=status_text)
+    # Lambda receives the new state value (1 arg) because inputs=None
+    workflow.change(lambda new_state_value: new_state_value, inputs=None, outputs=status_text)
 
     # Link Repo ID State variable change to UI status display
-    # Fix: The lambda here only needs the implicit new value of repo_id, remove inputs=[repo_id]
-    repo_id.change(lambda r: r if r else "None", inputs=None, outputs=repo_id_text)
+    # Lambda receives the new state value (1 arg) because inputs=None
+    # The warning about receiving 0 args persists, likely ignorable for this lambda
+    repo_id.change(lambda new_repo_id_value: new_repo_id_value if new_repo_id_value else "None", inputs=None, outputs=repo_id_text)
 
 
     # The main event handler for the Send button (generator)
-    # This call is direct, not in a .then() chain, so no implicit previous output is passed
+    # 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
+    # Ensure inputs and outputs match the ai_workflow_chat signature and yield tuple EXACTLY.
+    # This call is direct, not in a .then() chain, so it does NOT receive a prev_output arg.
+    # It receives args only from the inputs list.
     send_btn.click(
-        ai_workflow_chat, # The generator function to run (signature now handles potential args)
+        ai_workflow_chat, # The generator function to run (signature handles potential extra args, just in case)
         inputs=[
             user_input, chatbot, # UI component inputs (message, current chat history)
             hf_profile, hf_token, # HF State variables
@@ -1132,31 +1208,32 @@ 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)
-        show_profile, # Called directly by load, expects one arg (profile), should handle None initially
+        # show_profile expects 1 arg (profile) or None. It receives 1 from load. Correct.
+        show_profile,
         inputs=None,
         outputs=login_status # Updates UI. This output becomes the implicit first arg for the next .then()
     ).then(
         # Action 2: Configure Gemini using initial state
-        # Wrap configure_gemini because it's in a .then() chain and receives the output of the prior step (show_profile)
-        wrapper_configure_gemini, # Use the wrapper
+        # Use the then_inputs wrapper which expects (prev_output, api_key_val, model_name_val)
+        wrapper_configure_gemini_then,
         inputs=[gemini_api_key_state, gemini_model_state], # Explicitly pass the required states
         outputs=[gemini_status] # Update Gemini status display. This output becomes the implicit first arg for the next .then()
     ).then(
         # Action 3: After initial load checks, update the button state based on initial states
-        # check_send_button_ready signature now handles the implicit previous output + explicit inputs
-        check_send_button_ready,
-        inputs=send_button_interactive_binding_inputs,
+        # Use the then_inputs wrapper which expects (prev_output, *prereq_state_values)
+        wrapper_check_button_then,
+        inputs=send_button_interactive_binding_inputs, # Pass all 4 prerequisite states
         outputs=send_button_update_output, # Update the send button. This output becomes implicit first arg for next .then()
     ).then(
         # Action 4: Update the model description text based on the default selected model
-        # Use the wrapper because it's in a .then() chain
-        wrapper_get_model_description,
+        # Use the then_inputs wrapper which expects (prev_output, model_name_val)
+        wrapper_get_model_description_then,
         inputs=[gemini_model_state], # Get the default model name from state
-        outputs=[model_description_text] # Update description UI
+        outputs=[model_description_text] # Update description UI. This output becomes implicit first arg for next .then()
     ).then(
         # Action 5: Add the initial welcome message to the chat history
-        # Wrap greet because it's in a .then() chain and receives the output of the prior step (check_send_button_ready)
-        wrapper_greet, # Use the wrapper
+        # Use the prev_output wrapper which expects (prev_output,)
+        wrapper_greet_then,
         inputs=None, # Greet takes no explicit inputs
         outputs=chatbot # Updates the chatbot display
     )