app.py

import re, json, 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):
    """
    Write or overwrite any file in the Space. 
    The LLM can call this for app.py, requirements.txt, README.md, etc.
    """
    if not (profile and token):
        return "⚠️ Please log in first."
    if not repo_id:
        return "⚠️ Please create a Space first."
    # Save locally
    with open(path, "w") as f:
        f.write(content)
    # Upload to HF Space
    upload_file(path, path, repo_id,
                token=token.token, repo_type="space")
    return f"✅ Wrote and uploaded `{path}`"

def upload_file_to_space(file, path, repo_id, profile, token):
    if not (profile and 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(file.name, path, repo_id,
                token=token.token, repo_type="space")
    return f"✅ Uploaded `{path}` to `{repo_id}`"

def _fetch_space_logs_level(repo_id, level):
    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/{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: "):
                ev = json.loads(raw[6:].decode())
                lines.append(f"[{ev.get('timestamp','')}] {ev.get('data','')}")
    return "\n".join(lines)

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."
    return _fetch_space_logs_level(repo_id, "build")

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."
    return _fetch_space_logs_level(repo_id, "run")

# — GEMINI FUNCTION DECLARATIONS —

func_decls = [
    {
      "name":"create_space","description":"Create/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/overwrite a file in the Space",
      "parameters":{
        "type":"object",
        "properties":{
          "path":{"type":"string","description":"File path, e.g. app.py or requirements.txt"},
          "content":{"type":"string","description":"The full text content to write to the file."}
        },
        "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 the chat
    if session.get("chat") is None:
        client = genai.Client(api_key=gemini_key)
        cfg = types.GenerateContentConfig(
            system_instruction=(
                "You are a HF Spaces admin. You can create spaces, "
                "write files (app.py, requirements.txt, README.md...), "
                "list files, and fetch logs."
            ),
            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
        session["messages"]= []

    # Allow “call it NAME” in free chat
    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":f"✅ Created Space `{name}`. {log}"
            })
            return (session["messages"],
                    iframe, log, "", session)

    # Record user message
    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]

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

        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":
            fl = list_repo_files(
                session["repo_id"], token=token.token, repo_type="space"
            )
            result = {"files":"\n".join(fl)}

        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}"}

        # **Always** pass a JSON string here
        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
    else:
        assistant_text = part.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.get("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("#### Create a Space")
                    repo_m     = gr.Textbox(label="Name")
                    sdk_m      = gr.Radio(["gradio","streamlit"],
                                           value="gradio", label="SDK")
                    create_btn = gr.Button("Create Space")
                    sess_id    = gr.Textbox(visible=False)
                    log_c      = gr.Textbox(label="Log",
                                            interactive=False, lines=2)
                    preview    = gr.HTML("<p>No Space yet.</p>")

                    create_btn.click(create_space,
                                     inputs=[repo_m, sdk_m,
                                             profile_state, token_state],
                                     outputs=[sess_id, log_c, preview])

                    gr.Markdown("#### Upload File")
                    path        = gr.Textbox(label="Path", value="app.py")
                    file_u      = gr.File()
                    upload_btn  = gr.Button("Upload File")
                    log_u       = gr.Textbox(label="Log",
                                              interactive=False, lines=2)

                    upload_btn.click(upload_file_to_space,
                                     inputs=[file_u, path, sess_id,
                                             profile_state, token_state],
                                     outputs=[log_u])

                    gr.Markdown("#### Fetch Logs")
                    b_btn = gr.Button("Build Logs")
                    r_btn = gr.Button("Run Logs")
                    log_b = gr.Textbox(label="Build",
                                       interactive=False, lines=5)
                    log_r = gr.Textbox(label="Run",
                                       interactive=False, lines=5)

                    b_btn.click(get_build_logs,
                                inputs=[sess_id,
                                        profile_state, token_state],
                                outputs=[log_b])
                    r_btn.click(get_container_logs,
                                inputs=[sess_id,
                                        profile_state, token_state],
                                outputs=[log_r])

            # Persistent panels
            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()