app.py

import os
import json
import asyncio
import mimetypes
from typing import Optional

from fastapi import Request, UploadFile, File
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse, FileResponse
from gradio import Server
from gradio.oauth import attach_oauth, _get_valid_oauth_info_from_session
from openai import AsyncOpenAI
from pydantic import BaseModel

import agent

# Initialize gradio.Server (which is a subclass of FastAPI)
app = Server()

# Delegate Hugging Face OAuth to Gradio's battle-tested implementation.
# attach_oauth(app) registers /login/huggingface, /login/callback, /logout, AND
# the SessionMiddleware. After it runs, the user's token + profile are stored
# in request.session['oauth_info'] under the keys authlib returns
# (access_token, expires_at, userinfo, ...).
attach_oauth(app)

MAX_TOKENS_PER_TURN = 4096

# --------------------------------------------------------------------------
# Model provider selection
#
# Pick which GLM-5.2 API backend to use with the MODEL_PROVIDER Space
# secret/variable. Defaults to Hugging Face (unchanged behavior) if unset
# or unrecognized.
#
#   MODEL_PROVIDER=hf      Hugging Face Inference Router (default).
#                           Auth: sign in with HF (per-visitor OAuth) OR an
#                           HF_TOKEN secret. Same account that pays for it.
#   MODEL_PROVIDER=nvidia  NVIDIA NIM (build.nvidia.com). Free trial
#                           endpoint, no OAuth -- needs an NVIDIA_API_KEY
#                           secret (from build.nvidia.com, no card needed).
#   MODEL_PROVIDER=zai     Z.ai's own official API. Free trial credits then
#                           pay-as-you-go -- needs a ZAI_API_KEY secret
#                           (from z.ai/model-api).
#   MODEL_PROVIDER=zhipu   Zhipu AI's own China platform (open.bigmodel.cn,
#                           same company as Z.ai, separate portal). Huge
#                           free token grant on signup -- but that signup
#                           requires a mainland Chinese (+86) phone number
#                           and real-name ID verification, which most
#                           people outside China cannot complete. Needs a
#                           ZHIPU_API_KEY secret.
#
# Optionally set MODEL_ID to override the exact model string sent to
# whichever provider is active -- providers occasionally rename things, so
# if you see a "model not found" error, check your provider's own
# dashboard for the current id and set MODEL_ID to that instead of waiting
# for a code change.
#
# nvidia/zai/zhipu have no per-visitor login concept: whoever the Space
# owner is, their secret pays for every request, same as the "host token"
# path on hf.
PROVIDERS = {
    "hf": {
        "base_url": "https://router.huggingface.co/v1",
        "model": "zai-org/GLM-5.2:fireworks-ai",
        "api_key_env": "HF_TOKEN",
    },
    "nvidia": {
        "base_url": "https://integrate.api.nvidia.com/v1",
        "model": "z-ai/glm-5.2",
        "api_key_env": "NVIDIA_API_KEY",
    },
    "zai": {
        "base_url": "https://api.z.ai/api/paas/v4/",
        "model": "glm-5.2",
        "api_key_env": "ZAI_API_KEY",
    },
    "zhipu": {
        "base_url": "https://open.bigmodel.cn/api/paas/v4/",
        "model": "glm-5.2",
        "api_key_env": "ZHIPU_API_KEY",
    },
}

ACTIVE_PROVIDER = os.environ.get("MODEL_PROVIDER", "hf").strip().lower()
if ACTIVE_PROVIDER not in PROVIDERS:
    ACTIVE_PROVIDER = "hf"
PROVIDER = PROVIDERS[ACTIVE_PROVIDER]
MODEL_NAME = os.environ.get("MODEL_ID", "").strip() or PROVIDER["model"]


class ChatRequest(BaseModel):
    message: str
    temperature: float = 0.7
    system_prompt: Optional[str] = None


