Update app.py

app.py

@@ -5,11 +5,12 @@ 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()
@@ -30,14 +31,45 @@ class Config:
     WHITELIST_IPS = os.getenv("WHITELIST_IPS", "").split(",")
     ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "")
     CHUNK_SIZE = 1
+    MAX_RETRIES = 3  # 最大重试次数
+    COOLDOWN_DAYS = 30  # key冷却时间（天）
+
+# Key状态类
+class KeyStatus:
+    def __init__(self, key: str):
+        self.key = key
+        self.is_valid = True
+        self.last_check = datetime.now()
+        self.cooldown_until = None
+        self.usage_count = 0
+
+    def set_cooldown(self):
+        self.cooldown_until = datetime.now() + timedelta(days=Config.COOLDOWN_DAYS)
+        self.is_valid = False
+
+    def check_cooldown(self) -> bool:
+        if self.cooldown_until and datetime.now() > self.cooldown_until:
+            self.cooldown_until = None
+            self.is_valid = True
+            return True
+        return False
+
+    def to_dict(self) -> dict:
+        return {
+            "key": self.key,
+            "is_valid": self.is_valid,
+            "last_check": self.last_check.isoformat(),
+            "cooldown_until": self.cooldown_until.isoformat() if self.cooldown_until else None,
+            "usage_count": self.usage_count
+        }
 
 # 全局变量
-keys = []
-key_cycle = None
+key_statuses: Dict[str, KeyStatus] = {}
+chat_key_cycle = None
 
 # 初始化keys
 def init_keys():
-    global keys, key_cycle
+    global key_statuses, chat_key_cycle
     try:
         if Config.KEYS_URL:
             response = requests.get(Config.KEYS_URL)
@@ -46,13 +78,43 @@ def init_keys():
             with open("key.txt", "r") as f:
                 keys = [k.strip() for k in f.readlines() if k.strip()]
         
-        key_cycle = cycle(keys)
+        # 初始化key状态
+        key_statuses = {k: KeyStatus(k) for k in keys if k not in key_statuses}
+        
+        # 更新轮询器
+        update_key_cycle()
+        
         print(f"Loaded {len(keys)} API keys")
     except Exception as e:
         print(f"Error loading keys: {e}")
-        keys = []
-        key_cycle = cycle(keys)
+        key_statuses = {}
+        chat_key_cycle = None
+
+# 更新轮询器
+def update_key_cycle():
+    global chat_key_cycle
+    valid_keys = [k for k, status in key_statuses.items() 
+                 if status.is_valid and (not status.cooldown_until or status.check_cooldown())]
+    chat_key_cycle = cycle(valid_keys) if valid_keys else None
 
+# 获取下一个可用的key
+async def get_next_valid_key(for_chat: bool = False) -> str:
+    if not key_statuses:
+        raise HTTPException(status_code=500, detail="No API keys available")
+    
+    if for_chat:
+        # 聊天接口使用轮询
+        if not chat_key_cycle:
+            update_key_cycle()
+        if not chat_key_cycle:
+            raise HTTPException(status_code=500, detail="No valid API keys available")
+        return next(chat_key_cycle)
+    else:
+        # 其他接口使用第一个有效key
+        valid_keys = [k for k, status in key_statuses.items() if status.is_valid]
+        if not valid_keys:
+            raise HTTPException(status_code=500, detail="No valid API keys available")
+        return valid_keys[0]
 # 获取真实IP地址
 def get_client_ip(request: Request) -> str:
     forwarded_for = request.headers.get("x-forwarded-for")
@@ -70,21 +132,20 @@ def get_client_ip(request: Request) -> str:
 async def ip_whitelist(request: Request, call_next):
     if Config.WHITELIST_IPS and Config.WHITELIST_IPS[0]:
         client_ip = get_client_ip(request)
+        
+        # 检查是否是管理页面请求
+        is_admin_request = request.url.path.startswith("/admin") or request.url.path.startswith("/api/admin") or request.url.path.startswith("/api/keys")
+        
         print(f"Request from IP: {client_ip}")
         print(f"Allowed IPs: {Config.WHITELIST_IPS}")
-        print(f"Request headers: {dict(request.headers)}")
+        print(f"Is admin request: {is_admin_request}")
         
