dsProject/dsLightRag/Routes/MjRoute.py

import logging
import json
import time
import datetime
import requests
import uuid
import os
from typing import Optional, Dict, Any
from fastapi import APIRouter, HTTPException, Query, BackgroundTasks
from pydantic import BaseModel
from Midjourney.Txt2Img import Txt2Img
from Config import Config
import asyncio
import threading
from fastapi.responses import StreamingResponse
from fastapi import APIRouter, Request, HTTPException, BackgroundTasks

# 创建路由路由器
router = APIRouter(prefix="/api/mj", tags=["文生图"])

# 配置日志
logger = logging.getLogger(__name__)

class ImagineRequest(BaseModel):
    prompt: str
    base64_array: Optional[list] = None
    notify_hook: Optional[str] = None

class TaskStatusResponse(BaseModel):
    task_id: str
    status: str
    image_url: Optional[str] = None
    progress: Optional[int] = None
    error: Optional[str] = None

@router.post("/imagine", response_model=Dict[str, str])
async def submit_imagine(request: ImagineRequest):
    """
    提交文生图请求
    
    Args:
        request: 包含提示词、垫图和回调地址的请求体
    
    Returns:
        包含MJ服务器任务ID的字典
    """
    try:
        # 直接提交到Midjourney并返回MJ任务ID
        midjourney_task_id = Txt2Img.submit_imagine(
            prompt=request.prompt,
            base64_array=request.base64_array,
            notify_hook=request.notify_hook
        )
        
        logger.info(f"提交文生图请求成功，MJ任务ID: {midjourney_task_id}, 提示词: {request.prompt}")
        return {"task_id": midjourney_task_id}
    except Exception as e:
        logger.error(f"提交文生图请求失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"提交文生图请求失败: {str(e)}")

@router.get("/task_status")
async def get_task_status(task_id: str):
    """
    直接查询MJ服务器获取任务状态
    
    Args:
        task_id: MJ服务器返回的任务ID
    
    Returns:
        包含任务状态的字典
    """
    try:
        # 直接查询MJ服务器获取实时状态
        result = Txt2Img.query_task_status(task_id)
        
        # 处理查询结果
        if result.get("status") == "SUCCESS":
            return {
                "task_id": task_id,
                "status": "completed",
                "image_url": result.get("imageUrl"),
                "progress": 100,
                "error": None
            }
        elif result.get("status") == "FAILED":
            return {
                "task_id": task_id,
                "status": "failed",
                "image_url": None,
                "progress": 0,
                "error": result.get("errorMsg", "未知错误")
            }
        else:
            return {
                "task_id": task_id,
                "status": "processing",
                "progress": result.get("progress", 0),
                "error": None
            }
    except Exception as e:
        logger.error(f"查询任务状态失败: {str(e)}")
        return {
            "task_id": task_id,
            "status": "error",
            "message": f"查询任务状态失败: {str(e)}"
        }, 500