def _session_info(request: Request) -> dict:
    """Resolve everything downstream code needs from the current session in
    one place: who's asking, which token pays for it, and which private
    workspace folder is theirs.

    On the "hf" provider this distinguishes an OAuth-signed-in visitor
    (billed to their own HF account) from the Space owner's HF_TOKEN
    fallback. Every other provider has no per-visitor login concept -- the
    owner's own secret for that provider pays for every request, same as
    the "host" path on hf.
    """
    if ACTIVE_PROVIDER == "hf":
        oauth_info = _get_valid_oauth_info_from_session(request.session) or {}
        user_info = oauth_info.get("userinfo") or {}
        user_token = (oauth_info.get("access_token") or "").strip()
        host_token = os.environ.get("HF_TOKEN", "").strip()
        api_key = (user_token or host_token).strip()
        billed_to = "user" if user_token else ("host" if host_token else "none")
        # Everyone signed in via OAuth gets their own workspace; if the Space is
        # just running on the owner's own HF_TOKEN secret with nobody logged in,
        # treat that as a single personal workspace.
        user_key = agent.user_key_from_userinfo(user_info) if user_token else "_host"
        return {"user_info": user_info, "api_key": api_key, "billed_to": billed_to, "user_key": user_key}

    host_token = os.environ.get(PROVIDER["api_key_env"], "").strip()
    return {
        "user_info": {},
        "api_key": host_token,
        "billed_to": "host" if host_token else "none",
        "user_key": "_host",
    }


def _sse(payload: dict) -> str:
    return "data: " + json.dumps(payload, ensure_ascii=False) + "\n\n"


@app.get("/me")
async def me(request: Request):
    """Expose the current user's profile (or null) to the frontend, plus
    whether the active provider is usable without anyone signing in (either
    a non-hf provider with its secret set, or hf with an HF_TOKEN secret) --
    if so, chat still works for a visitor who isn't signed in via OAuth."""
    if ACTIVE_PROVIDER != "hf":
        host_available = bool(os.environ.get(PROVIDER["api_key_env"], "").strip())
        return JSONResponse({"user": None, "host_available": host_available, "provider": ACTIVE_PROVIDER})

    oauth_info = _get_valid_oauth_info_from_session(request.session) or {}
    user_info = oauth_info.get("userinfo") or {}
    host_available = bool(os.environ.get("HF_TOKEN", "").strip())
    if not oauth_info or not user_info:
        return JSONResponse({"user": None, "host_available": host_available, "provider": ACTIVE_PROVIDER})
    return JSONResponse({"user": user_info, "host_available": host_available, "provider": ACTIVE_PROVIDER})


@app.get("/api/history")
async def get_history(request: Request):
    info = _session_info(request)
    if not info["api_key"]:
        return JSONResponse({"history": [], "persistent": agent.IS_PERSISTENT})
    history = agent.load_history(info["user_key"])
    display_items = [h["display"] for h in history if h.get("display")]
    return JSONResponse({"history": display_items, "persistent": agent.IS_PERSISTENT})


@app.post("/api/history/clear")
async def clear_history_endpoint(request: Request):
    info = _session_info(request)
    if info["api_key"]:
        agent.clear_history(info["user_key"])
    return JSONResponse({"ok": True})


@app.get("/api/workspace")
async def workspace_listing(request: Request):
    info = _session_info(request)
    if not info["api_key"]:
        return JSONResponse({"files": [], "persistent": agent.IS_PERSISTENT})
    files_dir = agent.get_files_dir(info["user_key"])
    items = []
    for p in sorted(files_dir.rglob("*")):
        if p.is_file() and not p.name.startswith(".tmp_") and p.name != "chat_history.json":
            items.append({"name": str(p.relative_to(files_dir)), "size": p.stat().st_size})
    return JSONResponse({"files": items, "persistent": agent.IS_PERSISTENT})


