Update app.py

app.py

@@ -31,121 +31,107 @@ 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:
+# Key信息类
+class KeyInfo:
     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
+        self.cooling_until = None
+        self.last_used = None
+        self.error_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:
+    def to_dict(self):
         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
+            "cooling_until": self.cooling_until.isoformat() if self.cooling_until else None,
+            "last_used": self.last_used.isoformat() if self.last_used else None,
+            "error_count": self.error_count,
+            "status": self.get_status()
         }
 
+    def get_status(self):
+        if self.cooling_until and self.cooling_until > datetime.now():
+            return "cooling"
+        if not self.is_valid:
+            return "invalid"
+        return "valid"
+
 # 全局变量
-key_statuses: Dict[str, KeyStatus] = {}
-chat_key_cycle = None
+keys_info: Dict[str, KeyInfo] = {}
+chat_keys = []
+first_key = None
 
 # 初始化keys
 def init_keys():
-    global key_statuses, chat_key_cycle
+    global keys_info, chat_keys, first_key
     try:
         if Config.KEYS_URL:
             response = requests.get(Config.KEYS_URL)
-            keys = [k.strip() for k in response.text.splitlines() if k.strip()]
+            raw_keys = [k.strip() for k in response.text.splitlines() if k.strip()]
         else:
             with open("key.txt", "r") as f:
-                keys = [k.strip() for k in f.readlines() if k.strip()]
-        
-        # 初始化key状态
-        key_statuses = {k: KeyStatus(k) for k in keys if k not in key_statuses}
+                raw_keys = [k.strip() for k in f.readlines() if k.strip()]
         
-        # 更新轮询器
-        update_key_cycle()
-        
-        print(f"Loaded {len(keys)} API keys")
+        keys_info = {k: KeyInfo(k) for k in raw_keys}
+        chat_keys = list(raw_keys)
+        first_key = raw_keys[0] if raw_keys else None
+        print(f"Loaded {len(raw_keys)} API keys")
     except Exception as e:
         print(f"Error loading keys: {e}")
-        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
+        keys_info = {}
+        chat_keys = []
+        first_key = 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")
+# 获取可用的chat key
+def get_chat_key():
+    valid_keys = [k for k in chat_keys if is_key_available(k)]
+    if not valid_keys:
+        raise HTTPException(status_code=500, detail="No available API keys")
     
-    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]
+    # 简单轮询
+    key = valid_keys[0]
+    chat_keys.append(chat_keys.pop(0))
+    return key
+
+# 检查key是否可用
+def is_key_available(key: str) -> bool:
+    info = keys_info.get(key)
+    if not info or not info.is_valid:
+        return False
+    if info.cooling_until and info.cooling_until > datetime.now():
+        return False
+    return True
+
+# 设置key冷却
+def set_key_cooling(key: str, days: int = 30):
+    if key in keys_info:
+        keys_info[key].cooling_until = datetime.now() + timedelta(days=days)
+
 # 获取真实IP地址
 def get_client_ip(request: Request) -> str:
     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):
+async def access_control(request: Request, call_next):
+    path = request.url.path
+    
+    # 管理后台相关路径不检查IP
+    if path.startswith("/admin") or path.startswith("/api/admin") or path.startswith("/api/keys"):
+        return await call_next(request)
+        
+    # API调用检查白名单
     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"Is admin request: {is_admin_request}")
-        
-        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")
+        if client_ip not in Config.WHITELIST_IPS:
+            raise HTTPException(status_code=403, detail="IP not allowed")
             
     return await call_next(request)
-
 # 添加静态文件支持
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
@@ -171,16 +157,14 @@ 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": [status.to_dict() for status in key_statuses.values()]
-    }
+    return {"keys": [info.to_dict() for info in keys_info.values()]}
 
-# 检查单个key是否有效
+# 检查key是否有效
 async def check_key_valid(key: str) -> bool:
     headers = {
         "Authorization": f"Bearer {key}",
@@ -204,18 +188,16 @@ async def check_all_keys(password: str):
         raise HTTPException(status_code=401, detail="Invalid password")
     
     results = []
-    for key in list(key_statuses.keys()):
+    for key in keys_info:
         is_valid = await check_key_valid(key)
-        key_statuses[key].is_valid = is_valid
-        key_statuses[key].last_check = datetime.now()
+        keys_info[key].is_valid = is_valid
         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):
+async def delete_batch_keys(request: Request):
     try:
         body = await request.json()
         password = body.get("password")
@@ -224,14 +206,13 @@ async def delete_keys_batch(request: Request):
         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 in keys_info:
+                del keys_info[key]
+                if key in chat_keys:
+                    chat_keys.remove(key)
         
-        update_key_cycle()
-        return {"status": "success", "deleted_keys": deleted_keys}
+        return {"status": "success", "deleted_count": len(keys_to_delete)}
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
 
@@ -246,224 +227,126 @@ async def add_key(request: Request):
         if password != Config.ADMIN_PASSWORD:
             raise HTTPException(status_code=401, detail="Invalid password")
             
-        if key in key_statuses:
+        if key in keys_info:
             raise HTTPException(status_code=400, detail="Key already exists")
             
