Update app.py

app.py

@@ -22,7 +22,7 @@ def list_private_models(
     ]
     return "No models found." if not models else "Models:\n\n" + "\n - ".join(models)
 
-# — HELPERS FOR HF SPACE LOGS —
+# — LOG FETCHING —
 
 def _get_space_jwt(repo_id: str):
     url = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
@@ -45,7 +45,7 @@ def fetch_logs(repo_id: str, level: str):
                     continue
     return "\n".join(lines)
 
-# — CORE LOOP: ask LLM for JSON, write & deploy —
+# — CORE LOOP: ask LLM for structured JSON, write & deploy —
 
 def handle_user_message(
     history,               # list of {"role","content"} dicts
@@ -60,57 +60,75 @@ def handle_user_message(
 
     client = genai.Client(api_key=gemini_api_key)
 
-    # System prompt: instruct JSON output
+    # System prompt: ask for exact JSON per schema
     system_msg = {
         "role":"system",
         "content":(
-            "You are an AI assistant that generates all files needed for a HuggingFace Space.\n"
-            "Based on the user's prompt, return a **single JSON object** with keys:\n"
-            "  • repo_name: string – the name of the new Space repo (no username).\n"
-            "  • files: object mapping filenames to file contents (strings).  \n"
-            "    Required files: one code file (preferably app.py; if you choose a different name, update README.md),\n"
-            "    README.md frontmatter (with title, sdk, sdk_version, app_file), and requirements.txt.\n"
-            "Do **NOT** return any explanatory text—only that JSON object.\n"
+            "Generate **one JSON object** with keys:\n"
+            "  • repo_name: string (the Space name, no username)\n"
+            "  • files: object mapping filenames to file-contents strings\n\n"
+            "Include at least:\n"
+            "  - A code file (use app.py unless you explicitly choose another)\n"
+            "  - requirements.txt with all dependencies\n"
+            "  - README.md frontmatter (title, emoji, sdk, sdk_version, app_file, etc.)\n\n"
+            "Do NOT output any other text—only the JSON.\n"
+            f"The user wants a **{sdk_choice}** app with this prompt: {history[-1][1]}"
         )
     }
     chat = [system_msg] + history
 
+    # JSON schema for the response
+    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):
-        # dynamic version detection
-        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__
 
-        # ask the LLM
         tools = [Tool(google_search=GoogleSearch())] if grounding_enabled else []
-        cfg   = GenerateContentConfig(tools=tools, response_modalities=["TEXT"])
-        resp  = client.models.generate_content(
+        cfg = GenerateContentConfig(
+            tools=tools,
+            response_modalities=["TEXT"],
+            response_mime_type="application/json",
+            response_schema=schema
+        )
+
+        resp = client.models.generate_content(
             model="gemini-2.5-flash-preview-04-17",
             contents=[m["content"] for m in chat],
             config=cfg
         )
-        ai_text = resp.text
-        chat.append({"role":"assistant","content":ai_text})
+        ai_json = resp.text
 
-        # parse JSON
+        # parse JSON only
         try:
-            spec = json.loads(ai_text)
+            spec = json.loads(ai_json)
             repo_name = spec["repo_name"]
             files     = spec["files"]
-        except Exception as e:
-            # ask it to re-output valid JSON
-            chat.append({"role":"user","content":
-                         "Your last response wasn't valid JSON. "
-                         "Please reply with exactly the JSON object as specified."})
+        except Exception:
+            chat.append({
+                "role":"user",
+                "content":"The JSON was invalid. Please reply with exactly the JSON object per schema."
+            })
             continue
 
         repo_id = f"{profile.username}/{repo_name}"
-
-        # write & upload each file
+        # ensure Space exists
         create_repo(
             repo_id=repo_id,
             token=oauth_token.token,
@@ -118,9 +136,11 @@ def handle_user_message(
             repo_type="space",
             space_sdk=sdk_choice
         )
+
+        # write & upload all files
         for fn, content in files.items():
-            # fill in sdk_version into README if needed
-            if fn.lower() == "readme.md":
+            # inject the true sdk_version if placeholder present
+            if fn.lower()=="readme.md":
                 content = content.replace("<SDK_VERSION>", sdk_version)
             with open(fn, "w") as f:
                 f.write(content)
@@ -136,21 +156,20 @@ def handle_user_message(
         build_logs = fetch_logs(repo_id, "build")
         run_logs   = fetch_logs(repo_id, "run")
 
+        # break on clean build
         if "ERROR" not in build_logs.upper() and "ERROR" not in run_logs.upper():
             break
 
-        # feed errors back
         chat.append({
             "role":"user",
             "content":(
-                f"Deployment produced build logs:\n{build_logs}\n\n"
-                f"and run logs:\n{run_logs}\n\n"
-                "Please fix the code and output a new JSON spec."
+                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."
             )
         })
         time.sleep(2)
 
-    # prepare outputs
     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
@@ -159,14 +178,14 @@ 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 generate code, README, requirements, and iterate until successful.*\n\n---")
+                "1. Sign in  2. Prompt  3. Deploy & Debug\n"
+                "_The LLM controls filenames, title, README, requirements, and code._\n\n---")
 
     login_btn = gr.LoginButton("huggingface", size="lg")
     status_md = gr.Markdown("*Not logged in.*")
     models_md = gr.Markdown()
-    demo.load(show_profile,         None, status_md)
-    demo.load(list_private_models,  None, models_md)
+    demo.load(show_profile,        None, status_md)
+    demo.load(list_private_models, None, models_md)
     login_btn.click(show_profile,        None, status_md)
     login_btn.click(list_private_models, None, models_md)