app.py

import re
import json
import os
import gradio as gr
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

# — USER INFO & MODEL LISTING —

def show_profile(profile):
    return f"✅ Logged in as **{profile.username}**" if profile else "*Not logged in.*"

def list_private_models(profile, token):
    if not (profile and token):
        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=token.token)
    ]
    return "No models found." if not models else "Models:\n\n" + "\n - ".join(models)

# — BUTTON‑ENABLING HELPERS —

def enable_create(profile, token):
    return gr.update(interactive=bool(profile and token))

def enable_repo_actions(repo_id, profile, token):
    return gr.update(interactive=bool(repo_id and profile and token))

# — CORE ACTIONS —

def create_space(repo_name, sdk, profile, token):
    if not (profile and token):
        return "", "⚠️ Please log in first.", "<p>No Space created yet.</p>"
    rid = f"{profile.username}/{repo_name}"
    create_repo(
        rid,
        token=token.token,
        exist_ok=True,
        repo_type="space",
        space_sdk=sdk
    )
    url = f"https://huggingface.co/spaces/{rid}"
    log = f"✅ Space ready: {url} (SDK: {sdk})"
    iframe = f'<iframe src="{url}" width="100%" height="400px"></iframe>'
    return rid, log, iframe

def write_file(path, content, repo_id, profile, token):
    if not (profile and token):
        return "⚠️ Please log in first."
    if not repo_id:
        return "⚠️ Please create a Space first."
    with open(path, "w") as f:
        f.write(content)
    upload_file(path, path, repo_id, token=token.token, repo_type="space")
    return f"✅ Wrote and uploaded `{path}`"

def list_files(repo_id, profile, token):
    if not (profile and token and repo_id):
        return "⚠️ Please log in and create a Space first."
    files = list_repo_files(repo_id, token=token.token, repo_type="space")
    return "\n".join(files)

def get_build_logs(repo_id, profile, token):
    if not (profile and token and repo_id):
        return "⚠️ Please log in and create a Space first."
    r = get_session().get(
        f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt",
        headers=build_hf_headers()
    )
    hf_raise_for_status(r)
    jwt = r.json()["token"]
    url = f"https://api.hf.space/v1/{repo_id}/logs/build"
    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: "):
                ev = json.loads(raw[6:].decode())
                lines.append(f"[{ev.get('timestamp')}] {ev.get('data')}")
    return "\n".join(lines)

def get_container_logs(repo_id, profile, token):
    if not (profile and token and repo_id):
        return "⚠️ Please log in and create a Space first."
    r = get_session().get(
        f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt",
        headers=build_hf_headers()
    )
    hf_raise_for_status(r)
    jwt = r.json()["token"]
    url = f"https://api.hf.space/v1/{repo_id}/logs/run"
    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: "):
                ev = json.loads(raw[6:].decode())
                lines.append(f"[{ev.get('timestamp')}] {ev.get('data')}")
    return "\n".join(lines)

# — FUNCTION DECLARATIONS —

func_decls = [
    {
        "name": "create_space",
        "description": "Create or get a HF Space",
        "parameters": {
            "type": "object",
            "properties": {
                "repo_name": {"type": "string"},
                "sdk": {"type": "string", "enum": ["gradio", "streamlit"]}
            },
            "required": ["repo_name", "sdk"]
        }
    },
    {
        "name": "write_file",
        "description": "Write or overwrite a file in the Space",
        "parameters": {
            "type": "object",
            "properties": {
                "path": {"type": "string"},
                "content": {"type": "string"}
            },
            "required": ["path", "content"]
        }
    },
    {
        "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 build logs",
        "parameters": {
            "type": "object",
            "properties": {"repo_id": {"type": "string"}},
            "required": ["repo_id"]
        }
    },
    {
        "name": "get_run_logs",
        "description": "Fetch run logs",
        "parameters": {
            "type": "object",
            "properties": {"repo_id": {"type": "string"}},
            "required": ["repo_id"]
        }
    }
]

# — CHAT HANDLER —

