Update main.py

main.py

@@ -1,9 +1,12 @@
-from fastapi import FastAPI, HTTPException, Depends, Header
+from fastapi import FastAPI, HTTPException, Header, Depends, Request
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, JSONResponse
 import uvicorn
 import os
 import logging
 import httpx
+import json
+import asyncio
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -18,47 +21,56 @@ app.add_middleware(
 )
 
 # ──────────────────────────────────────────
-# 读取原项目固定命名的 Secrets
+# Provider 配置
 # ──────────────────────────────────────────
 
 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",
+    "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"
+            "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"
+            "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"
+            "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": {
@@ -68,13 +80,16 @@ PROVIDER_CONFIG = {
             "meta-llama/llama-3.2-3b-instruct:free",
             "google/gemma-3-1b-it:free",
             "deepseek/deepseek-r1:free",
+            "deepseek/deepseek-chat:free",
         ],
     },
     "mistral": {
         "base_url": "https://api.mistral.ai/v1",
         "models": [
-            "mistral-small-latest","mistral-large-latest",
-            "open-mistral-7b","open-mixtral-8x7b"
+            "mistral-small-latest",
+            "mistral-large-latest",
+            "open-mistral-7b",
+            "open-mixtral-8x7b",
         ],
     },
     "together": {
@@ -82,7 +97,8 @@ PROVIDER_CONFIG = {
         "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"
+            "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
         ],
     },
     "nvidia": {
@@ -90,46 +106,61 @@ PROVIDER_CONFIG = {
         "models": [
             "meta/llama-3.1-70b-instruct",
             "meta/llama-3.1-8b-instruct",
-            "mistralai/mixtral-8x7b-instruct"
+            "mistralai/mixtral-8x7b-instruct",
         ],
     },
     "cohere": {
         "base_url": "https://api.cohere.com/v2",
-        "models": ["command-r-plus","command-r","command"],
+        "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"
+            "mistralai/Mistral-7B-Instruct-v0.3",
         ],
     },
     "cerebras": {
         "base_url": "https://api.cerebras.ai/v1",
-        "models": ["llama3.1-8b","llama3.1-70b"],
+        "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"
+            "Meta-Llama-3.1-405B-Instruct",
         ],
     },
     "cloudflare": {
-        "base_url": "https://api.cloudflare.com/client/v4/accounts/{}/ai/v1",
+        "base_url": "https://api.cloudflare.com/client/v4/accounts/CF_ACCOUNT_ID/ai/v1",
         "models": [
             "@cf/meta/llama-3.1-8b-instruct",
-            "@cf/mistral/mistral-7b-instruct-v0.1"
+            "@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"],
+        "models": [
+            "glm-4-flash",
+            "glm-4",
+            "glm-3-turbo",
+        ],
     },
 }
 
 
+# ──────────────────────────────────────────
+# 启动加载
+# ──────────────────────────────────────────
+
 def load_config():
     raw_keys = os.getenv("API_KEYS", "")
     api_keys = set(k.strip() for k in raw_keys.split(",") if k.strip())
@@ -140,20 +171,34 @@ def load_config():
         if not key_value:
             continue
         cfg = PROVIDER_CONFIG.get(provider_name, {})
+        base_url = cfg.get("base_url", "")
+
+        # Cloudflare 需要 Account ID
+        if provider_name == "cloudflare":
+            cf_account = os.getenv("CLOUDFLARE_ACCOUNT_ID", "")
+            base_url = base_url.replace("CF_ACCOUNT_ID", cf_account)
+
         providers[provider_name] = {
             "api_key":  key_value,
-            "base_url": cfg.get("base_url", ""),
+            "base_url": base_url.rstrip("/"),
             "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())}")
+    logger.info(f"✅ API Keys 数量: {len(api_keys)}")
+    logger.info(f"✅ Provider 数量: {len(providers)} {list(providers.keys())}")
     return api_keys, providers
 
 
 API_KEYS, PROVIDERS = load_config()
 
+# model → provider 快速查找表
+MODEL_PROVIDER: dict = {}
+for _pname, _pcfg in PROVIDERS.items():
+    for _m in _pcfg["models"]:
+        MODEL_PROVIDER[_m] = _pname
+
+
 # ──────────────────────────────────────────
 # 鉴权
 # ──────────────────────────────────────────
