main.py

from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.middleware.cors import CORSMiddleware
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 Key 映射表
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 对应的 base_url 和支持的模型
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", "")
    api_keys = set(k.strip() for k in raw_keys.split(",") if k.strip())

    # ── 读取 Provider Keys ──
    providers = {}
    for env_name, provider_name in PROVIDER_MAP.items():
        key_value = os.getenv(env_name, "").strip()
        if not key_value:
            continue  # 没配置该 Provider，跳过

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

    logger.info(f"✅ 加载了 {len(api_keys)} 个 API Key")
    logger.info(f"✅ 加载了 {len(providers)} 个 Provider: {list(providers.keys())}")

    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("/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 with httpx.AsyncClient(timeout=60) as client:
        resp = await client.post(
            f"{provider['base_url']}/chat/completions",
            headers={"Authorization": f"Bearer {provider['api_key']}"},
            json=body,
        )
        return resp.json()


@app.get("/health")
async def health():
    return {"status": "ok"}


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