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