app.py

# app.py

import os
import time
import json
import requests

import gradio as gr
import google.generativeai as genai

from huggingface_hub import create_repo, list_models, upload_file, constants
from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status

# --- Helper functions for Hugging Face integration ---

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."
    try:
        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)
    except Exception as e:
        return f"Error listing models: {e}"

def create_space_action(repo_name: str, sdk: str, profile: gr.OAuthProfile, token: gr.OAuthToken):
    repo_id = f"{profile.username}/{repo_name}"
    create_repo(
        repo_id=repo_id,
        token=token.token,
        exist_ok=True,
        repo_type="space",
        space_sdk=sdk
    )
    url    = f"https://huggingface.co/spaces/{repo_id}"
    iframe = f'<iframe src="{url}" width="100%" height="500px"></iframe>'
    return repo_id, iframe

def upload_file_to_space_action(
    file_obj,
    path_in_repo: str,
    repo_id: str,
    profile: gr.OAuthProfile,
    token: gr.OAuthToken
) -> str:
    if not (profile and token and repo_id):
        return "⚠️ Please log in and create a Space first."
    try:
        upload_file(
            path_or_fileobj=file_obj,
            path_in_repo=path_in_repo,
            repo_id=repo_id,
            token=token.token,
            repo_type="space"
        )
        return f"✅ Uploaded `{path_in_repo}`"
    except Exception as e:
        return f"Error uploading file: {e}"

def _fetch_space_logs_level(repo_id: str, level: str, token: str) -> str:
    jwt_url  = f"{constants.ENDPOINT}/api/spaces/{repo_id}/jwt"
    r        = get_session().get(jwt_url, headers=build_hf_headers(token=token))
    hf_raise_for_status(r)
    jwt      = r.json()["token"]
    logs_url = f"https://api.hf.space/v1/{repo_id}/logs/{level}"
    lines, count = [], 0
    with get_session().get(logs_url, headers=build_hf_headers(token=jwt), stream=True, timeout=20) as resp:
        hf_raise_for_status(resp)
        for raw in resp.iter_lines():
            if count >= 200:
                lines.append("... truncated ...")
                break
            if not raw.startswith(b"data: "):
                continue
            payload = raw[len(b"data: "):]
            try:
                event = json.loads(payload.decode())
                ts    = event.get("timestamp", "")
                txt   = event.get("data", "").strip()
                if txt:
                    lines.append(f"[{ts}] {txt}")
                    count += 1
            except json.JSONDecodeError:
                continue
    return "\n".join(lines) if lines else f"No {level} logs found."

def get_build_logs_action(repo_id, profile, token):
    if not (repo_id and profile and token):
        return "⚠️ Please log in and create a Space first."
    return _fetch_space_logs_level(repo_id, "build", token.token)

def get_container_logs_action(repo_id, profile, token):
    if not (repo_id and profile and token):
        return "⚠️ Please log in and create a Space first."
    return _fetch_space_logs_level(repo_id, "run", token.token)


# --- Google Gemini integration with model selection ---

def configure_gemini(api_key: str | None, model_name: str | None) -> str:
    if not api_key:
        return "Gemini API key is not set."
    if not model_name:
        return "Please select a Gemini model."
    try:
        genai.configure(api_key=api_key)
        # Test using the selected model
        genai.GenerativeModel(model_name).generate_content("ping")
        return f"Gemini configured successfully with **{model_name}**."
    except Exception as e:
        return f"Error configuring Gemini: {e}"

def call_gemini(prompt: str, api_key: str, model_name: str) -> str:
    if not api_key or not model_name:
        return "Error: Gemini API key or model not provided."
    try:
        genai.configure(api_key=api_key)
        model    = genai.GenerativeModel(model_name)
        response = model.generate_content(prompt)
        return response.text or "Gemini returned an empty response."
    except Exception as e:
        return f"Error calling Gemini API with {model_name}: {e}"


# --- AI workflow logic (uses selected model) ---

def ai_workflow_chat(
    message: str,
    history: list[list[str | None]],
    hf_profile: gr.OAuthProfile | None,
    hf_token:   gr.OAuthToken   | None,
    gemini_api_key: str         | None,
    gemini_model:   str         | None,
    repo_id_state:  str         | None,
    workflow_state: str,
    space_sdk:      str,
    preview_html:   str,
    container_logs: str,
    build_logs:     str
) -> tuple[
    list[list[str | None]],
    str | None,
    str,
    str,
    str,
    str
]:
    # Append user message
    history.append([message, None])
    bot_message = ""
    new_repo_id = repo_id_state
    new_workflow = workflow_state
    updated_preview = preview_html
    updated_container = container_logs
    updated_build = build_logs

    # -- same workflow logic as before, but use call_gemini(prompt, gemini_api_key, gemini_model) --

    # example when generating code:
    # resp = call_gemini(prompt, gemini_api_key, gemini_model)

    # [Omitted for brevity; insert your existing logic here, replacing calls to
    # call_gemini(prompt, gemini_api_key) with call_gemini(prompt, gemini_api_key, gemini_model).]

    return history, new_repo_id, new_workflow, updated_preview, updated_container, updated_build


