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)

MODEL_NAME = "zai-org/GLM-5.2:fireworks-ai"
MAX_TOKENS_PER_TURN = 4096


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


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."""
    oauth_info = _get_valid_oauth_info_from_session(request.session) or {}
    user_info = oauth_info.get("userinfo") or {}
    if not oauth_info or not user_info:
        return JSONResponse({"user": None})
    return JSONResponse({"user": user_info})


@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:
        async def _need_login():
            yield _sse({
                "type": "error",
                "message": "Not signed in. Click the Hugging Face button in the sidebar to sign in "
                            "and run inference billed to your own HF account.",
            })
            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="https://huggingface.co/proxy/router.huggingface.co/v1", 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)