@@ -164,81 +209,167 @@ def verify_api_key(authorization: str = Header(...)):
         raise HTTPException(status_code=401, detail="Invalid API key")
     return token
 
+
 # ──────────────────────────────────────────
-# ★ 用 APIRouter 统一加 /v1 前缀
+# /v1 路由
 # ──────────────────────────────────────────
 
 from fastapi import APIRouter
 
-router = APIRouter(prefix="/v1")  # ← 关键！所有路由自动带 /v1
+v1 = APIRouter(prefix="/v1")
 
 
-@router.get("/models")
+# ── GET /v1 ── 访问 /v1 不报 404，返回说明
+@v1.get("")
+@v1.get("/")
+async def v1_root():
+    return {
+        "message": "FreeLLMAPI is running",
+        "endpoints": {
+            "models":           "/v1/models",
+            "chat_completions": "/v1/chat/completions",
+        },
+        "providers": list(PROVIDERS.keys()),
+        "total_models": len(MODEL_PROVIDER),
+    }
+
+
+# ── GET /v1/models ──
+@v1.get("/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"})
+    for provider_name, pcfg in PROVIDERS.items():
+        for m in pcfg["models"]:
+            data.append({
+                "id":       m,
+                "object":   "model",
+                "owned_by": provider_name,
+                "created":  0,
+            })
     return {"object": "list", "data": data}
 
 
-@router.post("/chat/completions")
+# ── POST /v1/chat/completions ──
+@v1.post("/chat/completions")
 async def chat_completions(
-    body: dict,
-    _: str = Depends(verify_api_key)
+    request: Request,
+    _: 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
+    body  = await request.json()
+    model = body.get("model", "")
+    stream = body.get("stream", False)
 
-    if not provider:
+    # 查找 provider
+    provider_name = MODEL_PROVIDER.get(model)
+    if not provider_name:
         raise HTTPException(
             status_code=404,
-            detail=f"没有 Provider 支持模型: {model}"
+            detail=f"模型 '{model}' 不存在，可用模型: {list(MODEL_PROVIDER.keys())}",
         )
 
-    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,
+    provider = PROVIDERS[provider_name]
+    url      = f"{provider['base_url']}/chat/completions"
+    headers  = {
+        "Authorization": f"Bearer {provider['api_key']}",
+        "Content-Type":  "application/json",
+    }
+
+    logger.info(f"转发请求 → {provider_name} | 模型: {model} | 流式: {stream}")
+
+    # ── 流式响应 ──
+    if stream:
+        async def event_stream():
+            async with httpx.AsyncClient(timeout=120) as client:
+                async with client.stream(
+                    "POST", url, headers=headers, json=body
+                ) as resp:
+                    async for chunk in resp.aiter_text():
+                        yield chunk
+
+        return StreamingResponse(
+            event_stream(),
+            media_type="text/event-stream",
+        )
+
+    # ── 普通响应 ──
+    async with httpx.AsyncClient(timeout=120) as client:
+        resp = await client.post(url, headers=headers, json=body)
+
+    if resp.status_code != 200:
+        logger.error(f"上游错误 {resp.status_code}: {resp.text[:300]}")
+        raise HTTPException(
+            status_code=resp.status_code,
+            detail=resp.text,
         )
-        return resp.json()
 
+    return resp.json()
 
-# ★ 把 router 注册到 app
-app.include_router(router)
+
+# ── POST /v1/embeddings（部分 Provider 支持）──
+@v1.post("/embeddings")
+async def embeddings(
+    request: Request,
+    _: str = Depends(verify_api_key),
+):
+    body  = await request.json()
+    model = body.get("model", "")
+
+    provider_name = MODEL_PROVIDER.get(model)
+    if not provider_name:
+        raise HTTPException(status_code=404, detail=f"模型 '{model}' 不存在")
+
+    provider = PROVIDERS[provider_name]
+    url      = f"{provider['base_url']}/embeddings"
+    headers  = {"Authorization": f"Bearer {provider['api_key']}"}
+
+    async with httpx.AsyncClient(timeout=60) as client:
+        resp = await client.post(url, headers=headers, json=body)
+
+    return resp.json()
+
+
+# 注册路由
+app.include_router(v1)
+
+
+# ──────────────────────────────────────────
+# 根路由
+# ──────────────────────────────────────────
+
+@app.get("/")
+async def root():
+    return {
+        "message":      "FreeLLMAPI",
+        "version":      "1.0.0",
+        "docs":         "/docs",
+        "health":       "/health",
+        "api_base":     "/v1",
+        "providers":    list(PROVIDERS.keys()),
+        "total_models": len(MODEL_PROVIDER),
+    }
 
 
-# 健康检查不需要 /v1 前缀，保持根路径
 @app.get("/health")
 async def health():
     return {
-        "status":    "ok",
-        "keys":      len(API_KEYS),
-        "providers": list(PROVIDERS.keys()),
+        "status":        "ok",
+        "api_keys":      len(API_KEYS),
+        "providers":     list(PROVIDERS.keys()),
+        "total_models":  len(MODEL_PROVIDER),
     }
+
+
 # ──────────────────────────────────────────
-# 调试路由（排错用）
+# 调试路由
 # ──────────────────────────────────────────
 
 @app.get("/debug")
 async def debug():
-    """查看所有配置和路由，排错用"""
-    
-    # 1. 检查环境变量
     env_check = {}
     for env_name in PROVIDER_MAP.keys():
         val = os.getenv(env_name, "")
         env_check[env_name] = "✅ 已设置" if val else "❌ 未设置"
-    
-    api_keys_raw = os.getenv("API_KEYS", "")
-    
-    # 2. 检查所有注册路由
+
     routes = []
     for route in app.routes:
         if hasattr(route, "methods"):
@@ -246,57 +377,52 @@ async def debug():
                 "path":    route.path,
                 "methods": list(route.methods),
             })