def process_message(profile, token, user_msg,
                    gemini_key, sidebar_repo, sidebar_sdk,
                    chat_history, session):
    # Initialize
    if session.get("chat") is None:
        client = genai.Client(api_key=gemini_key)
        cfg = types.GenerateContentConfig(
            system_instruction=(
                "You are a HF Spaces developer assistant. "
                "You MUST use function calls (create_space, write_file, list_files, "
                "get_build_logs, get_run_logs) to perform all actions—never return code inline."
            ),
            temperature=0,
            tools=[types.Tool(function_declarations=func_decls)],
            tool_config=types.ToolConfig(
                function_calling_config=types.FunctionCallingConfig(mode="ANY")
            )
        )
        session["chat"] = client.chats.create(
            model="gemini-2.0-flash", config=cfg
        )
        session["repo_id"] = None
        session["messages"] = []

    # Auto-create from “call it NAME”
    if session["repo_id"] is None:
        m = re.search(r"call (?:it )?([A-Za-z0-9_-]+)", user_msg, re.IGNORECASE)
        if m:
            name = m.group(1)
            rid, log, iframe = create_space(name, sidebar_sdk, profile, token)
            session["repo_id"] = rid
            session["iframe"] = iframe
            session["log"] = log
            session["files"] = ""
            session["messages"].append({
                "role": "assistant",
                "content": log
            })
            return session["messages"], iframe, log, "", session

    # Record user
    session["messages"].append({"role": "user", "content": user_msg})

    # Send to Gemini
    resp = session["chat"].send_message(user_msg)
    part = resp.candidates[0].content.parts[0]

    # Handle function call
    args = part.function_call.args
    if isinstance(args, str):
        args = json.loads(args)
    fn = part.function_call.name
    result = {}

    if fn == "create_space":
        rid, log, iframe = create_space(
            args["repo_name"], args["sdk"], profile, token
        )
        session["repo_id"] = rid
        result = {"log": log, "iframe": iframe}

    elif fn == "write_file":
        status = write_file(
            args["path"], args["content"],
            session["repo_id"], profile, token
        )
        result = {"status": status}

    elif fn == "list_files":
        result = {"files": list_files(
            session["repo_id"], profile, token
        )}

    elif fn == "get_build_logs":
        result = {"log": get_build_logs(
            session["repo_id"], profile, token
        )}

    elif fn == "get_run_logs":
        result = {"log": get_container_logs(
            session["repo_id"], profile, token
        )}

    else:
        result = {"log": f"⚠️ Unknown function {fn}"}

    # Respond via function_response
    session["chat"].send_message(
        types.Content(
            role="function",
            parts=[types.Part(
                function_call=part.function_call,
                function_response=json.dumps(result)
            )]
        )
    )
    assistant_text = session["chat"].get_history()[-1].parts[0].text

    # Record assistant
    session["messages"].append({
        "role": "assistant",
        "content": assistant_text
    })

    # Update panels
    if "iframe" in result: session["iframe"] = result["iframe"]
    if "log" in result:    session["log"] = result["log"]
    if "files" in result:  session["files"] = result["files"]

    return (
        session["messages"],
        session.get("iframe", ""),
        session.get("log", ""),
        session.get("files", ""),
        session
    )

# — SYNC MANUAL CHANGES —

def sync_manual(profile, token, session):
    if not (profile and token and session.get("repo_id")):
        return (
            session.get("iframe", ""),
            "⚠️ Cannot sync manual changes.",
            session.get("files", ""),
            session
        )
    fl = list_repo_files(
        session["repo_id"], token=token.token, repo_type="space"
    )
    session["files"] = "\n".join(fl)
    session["log"] = "🔄 Manual changes synced."
    return (
        session["iframe"],
        session["log"],
        session["files"],
        session
    )

# — BUILD THE UI —

with gr.Blocks(css="""
#sidebar { background:#f2f2f2; padding:1rem; border-right:1px solid #ccc; }
#main    { padding:1rem; }
""") as demo:

    with gr.Row():
        # Sidebar
        with gr.Column(elem_id="sidebar", scale=1):
            gr.Markdown("### 🔑 HF Login & Config")
            login_btn = gr.LoginButton(variant="huggingface", size="sm")
            profile_state = gr.State(None)
            token_state = gr.State(None)

            login_btn.click(None, [], [profile_state, token_state])

            status_md = gr.Markdown("*Not logged in.*")
            profile_state.change(
                show_profile, inputs=[profile_state], outputs=[status_md]
            )

            models_md = gr.Markdown()
            profile_state.change(
                list_private_models,
                inputs=[profile_state, token_state],
                outputs=[models_md]
            )

            gemini_key = gr.Textbox(
                label="Gemini API Key", type="password"
            )
            sidebar_repo = gr.Textbox(
                label="Space name", placeholder="my-space"
            )
            sidebar_sdk = gr.Radio(
                ["gradio", "streamlit"], value="gradio", label="SDK"
            )

            gr.Markdown("---")
            confirm_btn = gr.Button("🔄 Confirm Manual Changes")

        # Main area
        with gr.Column(elem_id="main", scale=3):
            tabs = gr.Tabs()
            with tabs:
                with gr.TabItem("💬 Chat"):
                    chatbox = gr.Chatbot(type="messages")
                    user_input = gr.Textbox(
                        show_label=False, placeholder="Ask the LLM…"
                    )
                    send_btn = gr.Button("Send")
                with gr.TabItem("🛠️ Manual"):
                    gr.Markdown("#### Manual Controls")
                    # (Repeat buttons for create/upload/logs as needed)

            gr.Markdown("---")
            iframe_out = gr.HTML(label="🖼️ Preview")
            log_out = gr.Textbox(label="📋 Latest Log", lines=4)
            files_out = gr.Textbox(label="📚 Files", lines=4)

    state = gr.State({})
    send_btn.click(
        process_message,
        inputs=[
            profile_state, token_state, user_input,
            gemini_key, sidebar_repo, sidebar_sdk,
            chatbox, state
        ],
        outputs=[chatbox, iframe_out, log_out, files_out, state]
    )

    confirm_btn.click(
        sync_manual,
        inputs=[profile_state, token_state, state],
        outputs=[iframe_out, log_out, files_out, state]
    )

if __name__ == "__main__":
    demo.launch()