Update main.py

main.py

@@ -6,8 +6,7 @@ import base64
 import re
 import os
 import argparse
-import time
-from datetime import datetime, timezone
+from datetime import datetime, timezone, timedelta
 from typing import List, Optional
 
 import httpx
@@ -26,21 +25,36 @@ from fastapi.staticfiles import StaticFiles
 
 from bearer_token import BearerTokenGenerator
 
-from fastapi import Depends, HTTPException, Security
-from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
-
-# 模型列表
+# 模型列表（根据需求，可自行调整）
 MODELS = ["gpt-4o", "gpt-4o-mini", "claude-3-5-sonnet", "claude"]
 
 # 默认端口
-INITIAL_PORT = 3000
+INITIAL_PORT = 8080
 
 # 外部API的URL
 EXTERNAL_API_URL = "https://api.chaton.ai/chats/stream"
 
+# 定义 images 目录
+IMAGES_DIR = "images"
+
+# 确保 images 目录存在
+def ensure_images_dir_exists(directory: str = IMAGES_DIR):
+    try:
+        os.makedirs(directory, exist_ok=True)
+        print(f"Directory '{directory}' is ready.")
+    except Exception as e:
+        print(f"Failed to create directory '{directory}': {e}")
+        sys.exit(1)
+
+# 在挂载静态文件之前确保 images 目录存在
+ensure_images_dir_exists()
+
 # 初始化FastAPI应用
 app = FastAPI()
 
