server-core/app/api/routers/relay.py

# -*- coding: utf-8 -*-
"""OpenAI 兼容模型中转端点。

PineAgents(8088) 的 ``pineagents`` provider 把 base_url 指向本服务的 ``/v1``，
请求进来后丢弃客户端任何 ``Authorization``，统一替换为服务端凭据。

模型调用统一经 **compute-engine**（loopback :3000）中转计量（引擎持有渠道 key，
前端不直连）；走 ``PINEAGENTS_COMPUTE_ADMIN_TOKEN`` 管理令牌。流式 SSE 原样透传。
"""
from __future__ import annotations

from typing import Any

import httpx
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.background import BackgroundTask

from ... import config
from ...services import compute_catalog

router = APIRouter(prefix="/v1", tags=["relay"])


def _auth_headers(request: Request | None = None) -> dict[str, str]:
    # 优先透传客户端 Authorization（用户算力 PAT → compute 归户计量）。
    # 无则回落服务端消费令牌（COMPUTE_RELAY_TOKEN）→ compute 计匿名账。
    if request is not None:
        client_auth = request.headers.get("Authorization", "")
        if client_auth:
            return {"Authorization": client_auth}
    token = config.COMPUTE_RELAY_TOKEN or config.COMPUTE_ADMIN_TOKEN
    if not token:
        return {}
    return {"Authorization": f"Bearer {token}"}


@router.post("/chat/completions")
async def relay_chat_completions(request: Request):
    """把 chat/completions 转发到 compute-engine，流式 SSE 原样透传。"""
    try:
        body: dict[str, Any] = await request.json()
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(status_code=400, detail="Invalid JSON body") from exc

    chat_url = f"{config.COMPUTE_BASE_URL.rstrip('/')}/v1/chat/completions"
    stream = bool(body.get("stream", False))
    headers = {**_auth_headers(request), "Content-Type": "application/json"}

    client = httpx.AsyncClient(timeout=None)
    upstream_request = client.build_request("POST", chat_url, json=body, headers=headers)

    if not stream:
        try:
            upstream = await client.send(upstream_request)
            payload = upstream.json() if upstream.content else None
        except Exception:  # noqa: BLE001
            await client.aclose()
            raise HTTPException(status_code=502, detail="Upstream relay failed") from None
        await client.aclose()
        return JSONResponse(status_code=upstream.status_code, content=payload)

    try:
        upstream = await client.send(upstream_request, stream=True)
    except Exception:  # noqa: BLE001
        await client.aclose()
        raise HTTPException(status_code=502, detail="Upstream relay failed") from None

    resp_headers: dict[str, str] = {}
    if "content-type" in upstream.headers:
        resp_headers["content-type"] = upstream.headers["content-type"]
    if "x-request-id" in upstream.headers:
        resp_headers["x-request-id"] = upstream.headers["x-request-id"]

    return StreamingResponse(
        upstream.aiter_raw(),
        status_code=upstream.status_code,
        media_type="text/event-stream",
        headers=resp_headers,
        background=BackgroundTask(client.aclose),
    )


@router.get("/models")
async def relay_models():
    """返回模型清单（= 算力中心 admin 模型，供 OpenAIProvider.fetch_models 使用）。"""
    ids = await compute_catalog.model_ids()
    return {
        "object": "list",
        "data": [{"id": m, "object": "model", "owned_by": "pineagents"} for m in ids],
    }