1.优化隔离工作目录2.新增图像生成和编辑功能3.生成模型替换为本地flux2 klein

2026-03-19 17:55:39 +08:00
parent b6ca7ae6ed
commit ac8a5e5a30
18 changed files with 1167 additions and 403 deletions
--- a/src/server/deep_agent/tools/generate_furniture_sketch.py
+++ b/src/server/deep_agent/tools/generate_furniture_sketch.py
@@ -1,15 +1,21 @@
 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
+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__)
 # 初始化全局凭证和客户端
@@ -27,47 +33,187 @@ client = genai.Client(
 )


-@tool
-async def generate_furniture(prompt: str) -> str:
-    """
-    使用 Gemini 图像生成模型根据详细的英文提示词生成家具设计草图。
-    """
-    print(f"\n[系统日志] 正在调用 Nano Banana (Gemini Image Gen) ...")
-
+def is_image_path_exist(image_path):
    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],
-            ),
-        )
+        return Path(image_path).exists()
+    except:
+        return False

-        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 "未能生成图像数据。"
-        object_name = f"furniture/sketches/{uuid.uuid4()}.png"
-        bucket = "fida-test"  # 替换为你的 bucket 名称
-        # 3. 调用你的上传函数
-        upload_res = oss_upload_image(
-            oss_client=minio_client,
-            bucket=bucket,
-            object_name=object_name,
-            image_bytes=image_bytes
-        )
+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 upload_res:
-            # 4. 构造访问链接 (如果是私有 bucket，需使用 presigned_get_object)
-            # 这里简单示例为直接访问地址
-            image_url = f"{bucket}/{object_name}"
-            return image_url
-        else:
-            return "图片生成成功，但上传至存储服务器失败。"
-    except Exception as e:
-        logger.warning(e)
-        return "绘图流程异常"
+            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