Update app.py

app.py

@@ -1,138 +1,194 @@
 import os, json
 import gradio as gr
-from huggingface_hub import create_repo, upload_file, list_repo_files, constants
+import requests
+from huggingface_hub import (
+    create_repo,
+    list_models,
+    upload_file,
+    list_repo_files,
+    constants,
+)
 from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status
 from google import genai
 from google.genai import types
 
-# ——— HELPERS FOR HF SPACES ———
+# — USER INFO & MODEL LISTING (UNCHANGED) —
 
-def create_space(repo_name: str, sdk: str, hf_token: str):
-    """Create or get a HF Space, return (repo_id, log, iframe)."""
-    token = hf_token or os.getenv("HF_TOKEN", "")
-    if not token:
-        return None, "⚠️ No HF token provided.", ""
-    os.environ["HF_TOKEN"] = token
-    repo_id = f"{os.getenv('HF_USERNAME', '')}/{repo_name}"
+def show_profile(profile: gr.OAuthProfile | None) -> str:
+    if profile is None:
+        return "*Not logged in.*"
+    return f"✅ Logged 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 log in to see your models."
+    models = [
+        f"{m.id} ({'private' if m.private else 'public'})"
+        for m in list_models(author=profile.username, token=oauth_token.token)
+    ]
+    return "No models found." if not models else "Models:\n\n" + "\n - ".join(models)
+
+def enable_create(
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None
+):
+    return gr.update(interactive=profile is not None)
+
+def enable_repo_actions(
+    repo_id: str,
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None
+):
+    return gr.update(interactive=bool(repo_id and profile and oauth_token))
+
+def create_space(
+    repo_name: str,
+    sdk: str,
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None
+) -> tuple[str, str, str]:
+    if not profile or not oauth_token:
+        return "", "⚠️ Please log in first.", "<p>No Space created yet.</p>"
+    repo_id = f"{profile.username}/{repo_name}"
     create_repo(
         repo_id=repo_id,
-        token=token,
+        token=oauth_token.token,
         exist_ok=True,
         repo_type="space",
-        space_sdk=sdk,
+        space_sdk=sdk
     )
-    url = f"https://huggingface.co/spaces/{repo_id}"
-    return repo_id, f"✅ Space ready: {url} (SDK={sdk})", f'<iframe src="{url}" width="100%" height="400px"></iframe>'
-
-def upload_file_to_space(repo_id: str, local_path: str, hf_token: str):
-    """Upload a local file to the Space."""
-    token = hf_token or os.getenv("HF_TOKEN", "")
-    if not (repo_id and token and os.path.exists(local_path)):
-        return "⚠️ Missing repo_id, token, or file not found."
-    os.environ["HF_TOKEN"] = token
+    url    = f"https://huggingface.co/spaces/{repo_id}"
+    logmsg = f"✅ Space ready: {url} (SDK: {sdk})"
+    iframe = f'<iframe src="{url}" width="100%" height="500px"></iframe>'
+    return repo_id, logmsg, iframe
+
+def upload_file_to_space(
+    file,
+    path_in_repo: str,
+    repo_id: str,
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None
+) -> str:
+    if not profile or not oauth_token:
+        return "⚠️ Please log in first."
+    if not repo_id:
+        return "⚠️ Please create a Space first."
+    if not file:
+        return "⚠️ No file selected."
     upload_file(
-        path_or_fileobj=local_path,
-        path_in_repo=os.path.basename(local_path),
+        path_or_fileobj=file.name,
+        path_in_repo=path_in_repo,
         repo_id=repo_id,
-        token=token,
-        repo_type="space",
+        token=oauth_token.token,
+        repo_type="space"
     )
-    return f"✅ Uploaded `{os.path.basename(local_path)}` to `{repo_id}`"
-
-def fetch_logs(repo_id: str, level: str, hf_token: str):
-    """Stream build or run logs."""
-    token = hf_token or os.getenv("HF_TOKEN", "")
-    if not (repo_id and token):
-        return "⚠️ Missing repo_id or HF token."
-    # 1) get JWT
+    return f"✅ Uploaded `{path_in_repo}` to `{repo_id}`"
+
+def _fetch_space_logs_level(repo_id: str, level: str) -> str:
     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"]
