app.py

import re
import json
import time
import importlib.metadata

import gradio as gr
from huggingface_hub import create_repo, upload_file, list_models, constants
from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status
from google import genai
from google.genai.types import Tool, GenerateContentConfig, GoogleSearch

# — USER INFO & MODEL LISTING —

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)

# — UTILITIES — 

def get_sdk_version(sdk_choice: str) -> str:
    pkg = "gradio" if sdk_choice == "gradio" else "streamlit"
    try:
        return importlib.metadata.version(pkg)
    except importlib.metadata.PackageNotFoundError:
        return "UNKNOWN"

def extract_code(text: str) -> str:
    """
    Extract the last ```…``` block. If none, return whole text.
    """
    blocks = re.findall(r"```(?:\w*\n)?([\s\S]*?)```", text)
    return blocks[-1].strip() if blocks else text.strip()

# — HF SPACE LOGGING — 

def _get_space_jwt(repo_id: str):
    url = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
    r = get_session().get(url, headers=build_hf_headers())
    hf_raise_for_status(r)
    return r.json()["token"]

def fetch_logs(repo_id: str, level: str) -> str:
    jwt = _get_space_jwt(repo_id)
    logs_url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
    lines = []
    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 raw.startswith(b"data: "):
                try:
                    ev = json.loads(raw[len(b"data: "):].decode())
                    ts  = ev.get("timestamp","")
                    txt = ev.get("data","")
                    lines.append(f"[{ts}] {txt}")
                except:
                    continue
    return "\n".join(lines)

# — CORE LOOP — 

def handle_user_message(
    history,
    sdk_choice: str,
    gemini_api_key: str,
    grounding_enabled: bool,
    space_suffix: str,
    profile: gr.OAuthProfile | None,
    oauth_token: gr.OAuthToken | None
):
    if profile is None or oauth_token is None:
        return (
            history + [{"role":"assistant","content":"⚠️ Please log in first."}],
            "",
            "",
            "<p>No Space yet.</p>"
        )

    username = profile.username
    repo_name = f"{username}-{space_suffix}"
    repo_id   = f"{username}/{repo_name}"

    client = genai.Client(api_key=gemini_api_key)
    system_msg = {
        "role":"system",
        "content":(
            f"You are an AI assistant writing a HuggingFace Space using the "
            f"{sdk_choice} SDK. After producing code, wait for logs; if errors appear, fix them."
        )
    }
    chat = [system_msg] + history

    code_fn   = "app.py" if sdk_choice=="gradio" else "streamlit_app.py"
    readme_fn = "README.md"
    reqs_fn   = "requirements.txt"

    # Try up to 5 times to generate & push working code
    for _ in range(5):
        tools = [Tool(google_search=GoogleSearch())] if grounding_enabled else []
        cfg   = GenerateContentConfig(tools=tools, response_modalities=["TEXT"])

        resp = client.models.generate_content(
            model="gemini-2.5-flash-preview-04-17",
            contents=[m["content"] for m in chat],
            config=cfg
        )

        raw_code = resp.text
        code = extract_code(raw_code)
        chat.append({"role":"assistant","content":code})

        # Write code file
        with open(code_fn, "w") as f:
            f.write(code)

        # Write README with dynamic SDK version
        sdk_version = get_sdk_version(sdk_choice)
        readme = f"""---
title: Wuhp Auto Space
emoji: 🐢
colorFrom: red
colorTo: pink
sdk: {sdk_choice}
sdk_version: {sdk_version}
app_file: {code_fn}
pinned: false
---

Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
"""
        with open(readme_fn, "w") as f:
            f.write(readme)

        # Write requirements
        base_reqs = "pandas\n"
        extra     = "streamlit\n" if sdk_choice=="streamlit" else "gradio\n"
        with open(reqs_fn, "w") as f:
            f.write(base_reqs + extra)

        # Push to Hugging Face
        create_repo(
            repo_id=repo_id,
            token=oauth_token.token,
            exist_ok=True,
            repo_type="space",
            space_sdk=sdk_choice
        )
        for fn in (code_fn, readme_fn, reqs_fn):
            upload_file(
                path_or_fileobj=fn,
                path_in_repo=fn,
                repo_id=repo_id,
                token=oauth_token.token,
                repo_type="space"
            )

        # If build succeeds without errors, break; otherwise feed logs back to LLM
        build = fetch_logs(repo_id, "build")
        run   = fetch_logs(repo_id, "run")
        if "ERROR" not in build.upper() and "ERROR" not in run.upper():
            break

        chat.append({
            "role":"user",
            "content":(
                f"Build logs:\n{build}\n\n"
                f"Run logs:\n{run}\n\n"
                "Please fix the code."
            )
        })
        time.sleep(2)

    # Prepare messages + immediate iframe preview
    messages = [{"role":m["role"],"content":m["content"]} for m in chat if m["role"]!="system"]
    preview_url = f"https://huggingface.co/spaces/{repo_id}"
    iframe = (
        f'<iframe src="{preview_url}" '
        f'width="100%" height="500px" frameborder="0"></iframe>'
    )

    placeholder = "✅ Space created! Click “Refresh Logs” to pull build/run logs."
    return messages, placeholder, placeholder, iframe

