DouBao2Api/app/main.py

import asyncio
import json
import os
import subprocess
import time
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from typing import AsyncGenerator, Optional

from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from loguru import logger

from .browser_manager import browser_manager
from .config import settings
from .doubao import DoubaoChat
from .models import (
    ChatCompletionChunk,
    ChatCompletionRequest,
    ChatCompletionResponse,
    Choice,
    ChoiceMessage,
    ChunkChoice,
    ChunkDelta,
    ModelInfo,
    ModelListResponse,
    SUPPORTED_MODELS,
)

DEFAULT_API_KEY = "sk-your-api-key-here"

VNC_MODE_FILE = Path("/app/data/.vnc_mode")


def _read_vnc_mode() -> str:
    if VNC_MODE_FILE.exists():
        mode = VNC_MODE_FILE.read_text().strip()
        return "viewonly" if mode in ("true", "viewonly") else "interactive"
    return "viewonly" if settings.vnc_view_only else "interactive"


def _write_vnc_mode(mode: str):
    VNC_MODE_FILE.parent.mkdir(parents=True, exist_ok=True)
    VNC_MODE_FILE.write_text(mode)


def _restart_x11vnc():
    try:
        subprocess.run(
            ["supervisorctl", "restart", "x11vnc"],
            capture_output=True, text=True, timeout=10,
        )
    except Exception as e:
        logger.warning(f"supervisorctl restart failed: {e}")


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("DouBao2Api starting up...")
    try:
        await browser_manager.start()
        logged_in = await browser_manager.check_login()
        if not logged_in:
            logger.warning(
                "Not logged in! Open noVNC at "
                f"http://localhost:{settings.novnc_port}/vnc.html to login."
            )
        else:
            logger.info("Login detected, ready to serve requests.")
    except Exception as e:
        logger.error(f"Browser startup failed: {e}")

    yield

    logger.info("DouBao2Api shutting down...")
    await browser_manager.stop()


