main.py

from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, HTMLResponse
import uvicorn
import os
import logging
import httpx

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="FreeLLMAPI")

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

# ──────────────────────────────────────────
# 读取原项目固定命名的 Secrets
# ──────────────────────────────────────────

PROVIDER_MAP = {
    "GOOGLE_API_KEY":        "google",
    "GROQ_API_KEY":          "groq",
    "GITHUB_TOKEN":          "github",
    "OPENROUTER_API_KEY":    "openrouter",
    "MISTRAL_API_KEY":       "mistral",
    "TOGETHER_API_KEY":      "together",
    "NVIDIA_API_KEY":        "nvidia",
    "COHERE_API_KEY":        "cohere",
    "HF_TOKEN":              "huggingface",
    "CEREBRAS_API_KEY":      "cerebras",
    "SAMBANOVA_API_KEY":     "sambanova",
    "CLOUDFLARE_API_TOKEN":  "cloudflare",
    "ZHIPU_API_KEY":         "zhipu",
}

PROVIDER_CONFIG = {
    "google": {
        "base_url": "https://generativelanguage.googleapis.com/v1beta/openai",
        "models": [
            "gemini-2.0-flash","gemini-2.0-flash-lite",
            "gemini-1.5-pro","gemini-1.5-flash","gemini-1.5-flash-8b"
        ],
    },
    "groq": {
        "base_url": "https://api.groq.com/openai/v1",
        "models": [
            "llama-3.3-70b-versatile","llama-3.1-8b-instant",
            "llama3-70b-8192","llama3-8b-8192",
            "mixtral-8x7b-32768","gemma2-9b-it"
        ],
    },
    "github": {
        "base_url": "https://models.inference.ai.azure.com",
        "models": [
            "gpt-4o","gpt-4o-mini",
            "Phi-3.5-mini-instruct","Phi-3.5-MoE-instruct",
            "Meta-Llama-3.1-70B-Instruct","Meta-Llama-3.1-405B-Instruct"
        ],
    },
    "openrouter": {
        "base_url": "https://openrouter.ai/api/v1",
        "models": [
            "mistralai/mistral-7b-instruct:free",
            "meta-llama/llama-3.2-3b-instruct:free",
            "google/gemma-3-1b-it:free",
            "deepseek/deepseek-r1:free",
        ],
    },
    "mistral": {
        "base_url": "https://api.mistral.ai/v1",
        "models": [
            "mistral-small-latest","mistral-large-latest",
            "open-mistral-7b","open-mixtral-8x7b"
        ],
    },
    "together": {
        "base_url": "https://api.together.xyz/v1",
        "models": [
            "meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo",
            "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
            "mistralai/Mixtral-8x7B-Instruct-v0.1"
        ],
    },
    "nvidia": {
        "base_url": "https://integrate.api.nvidia.com/v1",
        "models": [
            "meta/llama-3.1-70b-instruct",
            "meta/llama-3.1-8b-instruct",
            "mistralai/mixtral-8x7b-instruct"
        ],
    },
    "cohere": {
        "base_url": "https://api.cohere.com/v2",
        "models": ["command-r-plus","command-r","command"],
    },
    "huggingface": {
        "base_url": "https://api-inference.huggingface.co/v1",
        "models": [
            "meta-llama/Llama-3.2-3B-Instruct",
            "mistralai/Mistral-7B-Instruct-v0.3"
        ],
    },
    "cerebras": {
        "base_url": "https://api.cerebras.ai/v1",
        "models": ["llama3.1-8b","llama3.1-70b"],
    },
    "sambanova": {
        "base_url": "https://api.sambanova.ai/v1",
        "models": [
            "Meta-Llama-3.1-8B-Instruct",
            "Meta-Llama-3.1-70B-Instruct",
            "Meta-Llama-3.1-405B-Instruct"
        ],
    },
    "cloudflare": {
        "base_url": "https://api.cloudflare.com/client/v4/accounts/{}/ai/v1",
        "models": [
            "@cf/meta/llama-3.1-8b-instruct",
            "@cf/mistral/mistral-7b-instruct-v0.1"
        ],
    },
    "zhipu": {
        "base_url": "https://open.bigmodel.cn/api/paas/v4",
        "models": ["glm-4-flash","glm-4","glm-3-turbo"],
    },
}


