Update app.py

app.py

@@ -1,16 +1,14 @@
 from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse, HTMLResponse
-from fastapi.staticfiles import StaticFiles
+from fastapi.responses import StreamingResponse
 import httpx
 import os
 import json
-from typing import List, Optional, Dict
+from typing import List, Optional
 import requests
 from itertools import cycle
 import asyncio
 import time
-from datetime import datetime, timedelta
 
 # 创建FastAPI应用
 app = FastAPI()
@@ -29,107 +27,59 @@ 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", "")
-    CHUNK_SIZE = 1
-
-# Key信息类
-class KeyInfo:
-    def __init__(self, key: str):
-        self.key = key
-        self.is_valid = True
-        self.cooling_until = None
-        self.last_used = None
-        self.error_count = 0
-
-    def to_dict(self):
-        return {
-            "key": self.key,
-            "is_valid": self.is_valid,
-            "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"
 
 # 全局变量
-keys_info: Dict[str, KeyInfo] = {}
-chat_keys = []
+keys = []
+key_cycle = None
 first_key = None
+
 # 初始化keys
 def init_keys():
-    global keys_info, chat_keys, first_key
+    global keys, key_cycle, first_key
     try:
         if Config.KEYS_URL:
             response = requests.get(Config.KEYS_URL)
-            raw_keys = [k.strip() for k in response.text.splitlines() if k.strip()]
+            keys = [k.strip() for k in response.text.splitlines() if k.strip()]
         else:
             with open("key.txt", "r") as f:
-                raw_keys = [k.strip() for k in f.readlines() if k.strip()]
+                keys = [k.strip() for k in f.readlines() if k.strip()]
         
-        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")
+        if keys:
+            first_key = keys[0]
+            key_cycle = cycle(keys)
+            print(f"Loaded {len(keys)} API keys")
     except Exception as e:
         print(f"Error loading keys: {e}")
-        keys_info = {}
-        chat_keys = []
+        keys = []
+        key_cycle = None
         first_key = None
 
-# 获取可用的chat key
+# 获取聊天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")
-    
-    # 简单轮询
-    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)
-
+    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
 
 # IP白名单中间件
 @app.middleware("http")
-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调用检查白名单
+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)
 
 # 流式响应生成器
@@ -160,91 +110,11 @@ async def stream_generator(response):
     except Exception as e:
         print(f"Stream Error: {str(e)}")
         yield f"data: {json.dumps({'error': str(e)})}\n\n"
-
-
-# ���改 handle_api_request 函数中的超时设置
-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
-    
-    while current_try < max_retries:
-        try:
-            if for_chat:
-                key = get_chat_key()
-            else:
-                key = first_key
-                
-            if not key:
-                raise HTTPException(status_code=500, detail="No API keys available")
-                
-            headers["Authorization"] = f"Bearer {key}"
-            
-            # 减少连接超时，增加读取超时
-            timeout = httpx.Timeout(
-                connect=5.0,    # 连接超时5秒
-                read=60.0,      # 读取超时60秒
-                write=5.0,      # 写入超时5秒
-                pool=5.0        # 连接池超时5秒
-            )
-            
-            async with httpx.AsyncClient(timeout=timeout) as client:
-                response = await client.request(
-                    method=method,
-                    url=url,
-                    headers=headers,
-                    json=body
-                )
-                
-                # 检查配额不足
-                if response.status_code == 429 or "insufficient_quota" in response.text.lower():
-                    print(f"Key {key} quota exceeded, trying next key...")
-                    set_key_cooling(key)
-                    current_try += 1
-                    continue
-                    
-                return response
-                
-        except Exception as e:
-            print(f"Request error: {str(e)}")
-            current_try += 1
-            if current_try == max_retries:
-                raise HTTPException(status_code=500, detail=f"API request failed after {max_retries} retries")
-# 添加静态文件支持
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-# 管理页面路由
-@app.get("/admin", response_class=HTMLResponse)
-async def admin_page():
-    try:
-        with open("static/admin.html", "r", encoding="utf-8") as f:
-            return f.read()
-    except Exception as e:
-        print(f"Error loading admin page: {e}")
-        raise HTTPException(status_code=500, detail="Error loading admin page")
-
-# 验证管理密码
-@app.post("/api/admin/login")
-async def admin_login(request: Request):
-    try:
-        body = await request.json()
-        password = body.get("password")
-        if password == Config.ADMIN_PASSWORD:
-            return {"status": "success"}
-        raise HTTPException(status_code=401, detail="Invalid password")
-    except json.JSONDecodeError:
-        raise HTTPException(status_code=400, detail="Invalid JSON")
-
-# 获取所有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": [info.to_dict() for info in keys_info.values()]}
-
-# 检查key是否有效
-async def check_key_valid(key: str) -> bool:
+# 模型列表路由
+@app.get("/api/v1/models")
+async def list_models():
     headers = {
-        "Authorization": f"Bearer {key}",
+        "Authorization": f"Bearer {get_chat_key()}",
         "Content-Type": "application/json"
     }
     
