构建具有双上游自动容灾的 OpenAI 兼容网关

[客户端应用] 
       │ (OpenAI SDK / v1/chat/completions)
       ▼
[FastAPI 容灾网关] 
       │
       ├─► [主上游通道: SiliconFlow] (超时时间 < 8s?)
       │         │ (若发生故障或超时)
       │         ▼
       └─► [备用上游通道: DeepSeek 官方]
                 │
                 ▼
[SQLite 数据库] (记录 Token 消耗并校验可用额度)

import json
import time
import asyncio
import logging
from typing import AsyncGenerator
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.responses import StreamingResponse
import httpx
import aiosqlite

# 配置日志输出
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("llm-gateway")

app = FastAPI()

# 核心配置参数
DATABASE_PATH = "gateway.db"
PRIMARY_URL = "https://api.siliconflow.cn/v1/chat/completions"
PRIMARY_KEY = "sk-siliconflow-key-here"
FALLBACK_URL = "https://api.deepseek.com/v1/chat/completions"
FALLBACK_KEY = "sk-deepseek-key-here"
TIMEOUT_LIMIT = 8.0  # 超时阈值（秒）

# 初始化数据库表结构
async def init_db():
    async with aiosqlite.connect(DATABASE_PATH) as db:
        await db.execute("""
            CREATE TABLE IF NOT EXISTS users (
                api_key TEXT PRIMARY KEY,
                token_quota INTEGER,
                token_used INTEGER DEFAULT 0
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS usage_logs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                api_key TEXT,
                upstream TEXT,
                prompt_tokens INTEGER,
                completion_tokens INTEGER,
                timestamp REAL
            )
        """)
        # 预先插入一个本地测试 Key，赠送 50 万 Token 额度
        await db.execute(
            "INSERT OR IGNORE INTO users (api_key, token_quota) VALUES (?, ?)",
            ("sk-local-test-key", 500000)
        )
        await db.commit()

@app.on_event("startup")
async def startup_event():
    await init_db()

# 校验 API Key 及其剩余额度
async def verify_user(request: Request) -> str:
    auth_header = request.headers.get("Authorization")
    if not auth_header or not auth_header.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="API Key 格式错误或缺失")
    api_key = auth_header.split(" ")[1]
    
    async with aiosqlite.connect(DATABASE_PATH) as db:
        async with db.execute(
            "SELECT token_quota, token_used FROM users WHERE api_key = ?", 
            (api_key,)
        ) as cursor:
            row = await cursor.fetchone()
            if not row:
                raise HTTPException(status_code=401, detail="未授权的 API Key")
            quota, used = row
            if used >= quota:
                raise HTTPException(status_code=429, detail="额度已耗尽")
    return api_key

# 记录 Token 消耗并更新用户表
async def log_usage(api_key: str, upstream: str, prompt: int, completion: int):
    async with aiosqlite.connect(DATABASE_PATH) as db:
        await db.execute(
            """
            INSERT INTO usage_logs (api_key, upstream, prompt_tokens, completion_tokens, timestamp)
            VALUES (?, ?, ?, ?, ?)
            """,
            (api_key, upstream, prompt, completion, time.time())
        )
        await db.execute(
            "UPDATE users SET token_used = token_used + ? WHERE api_key = ?",
            (prompt + completion, api_key)
        )
        await db.commit()

# 异步迭代器：处理流式响应并实时解析 Token 数量
async def forward_stream(response: httpx.Response, api_key: str, upstream_name: str) -> AsyncGenerator[str, None]:
    total_prompt = 0
    total_completion = 0
    try:
        async for line in response.aiter_lines():
            if not line.strip():
                continue
            yield f"{line}\n\n"
            
            # 解析 SSE 数据流中的 usage 字段
            if line.startswith("data: "):
                data_str = line[6:]
                if data_str.strip() == "[DONE]":
                    continue
                try:
                    data_json = json.loads(data_str)
                    usage = data_json.get("usage")
                    if usage:
                        total_prompt = usage.get("prompt_tokens", 0)
                        total_completion = usage.get("completion_tokens", 0)
                except json.JSONDecodeError:
                    pass
    finally:
        await response.aclose()
        # 传输结束后写入日志
        if total_prompt > 0 or total_completion > 0:
            await log_usage(api_key, upstream_name, total_prompt, total_completion)

# 发送 HTTP 请求并捕获 5xx 错误
async def attempt_request(client: httpx.AsyncClient, url: str, key: str, payload: dict) -> httpx.Response:
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json"
    }
    response = await client.post(url, json=payload, headers=headers, timeout=TIMEOUT_LIMIT)
    if response.status_code >= 500:
        raise httpx.HTTPStatusError("上游服务器返回 5xx 错误", request=response.request, response=response)
    return response

@app.post("/v1/chat/completions")
async def chat_completions(request: Request, api_key: str = Depends(verify_user)):
    payload = await request.json()
    stream = payload.get("stream", False)
    
    # 配置主备上游通道
    upstreams = [
        ("SiliconFlow", PRIMARY_URL, PRIMARY_KEY),
        ("DeepSeek Official", FALLBACK_URL, FALLBACK_KEY)
    ]
    
    async with httpx.AsyncClient() as client:
        for name, url, key in upstreams:
            try:
                logger.info(f"正在尝试通过上游通道发送请求: {name}")
                response = await attempt_request(client, url, key, payload)
                
                if stream:
                    return StreamingResponse(
                        forward_stream(response, api_key, name),
                        media_type="text/event-stream"
                    )
                else:
                    data = response.json()
                    usage = data.get("usage", {})
                    prompt = usage.get("prompt_tokens", 0)
                    completion = usage.get("completion_tokens", 0)
                    await log_usage(api_key, name, prompt, completion)
                    return data
                    
            except (httpx.RequestError, httpx.HTTPStatusError, asyncio.TimeoutError) as e:
                logger.warning(f"上游通道 {name} 异常: {str(e)}。正在切换至下一个通道...")
                continue
                
        # 若所有上游通道均不可用，返回 502 错误
        raise HTTPException(status_code=502, detail="所有配置的上游 LLM 通道均无法访问或已超时。")