Update app.py

app.py

@@ -4,7 +4,7 @@ from fastapi import FastAPI, Request, HTTPException, Depends, status
 from fastapi.responses import Response # 导入Response
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials # 导入HTTPBearer和HTTPAuthorizationCredentials
 import httpx # 使用httpx替代requests，因为requests是同步的，而FastAPI是异步的
-import os
+import os, json
 from dotenv import load_dotenv
 from enum import Enum
 
@@ -12,7 +12,7 @@ from enum import Enum
 load_dotenv()
 
 # 从环境变量获取代理自身的API密钥
-PROXY_API_KEY = os.getenv("PROXY_API_KEY")
+proxy_api_key = os.getenv("PROXY_API_KEY")
 
 # 定义HTTPBearer安全方案
 security = HTTPBearer()
@@ -24,7 +24,7 @@ class ProtocolType(str, Enum):
 
 # 依赖函数：验证代理的API密钥
 def verify_proxy_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
-    if not PROXY_API_KEY or credentials.credentials != PROXY_API_KEY:
+    if not proxy_api_key or credentials.credentials != proxy_api_key:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Invalid or missing proxy API key",
@@ -43,27 +43,24 @@ def greet_json():
 async def health_check():
     return {"status": "ok", "message": "Proxy service is running."}
 
-# 1. 配置Gemini真实信息
-GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")  # 从环境变量获取密钥
-GEMINI_BASE_URL = os.getenv("GEMINI_BASE_URL")  # 从环境变量获取Gemini真实Base URL
-# print(f"Using Gemini Base URL: {GEMINI_BASE_URL}")
-# print(f"Using Gemini API Key: {GEMINI_API_KEY}")
+gemini_api_keys_str = os.getenv("GEMINI_API_KEYS_STR")
+try:
+    gemini_api_keys_arr = json.loads(gemini_api_keys_str)
+except json.JSONDecodeError as e:
+    print(f"JSON 解析错误：{e}")
+    gemini_api_keys_arr = []  # 解析失败时设置默认值
+except Exception as e:
+    print(f"其他错误：{e}")
+    gemini_api_keys_arr = []
 
-# 2. 通用转发接口（匹配Gemini的/predict等端点，路径动态匹配）
-@app.api_route("/v1/{protocol}/{host}/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
-async def proxy_gemini(request: Request, protocol: ProtocolType, host:str, path: str, proxy_api_key: str = Depends(verify_proxy_api_key)): # 添加代理认证依赖
-    # 拼接Gemini真实请求URL
-    gemini_url = f"{GEMINI_BASE_URL}/{path}"
-
-
-    print('gemini_url', gemini_url)
-
-    print('protocol.value', protocol.value)
-    print('host', host)
-    print('path', path)
+key_index = 0
+keys_count = len(gemini_api_keys_arr)
 
+# @app.api_route("/v1/{protocol}/{host}/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
+@app.api_route("/v1/{protocol}/{host}/{path:path}", methods=["POST"])
+async def proxy_gemini(request: Request, protocol: ProtocolType, host:str, path: str, proxy_api_key: str = Depends(verify_proxy_api_key)): # 添加代理认证依赖
+    global key_index # 声明使用全局变量
     real_url = f"{protocol.value}://{host}/{path}"
-    print('real_url', real_url)
     
     # 提取客户端请求的headers和body，透传给Gemini
     client_headers = dict(request.headers)
@@ -75,8 +72,16 @@ async def proxy_gemini(request: Request, protocol: ProtocolType, host:str, path:
     client_headers["User-Agent"] = "curl/8.7.1"
     
     # 添加Authorization头，用于Gemini的OpenAI兼容API
-    if GEMINI_API_KEY:
-        client_headers["Authorization"] = f"Bearer {GEMINI_API_KEY}"
+    if gemini_api_keys_arr: # 确保有可用的API密钥
+        print(f"key_index: {key_index}")
+        gemini_api_key = gemini_api_keys_arr[key_index]
+        key_index = (key_index + 1) % keys_count
+        print(f"next_key_index: {key_index}")
+        client_headers["Authorization"] = f"Bearer {gemini_api_key}"
+    elif gemini_api_key: # 如果没有API密钥列表，使用单个API密钥
+        client_headers["Authorization"] = f"Bearer {gemini_api_key}"
+    else:
+        raise HTTPException(status_code=500, detail="No Gemini API keys configured.")
     
     try:
         # 读取客户端请求体
@@ -86,7 +91,6 @@ async def proxy_gemini(request: Request, protocol: ProtocolType, host:str, path:
         async with httpx.AsyncClient() as client:
             response = await client.request(
                 method=request.method,
-                # url=gemini_url,
                 url=real_url,
                 headers=client_headers,
                 content=client_body, # httpx使用content参数传递请求体