@@ -254,77 +124,14 @@ async def check_key_valid(key: str) -> bool:
                 f"{Config.OPENAI_API_BASE}/models",
                 headers=headers
             )
-            return response.status_code == 200
-        except:
-            return False
-
-# 批量检查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 keys_info:
-        is_valid = await check_key_valid(key)
-        keys_info[key].is_valid = is_valid
-        results.append({"key": key, "valid": is_valid})
-    
-    return {"results": results}
-
-# 批量删除keys
-@app.post("/api/keys/delete-batch")
-async def delete_batch_keys(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")
-            
-        for key in keys_to_delete:
-            if key in keys_info:
-                del keys_info[key]
-                if key in chat_keys:
-                    chat_keys.remove(key)
-        
-        return {"status": "success", "deleted_count": len(keys_to_delete)}
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
+            return response.json()
+        except Exception as e:
+            print(f"Models Error: {str(e)}")
+            raise HTTPException(status_code=500, 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 keys_info:
-            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_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))
 # 聊天完成路由
 @app.post("/api/v1/chat/completions")
 async def chat_completions(request: Request):
-    global key_cycle  # 声明使用全局变量
-    
     try:
         # 获取请求体
         body = await request.body()
@@ -332,7 +139,7 @@ async def chat_completions(request: Request):
         
         # 获取headers
         headers = {
-            "Authorization": f"Bearer {next(key_cycle)}",
+            "Authorization": f"Bearer {get_chat_key()}",
             "Content-Type": "application/json",
             "Accept": "text/event-stream" if body_json.get("stream") else "application/json"
         }
@@ -376,7 +183,6 @@ 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):
@@ -385,36 +191,34 @@ async def proxy(path: str, request: Request):
         
     try:
         method = request.method
-        body = await request.json() if method in ["POST", "PUT"] else None
-        headers = {"Content-Type": "application/json"}
-        
-        response = await handle_api_request(
-            url=f"{Config.OPENAI_API_BASE}/{path}",
-            headers=headers,
-            method=method,
-            body=body,
-            for_chat=False
-        )
-        
-        return Response(
-            content=response.text,
-            status_code=response.status_code,
-            media_type=response.headers.get("content-type", "application/json")
-        )
+        body = await request.body() if method in ["POST", "PUT"] else None
+        headers = {
+            "Authorization": f"Bearer {get_chat_key()}",
+            "Content-Type": "application/json"
+        }
         
+        async with httpx.AsyncClient() as client:
+            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)}")
+        print(f"Proxy Error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
 # 健康检查路由
 @app.get("/api/health")
 async def health_check():
-    available_count = sum(1 for k in keys_info.values() if is_key_available(k.key))
-    return {
-        "status": "healthy",
-        "total_keys": len(keys_info),
-        "available_keys": available_count
-    }
+    return {"status": "healthy", "key_count": len(keys)}
 
 # 启动时初始化
 @app.on_event("startup")