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

import json
import logging
import os
import uuid
from pathlib import Path
from typing import Annotated

import httpx
from google.oauth2 import service_account
from langchain_core.tools import tool
from google import genai
from google.genai.types import GenerateContentConfig, Modality
from langgraph.prebuilt import ToolRuntime

from minio import Minio

from src.core.config import settings
from src.server.utils.new_oss_client import oss_upload_image, oss_get_image, is_minio_file_exist, oss_upload_image_file

logger = logging.getLogger(__name__)
# 初始化全局凭证和客户端
creds = service_account.Credentials.from_service_account_file(
    settings.GOOGLE_GENAI_USE_VERTEXAI,
    scopes=["https://www.googleapis.com/auth/cloud-platform"],
)

minio_client = Minio(settings.MINIO_URL, access_key=settings.MINIO_ACCESS, secret_key=settings.MINIO_SECRET, secure=settings.MINIO_SECURE)
client = genai.Client(
    credentials=creds,
    project=settings.GOOGLE_CLOUD_PROJECT,
    location=settings.GOOGLE_CLOUD_LOCATION,
    vertexai=True
)


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


def create_generate_furniture_tool(workspace_dir, width: int = 1024, height: int = 1024):
    @tool
    async def generate_furniture(prompt: str, runtime: ToolRuntime) -> str:
        """
        使用 Gemini 图像生成模型根据详细的英文提示词生成家具设计草图。
        """
        logger.info(f"\n[系统日志] 正在调用 generate_furniture ...")
        try:
            # 1. 生成图像 - local flux2-klein
            object_name = f"furniture/sketches/{uuid.uuid4()}.png"
            bucket_name = "fida-test"  # 替换为你的 bucket 名称
            request_data = {
                "prompt": prompt,
                "bucket_name": bucket_name,
                "object_name": object_name,
                "width": width,
                "height": height
            }
            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:
                filename = os.path.join(workspace_dir, image_url)
                # 2. 创建本地目录（确保目录存在）
                local_dir = os.path.dirname(filename)
                if not os.path.exists(local_dir):
                    os.makedirs(local_dir, exist_ok=True)

                img = oss_get_image(oss_client=minio_client, bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:])
                img.save(filename)

                return image_url
            else:
                return f"Image generation failed."

        except Exception as e:
            logger.warning(f"绘图流程异常：{e}")
            return "绘图流程异常"

    return generate_furniture


def create_edit_furniture_tool(workspace_dir, width: int = 1024, height: int = 1024):
    @tool
    async def edit_furniture(prompt: str, input_image_path) -> str:
        """
        使用图像生成模型根据详细的英文提示词编辑家具设计草图。
        """
        logger.info(f"\n[系统日志] 正在调用 edit_furniture ...")

        try:
            # 0. 编辑前先检查工作环境和minio上是否存在该图像
            input_image_path = input_image_path.lstrip('/')
            filename = os.path.join(workspace_dir, input_image_path)
            local_exist = is_image_path_exist(filename)
            minio_exist = is_minio_file_exist(minio_client=minio_client, bucket_name=input_image_path.split('/')[0], object_name=input_image_path.split('/')[0])

            if not local_exist and not minio_exist:
                # 两个地方都不存在 直接报错
                return f"Image generation failed."
            elif local_exist and not minio_exist:
                # 把本地的上传到minio
                oss_upload_image_file(oss_client=minio_client, bucket=input_image_path.split('/')[0], object_name=input_image_path.split('/')[0], file_path=filename)
            elif not local_exist and minio_exist:
                # minio的下载到本地
                img = oss_get_image(oss_client=minio_client, bucket=input_image_path.split('/')[0], object_name=input_image_path.split('/')[0], )
                img.save(filename)
            elif minio_exist and local_exist:
                # 两个地方都存在 直接跳过
                pass

            # 1. 生成图像 - local flux2-klein
            object_name = f"furniture/sketches/{uuid.uuid4()}.png"
            bucket_name = "fida-test"  # 替换为你的 bucket 名称
            request_data = {
                "input_image_paths": [input_image_path],
                "prompt": prompt,
                "bucket_name": bucket_name,
                "object_name": object_name,
                "width": width,
                "height": height
            }
            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:
                filename = os.path.join(workspace_dir, image_url)
                # 2. 创建本地目录（确保目录存在）
                local_dir = os.path.dirname(filename)
                if not os.path.exists(local_dir):
                    os.makedirs(local_dir, exist_ok=True)

                img = oss_get_image(oss_client=minio_client, bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:])
                img.save(filename)
                return image_url
            else:
                return f"Image generation failed."

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

    return edit_furniture

# 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