update whitelist, dalle

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
 from fastapi.responses import JSONResponse
@@ -16,10 +16,41 @@ import aiohttp
 import traceback
 import requests
 from openai import OpenAI
+from starlette.middleware.base import BaseHTTPMiddleware
+import uvicorn
+from collections import defaultdict
+from PIL import Image
 
 
 app = FastAPI()
 
+# 设置最大连接数
+MAX_CONNECTIONS = 10
+current_connections = 0
+
+# 设置优化设计接口的访问限制
+optimize_design_requests = defaultdict(int)  # 记录每个IP的请求次数
+optimize_design_timestamps = defaultdict(float)  # 记录每个IP的首次请求时间
+white_list = eval(os.getenv("WHITELIST"))
+
+class ConnectionLimitMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        global current_connections
+        if current_connections >= MAX_CONNECTIONS:
+            return JSONResponse(
+                status_code=503,
+                content={"detail": "已超过最大链接数，请稍后重试"}
+            )
+        current_connections += 1
+        try:
+            response = await call_next(request)
+            return response
+        finally:
+            current_connections -= 1
+
+# 添加中间件
+app.add_middleware(ConnectionLimitMiddleware)
+
 # 确保缓存目录存在
 CACHE_DIR = "cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
@@ -138,7 +169,11 @@ async def capture_screenshot(request: ScreenshotRequest):
 
 @app.get("/health")
 async def health_check():
-    return {"status": "ok"}
+    return {
+        "status": "ok",
+        "current_connections": current_connections,
+        "max_connections": MAX_CONNECTIONS
+    }
 
 
 @app.get("/")
@@ -169,10 +204,43 @@ async def analyze_feedback(request: AnalysisRequest):
         raise HTTPException(status_code=500, detail=f'Error: {e}, traceback: {traceback.format_exc()}')
 
 @app.post("/api/optimize-design")
-async def optimize_design(request: OptimizationRequest):
+async def optimize_design(request: OptimizationRequest, client_ip: str = None):
     try:
+        # 获取客户端IP（如果未提供，使用默认值）
+        if client_ip is None or client_ip not in white_list:
+            raise HTTPException(
+                status_code=503,
+                detail="当前用户无生图权限，请联系@王月(Phoebe)添加白名单后重试。"
+            )
+
+        user_rate_limit = white_list[client_ip]
+        
+        current_time = time.time()
+        
+        # 检查是否需要重置计数器（超过24小时）
+        if current_time - optimize_design_timestamps[client_ip] > 3600*24:
+            optimize_design_requests[client_ip] = 0
+            optimize_design_timestamps[client_ip] = current_time
+            
+        # 如果是首次请求，记录时间戳
+        if optimize_design_requests[client_ip] == 0:
+            optimize_design_timestamps[client_ip] = current_time
+            
+        # 检查是否超过限制
+        if optimize_design_requests[client_ip] >= user_rate_limit:
+            raise HTTPException(
+                status_code=503,
+                detail="用户当日改图接口访问已达上限，请24小时后重试"
+            )
+            
+        # 增加请求计数
+        optimize_design_requests[client_ip] += 1
+
+        # 提取设计类型
+        design_type = f"设计类型：{request.text.split()[0]}\n" if len(request.text.split()) > 1 else ""
+        
         # 构建图像生成提示词
-        prompt = f"基于以下设计反馈优化UI设计: {', '.join(request.suggestions)}"
+        prompt = f"{design_type}基于以下设计反馈优化UI设计: {', '.join(request.suggestions)}"
         
         # 处理图片数据
         image_data = request.image_data
@@ -185,9 +253,15 @@ async def optimize_design(request: OptimizationRequest):
         )
         
         return JSONResponse(response)
+    except HTTPException as he:
+        raise he
     except Exception as e:
         logging.error(f'Error: {e}, traceback: {traceback.format_exc()}')
         raise HTTPException(status_code=500, detail=f'Error: {e}, traceback: {traceback.format_exc()}')
+    finally:
+        # 如果发生异常，减少请求计数
+        if 'he' in locals() and isinstance(he, HTTPException):
+            optimize_design_requests[client_ip] -= 1
 
 @app.post("/api/optimize-text")
 async def optimize_text(request: TextOptimizationRequest):
@@ -324,11 +398,25 @@ async def call_openai_image_api(image_data: str, prompt: str, request_model_id='
         if image_data and 'base64,' in image_data:
             image_data = image_data.split('base64,')[1]
         
-        logging.log(logging.INFO, f"Processing image data (first 100 chars): {image_data[:100]}")
+        logging.log(logging.DEBUG, f"Processing image data (first 100 chars): {image_data[:100]}")
         
         # 将base64图片数据转换为文件对象
         image_bytes = base64.b64decode(image_data)
         image_file = BytesIO(image_bytes)
+        
+        # 如果是dall-e-2模型，需要将图片调整为800x800
+        if request_model_id == 'dall-e-2':
+            # 打开图片
+            img = Image.open(image_file)
+            # 调整图片大小为800x800，使用LANCZOS重采样方法以获得更好的质量
+            img = img.resize((800, 800), Image.Resampling.LANCZOS)
+            # 创建新的BytesIO对象
+            image_file = BytesIO()
+            # 保存调整后的图片
+            img.save(image_file, format='PNG')
+            # 将文件指针移到开始位置
+            image_file.seek(0)
+        
         image_file.name = "original-design.png"  # 设置文件名，与JS代码一致
         
         # 创建OpenAI客户端
@@ -338,7 +426,7 @@ async def call_openai_image_api(image_data: str, prompt: str, request_model_id='
         response = client.images.edit(
             model=request_model_id,
             image=image_file,
-            prompt=prompt # 明确要求返回base64格式
+            prompt=prompt
         )
         
         # 获取生成的图片数据
@@ -364,4 +452,10 @@ async def call_openai_image_api(image_data: str, prompt: str, request_model_id='
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(
+        app, 
+        host="0.0.0.0", 
+        port=7860,
+        limit_concurrency=MAX_CONNECTIONS,
+        limit_max_requests=0  # 0 means no limit on total requests
+    )

static/script.js

@@ -762,8 +762,8 @@ document.addEventListener('DOMContentLoaded', function() {
             }
 
             // 构建提示词
-            const prompt = `基于以下设计反馈优化UI设计: ${suggestions.join(', ')}`;
-            
+            const uploadedText = sessionStorage.getItem('uploadedText') || '';
+
             // 调用后端API
             const response = await fetch(`${BASE_URL}api/optimize-design`, {
                 method: 'POST',
@@ -771,7 +771,7 @@ document.addEventListener('DOMContentLoaded', function() {
                     'Content-Type': 'application/json'
                 },
                 body: JSON.stringify({
-                    text: analysisResult,
+                    text: uploadedText,
                     image_data: uploadedImage,
                     suggestions: suggestions,
                     request_model_id: 'dall-e-2'