aiData/Controller/HaiBaoController.py

import logging
import uuid
import asyncio
from datetime import datetime
from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from sqlalchemy.sql import text

from Util.BananaClient import BananaClient
from Util.LlmUtil import get_llm_response
from Model.HaiBaoModel import HaiBaoModel

router = APIRouter(prefix="/haibao")
logger = logging.getLogger("HaiBaoController")

class GenerateRequest(BaseModel):
    prompt: str
    width: int = 1024
    height: int = 1024

@router.on_event("startup")
async def startup_event():
    """初始化时逻辑"""
    logger.info("海报控制器启动成功")

class RefineRequest(BaseModel):
    prompt: str

@router.post("/refine")
async def refine_prompt(req: RefineRequest):
    """润色提示词"""
    async def generate():
        try:
            extra_requirements = ""
            if "春节" in req.prompt:
                extra_requirements = "7. 画面氛围要极其温馨喜庆，体现春节团圆主题。可以尝试在画面中自然融入“过年了”、“回家看看”等字样（如果模型支持文字生成），或者通过红灯笼、春联、全家福等元素强烈暗示这一主题。"

            refine_system_prompt = "你是一个资深的AI绘画提示词专家。你的任务是将用户简短的描述扩充为一段详细、高质量的画面描述提示词，用于生成宣传海报。"
            refine_user_prompt = f"""
            请根据以下主题，为充电企业“驿来特”设计一张宣传海报的画面描述。
            
            主题：{req.prompt}
            
            要求：
            1. 描述画面主体、背景、光影、色彩、构图。
            2. 风格要求：现代感、科技感、精美、3D渲染风格或高品质插画风格。
            3. 融入新能源、绿色环保、充电桩等元素。
            4. 关于品牌元素：画面中可自然融入品牌Logo的视觉元素（如配色、形状），能用多少就用多少，有元素体现即可，不必生搬硬套，保持画面自然和谐。
            5. 直接输出提示词内容，不要包含“好的”、“以下是”等无关废话。
            6. 字数在100-300字之间。
            {extra_requirements}
            """
            
            async for chunk in get_llm_response(query_text=refine_user_prompt, system_prompt=refine_system_prompt, stream=True):
                yield chunk
        except Exception as e:
            logger.error(f"提示词润色失败: {e}")
            yield f"Error: {str(e)}"

    return StreamingResponse(generate(), media_type="text/plain")