-        if client_ip not in Config.WHITELIST_IPS:
-            print(f"Access denied for IP: {client_ip}")
-            raise HTTPException(status_code=403, detail="IP not allowed")
+        if is_admin_request and client_ip not in Config.WHITELIST_IPS:
+            print(f"Admin access denied for IP: {client_ip}")
+            raise HTTPException(status_code=403, detail="IP not allowed for admin access")
+            
     return await call_next(request)
 
-# 获取下一个key
-def get_next_key():
-    if not keys:
-        raise HTTPException(status_code=500, detail="No API keys available")
-    return next(key_cycle)
-
 # 添加静态文件支持
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
@@ -110,13 +171,16 @@ async def admin_login(request: Request):
     except json.JSONDecodeError:
         raise HTTPException(status_code=400, detail="Invalid JSON")
 
-# 获取所有keys（需要密码）
+# 获取所有keys及其状态
 @app.get("/api/keys")
 async def get_keys(password: str):
     if password != Config.ADMIN_PASSWORD:
         raise HTTPException(status_code=401, detail="Invalid password")
-    return {"keys": list(keys)}
-# 检查key是否有效
+    return {
+        "keys": [status.to_dict() for status in key_statuses.values()]
+    }
+
+# 检查单个key是否有效
 async def check_key_valid(key: str) -> bool:
     headers = {
         "Authorization": f"Bearer {key}",
@@ -133,66 +197,102 @@ async def check_key_valid(key: str) -> bool:
         except:
             return False
 
-# 添加新key（需要密码）
-@app.post("/api/keys")
-async def add_key(key: str, password: str):
+# 批量检查keys
+@app.post("/api/keys/check-all")
+async def check_all_keys(password: str):
     if password != Config.ADMIN_PASSWORD:
         raise HTTPException(status_code=401, detail="Invalid password")
+    
+    results = []
+    for key in list(key_statuses.keys()):
+        is_valid = await check_key_valid(key)
+        key_statuses[key].is_valid = is_valid
+        key_statuses[key].last_check = datetime.now()
+        results.append({"key": key, "valid": is_valid})
+    
+    update_key_cycle()
+    return {"results": results}
+
+# 批量删除keys
+@app.post("/api/keys/delete-batch")
+async def delete_keys_batch(request: Request):
+    try:
+        body = await request.json()
+        password = body.get("password")
+        keys_to_delete = body.get("keys", [])
+        
+        if password != Config.ADMIN_PASSWORD:
+            raise HTTPException(status_code=401, detail="Invalid password")
+            
+        deleted_keys = []
+        for key in keys_to_delete:
+            if key in key_statuses:
+                del key_statuses[key]
+                deleted_keys.append(key)
         
-    if key not in keys:
+        update_key_cycle()
+        return {"status": "success", "deleted_keys": deleted_keys}
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+# 添加新key
+@app.post("/api/keys/add")
+async def add_key(request: Request):
+    try:
+        body = await request.json()
+        password = body.get("password")
+        key = body.get("key")
+        
+        if password != Config.ADMIN_PASSWORD:
+            raise HTTPException(status_code=401, detail="Invalid password")
+            
+        if key in key_statuses:
+            raise HTTPException(status_code=400, detail="Key already exists")
+            
         # 检查key是否有效
         is_valid = await check_key_valid(key)
         if not is_valid:
             raise HTTPException(status_code=400, detail="Invalid key")
             
-        keys.append(key)
-        global key_cycle
-        key_cycle = cycle(keys)
+        key_statuses[key] = KeyStatus(key)
+        update_key_cycle()
         return {"status": "success", "message": "Key added successfully"}
-    
-    raise HTTPException(status_code=400, detail="Key already exists")
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
 
-# 删除key（需要密码）
+# 删除单个key
 @app.delete("/api/keys/{key}")
 async def delete_key(key: str, password: str):
     if password != Config.ADMIN_PASSWORD:
         raise HTTPException(status_code=401, detail="Invalid password")
         
-    if key in keys:
-        keys.remove(key)
-        global key_cycle
-        key_cycle = cycle(keys)
+    if key in key_statuses:
+        del key_statuses[key]
+        update_key_cycle()
         return {"status": "success", "message": "Key deleted successfully"}
     
     raise HTTPException(status_code=404, detail="Key not found")
-
-# 检查key状态（需要密码）
-@app.get("/api/keys/check/{key}")
-async def check_key(key: str, password: str):
-    if password != Config.ADMIN_PASSWORD:
-        raise HTTPException(status_code=401, detail="Invalid password")
+# 处理API错误响应
+async def handle_api_error(response: httpx.Response, key: str) -> bool:
+    """
+    处理API错误响应，返回是否需要重试
+    """
+    try:
+        error_data = response.json()
+        error_message = error_data.get('error', {}).get('message', '').lower()
         
-    is_valid = await check_key_valid(key)
-    return {"valid": is_valid}
-
-# 模型列表路由
-@app.get("/api/v1/models")
-async def list_models():
-    headers = {
-        "Authorization": f"Bearer {get_next_key()}",
-        "Content-Type": "application/json"
-    }
-    
-    async with httpx.AsyncClient() as client:
-        try:
-            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))
+        # 检查是否是额度不足错误
+        if any(msg in error_message for msg in ['rate limit', 'quota exceeded', 'insufficient_quota']):
+            print(f"Key {key} quota exceeded, setting cooldown")
+            key_statuses[key].set_cooldown()
+            update_key_cycle()
+            return True
+            
+        return False
+    except:
+        return False
 
 # 流式响应生成器
 async def stream_generator(response):