app = FastAPI(title="DouBao2Api", version="1.0.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

_request_lock = asyncio.Lock()


async def verify_api_key(authorization: Optional[str] = Header(None)):
    if not settings.api_key or settings.api_key == DEFAULT_API_KEY:
        return
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing Authorization header. Expected: Bearer <key>")
    key = authorization[7:]
    if key != settings.api_key:
        raise HTTPException(status_code=401, detail="Invalid API key")


@app.get("/")
async def root():
    return {
        "service": "DouBao2Api",
        "version": "1.0.0",
        "browser_ready": browser_manager.ready,
        "endpoints": {
            "chat_completions": "POST /v1/chat/completions",
            "models": "GET /v1/models",
            "new_chat": "POST /chat/new",
            "login_status": "GET /login/status",
            "save_state": "POST /login/save",
            "inspect": "GET /inspect",
            "vnc_mode": "GET/POST /vnc/mode",
            "restart_browser": "POST /browser/restart",
        },
        "vnc_url": f"http://localhost:{settings.novnc_port}/vnc.html",
        "vnc_mode": _read_vnc_mode(),
    }


@app.get("/health")
async def health():
    return {"status": "ok", "browser_ready": browser_manager.ready}


@app.get("/v1/models")
async def list_models():
    return ModelListResponse(data=[ModelInfo(**m) for m in SUPPORTED_MODELS])


@app.get("/login/status")
async def login_status():
    logged_in = await browser_manager.check_login()
    return {
        "logged_in": logged_in,
        "vnc_url": f"http://localhost:{settings.novnc_port}/vnc.html",
        "instructions": (
            "Open the VNC URL in your browser, login to Doubao manually, "
            "then call POST /login/save to persist the session."
        )
        if not logged_in
        else "Session is active.",
    }


@app.post("/login/save")
async def save_state():
    await browser_manager.save_state()
    return {"status": "saved", "message": "Browser state persisted to disk."}


@app.get("/inspect")
async def inspect_dom():
    page = await browser_manager.get_page()
    doubao = DoubaoChat(page)
    result = await doubao.inspect()
    return result


@app.post("/chat/new")
async def new_chat(_: bool = Depends(verify_api_key)):
    page = await browser_manager.get_page()
    doubao = DoubaoChat(page)
    ok = await doubao.start_new_chat()
    return {"status": "new_chat_started" if ok else "failed", "ok": ok}


@app.get("/vnc/mode")
async def get_vnc_mode():
    mode = _read_vnc_mode()
    return {
        "mode": mode,
        "vnc_url": f"http://localhost:{settings.novnc_port}/vnc.html",
        "hint": "POST /vnc/mode with {'mode':'interactive'} to enable mouse/keyboard" if mode == "viewonly" else "VNC is interactive",
    }


@app.post("/vnc/mode")
async def set_vnc_mode(request: Request, _: bool = Depends(verify_api_key)):
    body = await request.json()
    mode = body.get("mode", "").lower()
    if mode not in ("viewonly", "interactive"):
        raise HTTPException(status_code=400, detail="mode must be 'viewonly' or 'interactive'")

    _write_vnc_mode(mode)
    _restart_x11vnc()
    logger.info(f"VNC mode set to {mode}, x11vnc restarted")
    return {"status": "ok", "mode": mode}


@app.post("/browser/restart")
async def restart_browser(_: bool = Depends(verify_api_key)):
    await browser_manager.restart()
    return {"status": "restarted", "browser_ready": browser_manager.ready}


@app.post("/v1/chat/completions")
async def chat_completions(
    request: ChatCompletionRequest,
    _=Depends(verify_api_key),
):
    async with _request_lock:
        if not browser_manager.ready:
            try:
                await browser_manager.start()
            except Exception as e:
                raise HTTPException(
                    status_code=503,
                    detail=f"Browser not ready: {e}. Please retry shortly.",
                )

        if not await browser_manager.check_login():
            raise HTTPException(
                status_code=401,
                detail=(
                    "Not logged in to Doubao. Open "
                    f"http://localhost:{settings.novnc_port}/vnc.html "
                    "to login manually, then call POST /login/save."
                ),
            )

        page = await browser_manager.get_page()
        doubao = DoubaoChat(page)

        if request.new_chat:
            await doubao.start_new_chat()
            logger.info("Started new chat session")

        prompt = DoubaoChat.messages_to_prompt(request.messages)
        logger.info(f"Sending prompt ({len(prompt)} chars), stream={request.stream}")

        try:
            old_text = await doubao.send_message(prompt)
        except RuntimeError as e:
            raise HTTPException(status_code=500, detail=str(e))
        except Exception as e:
            logger.error(f"Unexpected send error: {e}")
            raise HTTPException(status_code=500, detail=f"Failed to send message: {e}")

        if request.stream:
            return StreamingResponse(
                _stream_sse(doubao, request.model, old_text),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                },
            )

        try:
            full_response = await doubao.get_response_full(old_text)
        except Exception as e:
            logger.error(f"Response capture failed: {e}")
            raise HTTPException(status_code=502, detail=f"Failed to get response: {e}")

        await browser_manager.save_state()

        return ChatCompletionResponse(
            model=request.model,
            choices=[
                Choice(
                    index=0,
                    message=ChoiceMessage(role="assistant", content=full_response),
                    finish_reason="stop",
                )
            ],
        )


async def _stream_sse(
    doubao: DoubaoChat,
    model: str,
    old_text: str,
) -> AsyncGenerator[str, None]:
    chat_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created = int(time.time())

    first_chunk = ChatCompletionChunk(
        id=chat_id,
        created=created,
        model=model,
        choices=[ChunkChoice(index=0, delta=ChunkDelta(role="assistant"))],
    )
    yield f"data: {first_chunk.model_dump_json()}\n\n"

    try:
        async for chunk_text in doubao.stream_response(old_text):
            if not chunk_text:
                continue
            chunk = ChatCompletionChunk(
                id=chat_id,
                created=created,
                model=model,
                choices=[
                    ChunkChoice(index=0, delta=ChunkDelta(content=chunk_text))
                ],
            )
            yield f"data: {chunk.model_dump_json()}\n\n"
    except Exception as e:
        logger.error(f"Streaming error: {e}")
        err_chunk = ChatCompletionChunk(
            id=chat_id,
            created=created,
            model=model,
            choices=[
                ChunkChoice(
                    index=0,
                    delta=ChunkDelta(content=f"\n[Error: {e}]"),
                    finish_reason="stop",
                )
            ],
        )
        yield f"data: {err_chunk.model_dump_json()}\n\n"
    else:
        final_chunk = ChatCompletionChunk(
            id=chat_id,
            created=created,
            model=model,
            choices=[
                ChunkChoice(index=0, delta=ChunkDelta(), finish_reason="stop")
            ],
        )
        yield f"data: {final_chunk.model_dump_json()}\n\n"

    yield "data: [DONE]\n\n"

    try:
        await browser_manager.save_state()
    except Exception:
        pass


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=settings.api_port,
        log_level=settings.log_level.lower(),
    )