Update app.py

app.py

@@ -6,15 +6,30 @@ from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_st
 from google import genai
 from google.genai import types
 
-# --- Globals ---
-client = None
-chat = None
+# — USER INFO & MODEL LISTING —
+
+def show_profile(profile: gr.OAuthProfile | None) -> str:
+    if profile is None:
+        return "*Not signed in.*"
+    return f"✅ Signed in as **{profile.username}**"
+
+def list_private_models(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) -> str:
+    if profile is None or oauth_token is None:
+        return "Please sign in to HF."
+    models = list_models(author=profile.username, token=oauth_token.token)
+    if not models:
+        return "No models found."
+    return "Models:\n" + "\n".join(f"- {m.id} ({'private' if m.private else 'public'})" for m in models)
+
+# --- GLOBALS ---
+client: genai.Client | None = None
+chat: genai.chats.Chat | None = None
 
 # --- System instruction for Gemini ---
 system_instruction = (
     "You are a helpful assistant that writes, debugs, and pushes code to Hugging Face Spaces. "
-    "Treat Spaces as a sandbox: create, upload, debug. "
-    "Use function calling for logs and respond in JSON {success, data, message}."
+    "Treat Spaces as a sandbox: create spaces, upload code, and debug via function calling. "
+    "Always respond in JSON with {'success','data','message'}."
 )
 
 # --- Function declarations for logs ---
@@ -31,48 +46,46 @@ get_container_logs_decl = {
 tools = [ types.Tool(function_declarations=[get_build_logs_decl, get_container_logs_decl]) ]
 
 # --- HF helpers ---
-def create_space_backend(username, hf_token, repo_name, sdk):
+def create_space_backend(username: str, hf_token: str, repo_name: str, sdk: str) -> str:
     repo_id = f"{username}/{repo_name}"
     create_repo(repo_id=repo_id, token=hf_token, exist_ok=True, repo_type="space", space_sdk=sdk)
     return repo_id
 
-def fetch_logs(repo_id, level):
+def fetch_logs(repo_id: str, level: str) -> list[dict]:
     jwt_url = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
     r = get_session().get(jwt_url, headers=build_hf_headers())
     hf_raise_for_status(r)
     jwt = r.json()["token"]
-    url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
-    lines=[]
-    with get_session().get(url, headers=build_hf_headers(token=jwt), stream=True) as resp:
+    logs_url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
+    records = []
+    with get_session().get(logs_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())
-                    lines.append({"timestamp":ev.get("timestamp"),"message":ev.get("data")})
-                except: pass
-    return lines
+                    ev = json.loads(raw[len(b"data: "):].decode())
+                    records.append({"timestamp": ev.get("timestamp"), "message": ev.get("data")})
+                except:
+                    pass
+    return records
 
 # --- Chat init & respond ---
-def init_chat(repo_name, sdk, gemini_key, hf_profile, hf_token):
+def init_chat(repo_name: str, sdk: str, gemini_key: str, hf_profile: gr.OAuthProfile, hf_token: gr.OAuthToken):
     global client, chat
-    # Validate
     if hf_profile is None or hf_token is None:
-        return {"success":False,"data":None,"message":"Please sign in with HF."}, ""
+        return {"success": False, "data": None, "message": "Please sign in with Hugging Face."}, ""
     if not gemini_key:
-        return {"success":False,"data":None,"message":"Missing Gemini API key."}, ""
-    # create space
+        return {"success": False, "data": None, "message": "Missing Gemini API key."}, ""
     repo_id = create_space_backend(hf_profile.username, hf_token.token, repo_name, sdk)
     os.environ["HF_TOKEN"] = hf_token.token
-    # init Gemini
     client = genai.Client(api_key=gemini_key)
     chat = client.chats.create(
         model="gemini-2.5-flash-preview-04-17",
         config=types.GenerateContentConfig(system_instruction=system_instruction, tools=tools, temperature=0)
     )
-    return {"success":True,"data":None,"message":f"Sandbox ready: {repo_id}"}, repo_id
+    return {"success": True, "data": None, "message": f"Sandbox ready: {repo_id}"}, repo_id
 
-def chatbot_respond(message, history, repo_id, gemini_key):
+def chatbot_respond(message: str, history: list, repo_id: str, gemini_key: str):
     global chat
     if chat is None:
         history.append((None, "Error: chat not initialized."))
@@ -80,40 +93,43 @@ def chatbot_respond(message, history, repo_id, gemini_key):
     resp = chat.send_message(message)
     part = resp.candidates[0].content.parts[0]
     if part.function_call:
-        fn=part.function_call
-        args=json.loads(fn.args)
-        level = "build" if fn.name=="get_build_logs" else "run"
-        logs=fetch_logs(repo_id, level)
-        resp2 = chat.send_message("", function_response={fn.name:logs})
-        reply=resp2.candidates[0].content.parts[0].text
+        fn = part.function_call
+        args = json.loads(fn.args)
+        level = "build" if fn.name == "get_build_logs" else "run"
+        logs = fetch_logs(repo_id, level)
+        resp2 = chat.send_message("", function_response={fn.name: logs})
+        reply = resp2.candidates[0].content.parts[0].text
     else:
-        reply=part.text
+        reply = part.text
     history.append((message, reply))
     return history
 
 # --- UI ---
 with gr.Blocks() as demo:
     gr.Markdown("# HF Code Sandbox Chat")
-    # login
+
+    # --- Hugging Face Login ---
     login_btn = gr.LoginButton("Sign in with HF", variant="huggingface")
-    login_status = gr.Markdown("*Not signed in.*")
+    status_md = gr.Markdown("*Not signed in.*")
     models_md = gr.Markdown()
-    login_btn.click(lambda p: show_profile(p), inputs=[login_btn], outputs=[login_status])
-    login_btn.click(lambda p, t: list_private_models(p,t), inputs=[login_btn, login_btn.token], outputs=[models_md])
+    login_btn.click(show_profile, inputs=[login_btn], outputs=[status_md])
+    login_btn.click(list_private_models, inputs=[login_btn, login_btn.token], outputs=[models_md])
 
+    # --- Layout ---
     with gr.Row():
         with gr.Column(scale=2):
-            gr.Markdown("## Setup Sandbox")
+            gr.Markdown("### Setup Sandbox")
             gemini_key = gr.Textbox(label="Gemini API Key", type="password")
             repo_name = gr.Textbox(label="Space Name")
-            sdk = gr.Radio(choices=["gradio","streamlit"], label="SDK", value="gradio")
+            sdk = gr.Radio(choices=["gradio", "streamlit"], label="SDK", value="gradio")
             init_btn = gr.Button("Initialize Sandbox")
             init_status = gr.JSON()
             repo_store = gr.State("")
             init_btn.click(init_chat, inputs=[repo_name, sdk, gemini_key, login_btn, login_btn.token], outputs=[init_status, repo_store])
+
         with gr.Column(scale=8):
             chatbot = gr.Chatbot(type="messages")
-            user_input = gr.Textbox(show_label=False, placeholder="Ask to write/debug code...")
+            user_input = gr.Textbox(show_label=False, placeholder="Ask the sandbox to write/debug code...")
             user_input.submit(chatbot_respond, inputs=[user_input, chatbot, repo_store, gemini_key], outputs=[chatbot])
 
 if __name__ == "__main__":