# --- Build the Gradio UI ---

with gr.Blocks(title="AI-Powered HF Space App Builder") as ai_builder_tab:
    hf_profile      = gr.State(None)
    hf_token        = gr.State(None)
    gemini_key      = gr.State(None)
    gemini_model    = gr.State("gemini-2.5-pro-preview-03-25")
    repo_id         = gr.State(None)
    workflow        = gr.State("idle")
    sdk_state       = gr.State("gradio")

    with gr.Row():
        # Sidebar
        with gr.Column(scale=1, min_width=300):
            gr.Markdown("## Hugging Face Login")
            login_status = gr.Markdown("*Not logged in.*")
            login_btn    = gr.LoginButton(variant="huggingface")

            # init & update login status
            ai_builder_tab.load(show_profile, outputs=login_status)
            login_btn.click(show_profile, outputs=login_status)
            login_btn.click(lambda profile, token: (profile, token),
                            outputs=[hf_profile, hf_token])

            gr.Markdown("## Google AI Studio API Key")
            gemini_input  = gr.Textbox(label="API Key", type="password")
            gemini_status = gr.Markdown("")
            gemini_input.change(lambda k: k, inputs=gemini_input, outputs=gemini_key)

            gr.Markdown("## Gemini Model")
            model_selector = gr.Radio(
                choices=[
                    ("Gemini 2.5 Flash Preview 04-17", "gemini-2.5-flash-preview-04-17"),
                    ("Gemini 2.5 Pro Preview 03-25",   "gemini-2.5-pro-preview-03-25")
                ],
                value="gemini-2.5-pro-preview-03-25",
                label="Select model"
            )
            model_selector.change(lambda m: m, inputs=model_selector, outputs=gemini_model)

            # configure Gemini whenever key or model changes
            gr.Row().load(
                configure_gemini,
                inputs=[gemini_key, gemini_model],
                outputs=[gemini_status]
            )
            gemini_input.change(
                configure_gemini,
                inputs=[gemini_key, gemini_model],
                outputs=[gemini_status]
            )
            model_selector.change(
                configure_gemini,
                inputs=[gemini_key, gemini_model],
                outputs=[gemini_status]
            )

            gr.Markdown("## Space SDK")
            sdk_selector = gr.Radio(choices=["gradio","streamlit"], value="gradio", label="Template SDK")
            sdk_selector.change(lambda s: s, inputs=sdk_selector, outputs=sdk_state)

        # Main content
        with gr.Column(scale=3):
            chatbot    = gr.Chatbot()
            user_input = gr.Textbox(placeholder="Type your message…")
            send_btn   = gr.Button("Send", interactive=False)

            # enable send only when logged in & key & model selected
            ai_builder_tab.load(
                lambda p, k, m: gr.update(interactive=bool(p and k and m)),
                inputs=[hf_profile, gemini_key, gemini_model],
                outputs=[send_btn]
            )
            login_btn.click(
                lambda p, k, m: gr.update(interactive=bool(p and k and m)),
                inputs=[hf_profile, gemini_key, gemini_model],
                outputs=[send_btn]
            )
            gemini_input.change(
                lambda p, k, m: gr.update(interactive=bool(p and k and m)),
                inputs=[hf_profile, gemini_key, gemini_model],
                outputs=[send_btn]
            )
            model_selector.change(
                lambda p, k, m: gr.update(interactive=bool(p and k and m)),
                inputs=[hf_profile, gemini_key, gemini_model],
                outputs=[send_btn]
            )

            iframe    = gr.HTML("<p>No Space created yet.</p>")
            build_txt = gr.Textbox(label="Build Logs", lines=10, interactive=False)
            run_txt   = gr.Textbox(label="Container Logs", lines=10, interactive=False)

            def wrap_chat(msg, history, prof, tok, key, model, rid, wf, sdk, prev, run_l, build_l):
                out = ai_workflow_chat(
                    msg, history, prof, tok, key, model, rid, wf, sdk, prev, run_l, build_l
                )
                hist, new_rid, new_wf, new_prev, new_run, new_build = out
                return [(u or "", v or "") for u, v in hist], new_rid, new_wf, new_prev, new_run, new_build

            send_btn.click(
                wrap_chat,
                inputs=[
                    user_input, chatbot,
                    hf_profile, hf_token,
                    gemini_key, gemini_model,
                    repo_id, workflow, sdk_state,
                    iframe, run_txt, build_txt
                ],
                outputs=[
                    chatbot,
                    repo_id, workflow,
                    iframe, run_txt, build_txt
                ]
            )

with gr.Blocks(title="Manual Hugging Face Space Manager") as manual_control_tab:
    # ... (manual tab unchanged) ...

demo = gr.TabbedInterface(
    [ai_builder_tab, manual_control_tab],
    ["AI App Builder", "Manual Control"]
)

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