@app.get("/api/files/{filename:path}")
async def download_file(request: Request, filename: str):
    info = _session_info(request)
    if not info["api_key"]:
        return JSONResponse({"error": "Not signed in."}, status_code=401)
    files_dir = agent.get_files_dir(info["user_key"])
    try:
        target = agent.safe_join(files_dir, filename)
    except agent.UnsafePathError:
        return JSONResponse({"error": "Invalid path."}, status_code=400)
    if not target.exists() or not target.is_file():
        return JSONResponse({"error": "Not found."}, status_code=404)
    media_type, _ = mimetypes.guess_type(str(target))
    return FileResponse(str(target), media_type=media_type or "application/octet-stream", filename=target.name)


@app.post("/api/upload")
async def upload_file(request: Request, file: UploadFile = File(...)):
    info = _session_info(request)
    if not info["api_key"]:
        return JSONResponse({"error": "Not signed in."}, status_code=401)
    files_dir = agent.get_files_dir(info["user_key"])
    contents = await file.read()
    if len(contents) > agent.MAX_UPLOAD_BYTES:
        return JSONResponse({"error": f"File too large (max {agent.MAX_UPLOAD_BYTES} bytes)."}, status_code=400)
    safe_name = agent.safe_upload_name(file.filename or "upload.bin", files_dir)
    (files_dir / safe_name).write_bytes(contents)
    return JSONResponse({"name": safe_name, "size": len(contents)})