# — REFRESH LOGS — 

def refresh_logs(
    space_suffix: str,
    profile: gr.OAuthProfile | None,
    oauth_token: gr.OAuthToken | None
):
    if profile is None or oauth_token is None:
        return "⚠️ Please log in.", "⚠️ Please log in."
    repo_id = f"{profile.username}/{profile.username}-{space_suffix}"
    return fetch_logs(repo_id, "build"), fetch_logs(repo_id, "run")

# — BUILD THE UI — 

with gr.Blocks(title="HF Space Auto‑Builder") as demo:
    gr.Markdown("## Sign in + Auto‑Build Spaces\n\n"
                "1. Sign in  2. Enter your prompt  3. Watch code, README, requirements, logs, and preview\n\n---")

    # LOGIN controls
    login_btn = gr.LoginButton(variant="huggingface", size="lg")
    status_md = gr.Markdown("*Not logged in.*")
    models_md = gr.Markdown()
    demo.load(show_profile,        inputs=None, outputs=status_md)
    demo.load(list_private_models, inputs=None, outputs=models_md)
    login_btn.click(show_profile,        inputs=None, outputs=status_md)
    login_btn.click(list_private_models, inputs=None, outputs=models_md)

    # SETTINGS
    sdk_choice  = gr.Radio(["gradio","streamlit"], value="gradio", label="SDK template")
    api_key     = gr.Textbox(label="Gemini API Key", type="password")
    grounding   = gr.Checkbox(label="Enable grounding", value=False)
    space_suffix= gr.Textbox(label="Space suffix", value="auto-space",
                             info="E.g. 'auto-space', 'auto-space2', etc.")

    # CHAT + OUTPUTS
    chatbot    = gr.Chatbot(type="messages")
    user_in    = gr.Textbox(placeholder="Your prompt…", label="Prompt")
    send_btn   = gr.Button("Send")
    build_box  = gr.Textbox(label="Build logs", lines=5, interactive=False)
    run_box    = gr.Textbox(label="Run logs",   lines=5, interactive=False)
    preview    = gr.HTML("<p>No Space yet.</p>")

    send_btn.click(
        fn=handle_user_message,
        inputs=[chatbot, sdk_choice, api_key, grounding, space_suffix],
        outputs=[chatbot, build_box, run_box, preview]
    )

    # Manual log refresh
    refresh_btn = gr.Button("Refresh Logs")
    refresh_btn.click(
        fn=refresh_logs,
        inputs=[space_suffix],
        outputs=[build_box, run_box]
    )

    demo.launch(server_name="0.0.0.0", server_port=7860)