'commit'

dsLightRag/Routes/MjRoute.py

@@ -11,6 +11,7 @@ from pydantic import BaseModel
 from Midjourney.Txt2Img import Txt2Img
 from Config import Config
 import asyncio
+import threading
 from fastapi.responses import StreamingResponse
 
 # 创建路由路由器
@@ -71,6 +72,72 @@ async def submit_imagine(request: ImagineRequest, background_tasks: BackgroundTa
 
         # 存储Midjourney任务ID
         TASK_STATUS[task_id]["midjourney_task_id"] = midjourney_task_id
+
+        # 添加后台任务轮询状态
+        def poll_task_status_background(task_id, midjourney_task_id):
+            max_retries = 1000
+            retry_count = 0
+            retry_interval = 5  # 5秒
+
+            while retry_count < max_retries:
+                try:
+                    # 查询任务状态
+                    result = Txt2Img.query_task_status(midjourney_task_id)
+
+                    # 更新任务状态
+                    if result.get("status") == "SUCCESS":
+                        TASK_STATUS[task_id] = {
+                            "status": "completed",
+                            "image_url": result.get("imageUrl"),
+                            "progress": 100,
+                            "error": None,
+                            "midjourney_task_id": midjourney_task_id
+                        }
+                        logger.info(f"任务 {task_id} 完成，图片URL: {result.get('imageUrl')}")
+                        break
+                    elif result.get("status") == "FAILED":
+                        TASK_STATUS[task_id] = {
+                            "status": "failed",
+                            "image_url": None,
+                            "progress": 0,
+                            "error": result.get("errorMsg", "未知错误"),
+                            "midjourney_task_id": midjourney_task_id
+                        }
+                        logger.error(f"任务 {task_id} 失败: {result.get('errorMsg', '未知错误')}")
+                        break
+                    else:
+                        # 更新进度
+                        progress = result.get("progress", 0)
+                        TASK_STATUS[task_id]["progress"] = progress
+                        TASK_STATUS[task_id]["status"] = "processing"
+                        logger.info(f"任务 {task_id} 处理中，进度: {progress}%")
+
+                    # 增加重试计数
+                    retry_count += 1
+
+                    # 等待重试间隔
+                    time.sleep(retry_interval)
+
+                except Exception as e:
+                    logger.error(f"轮询任务 {task_id} 状态失败: {str(e)}")
+                    TASK_STATUS[task_id]["error"] = str(e)
+                    time.sleep(retry_interval)
+
+            if retry_count >= max_retries:
+                logger.error(f"任务 {task_id} 超时")
+                TASK_STATUS[task_id] = {
+                    "status": "failed",
+                    "image_url": None,
+                    "progress": 0,
+                    "error": "任务处理超时",
+                    "midjourney_task_id": midjourney_task_id
+                }
+
+        # 使用线程运行后台任务
+        thread = threading.Thread(target=poll_task_status_background, args=(task_id, midjourney_task_id))
+        thread.daemon = True
+        thread.start()
+
         return {"task_id": task_id}
     except Exception as e:
         logger.error(f"提交文生图请求失败: {str(e)}")