Update app.py

app.py

@@ -7,6 +7,17 @@ import requests
 import logging
 from typing import List, Dict, Any, Tuple, Optional, Literal, Generator, Union # Import Union for clarity
 
+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
+
+# --- Import for Hugging Face API Retries ---
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry # <-- Ensure this line is present and NOT commented out
+
 
 # --- Configure Logging ---
 # Replace print() statements with logging
@@ -1390,7 +1401,8 @@ def handle_complete(
 
 
 # --- Dispatch Table ---
-STATE_HANDLERS: Dict[WorkflowState, Any] = { # Use Any for type hint simplicity here
+# Using Union for the return type hint in the Dict value type for clarity
+STATE_HANDLERS: Dict[WorkflowState, Union[WorkflowOutputs, Generator[WorkflowOutputs, None, WorkflowOutputs]]] = {
     STATE_IDLE: handle_idle,
     STATE_AWAITING_REPO_NAME: handle_awaiting_repo_name,
     STATE_CREATING_SPACE: handle_creating_space,
@@ -1433,7 +1445,7 @@ def ai_workflow_chat(
     # Accept any extra args/kwargs passed by Gradio, common for generators
     *args,
     **kwargs
-) -> Any: # Use Any because it yields multiple times before returning the final value (None in this case)
+) -> Generator[WorkflowOutputs, None, None]: # Correct return type for the main generator
     """
     Generator function to handle the AI workflow state machine.
     Each 'yield' pauses execution and sends values to update Gradio outputs/state.