@@ -202,19 +302,16 @@ 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: '):
-                    data = event[6:]  # 移除 "data: " 前缀
+                    data = event[6:]
                     if data.strip() == '[DONE]':
                         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}")
@@ -226,62 +323,89 @@ async def stream_generator(response):
 # 聊天完成路由
 @app.post("/api/v1/chat/completions")
 async def chat_completions(request: Request):
+    body = await request.body()
+    body_json = json.loads(body)
+    
+    for attempt in range(Config.MAX_RETRIES):
+        try:
+            key = await get_next_valid_key(for_chat=True)
+            
+            headers = {
+                "Authorization": f"Bearer {key}",
+                "Content-Type": "application/json",
+                "Accept": "text/event-stream" if body_json.get("stream") else "application/json"
+            }
+            
+            url = f"{Config.OPENAI_API_BASE}/chat/completions"
+            
+            print(f"Chat request to: {url} (attempt {attempt + 1})")
+            print(f"Using key: {key}")
+            
+            async with httpx.AsyncClient(timeout=60.0) as client:
+                response = await client.post(
+                    url,
+                    headers=headers,
+                    json=body_json
+                )
+                
+                if response.status_code != 200:
+                    needs_retry = await handle_api_error(response, key)
+                    if needs_retry and attempt < Config.MAX_RETRIES - 1:
+                        continue
+                        
+                    return Response(
+                        content=response.text,
+                        status_code=response.status_code,
+                        media_type=response.headers.get("content-type", "application/json")
+                    )
+                
+                # 更新使用计数
+                key_statuses[key].usage_count += 1
+                
+                if body_json.get("stream"):
+                    return StreamingResponse(
+                        stream_generator(response),
+                        media_type="text/event-stream",
+                        headers={
+                            "Cache-Control": "no-cache",
+                            "Connection": "keep-alive",
+                            "Content-Type": "text/event-stream"
+                        }
+                    )
+                
+                return Response(
+                    content=response.text,
+                    media_type=response.headers.get("content-type", "application/json")
+                )
+                    
+        except Exception as e:
+            if attempt == Config.MAX_RETRIES - 1:
+                print(f"Chat Error: {str(e)}")
+                raise HTTPException(status_code=500, detail=str(e))
+
+# 模型列表路由
+@app.get("/api/v1/models")
+async def list_models():
     try:
-        # 获取请求体
-        body = await request.body()
-        body_json = json.loads(body)
+        key = await get_next_valid_key(for_chat=False)
         
-        # 获取headers
         headers = {
-            "Authorization": f"Bearer {get_next_key()}",
-            "Content-Type": "application/json",
-            "Accept": "text/event-stream" if body_json.get("stream") else "application/json"
+            "Authorization": f"Bearer {key}",
+            "Content-Type": "application/json"
         }
         
