Update app.py

app.py

@@ -12,9 +12,9 @@ from google.genai.types import Tool, GoogleSearch
 # -----------------------------------------------------------------------------
 # Configuration
 # -----------------------------------------------------------------------------
-MODEL_ID: str = "gemini-2.5-flash-preview-04-17"
-WORKSPACE_DIR: str = "workspace"
-SYSTEM_INSTRUCTION: str = (
+MODEL_ID = "gemini-2.5-flash-preview-04-17"
+WORKSPACE_DIR = "workspace"
+SYSTEM_INSTRUCTION = (
     "You are a helpful coding assistant that scaffolds a complete Hugging Face Space app. "
     "Based on the user's request, decide between Gradio or Streamlit (whichever fits best), "
     "and respond with exactly one JSON object with keys:\n"
@@ -24,27 +24,21 @@ SYSTEM_INSTRUCTION: str = (
     "Do not include extra text or markdown."
 )
 
-# In‑memory session store: maps session IDs -> state dicts
 state_store: Dict[str, Dict[str, Any]] = {}
 
-# -----------------------------------------------------------------------------
-# Helper functions
-# -----------------------------------------------------------------------------
-
-def start_app(gemini_key: str, hf_token: str, hf_username: str, repo_name: str) -> Dict[str, Any]:
-    """Initialise chat with Gemini + create a local workspace for the new app."""
+def start_app(
+    gemini_key: str,
+    hf_token: str,
+    hf_username: str,
+    repo_name: str
+) -> Dict[str, Any]:
     os.makedirs(WORKSPACE_DIR, exist_ok=True)
-
-    # Gemini client & chat session
     client = genai.Client(api_key=gemini_key)
     config = types.GenerateContentConfig(system_instruction=SYSTEM_INSTRUCTION)
     tools = [Tool(google_search=GoogleSearch())]
     chat = client.chats.create(model=MODEL_ID, config=config, tools=tools)
-
-    # Local project folder (used as upload snapshot source)
     local_path = os.path.join(WORKSPACE_DIR, repo_name)
     os.makedirs(local_path, exist_ok=True)
-
     return {
         "chat": chat,
         "hf_token": hf_token,
@@ -57,34 +51,24 @@ def start_app(gemini_key: str, hf_token: str, hf_username: str, repo_name: str)
         "logs": [f"Initialized workspace at {WORKSPACE_DIR}/{repo_name}."]
     }
 
-
 def handle_message(user_msg: str, state: Dict[str, Any]) -> Tuple[str, Dict[str, Any]]:
-    """Send *user_msg* to Gemini, act on the JSON response, and return assistant reply."""
     chat = state["chat"]
     logs = state.setdefault("logs", [])
-
     logs.append(f"> **User**: {user_msg}")
     resp = chat.send_message(user_msg)
     logs.append("Received response from Gemini.")
 
-    # ---------------------------------------------------------------------
-    # Parse Gemini JSON answer
-    # ---------------------------------------------------------------------
     try:
         data = json.loads(resp.text)
-        framework: str = data["framework"]
-        files: Dict[str, str] = data.get("files", {})
-        reply_msg: str = data.get("message", "")
+        framework = data["framework"]
+        files = data.get("files", {})
+        reply = data.get("message", "")
     except Exception:
         logs.append("⚠️ Failed to parse assistant JSON.\n" + resp.text)
         return "⚠️ Parsing error. Check logs.", state
 
-    # ---------------------------------------------------------------------
-    # Create the target Space on first run
-    # ---------------------------------------------------------------------
     if not state["created"]:
         full_repo = f"{state['hf_username']}/{state['repo_name']}"
-        logs.append(f"Creating HF Space **{full_repo}** (template '{framework}') …")
         create_repo(
             repo_id=full_repo,
             token=state["hf_token"],
@@ -98,9 +82,6 @@ def handle_message(user_msg: str, state: Dict[str, Any]) -> Tuple[str, Dict[str,
             "embed_url": f"https://huggingface.co/spaces/{full_repo}",
         })
 
-    # ---------------------------------------------------------------------
-    # Write / overwrite files
-    # ---------------------------------------------------------------------
     if files:
         logs.append(f"Writing {len(files)} file(s): {list(files)}")
         for relpath, content in files.items():
@@ -109,76 +90,54 @@ def handle_message(user_msg: str, state: Dict[str, Any]) -> Tuple[str, Dict[str,
             with open(dest, "w", encoding="utf‑8") as fp:
                 fp.write(content)
 
-    # ---------------------------------------------------------------------
-    # Push snapshot
-    # ---------------------------------------------------------------------
-    logs.append("Uploading snapshot to Hugging Face …")
     HfApi(token=state["hf_token"]).upload_folder(
         folder_path=state["local_path"],
         repo_id=state["repo_id"],
         repo_type="space",
     )
     logs.append("Snapshot upload complete.")
+    return reply, state
 
-    return reply_msg, state
-
-# -----------------------------------------------------------------------------
-# Gradio UI
-# -----------------------------------------------------------------------------
 with gr.Blocks(title="Gemini → HF Space scaffolder") as demo:
-    # -------------------------------------------------------------------------
-    # OAuth UI row (always visible)
-    # -------------------------------------------------------------------------
     with gr.Row():
         login_btn = gr.LoginButton(
-            logout_value="Logout ({username})",  # label once signed in
+            logout_value="Logout ({username})",
             variant="huggingface",
             size="lg"
         )
         status_md = gr.Markdown("Not logged in")
-
-    # Activate the login button’s routes (silences the warning)
     login_btn.activate()
 
-    # Callback to refresh the status text whenever OAuth state changes
     def show_profile(profile: gr.OAuthProfile | None):
         if profile is None:
             return "*Not logged in.*"
         return f"Logged in as **{profile.username}**"
-
     demo.load(fn=show_profile, inputs=None, outputs=status_md)
 
-    # -------------------------------------------------------------------------
-    # Main app controls (hidden behind OAuth)
-    # -------------------------------------------------------------------------
     with gr.Row():
         with gr.Column(scale=1):
-            gemini_key   = gr.Textbox(label="Gemini API Key", type="password")
-            hf_user      = gr.Textbox(label="HF Username")
-            repo_name    = gr.Textbox(label="New App (repo) name")
-            session_id   = gr.Textbox(value="", visible=False)
-            start_btn    = gr.Button("Start a new app")
-
+            gemini_key    = gr.Textbox(label="Gemini API Key", type="password")
+            hf_user       = gr.Textbox(label="HF Username")
+            repo_name     = gr.Textbox(label="New App (repo) name")
+            session_id    = gr.Textbox(value="", visible=False)
+            start_btn     = gr.Button("Start a new app")
         with gr.Column(scale=3):
-            chatbot        = gr.Chatbot(type="messages")
-            logs_display   = gr.Textbox(label="Operation Logs", interactive=False, lines=8)
-            preview_iframe = gr.HTML("<p>No deployed app yet.</p>")
-
-            user_msg = gr.Textbox(label="Your message")
-            send_btn = gr.Button("Send", interactive=False)
-
-    # ---------------------------------------------------------------------
-    # Callback: start a new app (requires OAuth token)
-    # ---------------------------------------------------------------------
-    def on_start(g_key: str, h_user: str, r_name: str, oauth_token: gr.OAuthToken | None):
+            chatbot       = gr.Chatbot(type="messages", value=[])
+            logs_display  = gr.Textbox(label="Operation Logs", interactive=False, lines=8)
+            preview_iframe= gr.HTML("<p>No deployed app yet.</p>")
+            user_msg      = gr.Textbox(label="Your message")
+            send_btn      = gr.Button("Send", interactive=False)
+
+    def on_start(
+        gemini_key: str,
+        hf_username: str,
+        repo_name: str,
+        oauth_token: gr.OAuthToken | None
+    ):
         if oauth_token is None:
-            return (
-                gr.Error("Please *Sign in with Hugging Face* first."),
-                "", "", gr.update(interactive=False)
-            )
-
+            return gr.Error("Please *Sign in with Hugging Face* first."), "", "", gr.update(interactive=False)
         new_id = str(uuid.uuid4())
-        state_store[new_id] = start_app(g_key, oauth_token.token, h_user, r_name)
+        state_store[new_id] = start_app(gemini_key, oauth_token.token, hf_username, repo_name)
         logs = "\n".join(state_store[new_id]["logs"])
         return new_id, logs, "<p>Awaiting first instruction…</p>", gr.update(interactive=True)
 
@@ -188,24 +147,23 @@ with gr.Blocks(title="Gemini → HF Space scaffolder") as demo:
         outputs=[session_id, logs_display, preview_iframe, send_btn],
     )
 
-    # ---------------------------------------------------------------------
-    # Callback: send a chat message
-    # ---------------------------------------------------------------------
-    def on_send(msg: str, chat_history: list[list[str]], sess_id: str):
+    def on_send(
+        msg: str,
+        chat_history: list[dict],
+        sess_id: str
+    ):
         if not sess_id or sess_id not in state_store:
-            err = "Error: No active session. Click *Start a new app* first."
-            return chat_history + [("", err)], sess_id, "", ""
-
+            return [], sess_id, "", ""
         reply, new_state = handle_message(msg, state_store[sess_id])
         state_store[sess_id] = new_state
-
-        chat_history.append((msg, reply))
-        logs  = "\n".join(new_state["logs"])
-        embed = (
-            f'<iframe src="{new_state["embed_url"]}" width="100%" height="500px"></iframe>'
-            if new_state.get("embed_url") else ""
-        )
-        return chat_history, sess_id, logs, embed
+        history = chat_history or []
+        history.append({"role": "user",      "content": msg})
+        history.append({"role": "assistant", "content": reply})
+        logs = "\n".join(new_state["logs"])
+        embed = ""
+        if new_state.get("embed_url"):
+            embed = f'<iframe src="{new_state["embed_url"]}" width="100%" height="500px"></iframe>'
+        return history, sess_id, logs, embed
 
     send_btn.click(
         on_send,
@@ -218,8 +176,5 @@ with gr.Blocks(title="Gemini → HF Space scaffolder") as demo:
         outputs=[chatbot, session_id, logs_display, preview_iframe],
     )
 
-# -----------------------------------------------------------------------------
-# Launch the Gradio app
-# -----------------------------------------------------------------------------
 if __name__ == "__main__":
     demo.launch()