@router.post("/generate")
async def generate_poster(req: GenerateRequest):
    """生成海报及文案"""
    try:
        # 并行执行生图和生文
        client = BananaClient()
        
        # 1. 构造生图任务 (包含智能润色判断)
        async def generate_image_task():
            final_prompt = req.prompt
            
            # 智能判断：如果提示词太短（少于50字），则认为用户未进行润色，自动执行润色
            # 如果用户使用了"一键扩写"功能，提示词通常会很长，这里就会跳过自动润色，尊重用户的修改
            if len(final_prompt) < 50:
                logger.info(f"提示词较短({len(final_prompt)}字)，执行自动润色...")
                
                extra_requirements = ""
                if "春节" in req.prompt:
                    extra_requirements = "7. 画面氛围要极其温馨喜庆，体现春节团圆主题。可以尝试在画面中自然融入“过年了”、“回家看看”等字样（如果模型支持文字生成），或者通过红灯笼、春联、全家福等元素强烈暗示这一主题。"

                refine_system_prompt = "你是一个资深的AI绘画提示词专家。你的任务是将用户简短的描述扩充为一段详细、高质量的画面描述提示词，用于生成宣传海报。"
                refine_user_prompt = f"""
                请根据以下主题，为充电企业“驿来特”设计一张宣传海报的画面描述。
                
                主题：{req.prompt}
                
                要求：
                1. 描述画面主体、背景、光影、色彩、构图。
                2. 风格要求：现代感、科技感、精美、3D渲染风格或高品质插画风格。
                3. 融入新能源、绿色环保、充电桩等元素。
                4. 关于品牌元素：画面中可自然融入品牌Logo的视觉元素（如配色、形状），能用多少就用多少，有元素体现即可，不必生搬硬套，保持画面自然和谐。
                5. 直接输出提示词内容，不要包含“好的”、“以下是”等无关废话。
                6. 字数在100-300字之间。
                {extra_requirements}
                """
                
                refined_prompt = ""
                try:
                    async for chunk in get_llm_response(query_text=refine_user_prompt, system_prompt=refine_system_prompt, stream=False):
                        refined_prompt += chunk
                    if refined_prompt and refined_prompt.strip():
                        final_prompt = refined_prompt
                except Exception as e:
                    logger.error(f"自动润色失败，使用原始提示词: {e}")
            
            logger.info(f"Final generation prompt: {final_prompt}")

            # 1.2 调用生图
            resp = await client.generate_image(prompt=final_prompt, size=f"{req.width}x{req.height}")
            # 定义Logo路径
            LOGO_PATH = r"d:\dsWork\aiData\static\Images\login_logo.png"
            obs_urls = await client.download_and_upload_to_obs(resp, overlay_logo_path=LOGO_PATH)
            if not obs_urls:
                raise Exception("未获取到有效的图片URL")
            return obs_urls[0]

        # 2. 构造生文任务
        async def generate_text_task():
            extra_text_req = ""
            if "春节" in req.prompt:
                extra_text_req = "6. 文案中必须包含“过年了，回家看看，记得提前充电噢”或类似的温馨提示。"

            scheme_prompt = f"""
            你是一个专业的社群运营专家。请为充电企业“驿来特”撰写一段发在微信群里的宣传文案。
            
            主题：{req.prompt}
            
            要求：
            1. 语气亲切、有吸引力，适合微信社群传播。
            2. 突出“驿来特”品牌，强调新能源、优惠、便利等特点（根据主题自由发挥）。
            3. 包含适当的emoji表情，增加趣味性。
            4. 字数控制在150字以内。
            5. 格式清晰，分段合理。
            {extra_text_req}
            """
            # get_llm_response 是一个异步生成器 (stream=True by default) 或者直接返回 (stream=False)
            # 这里我们强制 stream=False 获取完整文本
            text_response = ""
            # LlmUtil.get_llm_response 默认为 stream=True，我们需要修改调用方式或适配
            # 查看 LlmUtil 源码，如果 stream=False，它 yield 内容。
            # 所以我们需要迭代它
            async for chunk in get_llm_response(query_text=scheme_prompt, stream=False):
                text_response += chunk
            return text_response

        # 3. 并行执行
        image_url, scheme_content = await asyncio.gather(generate_image_task(), generate_text_task())
        
        # 4. 保存到数据库
        model = HaiBaoModel()
        record_id = str(uuid.uuid4())
        created_at = datetime.now()
        
        await model.insert_record(record_id, req.prompt, image_url, scheme_content, created_at)
                
        return {
            "id": record_id,
            "image_url": image_url,
            "prompt": req.prompt,
            "scheme_content": scheme_content,
            "created_at": created_at.strftime("%Y-%m-%d %H:%M:%S")
        }
        
    except Exception as e:
        logger.error(f"生成海报/文案失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/history")
async def get_history():
    """获取海报生成历史"""
    try:
        model = HaiBaoModel()
        result = await model.get_history(50)
        
        formatted_result = []
        for item in result:
            item_dict = dict(item) if not isinstance(item, dict) else item
            
            if isinstance(item_dict.get('created_at'), datetime):
                item_dict['created_at'] = item_dict['created_at'].strftime("%Y-%m-%d %H:%M:%S")
            
            # 确保 scheme_content 存在
            if 'scheme_content' not in item_dict:
                item_dict['scheme_content'] = ""
                
            formatted_result.append(item_dict)
                
        return formatted_result
    except Exception as e:
        logger.error(f"获取历史失败: {e}")
        return []