app/api/openai_compat.py

"""OpenAI 兼容三件套：/v1/chat/completions、/v1/embeddings、/v1/models。

完全兼容 OpenAI SDK 调用，头部支持 Authorization: Bearer 或 x-xtc-access-key。
"""
from __future__ import annotations

import json
from typing import Any, AsyncIterator, Optional

from fastapi import APIRouter, Depends, Header, Query, Request
from fastapi.responses import StreamingResponse
from sse_starlette.sse import EventSourceResponse

from ..adapters import gemini_api, openai as openai_adapter
from ..auth import require_access_key
from ..config import GEMINI_DEFAULT_MODEL, get_settings
from ..errors import HttpError
from ..providers.policy import is_model_allowed
from ..services.config_store import load_config
from ._common import (
    CORS_HEADERS,
    ok_with_cors,
    record_usage,
    select_provider_and_key,
)


router = APIRouter(prefix="/v1", tags=["openai-compat"])


def _provider_hint_header(provider_id: str, model: str, image_fix_mode: Optional[str] = None) -> dict:
    h = {"x-xtc-provider": provider_id, "x-xtc-model": model}
    if image_fix_mode:
        h["x-xtc-image-fix-mode"] = image_fix_mode
    return h


@router.get("/models")
async def list_models(_key: str = Depends(require_access_key)) -> dict:
    config = await load_config()
    # 聚合所有启用厂商的模型列表（取缓存或拉取）
    # 复用 xtc._fetch_provider_models：上游失败抛 HttpError、仅缓存非空结果，
    # 避免静默吞掉异常导致"该厂商无模型"的假象。
    from .xtc import _fetch_provider_models

    out: list[dict] = []
    for p in config.providers:
        if not p.enabled:
            continue
        try:
            models = await _fetch_provider_models(p)
        except HttpError as e:
            print(f"[models] provider {p.id} list failed: {e.message}")
            continue
        for m in models:
            mid = m.get("id") or ""
            if not mid:
                continue
            # 按厂商模型策略（黑/白名单）过滤，让前端选择器与实际可请求模型一致
            if not is_model_allowed(p, mid):
                continue
            full_id = f"{p.model_prefix}{mid}" if p.model_prefix else mid
            out.append({"id": full_id, "object": "model", "created": 0, "owned_by": p.id})
    return ok_with_cors({"object": "list", "data": out})


@router.post("/chat/completions")
async def chat_completions(
    request: Request,
    _key: str = Depends(require_access_key),
    x_provider: Optional[str] = Header(default=None, alias="x-provider"),
):
    body = await request.json()
    if not isinstance(body, dict):
        raise HttpError("invalid body", status=400, code="bad_request")

    model = body.get("model") or GEMINI_DEFAULT_MODEL
    messages = body.get("messages") or []
    if not isinstance(messages, list) or not messages:
        raise HttpError("messages is required", status=400, code="bad_request")
    stream = bool(body.get("stream"))

    config = await load_config()
    provider, api_key, clean_model = await select_provider_and_key(
        config=config,
        provider_id=x_provider,
        model=model,
        fallback_model=GEMINI_DEFAULT_MODEL,
    )

    if stream:
        return _stream_chat(provider, api_key, clean_model, body, _key)
    return await _nonstream_chat(provider, api_key, clean_model, body, messages, _key)


async def _nonstream_chat(
    provider,
    api_key: str,
    clean_model: str,
    body: dict,
    messages: list,
    access_key: str,
):
    try:
        if provider.type == "gemini":
            result = await gemini_api.chat_completions(
                api_key=api_key,
                model=clean_model,
                messages=messages,
                body=body,
                stream=False,
            )
        else:
            body_with_model = {**body, "model": clean_model}
            result = await openai_adapter.chat_completions(
                base_url=provider.base_url,
                api_key=api_key,
                body=body_with_model,
                stream=False,
            )
        usage = (result or {}).get("usage")
        record_usage(
            access_key=access_key,
            provider=provider.id,
            model=clean_model,
            usage=usage,
            ok=True,
        )
        return ok_with_cors(
            result,
            extra_headers=_provider_hint_header(provider.id, clean_model),
        )
    except HttpError as e:
        record_usage(
            access_key=access_key,
            provider=provider.id,
            model=clean_model,
            usage=None,
            ok=False,
            error_code=e.code,
        )
        raise


