Update app.py

app.py

@@ -1,136 +1,218 @@
-import os
-import json
+import os, json
 import gradio as gr
-from huggingface_hub import create_repo, list_models, upload_file, constants
+from huggingface_hub import create_repo, 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
 
-# — 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 spaces, upload code, and debug via function calling. "
-    "Always respond in JSON with {'success','data','message'}."
-)
-
-# --- Function declarations for logs ---
-get_build_logs_decl = {
-    "name": "get_build_logs",
-    "description": "Fetches build logs for a Space",
-    "parameters": {"type":"object","properties":{"repo_id":{"type":"string"}},"required":["repo_id"]}
-}
-get_container_logs_decl = {
-    "name": "get_container_logs",
-    "description": "Fetches container logs for a Space",
-    "parameters": {"type":"object","properties":{"repo_id":{"type":"string"}},"required":["repo_id"]}
-}
-tools = [ types.Tool(function_declarations=[get_build_logs_decl, get_container_logs_decl]) ]
-
-# --- HF helpers ---
-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: str, level: str) -> list[dict]:
+# ——— HELPERS FOR HF SPACES ———
+
+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}"
+    create_repo(
+        repo_id=repo_id,
+        token=token,
+        exist_ok=True,
+        repo_type="space",
+        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
+    upload_file(
+        path_or_fileobj=local_path,
+        path_in_repo=os.path.basename(local_path),
+        repo_id=repo_id,
+        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
     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"]
-    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:
+    # 2) stream
+    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:
         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())
-                    records.append({"timestamp": ev.get("timestamp"), "message": ev.get("data")})
-                except:
-                    pass
-    return records
-
-# --- Chat init & respond ---
-def init_chat(repo_name: str, sdk: str, gemini_key: str, hf_profile: gr.OAuthProfile, hf_token: gr.OAuthToken):
-    global client, chat
-    if hf_profile is None or hf_token is None:
-        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."}, ""
-    repo_id = create_space_backend(hf_profile.username, hf_token.token, repo_name, sdk)
-    os.environ["HF_TOKEN"] = hf_token.token
-    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
-
-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."))
-        return history
-    resp = chat.send_message(message)
+            if not raw.startswith(b"data: "): continue
+            ev = json.loads(raw[len(b"data: "):])
+            ts, txt = ev.get("timestamp",""), ev.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"))
+
+
+# ——— 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":"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_run_logs",
+        "description":"Fetch the run logs",
+        "parameters":{
+            "type":"object",
+            "properties":{"repo_id":{"type":"string"}},
+            "required":["repo_id"]
+        }
+    },
+]
+
+# ——— CHAT PROCESSING ———
+
+def process_message(user_msg, gemini_key, repo_name, sdk, hf_token, chat_history, session):
+    # init Gemini chat
+    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.",
+            temperature=0,
+            tools=[ types.Tool(function_declarations=func_decls) ]
+        )
+        session["chat"] = client.chats.create(model="gemini-2.0-flash", config=cfg)
+        session["repo_id"] = None
+
+    chat = session["chat"]
+    # record user
+    chat_history = chat_history + [(user_msg, None)]
+    resp = chat.send_message(user_msg)
     part = resp.candidates[0].content.parts[0]
+
+    # if function call requested
     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
+        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)
+            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)}
+        elif name=="get_build_logs":
+            result = {"log": fetch_logs(args["repo_id"], "build", hf_token)}
+        elif name=="get_run_logs":
+            result = {"log": fetch_logs(args["repo_id"], "run", hf_token)}
+        else:
+            result = {"log":f"⚠️ Unknown function {name}"}
+
+        # feed function result back to chat
+        chat.send_message(
+            types.Content(
+                role="function",
+                parts=[ types.Part(function_call=part.function_call, function_response=json.dumps(result)) ]
+            )
+        )
+        final = chat.get_history()[-1].parts[0].text
     else:
-        reply = part.text
-    history.append((message, reply))
-    return history
-
-# --- UI ---
-with gr.Blocks() as demo:
-    gr.Markdown("# HF Code Sandbox Chat")
-
-    # --- Hugging Face Login ---
-    login_btn = gr.LoginButton("Sign in with HF", variant="huggingface")
-    status_md = gr.Markdown("*Not signed in.*")
-    models_md = gr.Markdown()
-    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 ---
+        final = part.text
+
+    # update 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)
+
+    chat_history[-1] = (user_msg, final)
+    return chat_history, session["iframe"], session["log"], session["files"], session
+
+# ——— GRADIO INTERFACE ———
+
+with gr.Blocks(title="HF Spaces → Gemini Chat") as demo:
     with gr.Row():
-        with gr.Column(scale=2):
-            gr.Markdown("### Setup Sandbox")
+        with gr.Column(scale=1):
             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")
-            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 the sandbox to write/debug code...")
-            user_input.submit(chatbot_respond, inputs=[user_input, chatbot, repo_store, gemini_key], outputs=[chatbot])
+            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")
+        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")
+
+            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)
+
+    state = gr.State({})
+    send_btn.click(
+        fn=process_message,
+        inputs=[user_input, gemini_key, repo_name, sdk_choice, hf_token, chatbox, state],
+        outputs=[chatbox, iframe_out, log_out, files_out, state]
+    )
 
 if __name__ == "__main__":
     demo.launch()