'commit'

Controller/HaiBaoController.py

@@ -3,6 +3,7 @@ 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
 
@@ -64,33 +65,35 @@ class RefineRequest(BaseModel):
 @router.post("/refine")
 async def refine_prompt(req: RefineRequest):
     """润色提示词"""
-    try:
-        refine_system_prompt = "你是一个资深的AI绘画提示词专家。你的任务是将用户简短的描述扩充为一段详细、高质量的画面描述提示词，用于生成宣传海报。"
-        refine_user_prompt = f"""
-        请根据以下主题，为充电企业“驿来特”设计一张宣传海报的画面描述。
-        
-        主题：{req.prompt}
-        
-        要求：
-        1. 描述画面主体、背景、光影、色彩、构图。
-        2. 风格要求：现代感、科技感、精美、3D渲染风格或高品质插画风格。
-        3. 融入新能源、绿色环保、充电桩等元素。
-        4. 关于品牌元素：画面中可自然融入品牌Logo的视觉元素（如配色、形状），能用多少就用多少，有元素体现即可，不必生搬硬套，保持画面自然和谐。
-        5. 直接输出提示词内容，不要包含“好的”、“以下是”等无关废话。
-        6. 字数在100-300字之间。
-        """
-        
-        refined_prompt = ""
+    async def generate():
         try:
-            async for chunk in get_llm_response(query_text=refine_user_prompt, system_prompt=refine_system_prompt, stream=False):
-                refined_prompt += chunk
+            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}")
-            raise Exception("润色服务暂时不可用")
-            
-        return {"refined_prompt": refined_prompt}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+            yield f"Error: {str(e)}"
+
+    return StreamingResponse(generate(), media_type="text/plain")
 
 @router.post("/generate")
 async def generate_poster(req: GenerateRequest):
@@ -107,6 +110,11 @@ async def generate_poster(req: GenerateRequest):
             # 如果用户使用了"一键扩写"功能，提示词通常会很长，这里就会跳过自动润色，尊重用户的修改
             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"""
                 请根据以下主题，为充电企业“驿来特”设计一张宣传海报的画面描述。
@@ -120,6 +128,7 @@ async def generate_poster(req: GenerateRequest):
                 4. 关于品牌元素：画面中可自然融入品牌Logo的视觉元素（如配色、形状），能用多少就用多少，有元素体现即可，不必生搬硬套，保持画面自然和谐。
                 5. 直接输出提示词内容，不要包含“好的”、“以下是”等无关废话。
                 6. 字数在100-300字之间。
+                {extra_requirements}
                 """
                 
                 refined_prompt = ""
@@ -144,6 +153,10 @@ async def generate_poster(req: GenerateRequest):
 
         # 2. 构造生文任务
         async def generate_text_task():
+            extra_text_req = ""
+            if "春节" in req.prompt:
+                extra_text_req = "6. 文案中必须包含“过年了，回家看看，记得提前充电噢”或类似的温馨提示。"
+
             scheme_prompt = f"""
             你是一个专业的社群运营专家。请为充电企业“驿来特”撰写一段发在微信群里的宣传文案。
             
@@ -155,6 +168,7 @@ async def generate_poster(req: GenerateRequest):
             3. 包含适当的emoji表情，增加趣味性。
             4. 字数控制在150字以内。
             5. 格式清晰，分段合理。
+            {extra_text_req}
             """
             # get_llm_response 是一个异步生成器 (stream=True by default) 或者直接返回 (stream=False)
             # 这里我们强制 stream=False 获取完整文本
@@ -187,7 +201,7 @@ async def generate_poster(req: GenerateRequest):
             "created_at": created_at
         }
         
-        async with db.get_session() as session:
+        async with await db.get_session() as session:
             async with session.begin():
                 await session.execute(text(insert_sql), params)