Update main.py

main.py

@@ -172,60 +172,6 @@ async def chat_completions(
                 stream=True
             )
             
-            async def process_image_gen(job_id, prompt, negative):
-                """处理图片生成的异步函数"""
-                print(f"Starting image generation process for job_id: {job_id}")
-                
-                # 立即发送思考开始的消息
-                start_time = time.time()
-                think_msg = "<think>\n"
-                think_msg += "🎨 Generating image...\n\n"
-                think_msg += f"Prompt: {prompt}\n"
-                
-                # 发送思考开始消息
-                chunk1 = {
-                    "id": f"chatcmpl-{chat_id}",
-                    "object": "chat.completion.chunk",
-                    "created": int(time.time()),
-                    "model": data.get('model'),
-                    "choices": [{
-                        "delta": {"content": think_msg},
-                        "index": 0,
-                        "finish_reason": None
-                    }]
-                }
-                
-                # 检查图片状态和上传
-                image_base64 = await check_image_status(session, job_id, headers)
-                if image_base64:
-                    result = await upload_to_xinyew(image_base64, job_id)
-                else:
-                    result = None
-                
-                # 发送结束消息
-                elapsed_time = time.time() - start_time
-                end_msg = f"\n🤔 Thinking for {elapsed_time:.1f}s...\n"
-                end_msg += "</think>\n\n"
-                if result:
-                    end_msg += f"![Generated Image]({result})"
-                else:
-                    end_msg += "*Image generation or upload failed.*\n"
-                
-                # 发送结束消息
-                chunk2 = {
-                    "id": f"chatcmpl-{chat_id}",
-                    "object": "chat.completion.chunk",
-                    "created": int(time.time()),
-                    "model": data.get('model'),
-                    "choices": [{
-                        "delta": {"content": end_msg},
-                        "index": 0,
-                        "finish_reason": None
-                    }]
-                }
-                
-                return [chunk1, chunk2]
-            
             def generate():
                 content_buffer = ""
                 for line in response.iter_lines():
