Update app.py

app.py

@@ -1,9 +1,12 @@
 import re
 import json
 import time
+import requests
 import importlib.metadata
 import gradio as gr
-from huggingface_hub import create_repo, upload_file, list_models, constants
+from huggingface_hub import (
+    create_repo, upload_file, list_models, constants
+)
 from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status
 from google import genai
 from google.genai.types import Tool, GenerateContentConfig, GoogleSearch
@@ -11,15 +14,13 @@ from google.genai.types import Tool, GenerateContentConfig, GoogleSearch
 # — USER INFO & MODEL LISTING —
 
 def show_profile(profile: gr.OAuthProfile | None) -> str:
-    if profile is None:
-        return "*Not logged in.*"
-    return f"✅ Logged in as **{profile.username}**"
+    return f"✅ Logged in as **{profile.username}**" if profile else "*Not logged in.*"
 
 def list_private_models(
     profile: gr.OAuthProfile | None,
     oauth_token: gr.OAuthToken | None
 ) -> str:
-    if profile is None or oauth_token is None:
+    if not profile or not oauth_token:
         return "Please log in to see your models."
     models = [
         f"{m.id} ({'private' if m.private else 'public'})"
@@ -40,9 +41,20 @@ def extract_code(text: str) -> str:
     blocks = re.findall(r"```(?:\w*\n)?([\s\S]*?)```", text)
     return blocks[-1].strip() if blocks else text.strip()
 
+def classify_errors(logs: str) -> str:
+    errs = set()
+    for line in logs.splitlines():
+        if "SyntaxError" in line:
+            errs.add("syntax")
+        elif "ImportError" in line or "ModuleNotFoundError" in line:
+            errs.add("import")
+        elif "Traceback" in line or "Exception" in line:
+            errs.add("runtime")
+    return ", ".join(errs) or "unknown"
+
 # — HF SPACE LOGGING —
 
-def _get_space_jwt(repo_id: str):
+def _get_space_jwt(repo_id: str) -> str:
     url = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
     r = get_session().get(url, headers=build_hf_headers())
     hf_raise_for_status(r)
@@ -50,21 +62,27 @@ def _get_space_jwt(repo_id: str):
 
 def fetch_logs(repo_id: str, level: str) -> str:
     jwt = _get_space_jwt(repo_id)
-    logs_url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
+    url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
     lines = []
-    with get_session().get(logs_url, headers=build_hf_headers(token=jwt), stream=True) as resp:
+    with get_session().get(url, headers=build_hf_headers(token=jwt), stream=True) as resp:
         hf_raise_for_status(resp)
         for raw in resp.iter_lines():
             if raw.startswith(b"data: "):
                 try:
                     ev = json.loads(raw[len(b"data: "):].decode())
-                    ts  = ev.get("timestamp","")
-                    txt = ev.get("data","")
+                    ts, txt = ev.get("timestamp",""), ev.get("data","")
                     lines.append(f"[{ts}] {txt}")
                 except:
                     continue
     return "\n".join(lines)
 
+def check_iframe(url: str, timeout: int = 5) -> bool:
+    try:
+        r = requests.get(url, timeout=timeout)
+        return r.status_code == 200
+    except:
+        return False
+
 # — CORE LOOP —
 
 def handle_user_message(
@@ -72,18 +90,21 @@ def handle_user_message(
     sdk_choice: str,
     gemini_api_key: str,
     grounding_enabled: bool,
+    temperature: float,
+    max_output_tokens: int,
     profile: gr.OAuthProfile | None,
     oauth_token: gr.OAuthToken | None
 ):
-    if profile is None or oauth_token is None:
+    if not profile or not oauth_token:
         return history + [{"role":"assistant","content":"⚠️ Please log in first."}], "", "", "<p>No Space yet.</p>"
 
     client = genai.Client(api_key=gemini_api_key)
     system_msg = {
         "role":"system",
         "content":(
-            f"You are an AI assistant writing a HuggingFace Space using the "
-            f"{sdk_choice} SDK. After producing code, wait for logs; if errors appear, fix them."
+            f"You are an AI assistant that scaffolds a complete HuggingFace Space using the "
+            f"{sdk_choice} SDK. Generate working code, handle errors, and ensure the deployed "
+            "iframe loads successfully."
         )
     }
     chat = [system_msg] + history
@@ -92,120 +113,184 @@ def handle_user_message(
     readme_fn = "README.md"
     reqs_fn   = "requirements.txt"
     repo_id   = f"{profile.username}/{profile.username}-auto-space"
+    iframe_url = f"https://huggingface.co/spaces/{repo_id}"
 
     build_logs = run_logs = ""
-    for _ in range(5):
+    backoff = 1
+
+    for attempt in range(1, 7):
         tools = [Tool(google_search=GoogleSearch())] if grounding_enabled else []
-        cfg   = GenerateContentConfig(tools=tools, response_modalities=["TEXT"])
+        cfg = GenerateContentConfig(
+            tools=tools,
+            response_modalities=["TEXT"],
+            temperature=temperature,
+            max_output_tokens=max_output_tokens,
+        )
 
+        # generate code
         resp = client.models.generate_content(
             model="gemini-2.5-flash-preview-04-17",
             contents=[m["content"] for m in chat],
             config=cfg
         )
-
-        raw = resp.text
-        code = extract_code(raw)
+        code = extract_code(resp.text)
         chat.append({"role":"assistant","content":code})
 
-        # write code
-        with open(code_fn, "w") as f:
-            f.write(code)
-
-        # write dynamic README
+        # local syntax check
+        try:
+            compile(code, code_fn, "exec")
+        except SyntaxError as e:
+            chat.append({
+                "role":"user",
+                "content": f"SyntaxError caught locally: {e}. Please fix the code."
+            })
+            time.sleep(backoff)
+            backoff = min(backoff * 2, 30)
+            continue
+
+        # write code, README, requirements
         sdk_version = get_sdk_version(sdk_choice)
-        readme = f"""---
+        files = {
+            code_fn: code,
+            readme_fn: f"""---
 title: Wuhp Auto Space
 emoji: 🐢
-colorFrom: red
-colorTo: pink
 sdk: {sdk_choice}
 sdk_version: {sdk_version}
 app_file: {code_fn}
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
-"""
-        with open(readme_fn, "w") as f:
-            f.write(readme)
-
-        # write requirements
-        base_reqs = "pandas\n"
-        extra     = "streamlit\n" if sdk_choice=="streamlit" else "gradio\n"
-        with open(reqs_fn, "w") as f:
-            f.write(base_reqs + extra)
+See https://huggingface.co/docs/hub/spaces-config-reference
+""",
+            reqs_fn: "pandas\n" + ("streamlit\n" if sdk_choice=="streamlit" else "gradio\n")
+        }
+        for fn, content in files.items():
+            with open(fn, "w") as f:
+                f.write(content)
 
-        # push to HF
+        # push to HuggingFace
         create_repo(repo_id=repo_id, token=oauth_token.token,
                     exist_ok=True, repo_type="space", space_sdk=sdk_choice)
-        for fn in (code_fn, readme_fn, reqs_fn):
-            upload_file(path_or_fileobj=fn, path_in_repo=fn,
-                        repo_id=repo_id, token=oauth_token.token,
-                        repo_type="space")
+        for fn in files:
+            upload_file(fn, fn, repo_id=repo_id,
+                        token=oauth_token.token, repo_type="space")
 
+        # fetch logs
         build_logs = fetch_logs(repo_id, "build")
         run_logs   = fetch_logs(repo_id, "run")
-        if "ERROR" not in build_logs.upper() and "ERROR" not in run_logs.upper():
+        errors = classify_errors(build_logs + "\n" + run_logs)
+
+        # success criteria
+        if "ERROR" not in build_logs.upper() and \
+           "ERROR" not in run_logs.upper() and \
+           check_iframe(iframe_url):
             break
 
+        # ask for a fix
         chat.append({
             "role":"user",
             "content":(
+                f"Attempt {attempt} encountered {errors} errors.\n"
                 f"Build logs:\n{build_logs}\n\n"
                 f"Run logs:\n{run_logs}\n\n"
-                "Please fix the code."
+                f"Please fix the code and ensure the iframe at {iframe_url} returns HTTP 200."
             )
         })
-        time.sleep(2)
+        time.sleep(backoff)
+        backoff = min(backoff * 2, 30)
+
+    messages = [{"role":m["role"], "content":m["content"]} for m in chat if m["role"]!="system"]
+    iframe_html = (
+        f'<iframe src="{iframe_url}" width="100%" height="500px"></iframe>'
+        + ("" if check_iframe(iframe_url) else
+           "<p style='color:red;'>⚠️ iframe not responding.</p>")
+    )
 
-    messages = [{"role":m["role"],"content":m["content"]} for m in chat if m["role"]!="system"]
-    iframe   = f'<iframe src="https://huggingface.co/spaces/{repo_id}" width="100%" height="500px"></iframe>'
-    return messages, build_logs, run_logs, iframe
+    return messages, build_logs, run_logs, iframe_html
 
-# — BUILD THE UI —
+# — BUILD THE UI (Improved) —
 
-with gr.Blocks(title="HF Space Auto‑Builder") as demo:
-    gr.Markdown("## Sign in + Auto‑Build Spaces\n\n1. Sign in  2. Enter your prompt  3. Watch code, README, requirements, logs, and preview\n\n---")
+custom_css = """
+.gradio-container { max-width: 1200px; margin: auto; }
+.log-box { font-family: monospace; white-space: pre; }
+"""
 
-    # LOGIN & MODEL LISTING
-    login_btn = gr.LoginButton(variant="huggingface", size="lg")
-    status_md = gr.Markdown("*Not logged in.*")
-    models_md = gr.Markdown()
-    demo.load(show_profile,        inputs=None, outputs=status_md)
+with gr.Blocks(title="HF Space Auto‑Builder", css=custom_css) as demo:
+    with gr.Row():
+        # Sidebar
+        with gr.Column(scale=1, min_width=250):
+            login_btn = gr.LoginButton(variant="huggingface", size="sm")
+            status_md = gr.Markdown("*Not logged in.*")
+            models_md = gr.Markdown()
+            sdk_choice = gr.Radio(
+                ["gradio", "streamlit"],
+                value="gradio",
+                label="SDK template"
+            )
+            api_key = gr.Textbox(label="Gemini API Key", type="password")
+            grounding = gr.Checkbox(label="Enable grounding", value=False)
+            temp = gr.Slider(0, 1, value=0.2, label="LLM temperature")
+            max_tokens = gr.Slider(256, 4096, value=1024, step=256, label="Max tokens")
+            clear_btn = gr.Button("Clear Chat", variant="secondary")
+
+        # Main area
+        with gr.Column(scale=3):
+            gr.Markdown("## 🐢 Wuhp Auto‑Builder\nSign in, enter your prompt, and watch your Space live.")
+            chatbot = gr.Chatbot(elem_id="chatbot")
+            user_in = gr.Textbox(
+                placeholder="Type your prompt and hit Enter…",
+                label="Prompt", lines=1
+            )
+            send_btn = gr.Button("Send", variant="primary", show_progress=True)
+
+            with gr.Tabs():
+                with gr.TabItem("Chat"):
+                    pass
+                with gr.TabItem("Logs"):
+                    with gr.Accordion("🔨 Build Logs", open=True):
+                        build_box = gr.Textbox(
+                            lines=8, interactive=False, elem_classes="log-box"
+                        )
+                    with gr.Accordion("🚀 Run Logs", open=False):
+                        run_box = gr.Textbox(
+                            lines=8, interactive=False, elem_classes="log-box"
+                        )
+                with gr.TabItem("Preview"):
+                    preview = gr.HTML("<p>No Space yet.</p>")
+
+    # Wiring callbacks
+    demo.load(show_profile, inputs=None, outputs=status_md)
     demo.load(list_private_models, inputs=None, outputs=models_md)
-    login_btn.click(show_profile,        inputs=None, outputs=status_md)
+    login_btn.click(show_profile, inputs=None, outputs=status_md)
     login_btn.click(list_private_models, inputs=None, outputs=models_md)
 
-    # CONTROLS
-    sdk_choice = gr.Radio(["gradio","streamlit"], value="gradio", label="SDK template")
-    api_key    = gr.Textbox(label="Gemini API Key", type="password")
-    grounding  = gr.Checkbox(label="Enable grounding", value=False)
-
-    # CHAT + OUTPUTS
-    chatbot   = gr.Chatbot(type="messages")
-    user_in   = gr.Textbox(placeholder="Your prompt…", label="Prompt")
-    send_btn  = gr.Button("Send")
-
-    build_box = gr.Textbox(label="Build logs", lines=5, interactive=False)
-    run_box   = gr.Textbox(label="Run logs",   lines=5, interactive=False)
-    preview   = gr.HTML("<p>No Space yet.</p>")
-
     send_btn.click(
         fn=handle_user_message,
-        inputs=[chatbot, sdk_choice, api_key, grounding],
+        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, status_md, models_md],
+        outputs=[chatbot, build_box, run_box, preview]
+    )
+    user_in.submit(
+        fn=handle_user_message,
+        inputs=[chatbot, sdk_choice, api_key, grounding, temp, max_tokens, status_md, models_md],
+        outputs=[chatbot, build_box, run_box, preview]
+    )
+
+    def _clear_chat():
+        return [], "", "", "<p>No Space yet.</p>"
+    clear_btn.click(
+        fn=_clear_chat,
+        inputs=None,
         outputs=[chatbot, build_box, run_box, preview]
     )
 
-    # — Refresh Logs button —
-    def _refresh(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None):
+    # Refresh Logs
+    refresh_btn = gr.Button("Refresh Logs", variant="secondary", size="sm")
+    def _refresh(profile, oauth_token):
         if not profile or not oauth_token:
             return "", ""
         repo = f"{profile.username}/{profile.username}-auto-space"
         return fetch_logs(repo, "build"), fetch_logs(repo, "run")
-
-    refresh_btn = gr.Button("Refresh Logs")
-    # Gradio will auto‑inject `profile` and `oauth_token` here.
     refresh_btn.click(_refresh, inputs=None, outputs=[build_box, run_box])
 
     demo.launch(server_name="0.0.0.0", server_port=7860)