-    # 2) stream
-    url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
+    logs_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:
+    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 not raw.startswith(b"data: "): continue
-            ev = json.loads(raw[len(b"data: "):])
-            ts, txt = ev.get("timestamp",""), ev.get("data","")
+            if not raw.startswith(b"data: "):
+                continue
+            payload = raw[len(b"data: "):]
+            try:
+                event = json.loads(payload.decode())
+            except:
+                continue
+            ts  = event.get("timestamp", "")
+            txt = event.get("data", "")
             lines.append(f"[{ts}] {txt}")
     return "\n".join(lines)
 
-def list_files(repo_id: str, hf_token: str):
-    """List files in the Space repo."""
-    token = hf_token or os.getenv("HF_TOKEN", "")
-    if not (repo_id and token):
-        return "⚠️ Missing repo_id or HF token."
-    return "\n".join(list_repo_files(repo_id, token=token, repo_type="space"))
+def get_build_logs(
+    repo_id: str,
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None
+) -> str:
+    if not (profile and oauth_token and repo_id):
+        return "⚠️ Please log in and create a Space first."
+    return _fetch_space_logs_level(repo_id, "build")
+
+def get_container_logs(
+    repo_id: str,
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None
+) -> str:
+    if not (profile and oauth_token and repo_id):
+        return "⚠️ Please log in and create a Space first."
+    return _fetch_space_logs_level(repo_id, "run")
 
 
-# ——— GEMINI FUNCTION DECLARATIONS ———
+# — GEMINI FUNCTION DECLARATIONS — 
 
 func_decls = [
     {
-        "name": "create_space",
-        "description": "Create or get a Hugging Face Space",
-        "parameters": {
-            "type":"object",
-            "properties":{
-                "repo_name":{"type":"string"},
-                "sdk":{"type":"string","enum":["gradio","streamlit"]}
-            },
-            "required":["repo_name","sdk"]
-        }
-    },
-    {
-        "name":"upload_file",
-        "description":"Upload a file to the Space",
-        "parameters":{
-            "type":"object",
-            "properties":{
-                "repo_id":{"type":"string"},
-                "local_path":{"type":"string"}
-            },
-            "required":["repo_id","local_path"]
-        }
+      "name": "create_space",
+      "description": "Create or get a Hugging Face Space",
+      "parameters": {
+        "type": "object",
+        "properties": {
+          "repo_name": {"type":"string"},
+          "sdk":       {"type":"string","enum":["gradio","streamlit"]},
+        },
+        "required": ["repo_name","sdk"]
+      }
     },
     {
-        "name":"list_files",
-        "description":"List files in the Space",
-        "parameters":{
-            "type":"object",
-            "properties":{"repo_id":{"type":"string"}},
-            "required":["repo_id"]
-        }
+      "name":"list_files",
+      "description":"List files in the Space",
+      "parameters":{
+        "type":"object",
+        "properties":{"repo_id":{"type":"string"}},
+        "required":["repo_id"]
+      }
     },
     {
-        "name":"get_build_logs",
-        "description":"Fetch the build logs",
-        "parameters":{
-            "type":"object",
-            "properties":{"repo_id":{"type":"string"}},
-            "required":["repo_id"]
-        }
+      "name": "get_build_logs",
+      "description": "Fetch the build logs",
+      "parameters":{
+        "type":"object",
+        "properties":{"repo_id":{"type":"string"}},
+        "required":["repo_id"]
+      }
     },
     {
-        "name":"get_run_logs",
-        "description":"Fetch the run logs",
-        "parameters":{
-            "type":"object",
-            "properties":{"repo_id":{"type":"string"}},
-            "required":["repo_id"]
-        }
+      "name": "get_run_logs",
+      "description": "Fetch the run logs",
+      "parameters":{
+        "type":"object",
+        "properties":{"repo_id":{"type":"string"}},
+        "required":["repo_id"]
+      }
     },
 ]
 
-# ——— CHAT PROCESSING ———
+# — CHAT HANDLER — 
 
