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.
This commit is contained in:
parent
f1c843addc
commit
08ac557028
@ -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 # 代码最大行宽
|
||||
|
||||
@ -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
|
||||
|
||||
106
src/utils/minio_utils.py
Normal file
106
src/utils/minio_utils.py
Normal file
@ -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}")
|
||||
Loading…
Reference in New Issue
Block a user