From 08ac557028d92af8a668198ec46767956b5ca50d Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Tue, 2 Sep 2025 14:00:06 +0800 Subject: [PATCH] feat(agent): implement text-to-img tool with MinIO storage Adds a new agent tool `text_to_img` that generates a placeholder image from text and uploads it to a MinIO server, returning a public URL. Key changes: - Adds `minio` and `Pillow` to project dependencies. - Creates a reusable MinIO utility module in `src/utils/minio_utils.py` to handle file uploads and bucket management. - Refactors the agent tool to use this new utility, separating concerns and simplifying the tool's implementation. - Replaces the previous hardcoded URL with a real, functional image upload process. --- pyproject.toml | 2 + src/agents/chatbot/tools.py | 39 ++++++++++++- src/utils/minio_utils.py | 106 ++++++++++++++++++++++++++++++++++++ 3 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 src/utils/minio_utils.py diff --git a/pyproject.toml b/pyproject.toml index 6798d0d6..a2adac3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,8 @@ dependencies = [ "typer>=0.16.0", "mineru[core]>=2.1.6", "tabulate>=0.9.0", + "minio>=7.2.7", + "Pillow>=10.5.0", ] [tool.ruff] line-length = 120 # 代码最大行宽 diff --git a/src/agents/chatbot/tools.py b/src/agents/chatbot/tools.py index f4b921c2..a6e0d1a3 100644 --- a/src/agents/chatbot/tools.py +++ b/src/agents/chatbot/tools.py @@ -1,9 +1,12 @@ +from io import BytesIO from typing import Any from langchain_core.tools import tool +from PIL import Image, ImageDraw, ImageFont from src.agents.common.tools import get_buildin_tools from src.utils import logger +from src.utils.minio_utils import upload_image_to_minio @tool @@ -27,8 +30,42 @@ def calculator(a: float, b: float, operation: str) -> float: raise -def get_tools() -> dict[str, Any]: +@tool +async def text_to_img(text: str) -> str: + """ + 文生图函数，根据文本生成一张包含该文本的图片，并将其上传到文件服务器，最终返回图片的公开访问 URL。 + A text-to-image function that generates an image containing the given text, + uploads it to a file server, and returns the public URL of the image. + """ + logger.info(f"Generating image for text: {text}") + # 1. Simulate image generation using Pillow + try: + img = Image.new("RGB", (400, 100), color=(73, 109, 137)) + draw = ImageDraw.Draw(img) + try: + font = ImageFont.truetype("arial.ttf", 15) + except OSError: + font = ImageFont.load_default() + draw.text((10, 10), f"Generated from: {text}", fill=(255, 255, 0), font=font) + + img_bytes = BytesIO() + img.save(img_bytes, format="JPEG") + img_bytes.seek(0) + file_data = img_bytes.read() + logger.info("Image data generated successfully.") + except Exception as e: + logger.error(f"Failed to generate image with Pillow: {e}") + raise ValueError(f"Image generation failed: {e}") + + # 2. Upload to MinIO (Simplified) + image_url = upload_image_to_minio(data=file_data, file_extension="jpg") + logger.info(f"Image uploaded. URL: {image_url}") + return image_url + + +def get_tools() -> list[Any]: """获取所有可运行的工具（给大模型使用）""" tools = get_buildin_tools() tools.append(calculator) + tools.append(text_to_img) return tools diff --git a/src/utils/minio_utils.py b/src/utils/minio_utils.py new file mode 100644 index 00000000..dd03ffc4 --- /dev/null +++ b/src/utils/minio_utils.py @@ -0,0 +1,106 @@ +import json +import os +import uuid +from io import BytesIO + +from minio import Minio +from minio.error import S3Error + +from src.utils import logger + + +class _MinioClient: + """ + A private client for interacting with a MinIO server. + It reads connection details from environment variables and provides + methods for uploading files and ensuring buckets exist. + """ + + def __init__(self): + """ + Initializes the Minio client. + Reads configuration from environment variables set in docker-compose.yml. + """ + minio_uri = os.getenv("MINIO_URI", "http://milvus-minio:9000") + self.endpoint = minio_uri.split("://")[-1] + self.access_key = os.getenv("MINIO_ACCESS_KEY", "minioadmin") + self.secret_key = os.getenv("MINIO_SECRET_KEY", "minioadmin") + + if os.getenv("RUNNING_IN_DOCKER"): + self.public_endpoint = f"localhost:{self.endpoint.split(':')[-1]}" + else: + self.public_endpoint = self.endpoint + + try: + self.client = Minio(self.endpoint, access_key=self.access_key, secret_key=self.secret_key, secure=False) + except Exception as e: + logger.error(f"Failed to initialize MinIO client: {e}") + raise + + def ensure_bucket_exists(self, bucket_name: str): + try: + if not self.client.bucket_exists(bucket_name): + self.client.make_bucket(bucket_name) + logger.info(f"Bucket '{bucket_name}' created.") + policy = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": {"AWS": ["*"]}, + "Action": ["s3:GetObject"], + "Resource": [f"arn:aws:s3:::{bucket_name}/*"], + }, + ], + } + self.client.set_bucket_policy(bucket_name, json.dumps(policy)) + logger.info(f"Public read policy set for bucket '{bucket_name}'.") + except S3Error as e: + logger.error(f"Error with bucket '{bucket_name}': {e}") + raise + + def upload_file(self, bucket_name: str, file_name: str, data: bytes, content_type: str) -> str: + self.ensure_bucket_exists(bucket_name) + try: + data_stream = BytesIO(data) + self.client.put_object(bucket_name, file_name, data_stream, length=len(data), content_type=content_type) + logger.info(f"Successfully uploaded '{file_name}' to bucket '{bucket_name}'.") + return f"http://{self.public_endpoint}/{bucket_name}/{file_name}" + except S3Error as e: + logger.error(f"Failed to upload file '{file_name}' to MinIO: {e}") + raise + + +# Singleton instance of the private client +_minio_client = _MinioClient() + + +def upload_image_to_minio(data: bytes, file_extension: str = "jpg") -> str: + """ + Uploads image data to a predefined MinIO bucket and returns the public URL. + + This function abstracts away the details of MinIO client management, + bucket creation, and file naming. + + Args: + data (bytes): The raw image data. + file_extension (str): The file extension for the image (e.g., "jpg", "png"). + + Returns: + str: The public URL of the uploaded image. + + Raises: + ConnectionError: If the upload to the file server fails. + """ + try: + bucket_name = "generated-images" + file_name = f"{uuid.uuid4()}.{file_extension}" + content_type = f"image/{file_extension}" + + image_url = _minio_client.upload_file( + bucket_name=bucket_name, file_name=file_name, data=data, content_type=content_type + ) + return image_url + except Exception as e: + logger.error(f"High-level upload to MinIO failed: {e}") + raise ConnectionError(f"Failed to upload image to file server: {e}")