app.py

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

# ——— 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"]
    # 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 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:
        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:
        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=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")
        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()