-    
-    # 3. 检查已加载的 Provider
-    providers_loaded = {}
+
+    providers_info = {}
     for name, cfg in PROVIDERS.items():
-        providers_loaded[name] = {
+        providers_info[name] = {
             "base_url": cfg["base_url"],
             "models":   cfg["models"],
-            "api_key":  cfg["api_key"][:6] + "******",  # 只显示前6位
+            "api_key":  cfg["api_key"][:6] + "******",
         }
-    
+
     return {
-        "① 环境变量检查":   env_check,
-        "② API_KEYS设置":  "✅ 已设置" if api_keys_raw else "❌ 未设置（这是导致401的原因）",
-        "③ 已加载Provider": providers_loaded,
-        "④ 已注册路由":     routes,
-        "⑤ 内存中Key数量":  len(API_KEYS),
+        "① 环境变量":    env_check,
+        "② API_KEYS":   "✅ 已设置" if os.getenv("API_KEYS") else "❌ 未设置",
+        "③ Providers":  providers_info,
+        "④ 全部模型":    list(MODEL_PROVIDER.keys()),
+        "⑤ 注册路由":    routes,
+        "⑥ Key数量":    len(API_KEYS),
     }
 
 
-@app.get("/debug/test-provider/{provider_name}")
+@app.get("/debug/test/{provider_name}")
 async def test_provider(provider_name: str):
-    """测试某个 Provider 是否能正常请求"""
-    
     if provider_name not in PROVIDERS:
         return {
-            "status": "❌ 失败",
-            "reason": f"Provider '{provider_name}' 未加载",
-            "已加载的": list(PROVIDERS.keys()),
+            "status": "❌ 未加载",
+            "已加载":  list(PROVIDERS.keys()),
         }
-    
+
     cfg = PROVIDERS[provider_name]
-    
     try:
         async with httpx.AsyncClient(timeout=10) as client:
             resp = await client.get(
                 f"{cfg['base_url']}/models",
                 headers={"Authorization": f"Bearer {cfg['api_key']}"},
             )
-            return {
-                "status":      "✅ 成功" if resp.status_code == 200 else f"⚠️ 状态码 {resp.status_code}",
-                "provider":    provider_name,
-                "base_url":    cfg["base_url"],
-                "status_code": resp.status_code,
-                "response":    resp.json() if resp.status_code == 200 else resp.text[:200],
-            }
-    except Exception as e:
         return {
-            "status":   "❌ 请求失败",
-            "provider": provider_name,
-            "error":    str(e),
+            "status":      "✅ 连通" if resp.status_code == 200 else f"⚠️ {resp.status_code}",
+            "provider":    provider_name,
+            "status_code": resp.status_code,
         }
+    except Exception as e:
+        return {"status": "❌ 失败", "error": str(e)}
+
+
+# ──────────────────────────────────────────
+# 启动
+# ──────────────────────────────────────────
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)