支持并默认一次生成12张

src/server/deep_agent/tools/generate_furniture_sketch.py

@@ -1,73 +1,103 @@
-import uuid
-from typing import Optional
-
 import httpx
+import uuid
 import logging
 
 from langchain_core.runnables import RunnableConfig
 from minio import Minio
-# from pathlib import Path
-# from datetime import datetime
-from langchain_core.tools import tool
 from langgraph.prebuilt import ToolRuntime
-from src.core.config import settings, MONGO_URI
-# from src.server.deep_agent.utils.mongodb_util import ThreadImageMinIOStore
+from src.core.config import settings
 
 logger = logging.getLogger(__name__)
 minio_client = Minio(settings.MINIO_URL, access_key=settings.MINIO_ACCESS, secret_key=settings.MINIO_SECRET, secure=settings.MINIO_SECURE)
-# image_store = ThreadImageMinIOStore(MONGO_URI, "agent_tool_generate_db")
+
+from typing import List, Optional
+from langchain_core.tools import tool
+
+logger = logging.getLogger(__name__)
 
 
 @tool
-async def generate_furniture(runtime: ToolRuntime, prompts: list[str] = None, num_images: Optional[int] = 1):
+async def generate_furniture(runtime: ToolRuntime, prompts: List[str] = None, num_images: Optional[int] = 12, ):
     """
-    使用图像生成模型根据用户提供的详细英文提示词，从零生成一张全新的家具设计草图。
+    生成家具设计线稿草图（sketch / line drawing）。
 
     功能说明：
-    - 输入一段详细的英文描述，即可生成一张高品质的家具设计图片（可用于草图、效果图、渲染图等）。
-    - 生成后的图片会以 image_url 形式返回，自动加入对话上下文，后续 Agent 可以直接“看到”生成的家具图片并继续操作（描述、编辑、迭代等）。
+    - 默认生成 12 张家具设计线稿。
+    - 智能处理 prompts 数量与生成数量不一致的情况：
+      - 如果只有一个 prompt → 用该 prompt 生成全部 12 张（不同随机变体）。
+      - 如果有多个 prompt → 自动均匀分配生成数量（尽量让每个 prompt 生成相同数量）。
+    - 生成过程会一张一张进行，适合用户实时查看。
 
     参数说明：
-    - prompt (str): **必须是详细的英文提示词**，越详细越好，包含家具类型、风格、颜色、材质、尺寸比例、背景、视角、光影等具体要求。
-      示例："Generate a modern minimalist dining chair made of light oak wood and white leather, with slim metal legs, photographed in a bright Scandinavian living room with natural sunlight, high detail, 8k resolution."
-    - num_images (int, 可选): 要生成的图片数量，默认 1 张。最大只能是 4 张。如果输入超过 4，会自动限制为 4。
+    - prompts (list[str]):
+      必须是列表，即使只有一个提示词也要用 ["你的提示词"] 格式。
+      提供详细的英文提示词，描述越详细越好。
+    - num_images (int, 可选): 要生成的图片总数量，默认 12 张，最大限制为 12 张。
 
     返回值：
-    返回新生成家具图片的 image_url，后续对话中 Agent 可直接引用该图片进行描述、进一步编辑或分析。
-
-    使用场景：
-    - 从零创建新的家具设计方案
-    - 快速生成多种风格的家具概念图
-    - 室内设计初期灵感生成
-    - 家具产品可视化展示
-
-    注意：
-    - 生成的图片会自动携带到整个对话上下文中，支持后续使用 edit_furniture 等工具进行迭代修改。
-    - 如果需要生成多个方案，可以多次调用本工具或在 prompt 中明确要求生成不同变体。
+    返回 image_urls 列表，系统会自动依次展示生成的图片。
     """
+    # ====================== 参数安全处理 ======================
+    if prompts is None or len(prompts) == 0:
+        return "Error: prompts 参数不能为空。请至少提供一个详细的英文提示词。"
+
+    if not isinstance(prompts, list):
+        prompts = [str(prompts)]
+
+    # 数量限制
     if num_images is None or num_images < 1:
         num_images = 1
-    elif num_images > 4:
-        num_images = 4
-    # current_checkpoint_id = runtime.store.get(namespace=("image_history",), key="checkpoint_id", ).value.get("current_checkpoint_id")
+    elif num_images > 12:
+        num_images = 12
 
-    logger.info(f"\n[系统日志] 正在调用 generate_furniture ")
+    n_prompts = len(prompts)
 
