Update app.py

app.py

@@ -1,14 +1,16 @@
 from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
 import httpx
 import os
 import json
-from typing import List, Optional
+from typing import List, Optional, Dict
 import requests
 from itertools import cycle
 import asyncio
 import time
+from datetime import datetime, timedelta
 
 # 创建FastAPI应用
 app = FastAPI()
@@ -27,15 +29,56 @@ class Config:
     OPENAI_API_BASE = "https://api.x.ai/v1"
     KEYS_URL = os.getenv("KEYS_URL", "")
     WHITELIST_IPS = os.getenv("WHITELIST_IPS", "").split(",")
+    ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "admin")
+
+# 密钥状态类
+class KeyStatus:
+    def __init__(self, key: str):
+        self.key = key
+        self.status = "valid"  # valid, invalid, cooling
+        self.cooling_until = None
+        self.last_check = None
 
 # 全局变量
 keys = []
 key_cycle = None
 first_key = None
+key_status_map: Dict[str, KeyStatus] = {}
+
+# 静态文件挂载
+app.mount("/static", StaticFiles(directory="static"), name="static")
 
+# 管理页面路由
+@app.get("/admin")
+async def admin():
+    return FileResponse("static/admin.html")
+# 获取真实IP地址
+def get_client_ip(request: Request) -> str:
+    # 尝试从各种头部获取IP
+    forwarded_for = request.headers.get("x-forwarded-for")
+    if forwarded_for:
+        return forwarded_for.split(",")[0].strip()
+    
+    real_ip = request.headers.get("x-real-ip")
+    if real_ip:
+        return real_ip
+    
+    return request.client.host
+
+# IP白名单中间件
+@app.middleware("http")
+async def ip_whitelist(request: Request, call_next):
+    # 只对API接口启用白名单，排除管理后台
+    if "/api/" in request.url.path and "/api/admin/" not in request.url.path and "/api/keys" not in request.url.path:
+        if Config.WHITELIST_IPS and Config.WHITELIST_IPS[0]:
+            client_ip = get_client_ip(request)
+            if client_ip not in Config.WHITELIST_IPS:
+                raise HTTPException(status_code=403, detail="IP not allowed")
+    return await call_next(request)
+    
 # 初始化keys
 def init_keys():
-    global keys, key_cycle, first_key
+    global keys, key_cycle, first_key, key_status_map
     try:
         if Config.KEYS_URL:
             response = requests.get(Config.KEYS_URL)
@@ -47,40 +90,69 @@ def init_keys():
         if keys:
             first_key = keys[0]
             key_cycle = cycle(keys)
+            # 初始化所有key的状态
+            key_status_map = {key: KeyStatus(key) for key in keys}
             print(f"Loaded {len(keys)} API keys")
     except Exception as e:
         print(f"Error loading keys: {e}")
         keys = []
         key_cycle = None
         first_key = None
+        key_status_map = {}
 
-# 获取聊天key
-def get_chat_key():
+# 获取有效的key
+def get_valid_key():
     global key_cycle
     if not key_cycle:
         raise HTTPException(status_code=500, detail="No API keys available")
-    return next(key_cycle)
-# 获取真实IP地址
-def get_client_ip(request: Request) -> str:
-    # 尝试从各种头部获取IP
-    forwarded_for = request.headers.get("x-forwarded-for")
-    if forwarded_for:
-        return forwarded_for.split(",")[0].strip()
-    
-    real_ip = request.headers.get("x-real-ip")
-    if real_ip:
-        return real_ip
     
-    return request.client.host
+    # 尝试最多len(keys)次
+    for _ in range(len(keys)):
+        key = next(key_cycle)
+        key_info = key_status_map.get(key)
+        if not key_info:
+            key_info = KeyStatus(key)
+            key_status_map[key] = key_info
+            
+        # 检查key是否可用
+        if key_info.status == "valid":
+            return key
+        elif key_info.status == "cooling":
+            if key_info.cooling_until and datetime.now() > key_info.cooling_until:
+                key_info.status = "valid"
+                key_info.cooling_until = None
+                return key
+                
+    raise HTTPException(status_code=500, detail="No valid API keys available")
 