-        # 检查key是否有效
+        # 检查key有效性
         is_valid = await check_key_valid(key)
         if not is_valid:
             raise HTTPException(status_code=400, detail="Invalid key")
             
-        key_statuses[key] = KeyStatus(key)
-        update_key_cycle()
+        keys_info[key] = KeyInfo(key)
+        chat_keys.append(key)
+        
         return {"status": "success", "message": "Key added successfully"}
     except HTTPException:
         raise
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
 
-# 删除单个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 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")
-# 处理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()
-        
-        # 检查是否是额度不足错误
-        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):
-    buffer = ""
-    try:
-        async for chunk in response.aiter_bytes():
-            chunk_str = chunk.decode('utf-8')
-            buffer += chunk_str
-            
-            while '\n\n' in buffer:
-                event, buffer = buffer.split('\n\n', 1)
-                if event.startswith('data: '):
-                    data = event[6:]
-                    if data.strip() == '[DONE]':
-                        yield f"data: [DONE]\n\n"
-                    else:
-                        try:
-                            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"
-
-# 聊天完成路由
-@app.post("/api/v1/chat/completions")
-async def chat_completions(request: Request):
-    body = await request.body()
-    body_json = json.loads(body)
+# 处理API请求
+async def handle_api_request(url: str, headers: dict, method: str = "GET", body: dict = None, for_chat: bool = False):
+    max_retries = 3
+    current_try = 0
     
-    for attempt in range(Config.MAX_RETRIES):
+    while current_try < 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}")
+            # 获取key
+            key = get_chat_key() if for_chat else first_key
+            if not key:
+                raise HTTPException(status_code=500, detail="No API keys available")
+                
+            headers["Authorization"] = f"Bearer {key}"
             
-            async with httpx.AsyncClient(timeout=60.0) as client:
-                response = await client.post(
-                    url,
+            async with httpx.AsyncClient() as client:
+                response = await client.request(
+                    method=method,
+                    url=url,
                     headers=headers,
-                    json=body_json
+                    json=body
                 )
                 
-                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")
-                )
+                # 检查配额不足
+                if response.status_code == 429 or "insufficient_quota" in response.text:
+                    set_key_cooling(key)
+                    current_try += 1
+                    continue
                     
+                return response
+                
         except Exception as e:
-            if attempt == Config.MAX_RETRIES - 1:
-                print(f"Chat Error: {str(e)}")
-                raise HTTPException(status_code=500, detail=str(e))
+            current_try += 1
+            if current_try == max_retries:
+                raise e
 
-# 模型列表路由
-@app.get("/api/v1/models")
-async def list_models():
+# 聊天完成路由
+@app.post("/api/v1/chat/completions")
+async def chat_completions(request: Request):
     try:
-        key = await get_next_valid_key(for_chat=False)
-        
+        body = await request.json()
         headers = {
-            "Authorization": f"Bearer {key}",
-            "Content-Type": "application/json"
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream" if body.get("stream") else "application/json"
         }
         
-        async with httpx.AsyncClient() as client:
-            response = await client.get(
-                f"{Config.OPENAI_API_BASE}/models",
-                headers=headers
+        response = await handle_api_request(
+            url=f"{Config.OPENAI_API_BASE}/chat/completions",
+            headers=headers,
+            method="POST",
+            body=body,
+            for_chat=True
+        )
+        
+        if body.get("stream"):
+            return StreamingResponse(
+                stream_generator(response),
+                media_type="text/event-stream"
             )
-            
-            if response.status_code != 200:
-                await handle_api_error(response, key)
-            
-            return response.json()
+        
+        return Response(
+            content=response.text,
+            media_type="application/json"
+        )
+        
     except Exception as e:
-        print(f"Models Error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
-# 代理其他请求到X.AI
+# 代理其他请求
 @app.api_route("/api/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
 async def proxy(path: str, request: Request):
     if path == "chat/completions":
         return await chat_completions(request)
         
     try:
-        key = await get_next_valid_key(for_chat=False)
+        method = request.method
+        body = await request.json() if method in ["POST", "PUT"] else None
+        headers = {"Content-Type": "application/json"}
         
-        body = await request.body()
-        body_str = body.decode() if body else ""
+        response = await handle_api_request(
+            url=f"{Config.OPENAI_API_BASE}/{path}",
+            headers=headers,
+            method=method,
+            body=body,
+            for_chat=False
+        )
         
-        params = dict(request.query_params)
+        return Response(
+            content=response.text,
+            status_code=response.status_code,
+            media_type=response.headers.get("content-type", "application/json")
+        )
         
-        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,
-                params=params,
-                headers=headers,
-                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))
 
 # 健康检查路由
 @app.get("/api/health")
 async def health_check():
-    valid_keys = len([k for k, status in key_statuses.items() if status.is_valid])
-    total_keys = len(key_statuses)
+    available_count = sum(1 for k in keys_info.values() if is_key_available(k.key))
     return {
         "status": "healthy",
-        "total_keys": total_keys,
-        "valid_keys": valid_keys,
-        "keys_in_cooldown": total_keys - valid_keys
+        "total_keys": len(keys_info),
+        "available_keys": available_count
     }
 
 # 启动时初始化