dsProject/dsLightRag/Routes/JiMengRoute.py

import logging

import fastapi
from fastapi import APIRouter
from fastapi import HTTPException
from openai import AsyncOpenAI

from Config import Config
from JiMeng.Kit.JmTxt2ImgUtil import JmTxt2Img
from JiMeng.Kit.FenJingTouGenerator import FenJingTouGenerator  # 导入分镜头生成器

# 创建路由路由器
router = APIRouter(prefix="/api/jimeng", tags=["即梦"])

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

# 初始化异步 OpenAI 客户端
client = AsyncOpenAI(
    api_key=Config.ALY_LLM_API_KEY,
    base_url=Config.ALY_LLM_BASE_URL
)


@router.post("/prompt_input")
async def prompt_input(request: fastapi.Request):
    try:
        data = await request.json()
        prompt = data.get("prompt")
        if not prompt:
            raise HTTPException(status_code=400, detail="缺少提示词参数")

        logger.info(f"收到图片生成请求，提示词: {prompt}")
        # 调用 JmTxt2Img 生成图片
        image_url = JmTxt2Img.generate_image(prompt)
        logger.info(f"图片生成成功，URL: {image_url}")

        return {
            "code": 200,
            "message": "成功",
            "data": {
                "image_url": image_url
            }
        }
    except HTTPException as e:
        logger.error(f"请求参数错误: {str(e.detail)}")
        raise e
    except Exception as e:
        logger.error(f"图片生成失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"图片生成失败: {str(e)}")


@router.post("/generate_fenjingtou")
async def generate_fenjingtou(request: fastapi.Request):
    try:
        data = await request.json()
        image_url = data.get("image_url")
        prompt_text = data.get("prompt_text")

        if not image_url:
            raise HTTPException(status_code=400, detail="缺少图片地址参数")
        if not prompt_text:
            raise HTTPException(status_code=400, detail="缺少提示词参数")

        logger.info(f"收到分镜头生成请求，图片地址: {image_url}，提示词: {prompt_text}")

        # 创建分镜头生成器实例并调用方法
        generator = FenJingTouGenerator()
        fenjingtou_content = generator.generate_fen_jing_tou(image_url, prompt_text)

        if not fenjingtou_content:
            raise HTTPException(status_code=500, detail="分镜头脚本生成失败")

        logger.info(f"分镜头脚本生成成功")

        return {
            "code": 200,
            "message": "成功",
            "data": {
                "fenjingtou_content": fenjingtou_content
            }
        }
    except HTTPException as e:
        logger.error(f"请求参数错误: {str(e.detail)}")
        raise e
    except Exception as e:
        logger.error(f"分镜头脚本生成失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"分镜头脚本生成失败: {str(e)}")