'commit'

dsLightRag/Midjourney/Txt2Img.py

@@ -147,5 +147,52 @@ class Txt2Img(MjCommon):
         except Exception as e:
             log.error(f"程序执行异常: {str(e)}")
 
+@staticmethod
+def query_task_status(task_id):
+    """
+    查询任务状态
+
+    :param task_id: 任务ID
+    :return: 任务状态信息
+    :raises Exception: 异常信息
+    """
+    try:
+        # 创建请求
+        url = f"{MjCommon.BASE_URL}/mj/task/status"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {MjCommon.ak}"
+        }
+        params = {
+            "taskId": task_id
+        }
+
+        # 发送请求并获取响应
+        log.info(f"查询Midjourney任务状态: {task_id}")
+        response = requests.get(url, headers=headers, params=params, timeout=30)
+
+        # 检查响应状态
+        if response.status_code != 200:
+            error_msg = f"Midjourney API请求失败，状态码: {response.status_code}"
+            log.error(error_msg)
+            raise Exception(error_msg)
+
+        # 解析响应
+        response_body = response.text
+        log.info(f"Midjourney任务状态响应: {response_body}")
+
+        response_json = json.loads(response_body)
+
+        # 检查响应状态 - code=1 表示成功
+        if response_json.get("code") != 1:
+            error_msg = f"Midjourney任务状态查询失败: {response_json.get('description', '未知错误')}"
+            log.error(error_msg)
+            raise Exception(error_msg)
+
+        # 返回任务状态信息
+        return response_json.get("result", {})
+    except Exception as e:
+        log.error(f"查询任务状态异常: {str(e)}")
+        raise
 if __name__ == "__main__":
     Txt2Img.main()