-# IP白名单中间件
-@app.middleware("http")
-async def ip_whitelist(request: Request, call_next):
-    if Config.WHITELIST_IPS and Config.WHITELIST_IPS[0]:
-        client_ip = get_client_ip(request)
-        if client_ip not in Config.WHITELIST_IPS:
-            raise HTTPException(status_code=403, detail="IP not allowed")
-    return await call_next(request)
+# 标记key为冷却状态
+def mark_key_cooling(key: str):
+    if key in key_status_map:
+        key_status_map[key].status = "cooling"
+        key_status_map[key].cooling_until = datetime.now() + timedelta(days=30)
+
+# 检查key状态
+async def check_key_status(key: str) -> bool:
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.get(
+                f"{Config.OPENAI_API_BASE}/models",
+                headers={"Authorization": f"Bearer {key}"}
+            )
+            is_valid = response.status_code == 200
+            key_info = key_status_map.get(key)
+            if key_info:
+                if not is_valid:
+                    key_info.status = "cooling"
+                    key_info.cooling_until = datetime.now() + timedelta(days=30)
+                else:
+                    key_info.status = "valid"
+                    key_info.cooling_until = None
+                key_info.last_check = datetime.now()
+            return is_valid
+    except Exception as e:
+        print(f"Error checking key {key}: {e}")
+        return False
 
 # 流式响应生成器
 async def stream_generator(response):
@@ -90,7 +162,6 @@ async def stream_generator(response):
             chunk_str = chunk.decode('utf-8')
             buffer += chunk_str
             
-            # 处理buffer中的完整事件
             while '\n\n' in buffer:
                 event, buffer = buffer.split('\n\n', 1)
                 if event.startswith('data: '):
@@ -99,35 +170,165 @@ async def stream_generator(response):
                         yield f"data: [DONE]\n\n"
                     else:
                         try:
-                            # 解析JSON并重新格式化
                             json_data = json.loads(data)
                             yield f"data: {json.dumps(json_data)}\n\n"
-                            # 添加小延迟使流更平滑
-                            await asyncio.sleep(0.01)
                         except json.JSONDecodeError:
                             print(f"JSON decode error for data: {data}")
                             continue
     except Exception as e:
         print(f"Stream Error: {str(e)}")
         yield f"data: {json.dumps({'error': str(e)})}\n\n"
+# 验证管理员密码
+def verify_admin(password: str):
+    if not password or password != Config.ADMIN_PASSWORD:
+        raise HTTPException(status_code=403, detail="Invalid admin password")
+
+# 管理员登录
+@app.post("/api/admin/login")
+async def admin_login(request: Request):
+    data = await request.json()
+    password = data.get("password")
+    verify_admin(password)
+    return {"status": "success"}
+
+# 获取所有密钥状态
+@app.get("/api/keys")
+async def list_keys(password: str):
+    verify_admin(password)
+    return {
+        "keys": [
+            {
+                "key": k,
+                "status": key_status_map[k].status if k in key_status_map else "valid",
+                "cooling_until": key_status_map[k].cooling_until.isoformat() if k in key_status_map and key_status_map[k].cooling_until else None,
+                "last_check": key_status_map[k].last_check.isoformat() if k in key_status_map and key_status_map[k].last_check else None
+            }
+            for k in keys
+        ]
+    }
+
+# 添加新密钥
+@app.post("/api/keys/add")
+async def add_key(request: Request):
+    data = await request.json()
+    verify_admin(data.get("password"))
+    new_key = data.get("key", "").strip()
+    
+    if not new_key:
+        raise HTTPException(status_code=400, detail="Key is required")
+    if new_key in keys:
+        raise HTTPException(status_code=400, detail="Key already exists")
+        
+    keys.append(new_key)
+    key_status_map[new_key] = KeyStatus(new_key)
+    
+    # 重新初始化key_cycle
+    global key_cycle
+    key_cycle = cycle(keys)
+    
+    # 保存到文件
+    if not Config.KEYS_URL:
+        with open("key.txt", "w") as f:
+            f.write("\n".join(keys))
+            
+    return {"status": "success"}
+
+# 删除密钥
+@app.delete("/api/keys/{key}")
+async def delete_key(key: str, password: str):
+    verify_admin(password)
+    if key in keys:
+        keys.remove(key)
+        if key in key_status_map:
+            del key_status_map[key]
+            
+        # 重新初始化key_cycle
+        global key_cycle
+        key_cycle = cycle(keys)
+        
+        # 保存到文件
+        if not Config.KEYS_URL:
+            with open("key.txt", "w") as f:
+                f.write("\n".join(keys))
+                
+    return {"status": "success"}
+
+# 批量删除密钥
+@app.post("/api/keys/delete-batch")
+async def delete_keys_batch(request: Request):
+    data = await request.json()
+    verify_admin(data.get("password"))
+    keys_to_delete = data.get("keys", [])
+    
+    for key in keys_to_delete:
+        if key in keys:
+            keys.remove(key)
+            if key in key_status_map:
+                del key_status_map[key]
+                
+    # 重新初始化key_cycle
+    global key_cycle
+    key_cycle = cycle(keys)
+    
+    # 保存到文件
+    if not Config.KEYS_URL:
+        with open("key.txt", "w") as f:
+            f.write("\n".join(keys))
+            
+    return {"status": "success"}
+
+# 检查单个密钥
+@app.get("/api/keys/check/{key}")
+async def check_single_key(key: str, password: str):
+    verify_admin(password)
+    if key not in keys:
+        raise HTTPException(status_code=404, detail="Key not found")
+    
+    is_valid = await check_key_status(key)
+    return {"status": "success", "valid": is_valid}
+
+# 检查所有密钥
+@app.post("/api/keys/check-all")
+async def check_all_keys(password: str):
+    verify_admin(password)
+    for key in keys:
+        await check_key_status(key)
+    return {"status": "success"}
 # 模型列表路由
 @app.get("/api/v1/models")
 async def list_models():