def load_config():
    """从环境变量读取所有配置"""
    # ── 读取用户 API Keys ──
    raw_keys = os.getenv("API_KEYS", "").strip()
    api_keys = set(k.strip() for k in raw_keys.split(",") if k.strip())

    # ── 读取 Provider Keys ──
    providers = {}
    cf_account = os.getenv("CLOUDFLARE_ACCOUNT_ID", "default_id")
    
    for env_name, provider_name in PROVIDER_MAP.items():
        key_value = os.getenv(env_name, "").strip()
        if not key_value:
            continue

        cfg = PROVIDER_CONFIG.get(provider_name, {})
        base_url = cfg.get("base_url", "")
        
        # 安全修复：动态组装 Cloudflare 的账户 ID 路径
        if provider_name == "cloudflare" and "{}" in base_url:
            base_url = base_url.format(cf_account)

        providers[provider_name] = {
            "api_key":  key_value,
            "base_url": base_url,
            "models":   cfg.get("models", []),
        }
        logger.info(f"✅ 加载 Provider: {provider_name}")

    logger.info(f"🚀 加载成功！用户授权 Key 数量: {len(api_keys)} | 通道网关数量: {len(providers)}")
    return api_keys, providers


API_KEYS, PROVIDERS = load_config()

# ──────────────────────────────────────────
# 鉴权
# ──────────────────────────────────────────

def verify_api_key(authorization: str = Header(...)):
    token = authorization.removeprefix("Bearer ").strip()
    if token not in API_KEYS:
        raise HTTPException(status_code=401, detail="Invalid API key")
    return token

# ──────────────────────────────────────────
# 路由映射
# ──────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def index():
    return """
    <html>
        <head><title>FreeLLMAPI Gateway</title></head>
        <body style="font-family: Arial, sans-serif; text-align: center; margin-top: 100px; background-color: #f4f4f9;">
            <h1 style="color: #4f46e5;">🚀 FreeLLMAPI 环境变量极速版</h1>
            <p style="color: #6b7280;">云端轻量化无持久化转发网关运行中...</p>
            <div style="margin-top: 20px;">
                <span style="background-color: #10b981; color: white; padding: 6px 16px; border-radius: 20px; font-size: 0.9em; font-weight: bold;">STATUS: ONLINE</span>
            </div>
        </body>
    </html>
    """

@app.get("/health")
async def health():
    # 完美合并后的检查路由
    return {
        "status":    "ok",
        "keys":      len(API_KEYS),
        "providers": list(PROVIDERS.keys()),
    }


@app.get("/v1/models")
async def list_models(_: str = Depends(verify_api_key)):
    data = []
    for p_cfg in PROVIDERS.values():
        for m in p_cfg["models"]:
            data.append({"id": m, "object": "model"})
    return {"object": "list", "data": data}


@app.post("/v1/chat/completions")
async def chat_completions(body: dict, _: str = Depends(verify_api_key)):
    model = body.get("model", "")
    provider = None

    for p_cfg in PROVIDERS.values():
        if model in p_cfg["models"]:
            provider = p_cfg
            break

    if not provider:
        raise HTTPException(
            status_code=404,
            detail=f"没有 Provider 支持模型: {model}"
        )

    # 完美支持流式输出与普通响应的核心中转逻辑
    async def stream_generator(target_url, headers, payload):
        async with httpx.AsyncClient(timeout=60.0) as client:
            async with client.stream("POST", target_url, headers=headers, json=payload) as response:
                async for chunk in response.aiter_bytes():
                    yield chunk

    target_url = f"{provider['base_url']}/chat/completions"
    headers = {"Authorization": f"Bearer {provider['api_key']}"}

    if body.get("stream", False):
        return StreamingResponse(
            stream_generator(target_url, headers, body),
            media_type="text/event-stream"
        )
    else:
        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.post(target_url, headers=headers, json=body)
            return resp.json()


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)