@@ -245,22 +191,24 @@ async def chat_completions(
                             # 在处理消息时先判断模型类型
                             if data.get('model') == 'AkashGen' and "<image_generation>" in msg_data:
                                 # 图片生成模型的特殊处理
-                                match = re.search(r"jobId='([^']+)' prompt='([^']+)' negative='([^']*)'", msg_data)
-                                if match:
-                                    job_id, prompt, negative = match.groups()
-                                    
-                                    # 创建新的事件循环
-                                    loop = asyncio.new_event_loop()
-                                    asyncio.set_event_loop(loop)
-                                    try:
-                                        chunks = loop.run_until_complete(process_image_gen(job_id, prompt, negative))
-                                    finally:
-                                        loop.close()
-                                    
-                                    # 发送所有消息块
-                                    for chunk in chunks:
-                                        yield f"data: {json.dumps(chunk)}\n\n"
-                                    continue
+                                async def process_and_send():
+                                    messages = await process_image_generation(msg_data, session, headers, chat_id)
+                                    if messages:
+                                        return messages
+                                    return None
+
+                                # 创建新的事件循环
+                                loop = asyncio.new_event_loop()
+                                asyncio.set_event_loop(loop)
+                                try:
+                                    result_messages = loop.run_until_complete(process_and_send())
+                                finally:
+                                    loop.close()
+                                
+                                if result_messages:
+                                    for message in result_messages:
+                                        yield f"data: {json.dumps(message)}\n\n"
+                                continue
                             
                             content_buffer += msg_data
                             
@@ -282,7 +230,7 @@ async def chat_completions(
                                 "id": f"chatcmpl-{chat_id}",
                                 "object": "chat.completion.chunk",
                                 "created": int(time.time()),
-                                "model": data.get('model'),
+                                "model": data.get('model'),  # 使用请求中指定的模型
                                 "choices": [{
                                     "delta": {},
                                     "index": 0,
@@ -459,104 +407,59 @@ async def process_image_generation(msg_data: str, session: requests.Session, hea
         job_id, prompt, negative = match.groups()
         print(f"Starting image generation process for job_id: {job_id}")
         
-        # 返回多个消息块
-        messages = []
-        
-        # 发送思考开始标记
-        messages.append({
-            "id": f"chatcmpl-{chat_id}-1",
-            "object": "chat.completion.chunk",
-            "created": int(time.time()),
-            "model": "AkashGen",
-            "choices": [{
-                "delta": {"content": "<think>\n"},
-                "index": 0,
-                "finish_reason": None
-            }]
-        })
-        
-        # 发送生成图片的提示
-        messages.append({
-            "id": f"chatcmpl-{chat_id}-2",
-            "object": "chat.completion.chunk",
-            "created": int(time.time()),
-            "model": "AkashGen",
-            "choices": [{
-                "delta": {"content": "🎨 Generating image...\n\n"},
-                "index": 0,
-                "finish_reason": None
-            }]
-        })
-        
-        # 发送提示词
-        messages.append({
-            "id": f"chatcmpl-{chat_id}-3",
-            "object": "chat.completion.chunk",
-            "created": int(time.time()),
-            "model": "AkashGen",
-            "choices": [{
-                "delta": {"content": f"Prompt: {prompt}\n\n"},
-                "index": 0,
-                "finish_reason": None
-            }]
-        })
-        
-        # 记录开始时间
+        # 发送思考开始的消息
         start_time = time.time()
+        think_msg = "<think>\n"
+        think_msg += "🎨 Generating image...\n\n"
+        think_msg += f"Prompt: {prompt}\n"
         
         # 检查图片状态和上传
         result = await check_image_status(session, job_id, headers)
         
-        # 计算实际经过的时间
+        # 完成思考部分
         elapsed_time = time.time() - start_time
+        think_msg += f"\n🤔 Thinking for {elapsed_time:.1f}s...\n"
+        think_msg += "</think>"
         
-        # 发送思考时间
-        messages.append({
-            "id": f"chatcmpl-{chat_id}-4",
-            "object": "chat.completion.chunk",
-            "created": int(time.time()),
-            "model": "AkashGen",
-            "choices": [{
-                "delta": {"content": f"🤔 Thinking for {elapsed_time:.1f}s...\n"},
-                "index": 0,
-                "finish_reason": None
-            }]
-        })
+        # 返回两个独立的消息块
+        messages = []
         
-        # 发送思考结束标记
+        # 第一个消息块：思考过程
         messages.append({
-            "id": f"chatcmpl-{chat_id}-5",
+            "id": f"chatcmpl-{chat_id}-think",
             "object": "chat.completion.chunk",
             "created": int(time.time()),
             "model": "AkashGen",
             "choices": [{
-                "delta": {"content": "</think>\n\n"},
+                "delta": {"content": think_msg},
                 "index": 0,
                 "finish_reason": None
             }]
         })
         
-        # 发送图片结果
+        # 第二个消息块：图片结果
         if result:
+            image_msg = f"\n\n![Generated Image]({result})"
             messages.append({
-                "id": f"chatcmpl-{chat_id}-6",
+                "id": f"chatcmpl-{chat_id}-image",
                 "object": "chat.completion.chunk",
                 "created": int(time.time()),
                 "model": "AkashGen",
                 "choices": [{
-                    "delta": {"content": f"![Generated Image]({result})"},
+                    "delta": {"content": image_msg},
                     "index": 0,
                     "finish_reason": None
                 }]
             })
         else:
+            fail_msg = "\n\n*Image generation or upload failed.*"
             messages.append({
-                "id": f"chatcmpl-{chat_id}-6",
+                "id": f"chatcmpl-{chat_id}-fail",
                 "object": "chat.completion.chunk",
                 "created": int(time.time()),
                 "model": "AkashGen",
                 "choices": [{
-                    "delta": {"content": "*Image generation or upload failed.*"},
+                    "delta": {"content": fail_msg},
                     "index": 0,
                     "finish_reason": None
                 }]