+# 挂载静态文件路由以提供 images 目录的内容
+app.mount("/images", StaticFiles(directory=IMAGES_DIR), name="images")
+
 # 添加CORS中间件
 app.add_middleware(
     CORSMiddleware,
@@ -50,9 +64,6 @@ app.add_middleware(
     allow_headers=["Content-Type", "Authorization"],  # 允许的头部
 )
 
-# 挂载静态文件路由以提供 images 目录的内容
-app.mount("/images", StaticFiles(directory="images"), name="images")
-
 # 辅助函数
 def send_error_response(message: str, status_code: int = 400):
     """构建错误响应，并确保包含CORS头"""
@@ -106,17 +117,50 @@ async def download_image(image_url: str) -> Optional[bytes]:
             print(f"Error downloading image: {e}")
             return None
 
-def save_base64_image(base64_str: str, images_dir: str = "images") -> str:
+def cleanup_images(images_dir: str = IMAGES_DIR, age_seconds: int = 60):
     """
-    将Base64编码的图片保存到images目录，返回文件名
+    清理 images 目录中创建时间超过指定秒数的图片
     """
+    now = datetime.now(timezone.utc)
+    cutoff_time = now - timedelta(seconds=age_seconds)
+
     if not os.path.exists(images_dir):
-        os.makedirs(images_dir)
-    image_data = base64.b64decode(base64_str)
+        return
+
+    for filename in os.listdir(images_dir):
+        file_path = os.path.join(images_dir, filename)
+        if os.path.isfile(file_path):
+            try:
+                file_creation_time = datetime.fromtimestamp(os.path.getctime(file_path), timezone.utc)
+                if file_creation_time < cutoff_time:
+                    os.remove(file_path)
+                    print(f"已删除旧图片: {filename}")
+            except Exception as e:
+                print(f"无法删除文件 {filename}: {e}")
+
+def save_base64_image(base64_str: str, images_dir: str = IMAGES_DIR) -> str:
+    """
+    将Base64编码的图片保存到images目录，返回文件名
+    """
+    # 先清理1分钟前的所有图片
+    cleanup_images(images_dir, age_seconds=60)
+
+    try:
+        image_data = base64.b64decode(base64_str)
+    except base64.binascii.Error as e:
+        print(f"Base64解码失败: {e}")
+        raise ValueError("Invalid base64 image data")
+
     filename = f"{uuid.uuid4()}.png"  # 默认保存为png格式
     file_path = os.path.join(images_dir, filename)
-    with open(file_path, "wb") as f:
-        f.write(image_data)
+    try:
+        with open(file_path, "wb") as f:
+            f.write(image_data)
+        print(f"保存图片: {filename}")
+    except Exception as e:
+        print(f"保存图片失败: {e}")
+        raise
+
     return filename
 
 def is_base64_image(url: str) -> bool:
@@ -125,73 +169,26 @@ def is_base64_image(url: str) -> bool:
     """
     return url.startswith("data:image/")
 
-# 添加 HTTPBearer 实例
-security = HTTPBearer()
-
-# 添加 API_KEY 验证函数
-def verify_api_key(credentials: HTTPAuthorizationCredentials = Security(security)):
-    api_key = os.environ.get("API_KEY")
-    if api_key is None:
-        raise HTTPException(status_code=500, detail="API_KEY not set in environment variables")
-    if credentials.credentials != api_key:
-        raise HTTPException(status_code=401, detail="Invalid API key")
-    return credentials.credentials
-
 # 根路径GET请求处理
-@app.get("/")
-async def root():
-    return JSONResponse(content={
-        "service": "AI Chat Completion Proxy",
-        "usage": {
-            "endpoint": "/ai/v1/chat/completions",
-            "method": "POST",
-            "headers": {
-                "Content-Type": "application/json",
-                "Authorization": "Bearer YOUR_API_KEY"
-            },
-            "body": {
-                "model": "One of: " + ", ".join(MODELS),
-                "messages": [
-                    {"role": "system", "content": "You are a helpful assistant."},
-                    {"role": "user", "content": "Hello, who are you?"}
-                ],
-                "stream": False,
-                "temperature": 0.7,
-                "max_tokens": 8000
-            }
-        },
-        "availableModels": MODELS,
-        "endpoints": {
-            "/ai/v1/chat/completions": "Chat completion endpoint",
-            "/ai/v1/images/generations": "Image generation endpoint",
-            "/ai/v1/models": "List available models"
-        },
-        "note": "Replace YOUR_API_KEY with your actual API key."
-    })
-
-# 返回模型列表
-@app.get("/ai/v1/models")
-async def list_models(api_key: str = Depends(verify_api_key)):
-    """返回可用模型列表。"""
-    models = [
-        {
-            "id": model,
-            "object": "model",
-            "created": int(time.time()),
-            "owned_by": "chaton",
-            "permission": [],
-            "root": model,
-            "parent": None,
-        } for model in MODELS
-    ]
-    return JSONResponse(content={
-        "object": "list",
-        "data": models
-    })
+@app.get("/", response_class=HTMLResponse)
+async def read_root():
+    """返回欢迎页面"""
+    html_content = """
+    <html>
+        <head>
+            <title>Welcome to API</title>
+        </head>
+        <body>
+            <h1>Welcome to API</h1>
+            <p>You can send messages to the model and receive responses.</p>
+        </body>
+    </html>
+    """
+    return HTMLResponse(content=html_content, status_code=200)
 
-# 聊天完成处理
-@app.post("/ai/v1/chat/completions")
-async def chat_completions(request: Request, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
+# 聊天完成处理（保留原有逻辑，未修改）
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request, background_tasks: BackgroundTasks):
     """
     处理聊天完成请求
     """
@@ -206,11 +203,15 @@ async def chat_completions(request: Request, background_tasks: BackgroundTasks,
     # 处理消息内容
     messages = request_body.get("messages", [])
     temperature = request_body.get("temperature", 1.0)
-    top_p = request_body.get("top_p", 1.0)
+    #top_p = request_body.get("top_p", 1.0)
     max_tokens = request_body.get("max_tokens", 8000)
     model = request_body.get("model", "gpt-4o")
     is_stream = request_body.get("stream", False)  # 获取 stream 字段
 
+    # 验证模型
+    if model not in MODELS:
+        raise HTTPException(status_code=400, detail=f"无效的 model: {model}. 可用的模型有: {', '.join(MODELS)}")
+
     has_image = False
     has_text = False
 
@@ -230,11 +231,15 @@ async def chat_completions(request: Request, background_tasks: BackgroundTasks,
                     url = image_info.get("url", "")
                     if is_base64_image(url):
                         # 解码并保存图片
-                        base64_str = url.split(",")[1]
-                        filename = save_base64_image(base64_str)
-                        base_url = app.state.base_url
-                        image_url = f"{base_url}/images/{filename}"
-                        images.append({"data": image_url})
+                        try:
+                            base64_str = url.split(",")[1]
+                            filename = save_base64_image(base64_str)
+                            base_url = app.state.base_url
+                            image_url = f"{base_url}/images/{filename}"
+                            images.append({"data": image_url})
+                        except (IndexError, ValueError) as e:
+                            print(f"处理Base64图片失败: {e}")
+                            continue
                     else:
                         images.append({"data": url})
             extracted_content = " ".join(text_parts).strip()
@@ -269,19 +274,15 @@ async def chat_completions(request: Request, background_tasks: BackgroundTasks,
     if not cleaned_messages:
         raise HTTPException(status_code=400, detail="所有消息的内容均为空。")
 
-    # 验证模型
-    if model not in MODELS:
-        model = "gpt-4o"
-
     # 构建新的请求JSON
     new_request_json = {
         "function_image_gen": False,
         "function_web_search": True,
         "max_tokens": max_tokens,
         "model": model,
-        "source": "chat/free",
+        "source": "chat/pro",
         "temperature": temperature,
-        "top_p": top_p,
+        #"top_p": top_p,
         "messages": cleaned_messages,
     }
 
@@ -299,13 +300,32 @@ async def chat_completions(request: Request, background_tasks: BackgroundTasks,
         "Date": formatted_date,
         "Client-time-zone": "-05:00",
         "Authorization": bearer_token,
-        "User-Agent": "ChatOn_Android/1.53.502",
+        "User-Agent": "ChatOn_Android/1.55.488",
         "Accept-Language": "en-US",
         "X-Cl-Options": "hb",
         "Content-Type": "application/json; charset=UTF-8",
     }
 
     if is_stream:
+        import uuid
+        from datetime import datetime, timezone
+
+        # 定义 should_filter_out 函数
+        def should_filter_out(json_data):
+            if 'ping' in json_data:
+                return True
+            if 'data' in json_data:
+                data = json_data['data']
+                if 'analytics' in data:
+                    return True
+                if 'operation' in data and 'message' in data:
+                    return True
+            return False
+
+        # 定义 generate_id 函数
+        def generate_id():
+            return uuid.uuid4().hex[:24]
+
         # 流式响应处理
         async def event_generator():
             async with httpx.AsyncClient(timeout=None) as client_stream:
@@ -320,32 +340,54 @@ async def chat_completions(request: Request, background_tasks: BackgroundTasks,
                                     break
                                 try:
                                     sse_json = json.loads(data)
-                                    if "choices" in sse_json:
-                                        for choice in sse_json["choices"]:
-                                            delta = choice.get("delta", {})
-                                            content = delta.get("content")
-                                            if content:
-                                                new_sse_json = {
-                                                    "choices": [
-                                                        {
-                                                            "index": choice.get("index", 0),
-                                                            "delta": {"content": content},
-                                                        }
-                                                    ],
-                                                    "created": sse_json.get(
-                                                        "created", int(datetime.now(timezone.utc).timestamp())
-                                                    ),
-                                                    "id": sse_json.get(
-                                                        "id", str(uuid.uuid4())
-                                                    ),
-                                                    "model": sse_json.get("model", "gpt-4o"),
-                                                    "system_fingerprint": f"fp_{uuid.uuid4().hex[:12]}",
-                                                }
-                                                new_sse_line = f"data: {json.dumps(new_sse_json, ensure_ascii=False)}\n\n"
-                                                yield new_sse_line
-                                except json.JSONDecodeError:
-                                    print("JSON解析错误")
+                                    
+                                    # 判断是否需要过滤
+                                    if should_filter_out(sse_json):
+                                        continue
+
+                                    # 处理包含 web sources 的消息
+                                    if 'data' in sse_json and 'web' in sse_json['data']:
+                                        web_data = sse_json['data']['web']
+                                        if 'sources' in web_data:
+                                            sources = web_data['sources']
+                                            urls_list = []
+                                            for source in sources:
+                                                if 'url' in source:
+                                                    urls_list.append(source['url'])
+                                            urls_content = '\n\n'.join(urls_list)
+                                            print(f"从 API 接收到的内容: {urls_content}")
+                                            # 构造新的 SSE 消息，填入 content 字段
+                                            new_sse_json = {
+                                                "id": generate_id(),
+                                                "object": "chat.completion.chunk",
+                                                "created": int(datetime.now(timezone.utc).timestamp()),
+                                                "model": sse_json.get("model", "gpt-4o"),
+                                                "choices": [
+                                                    {
+                                                        "delta": {"content": "\n" + urls_content + "\n"},
+                                                        "index": 0,
+                                                        "finish_reason": None
+                                                    }
+                                                ]
+                                            }
+                                            new_sse_line = f"data: {json.dumps(new_sse_json, ensure_ascii=False)}\n\n"
+                                            yield new_sse_line
+                                    else:
+                                        # 尝试打印内容
+                                        if 'choices' in sse_json:
+                                            for choice in sse_json['choices']:
+                                                delta = choice.get('delta', {})
+                                                content = delta.get('content')
+                                                if content:
+                                                    print(content, end='')
+                                        # 直接转发其他消息
+                                        yield f"data: {data}\n\n"
+                                except json.JSONDecodeError as e:
+                                    print(f"JSON解析错误: {e}")
                                     continue
+                            else:
+                                # 忽略不以 "data: " 开头的行
+                                continue
                 except httpx.RequestError as exc:
                     print(f"外部API请求失败: {exc}")
                     yield f"data: {{\"error\": \"外部API请求失败: {str(exc)}\"}}\n\n"
@@ -401,7 +443,7 @@ async def chat_completions(request: Request, background_tasks: BackgroundTasks,
                     "id": f"chatcmpl-{uuid.uuid4()}",
                     "object": "chat.completion",
                     "created": int(datetime.now(timezone.utc).timestamp()),
-                    "model": model,
+                    "model": model,  # 使用用户传入的model
                     "choices": [
                         {
                             "index": 0,
@@ -430,8 +472,8 @@ async def chat_completions(request: Request, background_tasks: BackgroundTasks,
                 raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(exc)}")
 
 # 图像生成处理
-@app.post("/ai/v1/images/generations")
-async def images_generations(request: Request, api_key: str = Depends(verify_api_key)):
+@app.post("/v1/images/generations")
+async def images_generations(request: Request):
     """
     处理图像生成请求
     """
@@ -447,177 +489,238 @@ async def images_generations(request: Request, api_key: str = Depends(verify_api
         return send_error_response("缺少必需的字段: prompt", status_code=400)
 
     user_prompt = request_body.get("prompt", "").strip()
-    response_format = request_body.get("response_format", "b64_json").strip()
+    response_format = request_body.get("response_format", "").strip().lower()
+    model = request_body.get("model", "gpt-4o")
+    n = request_body.get("n", 1)  # 生成图像数量，默认1
+    size = request_body.get("size", "1024x1024")  # 图片尺寸，默认1024x1024
+
+    is_base64_response = response_format == "b64_json"
 
     if not user_prompt:
         return send_error_response("Prompt 不能为空。", status_code=400)
 
     print(f"Prompt: {user_prompt}")
+    print(f"Response Format: {response_format}")
+    print(f"Number of images to generate (n): {n}")
+    print(f"Size: {size}")
+    
+    # 设置最大尝试次数为 2 * n
+    max_attempts = 2 * n
+    print(f"Max Attempts: {max_attempts}")
+
+    # 初始化用于存储多个 URL 的线程安全列表
+    final_download_urls: List[str] = []
+
+    async def attempt_generate_image(attempt: int) -> Optional[str]:
+        """
+        尝试生成单张图像，带有重试机制。
+        """
+        try:
+            # 构建新的 TextToImage JSON 请求体
+            text_to_image_json = {
+                "function_image_gen": True,
+                "function_web_search": True,
+                "image_aspect_ratio": "1:1",  # 图片比例可选：1:1/9:19/16:9/4:3
+                "image_style": "photographic",  # 固定 image_style，可根据需要调整
+                "max_tokens": 8000,
+                "messages": [
+                    {
+                        "content": "You are a helpful artist, please draw a picture. Based on imagination, draw a picture with user message.",
+                        "role": "system"
+                    },
+                    {
+                        "content": "Draw: " + user_prompt,
+                        "role": "user"
+                    }
+                ],
+                "model": "gpt-4o",
+                "source": "chat/free"  # 固定 source
+            }
 
-    # 构建新的 TextToImage JSON 请求体
-    text_to_image_json = {
-        "function_image_gen": True,
-        "function_web_search": True,
-        "image_aspect_ratio": "1:1",
-        "image_style": "photographic",  # 暂时固定 image_style
-        "max_tokens": 8000,
-        "messages": [
-            {
-                "content": "You are a helpful artist, please based on imagination draw a picture.",
-                "role": "system"
-            },
-            {
-                "content": "Draw: " + user_prompt,
-                "role": "user"
+            modified_request_body = json.dumps(text_to_image_json, ensure_ascii=False)
+            print(f"Attempt {attempt} - Modified Request JSON: {modified_request_body}")
+
+            # 获取Bearer Token
+            tmp_token = BearerTokenGenerator.get_bearer(modified_request_body, path="/chats/stream")
+            if not tmp_token:
+                print(f"Attempt {attempt} - 无法生成 Bearer Token")
+                return None
+
+            bearer_token, formatted_date = tmp_token
+
+            headers = {
+                "Date": formatted_date,
+                "Client-time-zone": "-05:00",
+                "Authorization": bearer_token,
+                "User-Agent": "ChatOn_Android/1.53.502",
+                "Accept-Language": "en-US",
+                "X-Cl-Options": "hb",
+                "Content-Type": "application/json; charset=UTF-8",
             }
-        ],
-        "model": "gpt-4o",  # 固定 model，只能gpt-4o或gpt-4o-mini
-        "source": "chat/pro_image"  # 固定 source
-    }
 
-    modified_request_body = json.dumps(text_to_image_json, ensure_ascii=False)
-    print("Modified Request JSON:", modified_request_body)
+            async with httpx.AsyncClient(timeout=None) as client:
+                response = await client.post(
+                    EXTERNAL_API_URL,
+                    headers=headers,
+                    content=modified_request_body,
+                    timeout=None
+                )
 
-    # 获取Bearer Token
-    tmp_token = BearerTokenGenerator.get_bearer(modified_request_body, path="/chats/stream")
-    if not tmp_token:
-        return send_error_response("无法生成 Bearer Token", status_code=500)
+                if response.status_code != 200:
+                    print(f"Attempt {attempt} - API 错误: {response.status_code}")
+                    return None
 
-    bearer_token, formatted_date = tmp_token
+                # 读取 SSE 流并提取图像URL
+                sse_lines = response.text.splitlines()
+                image_markdown = ""
 
-    headers = {
-        "Date": formatted_date,
-        "Client-time-zone": "-05:00",
-        "Authorization": bearer_token,
-        "User-Agent": "ChatOn_Android/1.53.502",
-        "Accept-Language": "en-US",
-        "X-Cl-Options": "hb",
-        "Content-Type": "application/json; charset=UTF-8",
-    }
+                for line in sse_lines:
+                    if line.startswith("data: "):
+                        data = line[6:].strip()
+                        if data == "[DONE]":
+                            break
+                        try:
+                            sse_json = json.loads(data)
+                            if "choices" in sse_json:
+                                for choice in sse_json["choices"]:
+                                    delta = choice.get("delta", {})
+                                    content = delta.get("content")
+                                    if content:
+                                        image_markdown += content
+                        except json.JSONDecodeError:
+                            print(f"Attempt {attempt} - JSON解析错误")
+                            continue
 
-    async with httpx.AsyncClient(timeout=None) as client:
-        try:
-            response = await client.post(
-                EXTERNAL_API_URL, headers=headers, content=modified_request_body, timeout=None
-            )
-            if response.status_code != 200:
-                return send_error_response(f"API 错误: {response.status_code}", status_code=500)
-
-            # 初始化用于拼接 URL 的字符串
-            url_builder = ""
-
-            # 读取 SSE 流并拼接 URL
-            async for line in response.aiter_lines():
-                if line.startswith("data: "):
-                    data = line[6:].strip()
-                    if data == "[DONE]":
-                        break
-                    try:
-                        sse_json = json.loads(data)
-                        if "choices" in sse_json:
-                            for choice in sse_json["choices"]:
-                                delta = choice.get("delta", {})
-                                content = delta.get("content")
-                                if content:
-                                    url_builder += content
-                    except json.JSONDecodeError:
-                        print("JSON解析错误")
-                        continue
-
-            image_markdown = url_builder
-            # Step 1: 检查Markdown文本是否为空
-            if not image_markdown:
-                print("无法从 SSE 流中构建图像 Markdown。")
-                return send_error_response("无法从 SSE 流中构建图像 Markdown。", status_code=500)
-
-            # Step 2, 3, 4, 5: 处理图像
-            extracted_path = extract_path_from_markdown(image_markdown)
-            if not extracted_path:
-                print("无法从 Markdown 中提取路径。")
-                return send_error_response("无法从 Markdown 中提取路径。", status_code=500)
-
-            print(f"提��的路径: {extracted_path}")
-
-            # Step 5: 拼接最终的存储URL
-            storage_url = f"https://api.chaton.ai/storage/{extracted_path}"
-            print(f"存储URL: {storage_url}")
-
-            # 获取最终下载URL
-            final_download_url = await fetch_get_url_from_storage(storage_url)
-            if not final_download_url:
-                return send_error_response("无法从 storage URL 获取最终下载链接。", status_code=500)
-
-            print(f"Final Download URL: {final_download_url}")
-
-            # 下载图像
-            image_bytes = await download_image(final_download_url)
-            if not image_bytes:
-                return send_error_response("无法从 URL 下载图像。", status_code=500)
-
-            # 转换为 Base64
-            image_base64 = base64.b64encode(image_bytes).decode('utf-8')
-
-            # 将图片保存到images目录并构建可访问的URL
-            filename = save_base64_image(image_base64)
-            base_url = app.state.base_url
-            accessible_url = f"{base_url}/images/{filename}"
-
-            # 根据 response_format 返回相应的响应
-            if response_format.lower() == "b64_json":
-                response_json = {
-                    "data": [
-                        {
-                            "b64_json": image_base64
-                        }
-                    ]
-                }
-                return JSONResponse(content=response_json, status_code=200)
-            else:
-                # 构建包含可访问URL的响应
-                response_json = {
-                    "data": [
-                        {
-                            "url": accessible_url
-                        }
-                    ]
-                }
-                return JSONResponse(content=response_json, status_code=200)
-        except httpx.RequestError as exc:
-            print(f"请求失败: {exc}")
-            return send_error_response(f"请求失败: {str(exc)}", status_code=500)
-        except Exception as exc:
-            print(f"内部服务器错误: {exc}")
-            return send_error_response(f"内部服务器错误: {str(exc)}", status_code=500)
-
-# 运行服务器
-def main():
-    parser = argparse.ArgumentParser(description="启动ChatOn API服务器")
-    parser.add_argument('--base_url', type=str, default='http://localhost', help='Base URL for accessing images')
-    parser.add_argument('--port', type=int, default=INITIAL_PORT, help='服务器监听端口')
-    args = parser.parse_args()
-    base_url = args.base_url
-    port = args.port
+                # 检查Markdown文本是否为空
+                if not image_markdown:
+                    print(f"Attempt {attempt} - 无法从 SSE 流中构建图像 Markdown。")
+                    return None
 
-    # 检查 API_KEY 是否设置
-    if not os.environ.get("API_KEY"):
-        print("警告: API_KEY 环境变量未设置。客户端验证将无法正常工作。")
+                # 从Markdown中提取图像路径
+                extracted_path = extract_path_from_markdown(image_markdown)
+                if not extracted_path:
+                    print(f"Attempt {attempt} - 无法从 Markdown 中提取路径。")
+                    return None
 
-    # 确保 images 目录存在
-    if not os.path.exists("images"):
-        os.makedirs("images")
+                print(f"Attempt {attempt} - 提取的路径: {extracted_path}")
 
-    # 设置 FastAPI 应用的 state
-    app.state.base_url = base_url
+                # 拼接最终的存储URL
+                storage_url = f"https://api.chaton.ai/storage/{extracted_path}"
+                print(f"Attempt {attempt} - 存储URL: {storage_url}")
 
-    print(f"Server started on port {port} with base_url: {base_url}")
+                # 获取最终下载URL
+                final_download_url = await fetch_get_url_from_storage(storage_url)
+                if not final_download_url:
+                    print(f"Attempt {attempt} - 无法从 storage URL 获取最终下载链接。")
+                    return None
 
-    # 运行FastAPI应用
-    uvicorn.run(app, host="0.0.0.0", port=port)
+                print(f"Attempt {attempt} - Final Download URL: {final_download_url}")
+
+                return final_download_url
+        except Exception as e:
+            print(f"Attempt {attempt} - 处理响应时发生错误: {e}")
+            return None
+
+    # 定义一个异步任务池，限制并发数量
+    semaphore = asyncio.Semaphore(10)  # 限制同时进行的任务数为10
+
+    async def generate_with_retries(attempt: int) -> Optional[str]:
+        async with semaphore:
+            return await attempt_generate_image(attempt)
+
+    # 开始尝试生成图像
+    for attempt in range(1, max_attempts + 1):
+        needed = n - len(final_download_urls)
+        if needed <= 0:
+            break
+
+        print(f"Attempt {attempt} - 需要生成的图像数量: {needed}")
+
+        # 创建多个任务同时生成所需数量的图��
+        tasks = [asyncio.create_task(generate_with_retries(attempt)) for _ in range(needed)]
+
+        # 等待所有任务完成
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        for result in results:
+            if isinstance(result, Exception):
+                print(f"Attempt {attempt} - 任务发生异常: {result}")
+                continue
+            if result:
+                final_download_urls.append(result)
+                print(f"Attempt {attempt} - 收集到下载链接: {result}")
+
+        # 检查是否已经收集到足够的下载链接
+        if len(final_download_urls) >= n:
+            break
+
+    # 检查是否收集到足够的链接
+    if len(final_download_urls) < n:
+        print("已达到最大尝试次数，仍未收集到足够数量的下载链接。")
+        return send_error_response("无法生成足够数量的图像。", status_code=500)
+
+    # 根据 response_format 返回相应的响应
+    data_array = []
+
+    if is_base64_response:
+        for download_url in final_download_urls[:n]:
+            try:
+                image_bytes = await download_image(download_url)
+                if not image_bytes:
+                    print(f"无法从 URL 下载图像: {download_url}")
+                    continue
+
+                image_base64 = base64.b64encode(image_bytes).decode('utf-8')
+                data_array.append({
+                    "b64_json": image_base64
+                })
+            except Exception as e:
+                print(f"处理图像时发生错误: {e}")
+                continue
+    else:
+        for download_url in final_download_urls[:n]:
+            data_array.append({
+                "url": download_url
+            })
+
+    # 如果收集的 URL 数量不足 n，则通过复制现有的 URL 来填充
+    while len(data_array) < n and len(data_array) > 0:
+        for item in data_array.copy():
+            if len(data_array) >= n:
+                break
+            data_array.append(item)
+
+    # 构建最终响应
+    response_json = {
+        "created": int(datetime.now(timezone.utc).timestamp()),
+        "data": data_array
+    }
+
+    # 如果data_array为空，返回错误
+    if not data_array:
+        return send_error_response("无法生成图像。", status_code=500)
+
+    # 返回响应
+    return JSONResponse(content=response_json, status_code=200)
+
+@app.get("/v1/models", response_class=JSONResponse)
+async def get_models():
+    models_data = {
+        "object": "list",
+        "data": [
+            {"id": "gpt-4o", "object": "model"},
+            {"id": "gpt-4o-mini", "object": "model"},
+            {"id": "claude-3-5-sonnet", "object": "model"},
+            {"id": "claude", "object": "model"}
+        ]
+    }
+    return models_data
 
 async def get_available_port(start_port: int = INITIAL_PORT, end_port: int = 65535) -> int:
     """查找可用的端口号"""
     for port in range(start_port, end_port + 1):
         try:
+            # 尝试绑定端口
             server = await asyncio.start_server(lambda r, w: None, host="0.0.0.0", port=port)
             server.close()
             await server.wait_closed()
@@ -626,5 +729,31 @@ async def get_available_port(start_port: int = INITIAL_PORT, end_port: int = 655
             continue
     raise RuntimeError(f"No available ports between {start_port} and {end_port}")
 
+def main():
+    parser = argparse.ArgumentParser(description="启动ChatOn API服务器")
+    parser.add_argument('--base_url', type=str, default='http://localhost', help='Base URL for accessing images')
+    parser.add_argument('--port', type=int, default=INITIAL_PORT, help='服务器监听端口')
+    args = parser.parse_args()
+
+    base_url = args.base_url
+    port = args.port
+
+    # 设置 FastAPI 应用的 state
+    app.state.base_url = base_url
+
+    print(f"Starting server on port {port} with base_url: {base_url}")
+
+    # 检查端口可用性
+    try:
+        port = asyncio.run(get_available_port(start_port=port))
+    except RuntimeError as e:
+        print(e)
+        return
+
+    print(f"Server running on available port: {port}")
+
+    # 运行FastAPI应用
+    uvicorn.run(app, host="0.0.0.0", port=port)
+
 if __name__ == "__main__":
     main()