-        # 构建目标URL
-        url = f"{Config.OPENAI_API_BASE}/chat/completions"
-        
-        print(f"Chat request to: {url}")
-        print(f"Request from IP: {get_client_ip(request)}")
-        print(f"Headers: {headers}")
-        print(f"Body: {json.dumps(body_json)}")
-        
-        async with httpx.AsyncClient(timeout=60.0) as client:
-            response = await client.post(
-                url,
-                headers=headers,
-                json=body_json
+        async with httpx.AsyncClient() as client:
+            response = await client.get(
+                f"{Config.OPENAI_API_BASE}/models",
+                headers=headers
             )
             
-            # 检查响应状态
             if response.status_code != 200:
-                print(f"Error response: {response.status_code} - {response.text}")
-                return Response(
-                    content=response.text,
-                    status_code=response.status_code,
-                    media_type=response.headers.get("content-type", "application/json")
-                )
+                await handle_api_error(response, key)
             
-            # 处理流式响应
-            if body_json.get("stream"):
-                return StreamingResponse(
-                    stream_generator(response),
-                    media_type="text/event-stream",
-                    headers={
-                        "Cache-Control": "no-cache",
-                        "Connection": "keep-alive",
-                        "Content-Type": "text/event-stream"
-                    }
-                )
-            
-            # 处理普通响应
-            return Response(
-                content=response.text,
-                media_type=response.headers.get("content-type", "application/json")
-            )
-                
+            return response.json()
     except Exception as e:
-        print(f"Chat Error: {str(e)}")
+        print(f"Models Error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
 # 代理其他请求到X.AI
@@ -290,27 +414,25 @@ async def proxy(path: str, request: Request):
     if path == "chat/completions":
         return await chat_completions(request)
         
-    # 获取请求体
-    body = await request.body()
-    body_str = body.decode() if body else ""
-    
-    # 获取查询参数
-    params = dict(request.query_params)
-    
-    # 获取并处理headers
-    headers = dict(request.headers)
-    headers.pop("host", None)
-    headers["Authorization"] = f"Bearer {get_next_key()}"
-    headers["Content-Type"] = "application/json"
-    
-    # 构建目标URL
-    url = f"{Config.OPENAI_API_BASE}/{path}"
-    
-    print(f"Proxy request to: {url}")
-    print(f"Request from IP: {get_client_ip(request)}")
-    
-    async with httpx.AsyncClient() as client:
-        try:
+    try:
+        key = await get_next_valid_key(for_chat=False)
+        
+        body = await request.body()
+        body_str = body.decode() if body else ""
+        
+        params = dict(request.query_params)
+        
+        headers = dict(request.headers)
+        headers.pop("host", None)
+        headers["Authorization"] = f"Bearer {key}"
+        headers["Content-Type"] = "application/json"
+        
+        url = f"{Config.OPENAI_API_BASE}/{path}"
+        
+        print(f"Proxy request to: {url}")
+        print(f"Using key: {key}")
+        
+        async with httpx.AsyncClient() as client:
             response = await client.request(
                 method=request.method,
                 url=url,
@@ -319,20 +441,30 @@ async def proxy(path: str, request: Request):
                 content=body_str if body_str else None
             )
             
+            if response.status_code != 200:
+                await handle_api_error(response, key)
+            
             return Response(
                 content=response.text,
                 status_code=response.status_code,
                 headers=dict(response.headers)
             )
                 
-        except Exception as e:
-            print(f"Proxy Error: {str(e)}")
-            raise HTTPException(status_code=500, detail=str(e))
+    except Exception as e:
+        print(f"Proxy Error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
 
 # 健康检查路由
 @app.get("/api/health")
 async def health_check():
-    return {"status": "healthy", "key_count": len(keys)}
+    valid_keys = len([k for k, status in key_statuses.items() if status.is_valid])
+    total_keys = len(key_statuses)
+    return {
+        "status": "healthy",
+        "total_keys": total_keys,
+        "valid_keys": valid_keys,
+        "keys_in_cooldown": total_keys - valid_keys
+    }
 
 # 启动时初始化
 @app.on_event("startup")