Building an OpenAI Compatible Gateway with Automatic Dual Upstream Failover

[Client Application] 
       │ (OpenAI SDK / v1/chat/completions)
       ▼
[FastAPI Gateway] 
       │
       ├─► [Primary Upstream: SiliconFlow] (Timeout < 8s?)
       │         │ (If Fail / Timeout)
       │         ▼
       └─► [Fallback Upstream: DeepSeek Official]
                 │
                 ▼
[SQLite Database] (Log Usage & Verify Token Quota)

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

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("llm-gateway")

app = FastAPI()

# Configuration Constants
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  # seconds

# Initialize SQLite database schema
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
            )
        """)
        # Seed a test user key if not exists
        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()

# Helper to authenticate and check quotas
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="Invalid or missing 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="Unauthorized key")
            quota, used = row
            if used >= quota:
                raise HTTPException(status_code=429, detail="Quota exceeded")
    return api_key

# Helper to log token consumption
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()

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"
            
            # Attempt to parse chunk for token usage tracking
            if line.startswith("data: "):
                data_str = line[6:]
                if data_str.strip() == "[DONE]":
                    continue
                try:
                    data_json = json.loads(data_str)
                    # Extract usage metadata if provided by the upstream
                    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()
        # Log final usage metrics
        if total_prompt > 0 or total_completion > 0:
            await log_usage(api_key, upstream_name, total_prompt, total_completion)

async def attempt_request(client: httpx.AsyncClient, url: str, key: str, payload: dict) -> httpx.Response:
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json"
    }
    # Send request with a strict timeout limit
    response = await client.post(url, json=payload, headers=headers, timeout=TIMEOUT_LIMIT)
    if response.status_code >= 500:
        raise httpx.HTTPStatusError("Upstream server error", 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"Attempting request using upstream: {name}")
                response = await attempt_request(client, url, key, payload)
                
                if stream:
                    # Return custom streaming response iterator
                    return StreamingResponse(
                        forward_stream(response, api_key, name),
                        media_type="text/event-stream"
                    )
                else:
                    # Handle standard non-streaming response
                    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"Upstream {name} failed with error: {str(e)}. Attempting failover...")
                continue
                
        # If all upstreams are exhausted, raise 502
        raise HTTPException(status_code=502, detail="All upstream providers failed or timed out.")