Spaces:
Sleeping
Sleeping
Junhui Ji
commited on
Commit
·
df5d8bf
1
Parent(s):
d22b2fa
update logging, update key popup position
Browse files
main.py
CHANGED
|
@@ -32,7 +32,7 @@ current_connections = 0
|
|
| 32 |
optimize_design_requests = defaultdict(int) # 记录每个IP的请求次数
|
| 33 |
optimize_design_timestamps = defaultdict(float) # 记录每个IP的首次请求时间
|
| 34 |
white_list = eval(os.getenv("WHITELIST"))
|
| 35 |
-
|
| 36 |
|
| 37 |
class ConnectionLimitMiddleware(BaseHTTPMiddleware):
|
| 38 |
async def dispatch(self, request: Request, call_next):
|
|
@@ -65,9 +65,6 @@ OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
|
| 65 |
OPENAI_API_IMAGE_EDIT_KEY = os.getenv("OPENAI_API_IMAGE_EDIT_KEY")
|
| 66 |
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
|
| 67 |
SEARCH_ENGINE_ID = os.getenv("SEARCH_ENGINE_ID", "27acb0d55ad504716")
|
| 68 |
-
print(OPENAI_API_KEY)
|
| 69 |
-
print(OPENAI_API_IMAGE_EDIT_KEY)
|
| 70 |
-
print(GOOGLE_API_KEY)
|
| 71 |
|
| 72 |
class ScreenshotRequest(BaseModel):
|
| 73 |
url: str
|
|
@@ -113,7 +110,7 @@ async def capture_screenshot(request: ScreenshotRequest):
|
|
| 113 |
filename = f"{domain}_{timestamp}.{request.format}"
|
| 114 |
filepath = os.path.join(CACHE_DIR, filename)
|
| 115 |
|
| 116 |
-
|
| 117 |
|
| 118 |
# 默认请求头
|
| 119 |
default_headers = {
|
|
@@ -153,7 +150,7 @@ async def capture_screenshot(request: ScreenshotRequest):
|
|
| 153 |
|
| 154 |
await browser.close()
|
| 155 |
|
| 156 |
-
|
| 157 |
|
| 158 |
# 返回截图URL
|
| 159 |
screenshot_url = f"/screenshots/{filename}"
|
|
@@ -210,7 +207,7 @@ async def optimize_design(request: OptimizationRequest, client_ip: str = None):
|
|
| 210 |
try:
|
| 211 |
# 获取客户端IP(如果未提供,使用默认值)
|
| 212 |
if client_ip is None or client_ip not in white_list:
|
| 213 |
-
|
| 214 |
raise HTTPException(
|
| 215 |
status_code=503,
|
| 216 |
detail="当前用户无生图权限,请点击'想使用自己的OpenAI API Key?'输入您的OpenAI API Key或联系@wangyue161添加白名单后重试。"
|
|
@@ -230,7 +227,7 @@ async def optimize_design(request: OptimizationRequest, client_ip: str = None):
|
|
| 230 |
optimize_design_timestamps[client_ip] = current_time
|
| 231 |
|
| 232 |
# 检查是否超过限制
|
| 233 |
-
|
| 234 |
if optimize_design_requests[client_ip] >= user_rate_limit:
|
| 235 |
raise HTTPException(
|
| 236 |
status_code=503,
|
|
|
|
| 32 |
optimize_design_requests = defaultdict(int) # 记录每个IP的请求次数
|
| 33 |
optimize_design_timestamps = defaultdict(float) # 记录每个IP的首次请求时间
|
| 34 |
white_list = eval(os.getenv("WHITELIST"))
|
| 35 |
+
print(logging.INFO, white_list)
|
| 36 |
|
| 37 |
class ConnectionLimitMiddleware(BaseHTTPMiddleware):
|
| 38 |
async def dispatch(self, request: Request, call_next):
|
|
|
|
| 65 |
OPENAI_API_IMAGE_EDIT_KEY = os.getenv("OPENAI_API_IMAGE_EDIT_KEY")
|
| 66 |
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
|
| 67 |
SEARCH_ENGINE_ID = os.getenv("SEARCH_ENGINE_ID", "27acb0d55ad504716")
|
|
|
|
|
|
|
|
|
|
| 68 |
|
| 69 |
class ScreenshotRequest(BaseModel):
|
| 70 |
url: str
|
|
|
|
| 110 |
filename = f"{domain}_{timestamp}.{request.format}"
|
| 111 |
filepath = os.path.join(CACHE_DIR, filename)
|
| 112 |
|
| 113 |
+
print(f"开始为 {request.url} 生成截图...")
|
| 114 |
|
| 115 |
# 默认请求头
|
| 116 |
default_headers = {
|
|
|
|
| 150 |
|
| 151 |
await browser.close()
|
| 152 |
|
| 153 |
+
print(f"截图完成: {filepath}")
|
| 154 |
|
| 155 |
# 返回截图URL
|
| 156 |
screenshot_url = f"/screenshots/{filename}"
|
|
|
|
| 207 |
try:
|
| 208 |
# 获取客户端IP(如果未提供,使用默认值)
|
| 209 |
if client_ip is None or client_ip not in white_list:
|
| 210 |
+
print(f'user_ip {client_ip} not in whitelist')
|
| 211 |
raise HTTPException(
|
| 212 |
status_code=503,
|
| 213 |
detail="当前用户无生图权限,请点击'想使用自己的OpenAI API Key?'输入您的OpenAI API Key或联系@wangyue161添加白名单后重试。"
|
|
|
|
| 227 |
optimize_design_timestamps[client_ip] = current_time
|
| 228 |
|
| 229 |
# 检查是否超过限制
|
| 230 |
+
print(f'user_ip {client_ip}, total requests in routine: {optimize_design_requests[client_ip]}, rate limit: {user_rate_limit}')
|
| 231 |
if optimize_design_requests[client_ip] >= user_rate_limit:
|
| 232 |
raise HTTPException(
|
| 233 |
status_code=503,
|