@app.post("/api/chat")
async def chat_endpoint(request: Request, payload: ChatRequest):
    info = _session_info(request)
    api_key = info["api_key"]
    billed_to = info["billed_to"]
    user_key = info["user_key"]

    headers = {
        "Cache-Control": "no-cache",
        "Connection": "keep-alive",
        "X-Accel-Buffering": "no",
    }

    if not api_key:
        if ACTIVE_PROVIDER == "hf":
            not_ready_message = (
                "Not signed in. Click the Hugging Face button in the sidebar to sign in "
                "and run inference billed to your own HF account."
            )
        else:
            not_ready_message = (
                f"This Space isn't configured yet: ask the owner to add a "
                f"{PROVIDER['api_key_env']} secret in the Space settings "
                f"(Settings -> Variables and secrets)."
            )

        async def _need_login():
            yield _sse({"type": "error", "message": not_ready_message})
            yield _sse({"type": "done", "billedTo": "none"})
        return StreamingResponse(_need_login(), media_type="text/event-stream", headers=headers)

    user_message = (payload.message or "").strip()
    if not user_message:
        async def _empty():
            yield _sse({"type": "error", "message": "Empty message."})
            yield _sse({"type": "done", "billedTo": billed_to})
        return StreamingResponse(_empty(), media_type="text/event-stream", headers=headers)

    client = AsyncOpenAI(base_url=PROVIDER["base_url"], api_key=api_key)
    files_dir = agent.get_files_dir(user_key)
    system_prompt = agent.build_system_prompt(payload.system_prompt)

    history = agent.load_history(user_key)
    history.append({
        "role": "user",
        "content": user_message,
        "display": {"type": "user", "text": user_message},
    })
    agent.save_history(user_key, history)

    async def event_generator():
        try:
            finished = False

            for _step in range(agent.MAX_AGENT_STEPS):
                yield _sse({"type": "turn_start"})
                api_messages = [{"role": "system", "content": system_prompt}] + agent.history_to_api_messages(history)
                turn_text = ""
                try:
                    stream = await client.chat.completions.create(
                        model=MODEL_NAME,
                        messages=api_messages,
                        temperature=payload.temperature,
                        max_tokens=MAX_TOKENS_PER_TURN,
                        stream=True,
                    )
                    async for chunk in stream:
                        if chunk.choices:
                            delta = chunk.choices[0].delta.content
                            if delta:
                                turn_text += delta
                                yield _sse({"type": "token", "content": delta})
                except Exception as e:
                    yield _sse({"type": "error", "message": f"Model call failed: {e}"})
                    yield _sse({"type": "done", "billedTo": billed_to})
                    return

                call, parse_error, display_text = agent.find_tool_call(turn_text)

                history.append({
                    "role": "assistant",
                    "content": turn_text,
                    "display": {"type": "assistant_text", "text": display_text, "tool_call": call},
                })
                agent.save_history(user_key, history)

                if parse_error:
                    yield _sse({"type": "tool_result", "tool": None, "success": False,
                                "stdout": "", "stderr": parse_error, "exit_code": None, "file": None,
                                "text": display_text})
                    history.append({
                        "role": "user",
                        "content": f"[SYSTEM]: {parse_error}",
                        "display": {"type": "tool_result", "tool": None, "success": False,
                                    "stdout": "", "stderr": parse_error, "exit_code": None, "file": None},
                    })
                    agent.save_history(user_key, history)
                    continue

                if not call:
                    yield _sse({"type": "done", "billedTo": billed_to})
                    finished = True
                    break

                tool_name = call.get("tool")
                tool_args = call.get("args") or {}
                yield _sse({"type": "tool_call", "tool": tool_name, "args": tool_args, "text": display_text})

                result = await asyncio.to_thread(agent.execute_tool, tool_name, tool_args, files_dir)
                yield _sse({
                    "type": "tool_result",
                    "tool": tool_name,
                    "success": bool(result.get("success")),
                    "stdout": result.get("stdout") or "",
                    "stderr": result.get("stderr") or "",
                    "exit_code": result.get("exit_code"),
                    "file": result.get("file") if result.get("success") else None,
                })

                produced_file = result.get("file") if result.get("success") else None
                if produced_file:
                    try:
                        size = (files_dir / produced_file).stat().st_size
                    except OSError:
                        size = None
                    yield _sse({"type": "file", "name": produced_file, "url": f"/api/files/{produced_file}", "size": size})

                history.append({
                    "role": "user",
                    "content": agent.format_tool_result_message(tool_name, result),
                    "display": {
                        "type": "tool_result",
                        "tool": tool_name,
                        "success": bool(result.get("success")),
                        "stdout": result.get("stdout") or "",
                        "stderr": result.get("stderr") or "",
                        "exit_code": result.get("exit_code"),
                        "file": produced_file,
                    },
                })
                agent.save_history(user_key, history)

            if not finished:
                # Ran out of steps: force exactly one more plain-text round
                # with no tool parsing, then stop no matter what.
                api_messages = [{"role": "system", "content": system_prompt}] + agent.history_to_api_messages(history)
                api_messages.append({
                    "role": "user",
                    "content": "[SYSTEM]: You're out of tool calls for this message. "
                               "Give your best final answer now in plain text -- no tool_call block.",
                })
                yield _sse({"type": "turn_start"})
                turn_text = ""
                try:
                    stream = await client.chat.completions.create(
                        model=MODEL_NAME,
                        messages=api_messages,
                        temperature=payload.temperature,
                        max_tokens=MAX_TOKENS_PER_TURN,
                        stream=True,
                    )
                    async for chunk in stream:
                        if chunk.choices:
                            delta = chunk.choices[0].delta.content
                            if delta:
                                turn_text += delta
                                yield _sse({"type": "token", "content": delta})
                except Exception as e:
                    yield _sse({"type": "error", "message": f"Model call failed: {e}"})

                history.append({
                    "role": "assistant",
                    "content": turn_text,
                    "display": {"type": "assistant_text", "text": turn_text.strip(), "tool_call": None},
                })
                agent.save_history(user_key, history)
                yield _sse({"type": "done", "billedTo": billed_to})

        except Exception as e:
            yield _sse({"type": "error", "message": str(e)})
            yield _sse({"type": "done", "billedTo": billed_to})

    return StreamingResponse(event_generator(), media_type="text/event-stream", headers=headers)


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    if not os.path.exists(html_path):
        return HTMLResponse("<h3>index.html not found. Please ensure it is created in the workspace directory.</h3>", status_code=404)
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(f.read())


if __name__ == "__main__":
    # Launch Gradio Server (which binds to the FastAPI app underneath)
    app.launch(show_error=True)