-    headers = {
-        "Authorization": f"Bearer {get_chat_key()}",
-        "Content-Type": "application/json"
-    }
-    
-    async with httpx.AsyncClient() as client:
-        try:
+    # 对于模型列表，使用第一个可用的key
+    try:
+        key = first_key
+        if key in key_status_map and key_status_map[key].status == "cooling":
+            key = get_valid_key()  # 如果第一个key在冷却，则使用轮询获取
+            
+        headers = {
+            "Authorization": f"Bearer {key}",
+            "Content-Type": "application/json"
+        }
+        
+        async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{Config.OPENAI_API_BASE}/models",
                 headers=headers
             )
+            
+            if response.status_code == 429:  # 如果遇到限流
+                mark_key_cooling(key)
+                # 重试一次，使用轮询的key
+                key = get_valid_key()
+                headers["Authorization"] = f"Bearer {key}"
+                response = await client.get(
+                    f"{Config.OPENAI_API_BASE}/models",
+                    headers=headers
+                )
+                
             return response.json()
-        except Exception as e:
-            print(f"Models Error: {str(e)}")
-            raise HTTPException(status_code=500, detail=str(e))
+            
+    except Exception as e:
+        print(f"Models Error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
 
 # 聊天完成路由
 @app.post("/api/v1/chat/completions")
@@ -137,9 +338,12 @@ async def chat_completions(request: Request):
         body = await request.body()
         body_json = json.loads(body)
         
+        # 使用轮询获取key
+        key = get_valid_key()
+        
         # 获取headers
         headers = {
-            "Authorization": f"Bearer {get_chat_key()}",
+            "Authorization": f"Bearer {key}",
             "Content-Type": "application/json",
             "Accept": "text/event-stream" if body_json.get("stream") else "application/json"
         }
@@ -154,6 +358,18 @@ async def chat_completions(request: Request):
                 json=body_json
             )
             
+            # 如果遇到限流，标记key为冷却状态并重试
+            if response.status_code == 429:
+                mark_key_cooling(key)
+                # 重试一次
+                key = get_valid_key()
+                headers["Authorization"] = f"Bearer {key}"
+                response = await client.post(
+                    url,
+                    headers=headers,
+                    json=body_json
+                )
+            
             # 检查响应状态
             if response.status_code != 200:
                 return Response(
@@ -183,6 +399,7 @@ async def chat_completions(request: Request):
     except Exception as e:
         print(f"Chat Error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
+
 # 代理其他请求
 @app.api_route("/api/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
 async def proxy(path: str, request: Request):
@@ -192,8 +409,14 @@ async def proxy(path: str, request: Request):
     try:
         method = request.method
         body = await request.body() if method in ["POST", "PUT"] else None
+        
+        # 使用第一个可用的key
+        key = first_key
+        if key in key_status_map and key_status_map[key].status == "cooling":
+            key = get_valid_key()
+            
         headers = {
-            "Authorization": f"Bearer {get_chat_key()}",
+            "Authorization": f"Bearer {key}",
             "Content-Type": "application/json"
         }
         
@@ -205,12 +428,24 @@ async def proxy(path: str, request: Request):
                 content=body
             )
             
+            if response.status_code == 429:  # 如果遇到限流
+                mark_key_cooling(key)
+                # 重试一次，使用轮询的key
+                key = get_valid_key()
+                headers["Authorization"] = f"Bearer {key}"
+                response = await client.request(
+                    method=method,
+                    url=f"{Config.OPENAI_API_BASE}/{path}",
+                    headers=headers,
+                    content=body
+                )
+            
             return Response(
                 content=response.text,
                 status_code=response.status_code,
                 media_type=response.headers.get("content-type", "application/json")
             )
-            
+                
     except Exception as e:
         print(f"Proxy Error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))