+    logger.info(f"[generate_furniture] 开始生成 | prompts数量={n_prompts} | num_images={num_images}（默认12）")
+
+    # ====================== 均匀分配 prompts（核心逻辑） ======================
+    if n_prompts == 0:
+        return "Error: prompts 列表为空"
+
+    # 计算每个 prompt 应该生成的张数
+    base_count = num_images // n_prompts
+    remainder = num_images % n_prompts
+
+    images_per_prompt = [base_count] * n_prompts
+    for i in range(remainder):
+        images_per_prompt[i] += 1
+
+    # 构建实际使用的 prompt 列表
+    expanded_prompts: List[str] = []
+    for i, count in enumerate(images_per_prompt):
+        expanded_prompts.extend([prompts[i]] * count)
+
+    logger.info(f"[generate_furniture] 分配完成: {images_per_prompt} （每个prompt生成张数）")
+
+    # ====================== 生成图片 ======================
     try:
         bucket_name = "fida-public-bucket"
-        object_name = f"furniture/sketches/{uuid.uuid4()}"
+        base_object_name = f"furniture/sketches/{uuid.uuid4()}"
         image_urls = []
-        for i in range(num_images):
-            image_urls.append(await generate_or_edit_image(prompt=prompts[i], bucket_name=bucket_name, object_name=f"{object_name}-{i}.png"))
 
-        # if image_urls:
-        #     image_store.save_image_path(thread_id=current_checkpoint_id, object_path=image_urls, metadata={"prompt": prompt, "generated_at": str(datetime.now())})
+        for i in range(num_images):
+            prompt = expanded_prompts[i]
+            object_name = f"{base_object_name}-{i:02d}.png"
+
+            image_url = await generate_or_edit_image(
+                prompt=prompt,
+                bucket_name=bucket_name,
+                object_name=object_name
+            )
+            image_urls.append(image_url)
+
+            logger.info(f"[generate_furniture] 已生成第 {i + 1}/{num_images} 张")
+
+        logger.info(f"[generate_furniture] 成功生成 {len(image_urls)} 张图片")
         return image_urls
-        # else:
-        #     return "Image generation failed."
+
     except Exception as e:
-        logger.warning(f"绘图流程异常：{e}")
-        return "generate furniture error"
+        logger.error(f"generate_furniture 执行异常: {e}", exc_info=True)
+        return f"generate furniture error: {str(e)}"
 
 
 @tool
@@ -112,26 +142,7 @@ async def edit_furniture(runtime: ToolRuntime, config: RunnableConfig, input_ima
         "Change the chair to a sleek modern design with black leather and chrome legs."
     ]
     """
-
-    # image_history = runtime.store.get(namespace=("image_history",), key="checkpoint_id", )
-    # last_checkpoint_id = image_history.value.get("last_checkpoint_id")
-    # current_checkpoint_id = image_history.value.get("current_checkpoint_id")
-    #
-    # logger.info(f"\n[系统日志] 正在调用 edit_furniture ...current_checkpoint_id={current_checkpoint_id}  ---  last_checkpoint_id={last_checkpoint_id}")
-    #
-    # if image_store.get_image_path(last_checkpoint_id):
-    #     current_image_path = image_store.get_image_path(last_checkpoint_id).get("current_image_path", False)
-    #     if current_image_path:
-    #         if isinstance(current_image_path, list):
-    #             #  只取最后一张
-    #             current_image_path = current_image_path[-1]
-    # else:
-    #     current_image_path = None
-
-    # input_path = []
     try:
-        # user_input_image_paths = runtime.state.get("files").get("input_image", [])
-        # user_quote_image_path = runtime.state.get("files").get("quote_image", "")
         result = []
         if len(input_image_paths):
             for i in range(len(input_image_paths)):
@@ -139,17 +150,7 @@ async def edit_furniture(runtime: ToolRuntime, config: RunnableConfig, input_ima
                 object_name = f"furniture/sketches/{uuid.uuid4()}.png"
                 image_url = await generate_or_edit_image(input_path=[input_image_paths[i]], prompt=prompts[i], bucket_name=bucket_name, object_name=f"{object_name}-{i}.png")
                 result.append(image_url)
-                # image_url = await generate_or_edit_image(input_path=input_path, prompt=prompt, bucket_name=bucket_name, object_name=object_name)
-                # if image_url:
-                #     image_store.save_image_path(thread_id=current_checkpoint_id, object_path=[image_url], metadata={"prompt": prompt, "generated_at": str(datetime.now())})
         return result
-        # else:
-        #     return "Image generation failed."
-        # else:
-        #     return "The picture to be edited does not exist."
-        # else:
-        #     return "No recent image found, please upload or cite it"
-
     except Exception as e:
         logger.warning(f"edit_furniture error ：{e}")
         return "edit_furniture error"