-def process_message(user_msg, gemini_key, repo_name, sdk, hf_token, chat_history, session):
-    # init Gemini chat
+def process_message(
+    user_msg: str,
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None,
+    gemini_key: str,
+    repo_name: str,
+    sdk: str,
+    chat_history,
+    session
+):
+    # init
     if session.get("chat") is None:
         client = genai.Client(api_key=gemini_key)
         cfg = types.GenerateContentConfig(
-            system_instruction="You are a HF Spaces assistant. Use functions to manage spaces.",
+            system_instruction="You are an admin for HF Spaces: use functions to create spaces, list files, and fetch logs.",
             temperature=0,
             tools=[ types.Tool(function_declarations=func_decls) ]
         )
@@ -140,32 +196,31 @@ def process_message(user_msg, gemini_key, repo_name, sdk, hf_token, chat_history
         session["repo_id"] = None
 
     chat = session["chat"]
-    # record user
+    # user turn
     chat_history = chat_history + [(user_msg, None)]
     resp = chat.send_message(user_msg)
     part = resp.candidates[0].content.parts[0]
 
-    # if function call requested
+    # function call?
     if part.function_call:
         name = part.function_call.name
         args = json.loads(part.function_call.args)
         # dispatch
         if name=="create_space":
-            rid, log, iframe = create_space(args["repo_name"], args["sdk"], hf_token)
+            rid, log, iframe = create_space(args["repo_name"], args["sdk"], profile, oauth_token)
             session["repo_id"] = rid
             result = {"repo_id":rid,"log":log,"iframe":iframe}
-        elif name=="upload_file":
-            result = {"log": upload_file_to_space(args["repo_id"], args["local_path"], hf_token)}
         elif name=="list_files":
-            result = {"files": list_files(args["repo_id"], hf_token)}
+            files = list_repo_files(session["repo_id"], token=oauth_token.token, repo_type="space") if session["repo_id"] else []
+            result = {"files":"\n".join(files)}
         elif name=="get_build_logs":
-            result = {"log": fetch_logs(args["repo_id"], "build", hf_token)}
+            result = {"log":get_build_logs(session["repo_id"], profile, oauth_token)}
         elif name=="get_run_logs":
-            result = {"log": fetch_logs(args["repo_id"], "run", hf_token)}
+            result = {"log":get_container_logs(session["repo_id"], profile, oauth_token)}
         else:
             result = {"log":f"⚠️ Unknown function {name}"}
 
-        # feed function result back to chat
+        # send back
         chat.send_message(
             types.Content(
                 role="function",
@@ -176,41 +231,94 @@ def process_message(user_msg, gemini_key, repo_name, sdk, hf_token, chat_history
     else:
         final = part.text
 
-    # update panels
+    # pick up panels
     iframe = session.get("iframe", "")
     log    = session.get("log", "")
     files  = session.get("files", "")
 
     if part.function_call:
         session["iframe"] = result.get("iframe", iframe)
-        session["log"]    = result.get("log", log)
-        session["files"]  = result.get("files", files)
+        session["log"]    = result.get("log",    log)
+        session["files"]  = result.get("files",  files)
 
     chat_history[-1] = (user_msg, final)
     return chat_history, session["iframe"], session["log"], session["files"], session
 
-# ——— GRADIO INTERFACE ———
+# — BUILD THE UI — 
 
-with gr.Blocks(title="HF Spaces → Gemini Chat") as demo:
+with gr.Blocks(title="HF Spaces + Gemini Chat") as demo:
     with gr.Row():
+        # ◀ Sidebar: login, models, chat config
         with gr.Column(scale=1):
-            gemini_key = gr.Textbox(label="Gemini API Key", type="password")
-            repo_name  = gr.Textbox(label="Space name", placeholder="my-space")
-            sdk_choice = gr.Radio(["gradio","streamlit"], value="gradio", label="SDK")
-            hf_token   = gr.Textbox(label="HF Token (opt)", type="password")
+            gr.Markdown("## 🔑 Sign in & Chat Config")
+            login_btn   = gr.LoginButton(variant="huggingface", size="sm")
+            status_md   = gr.Markdown("*Not logged in.*")
+            models_md   = gr.Markdown()
+            demo.load(show_profile, inputs=None, outputs=status_md)
+            login_btn.click(show_profile, inputs=None, outputs=status_md)
+            demo.load(list_private_models, inputs=None, outputs=models_md)
+            login_btn.click(list_private_models, inputs=None, outputs=models_md)
+
+            gemini_key  = gr.Textbox(label="Gemini API Key", type="password")
+            repo_name   = gr.Textbox(label="Space name", placeholder="e.g. my-space")
+            sdk_choice  = gr.Radio(["gradio","streamlit"], value="gradio", label="SDK")
+        # ▶ Main: chat + panels + manual UI
         with gr.Column(scale=3):
-            chatbox    = gr.Chatbot()
-            user_input = gr.Textbox(show_label=False, placeholder="e.g. Create Space")
-            send_btn   = gr.Button("Send")
+            gr.Markdown("## 💬 Chat to manage your Space")
+            chatbox   = gr.Chatbot()
+            user_in   = gr.Textbox(show_label=False, placeholder="e.g. Create Space")
+            send_btn  = gr.Button("Send")
+
+            iframe_out = gr.HTML(label="🖼️ Preview")
+            log_out    = gr.Textbox(label="📋 Latest Log", lines=8)
+            files_out  = gr.Textbox(label="📚 Files in Space", lines=4)
+
+            gr.Markdown("---\n## 🛠️ Manual Controls (unchanged)")
+            # — Manual CREATE SPACE —
+            repo_name_m   = gr.Textbox(label="New Space name", placeholder="my-space")
+            sdk_selector  = gr.Radio(["gradio","streamlit"], value="gradio", label="SDK for manual")
+            create_btn    = gr.Button("Create Space", interactive=False)
+            session_id    = gr.Textbox(visible=False)
+            create_logs   = gr.Textbox(label="Create Logs", interactive=False, lines=3)
+            preview_m     = gr.HTML("<p>No Space created yet.</p>")
+
+            demo.load(enable_create, inputs=None, outputs=[create_btn])
+            login_btn.click(enable_create, inputs=None, outputs=[create_btn])
+            create_btn.click(
+                fn=create_space,
+                inputs=[repo_name_m, sdk_selector, status_md, login_btn.token],
+                outputs=[session_id, create_logs, preview_m]
+            )
+
+            # — Manual UPLOAD FILES —
+            path_in_repo   = gr.Textbox(label="Path in Space", value="app.py")
+            file_uploader  = gr.File(label="Select file")
+            upload_btn     = gr.Button("Upload File", interactive=False)
+            upload_logs    = gr.Textbox(label="Upload Logs", interactive=False, lines=2)
+
+            demo.load(enable_repo_actions, inputs=[session_id, status_md, login_btn.token], outputs=[upload_btn])
+            upload_btn.click(
+                fn=upload_file_to_space,
+                inputs=[file_uploader, path_in_repo, session_id, status_md, login_btn.token],
+                outputs=[upload_logs]
+            )
+
+            # — Manual FETCH LOGS —
+            build_btn      = gr.Button("Get Build Logs", interactive=False)
+            run_btn        = gr.Button("Get Container Logs", interactive=False)
+            build_logs_md  = gr.Textbox(label="Build Logs", interactive=False, lines=10)
+            run_logs_md    = gr.Textbox(label="Container Logs", interactive=False, lines=10)
+
+            for b in (build_btn, run_btn):
+                demo.load(enable_repo_actions, inputs=[session_id, status_md, login_btn.token], outputs=[b])
 
-            iframe_out = gr.HTML(label="Space Preview")
-            log_out    = gr.Textbox(label="Latest Log", lines=8)
-            files_out  = gr.Textbox(label="Files in Space", lines=4)
+            build_btn.click(fn=get_build_logs, inputs=[session_id, status_md, login_btn.token], outputs=[build_logs_md])
+            run_btn.click(  fn=get_container_logs, inputs=[session_id, status_md, login_btn.token], outputs=[run_logs_md])
 
     state = gr.State({})
     send_btn.click(
         fn=process_message,
-        inputs=[user_input, gemini_key, repo_name, sdk_choice, hf_token, chatbox, state],
+        inputs=[user_in, status_md, login_btn.token, gemini_key, repo_name, sdk_choice, chatbox, state],
         outputs=[chatbox, iframe_out, log_out, files_out, state]
     )