Update app.py

app.py

@@ -1,10 +1,17 @@
 import gradio as gr
 import json, time
+from pydantic import BaseModel
 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
 
+# — JSON spec model —
+
+class RepoSpec(BaseModel):
+    repo_name: str
+    files: dict[str, str]
+
 # — USER INFO & MODEL LISTING —
 
 def show_profile(profile: gr.OAuthProfile | None) -> str:
@@ -66,41 +73,26 @@ def handle_user_message(
     system_msg = {
         "role":"system",
         "content":(
-            "You are an AI assistant that generates all files needed for a HuggingFace Space, "
-            "and returns exactly one JSON object with:\n"
-            "  • repo_name: the new Space name (no username)\n"
-            "  • files: a map of filename→file-content (strings)\n\n"
+            "Return exactly one JSON object matching this schema:\n"
+            "  • repo_name (string)\n"
+            "  • files (object mapping filename→file-content)\n\n"
             "Files must include at least:\n"
             "  - A code file (default name: app.py unless you choose otherwise)\n"
             "  - requirements.txt with dependencies\n"
-            "  - README.md frontmatter (title, emoji, sdk, sdk_version, app_file, etc.)\n\n"
-            "Do NOT output any extra text—only the JSON.  "
-            "The user’s request is provided as the next message."
+            "  - README.md with frontmatter (title, emoji, sdk, sdk_version, app_file)\n\n"
+            "Do NOT output any extra text—only the JSON object."
         )
     }
 
     # build the chat context
     chat = [system_msg] + history + [{"role":"user", "content":user_prompt}]
 
-    # strict JSON schema
-    schema = {
-      "type": "object",
-      "properties": {
-        "repo_name": { "type": "string" },
-        "files": {
-          "type": "object",
-          "additionalProperties": { "type": "string" }
-        }
-      },
-      "required": ["repo_name","files"]
-    }
-
     repo_id = None
     build_logs = run_logs = ""
 
     for _ in range(5):
-        # detect sdk_version
-        if sdk_choice=="gradio":
+        # detect sdk_version at runtime
+        if sdk_choice == "gradio":
             import gradio as _gr; sdk_version = _gr.__version__
         else:
             import streamlit as _st; sdk_version = _st.__version__
@@ -110,7 +102,7 @@ def handle_user_message(
             tools=tools,
             response_modalities=["TEXT"],
             response_mime_type="application/json",
-            response_schema=schema
+            response_schema=RepoSpec
         )
 
         # call the LLM
@@ -119,16 +111,9 @@ def handle_user_message(
             contents=[m["content"] for m in chat],
             config=cfg
         )
-        ai_json = resp.text
-
-        # parse JSON
-        try:
-            spec = json.loads(ai_json)
-            repo_name = spec["repo_name"]
-            files     = spec["files"]
-        except Exception:
-            chat.append({"role":"user","content":"Invalid JSON—please output exactly the JSON object per schema."})
-            continue
+        spec = RepoSpec.model_validate_json(resp.text)
+        repo_name = spec.repo_name
+        files     = spec.files
 
         # prepare repo
         repo_id = f"{profile.username}/{repo_name}"
@@ -142,8 +127,8 @@ def handle_user_message(
 
         # write & upload all files
         for fn, content in files.items():
-            # if README uses placeholder <SDK_VERSION>, replace it
-            if fn.lower()=="readme.md":
+            # replace placeholder in README
+            if fn.lower() == "readme.md":
                 content = content.replace("<SDK_VERSION>", sdk_version)
             with open(fn, "w") as f:
                 f.write(content)
@@ -158,7 +143,6 @@ def handle_user_message(
         # 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():
             break
 
@@ -168,12 +152,12 @@ def handle_user_message(
             "content":(
                 f"Build logs:\n{build_logs}\n\n"
                 f"Run logs:\n{run_logs}\n\n"
-                "Please fix the JSON spec and return updated JSON."
+                "Please fix the JSON spec and return updated JSON only."
             )
         })
         time.sleep(2)
 
-    messages = [{"role":m["role"],"content":m["content"]} for m in chat if m["role"]!="system"]
+    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
 
@@ -182,7 +166,7 @@ def handle_user_message(
 with gr.Blocks(title="HF Space Auto‑Builder") as demo:
     gr.Markdown("## Sign in + Auto‑Build Spaces\n\n"
                 "1. Sign in  2. Prompt  3. Deploy & Debug\n\n"
-                "_LLM will choose filenames, title, code, README, requirements, and iterate._\n\n---")
+                "_LLM controls filenames, code, README, requirements, and iterates until successful._\n\n---")
 
     login_btn = gr.LoginButton("huggingface", size="lg")
     status_md = gr.Markdown("*Not logged in.*")
@@ -203,11 +187,11 @@ with gr.Blocks(title="HF Space Auto‑Builder") as demo:
     run_box   = gr.Textbox(label="Run logs",   lines=5)
     preview   = gr.HTML("<p>No Space yet.</p>")
 
-    # Now pass user_in explicitly, so history isn’t empty
     send_btn.click(
         fn=handle_user_message,
         inputs=[chatbot, user_in, sdk_choice, api_key, grounding],
         outputs=[chatbot, build_box, run_box, preview]
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()