async def _stream_chat(provider, api_key: str, clean_model: str, body: dict, access_key: str):
    """流式响应：Gemini 转 OpenAI chunk，OpenAI 直接透传。"""
    from ..utils.sse import sse_data, sse_done, passthrough_with_usage

    # 让上游在最后一个流式 chunk 里返回 usage，否则 token 计数永远为 0。
    body = {**body, "stream_options": {**(body.get("stream_options") or {}), "include_usage": True}}

    if provider.type == "gemini":
        # Gemini 转 OpenAI chunk 流
        captured: dict = {}

        async def gen() -> AsyncIterator[str]:
            try:
                chunk_iter = await gemini_api.chat_completions(
                    api_key=api_key,
                    model=clean_model,
                    messages=body.get("messages") or [],
                    body=body,
                    stream=True,
                )
                async for chunk in chunk_iter:
                    if isinstance(chunk, dict) and chunk.get("usage"):
                        captured["usage"] = chunk["usage"]
                    yield sse_data(chunk)
                yield sse_done()
            except HttpError as e:
                # 错误以 SSE event 形式返回
                record_usage(
                    access_key=access_key,
                    provider=provider.id,
                    model=clean_model,
                    usage=None,
                    ok=False,
                    error_code=e.code,
                )
                yield sse_data(
                    {
                        "error": {
                            "code": e.code,
                            "message": e.message,
                            "status": e.status,
                        }
                    }
                )
                yield sse_done()
            else:
                record_usage(
                    access_key=access_key,
                    provider=provider.id,
                    model=clean_model,
                    usage=captured.get("usage"),
                    ok=True,
                )

        return EventSourceResponse(
            gen(),
            headers={
                **CORS_HEADERS,
                **_provider_hint_header(provider.id, clean_model),
                "Cache-Control": "no-cache",
                "x-accel-buffering": "no",
            },
        )

    # OpenAI 透传
    body_with_model = {**body, "model": clean_model}
    captured_pt: dict = {}

    async def gen_passthrough() -> AsyncIterator[bytes]:
        try:
            chunk_iter = await openai_adapter.chat_completions(
                base_url=provider.base_url,
                api_key=api_key,
                body=body_with_model,
                stream=True,
            )
            async for chunk in passthrough_with_usage(chunk_iter, captured_pt):
                yield chunk
        except HttpError as e:
            record_usage(
                access_key=access_key,
                provider=provider.id,
                model=clean_model,
                usage=None,
                ok=False,
                error_code=e.code,
            )
            err_payload = sse_data(
                {"error": {"code": e.code, "message": e.message, "status": e.status}}
            )
            yield err_payload.encode("utf-8")
            yield sse_done().encode("utf-8")
        else:
            record_usage(
                access_key=access_key,
                provider=provider.id,
                model=clean_model,
                usage=captured_pt.get("usage"),
                ok=True,
            )

    return StreamingResponse(
        gen_passthrough(),
        media_type="text/event-stream",
        headers={
            **CORS_HEADERS,
            **_provider_hint_header(provider.id, clean_model),
            "Cache-Control": "no-cache",
            "x-accel-buffering": "no",
        },
    )


@router.post("/embeddings")
async def embeddings(
    request: Request,
    _key: str = Depends(require_access_key),
    x_provider: Optional[str] = Header(default=None, alias="x-provider"),
):
    body = await request.json()
    if not isinstance(body, dict):
        raise HttpError("invalid body", status=400, code="bad_request")
    input_data = body.get("input")
    if input_data is None:
        raise HttpError("input is required", status=400, code="bad_request")

    config = await load_config()
    provider, api_key, clean_model = await select_provider_and_key(
        config=config,
        provider_id=x_provider,
        model=body.get("model"),
        fallback_model="text-embedding-3-small",
    )

    try:
        if provider.type == "gemini":
            result = await gemini_api.embeddings(
                api_key=api_key,
                model=clean_model,
                input_data=input_data,
            )
        else:
            body_with_model = {**body, "model": clean_model}
            result = await openai_adapter.embeddings(
                base_url=provider.base_url,
                api_key=api_key,
                body=body_with_model,
            )
        record_usage(
            access_key=_key,
            provider=provider.id,
            model=clean_model,
            usage=(result or {}).get("usage"),
            ok=True,
        )
        return ok_with_cors(
            result, extra_headers=_provider_hint_header(provider.id, clean_model)
        )
    except HttpError as e:
        record_usage(
            access_key=_key,
            provider=provider.id,
            model=clean_model,
            usage=None,
            ok=False,
            error_code=e.code,
        )
        raise