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

import uuid
import httpx
import logging

from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.stores import BaseStore
from minio import Minio
from google import genai
from pathlib import Path
from datetime import datetime
from langchain_core.tools import tool
from google.oauth2 import service_account
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.server.utils.new_oss_client import get_presigned_url, check_and_extract_minio_image

# from google.genai.types import GenerateContentConfig, Modality
# from langgraph.config import get_stream_writer
# from src.server.utils.new_oss_client import oss_upload_image, oss_get_image, is_minio_file_exist, oss_upload_image_file

# 初始化全局凭证和客户端
# creds = service_account.Credentials.from_service_account_file(
#     settings.GOOGLE_GENAI_USE_VERTEXAI,
#     scopes=["https://www.googleapis.com/auth/cloud-platform"],
# )
# client = genai.Client(
#     credentials=creds,
#     project=settings.GOOGLE_CLOUD_PROJECT,
#     location=settings.GOOGLE_CLOUD_LOCATION,
#     vertexai=True
# )

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


def is_image_path_exist(image_path):
    try:
        return Path(image_path).exists()
    except:
        return False


@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 中明确要求生成不同变体。
    """
    logger.info(f"\n[系统日志] 正在调用 generate_furniture ...")
    # thread_id = runtime.config.get("configurable").get("thread_id")
    try:
        # 1. 生成图像 - local flux2-klein
        object_name = f"furniture/sketches/{uuid.uuid4()}.png"
        bucket_name = "fida-public-bucket"  # 替换为你的 bucket 名称
        request_data = {
            "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)
        # image_presigned_url = get_presigned_url(oss_client=minio_client, bucket=bucket_name, object_name=object_name)
        if image_url:
            # image_store.save_image_path(thread_id=thread_id, object_path=image_url, metadata={"prompt": prompt, "generated_at": str(datetime.now())})
            # image_id = str(uuid.uuid4())
            # runtime.store.put(
            #     namespace=("images", thread_id),  # e.g. ("images", thread_id)
            #     key=image_id,
            #     value={
            #         "url": image_url,
            #         "prompt": prompt,
            #         "timestamp": str(uuid.uuid4()),  # 或用 datetime
            #         "version": "v1"
            #     }
            # )
            return f"Image has been generated: https://minio-api.aida.com.hk/{image_url}"
        else:
            return "Image generation failed."
    except Exception as e:
        logger.warning(f"绘图流程异常：{e}")
        return "generate furniture error"


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

    功能说明：
    - 输入一张家具图片（草图/效果图），根据用户提供的**详细英文提示词**，生成修改后的新家具图片。
    - 修改后的图片会以 image_url 形式返回，自动加入对话上下文，后续 Agent 可以直接“看到”编辑结果并继续操作。

    参数说明：
    - image_url (str): 原始家具图片的 URL（支持公开 http/https 链接或 data:image/...;base64 格式）。
    - 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_url，后续对话中 Agent 可直接引用该图片进行描述、进一步编辑或分析。

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

    注意：如果需要多次迭代编辑，直接在下一次调用时传入上一次返回的新 image_url 即可。
    """
    logger.info(f"\n[系统日志] 正在调用 edit_furniture ...")
    thread_id = runtime.config.get("configurable").get("thread_id")
    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", "")

        extract_result = check_and_extract_minio_image(url=image_url)
        if extract_result['state']:
            current_image_path = extract_result['data']
            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]
                object_name = f"furniture/sketches/{uuid.uuid4()}.png"
                bucket_name = "fida-public-bucket"  # 替换为你的 bucket 名称
                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)

                if image_url:
                    # image_store.save_image_path(thread_id=thread_id, object_path=image_url, metadata={"prompt": prompt, "generated_at": str(datetime.now())})
                    return f"Image has been generated: https://minio-api.aida.com.hk/{image_url}"
                else:
                    return "Image generation failed."
            else:
                return "The picture to be edited does not exist."
        else:
            return extract_result['message']

    except Exception as e:
        logger.warning(f"edit_furniture error ：{e}")
        return "edit_furniture error"

# def create_generate_furniture_tool(workspace_dir):
#     @tool
#     async def generate_furniture(prompt: str) -> str:
#         """
#         使用 Gemini 图像生成模型根据详细的英文提示词生成家具设计草图。
#         """
#         print(f"\n[系统日志] 正在调用 Nano Banana (Gemini Image Gen) ...")
#
#         try:
#             response = client.models.generate_content(
#                 model="gemini-2.5-flash-image",
#                 contents=(f"Generate a professional furniture design sketch: {prompt}"),
#                 config=GenerateContentConfig(
#                     response_modalities=[Modality.TEXT, Modality.IMAGE],
#                 ),
#             )
#
#             image_bytes = None
#             for part in response.candidates[0].content.parts:
#                 if part.inline_data:
#                     image_bytes = part.inline_data.data
#                     break
#
#             if not image_bytes:
#                 return "未能生成图像数据。"
#             # 1. 定义OSS存储路径和本地保存路径
#             object_name = f"furniture/sketches/{uuid.uuid4()}.png"
#             bucket = "fida-test"  # 替换为你的 bucket 名称
#             filename = os.path.join(workspace_dir, f"{bucket}/{object_name}")
#
#             # 2. 创建本地目录（确保目录存在）
#             local_dir = os.path.dirname(filename)
#             if not os.path.exists(local_dir):
#                 os.makedirs(local_dir, exist_ok=True)
#
#             # 3. 保存图片到本地文件（新增核心逻辑）
#             try:
#                 with open(filename, "wb") as f:
#                     f.write(image_bytes)
#                 print(f"[系统日志] 图片已保存到本地：{filename}")
#             except Exception as save_e:
#                 logger.warning(f"保存图片到本地失败：{save_e}")
#                 # 本地保存失败不中断上传流程，仅记录日志
#
#             # 4. 上传图片到OSS（原有逻辑）
#             upload_res = oss_upload_image(
#                 oss_client=minio_client,
#                 bucket=bucket,
#                 object_name=object_name,
#                 image_bytes=image_bytes
#             )
#
#             if upload_res:
#                 image_url = f"{bucket}/{object_name}"
#                 return image_url
#             else:
#                 return f"图片生成成功（本地路径：{filename}），但上传至存储服务器失败。"
#
#         except Exception as e:
#             logger.warning(f"绘图流程异常：{e}")
#             return "绘图流程异常"
#
#     return generate_furniture