FiDA_Python/src/server/deep_agent/tools/generate_furniture_sketch.py

import uuid
import httpx
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

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")


@tool
async def generate_furniture(prompt: str, runtime: ToolRuntime):
    """
    使用图像生成模型根据用户提供的详细英文提示词，从零生成一张全新的家具设计草图。

    功能说明：
    - 输入一段详细的英文描述，即可生成一张高品质的家具设计图片（可用于草图、效果图、渲染图等）。
    - 生成后的图片会以 image_url 形式返回，自动加入对话上下文，后续 Agent 可以直接“看到”生成的家具图片并继续操作（描述、编辑、迭代等）。

    参数说明：
    - 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."

    返回值：
    返回新生成家具图片的 image_url，后续对话中 Agent 可直接引用该图片进行描述、进一步编辑或分析。

    使用场景：
    - 从零创建新的家具设计方案
    - 快速生成多种风格的家具概念图
    - 室内设计初期灵感生成
    - 家具产品可视化展示

    注意：
    - 生成的图片会自动携带到整个对话上下文中，支持后续使用 edit_furniture 等工具进行迭代修改。
    - 如果需要生成多个方案，可以多次调用本工具或在 prompt 中明确要求生成不同变体。
    """
    current_checkpoint_id = runtime.store.get(namespace=("image_history",), key="checkpoint_id", ).value.get("current_checkpoint_id")

    logger.info(f"\n[系统日志] 正在调用 generate_furniture ...当前checkpoint_id={current_checkpoint_id}")

    try:
        image_url = await generate_or_edit_image(prompt=prompt)

        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 image_url
        else:
            return "Image generation failed."
    except Exception as e:
        logger.warning(f"绘图流程异常：{e}")
        return "generate furniture error"


@tool
async def edit_furniture(prompt: str, runtime: ToolRuntime, config: RunnableConfig):
    """
    使用先进的图像编辑模型（image editing model）对家具设计草图进行精准修改。

    功能说明：
    - 根据用户提供的**详细英文提示词**，生成修改后的新家具图片。

    参数说明：
    - prompt (str): **必须是详细的英文提示词**，描述想要的具体修改（风格、颜色、材质、形状、添加/删除元素、比例等）。
      示例："Change the sofa to a modern minimalist style with dark gray fabric and metal legs, add a matching coffee table, make the overall lighting warmer and more luxurious."

    使用场景：
    - 家具设计迭代
    - 室内设计方案修改
    - 风格转换（现代/北欧/工业风等）
    - 材质/颜色调整

    """
    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}")

    input_path = []
    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)
    else:
        current_image_path = None

    try:
        user_input_image_paths = runtime.state.get("files").get("input_image", [])
        user_quote_image_path = runtime.state.get("files").get("quote_image", "")

        if current_image_path:
            if len(user_input_image_paths) or current_image_path:
                for path in user_input_image_paths:
                    input_path.append(path)
                if user_quote_image_path:
                    input_path.append(user_quote_image_path)
                if not len(user_input_image_paths) and not user_quote_image_path:
                    input_path = [current_image_path]

                bucket_name = "fida-public-bucket"
                object_name = f"furniture/sketches/{uuid.uuid4()}.png"
                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 image_url
                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"


async def generate_or_edit_image(input_path=None, bucket_name="fida-public-bucket",
                                 object_name=f"furniture/sketches/{uuid.uuid4()}.png",
                                 prompt="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."):
    if input_path is None:
        input_path = []
    request_data = {
        "input_image_paths": input_path,
        "prompt": prompt,
        "bucket_name": bucket_name,
        "object_name": object_name,
        "width": 1024,
        "height": 1024
    }
    async with httpx.AsyncClient(timeout=120) as client:
        resp = await client.post(
            f"http://{settings.FLUX2_GEN_IMG_MODEL_URL}/predict",
            json=request_data,
        )
    result = resp.json()
    image_url = result.get("output_path", None)
    return image_url