bobocup commited on
Commit
3c0a2c4
·
verified ·
1 Parent(s): a39b29c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +98 -0
app.py ADDED
@@ -0,0 +1,98 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, HTTPException, Request
2
+ from fastapi.middleware.cors import CORSMiddleware
3
+ import httpx
4
+ import os
5
+ import json
6
+ from typing import List, Optional
7
+ import requests
8
+ from itertools import cycle
9
+
10
+ # 创建FastAPI应用,设置API前缀
11
+ app = FastAPI(openapi_url="/api/openapi.json", docs_url="/api/docs")
12
+
13
+ # CORS设置
14
+ app.add_middleware(
15
+ CORSMiddleware,
16
+ allow_origins=["*"],
17
+ allow_credentials=True,
18
+ allow_methods=["*"],
19
+ allow_headers=["*"],
20
+ )
21
+
22
+ # 配置
23
+ class Config:
24
+ OPENAI_API_BASE = "https://api.x.ai/v1"
25
+ KEYS_URL = os.getenv("KEYS_URL", "")
26
+ WHITELIST_IPS = os.getenv("WHITELIST_IPS", "").split(",")
27
+
28
+ # 全局变量
29
+ keys = []
30
+ key_cycle = None
31
+
32
+ # 初始化keys
33
+ def init_keys():
34
+ global keys, key_cycle
35
+ try:
36
+ if Config.KEYS_URL:
37
+ response = requests.get(Config.KEYS_URL)
38
+ keys = [k.strip() for k in response.text.splitlines() if k.strip()]
39
+ else:
40
+ with open("key.txt", "r") as f:
41
+ keys = [k.strip() for k in f.readlines() if k.strip()]
42
+
43
+ key_cycle = cycle(keys)
44
+ print(f"Loaded {len(keys)} API keys")
45
+ except Exception as e:
46
+ print(f"Error loading keys: {e}")
47
+ keys = []
48
+ key_cycle = cycle(keys)
49
+
50
+ # IP白名单中间件
51
+ @app.middleware("http")
52
+ async def ip_whitelist(request: Request, call_next):
53
+ if Config.WHITELIST_IPS and Config.WHITELIST_IPS[0]:
54
+ client_ip = request.client.host
55
+ if client_ip not in Config.WHITELIST_IPS:
56
+ raise HTTPException(status_code=403, detail="IP not allowed")
57
+ return await call_next(request)
58
+
59
+ # 获取下一个key
60
+ def get_next_key():
61
+ if not keys:
62
+ raise HTTPException(status_code=500, detail="No API keys available")
63
+ return next(key_cycle)
64
+
65
+ # 代理请求到X.AI,注意路径前缀改为/api
66
+ @app.api_route("/api/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
67
+ async def proxy(path: str, request: Request):
68
+ # 获取请求体
69
+ body = await request.body()
70
+ # 获取查询参数
71
+ params = dict(request.query_params)
72
+ # 获取headers
73
+ headers = dict(request.headers)
74
+ headers.pop("host", None)
75
+
76
+ # 设置API key
77
+ headers["Authorization"] = f"Bearer {get_next_key()}"
78
+
79
+ # 构建目标URL
80
+ url = f"{Config.OPENAI_API_BASE}/{path}"
81
+
82
+ async with httpx.AsyncClient() as client:
83
+ try:
84
+ response = await client.request(
85
+ method=request.method,
86
+ url=url,
87
+ params=params,
88
+ headers=headers,
89
+ content=body
90
+ )
91
+ return response.json()
92
+ except Exception as e:
93
+ raise HTTPException(status_code=500, detail=str(e))
94
+
95
+ # 启动时初始化
96
+ @app.on_event("startup")
97
+ async def startup_event():
98
+ init_keys()