image_generator_tool.py

from crewai.tools import BaseTool
from pydantic import BaseModel, Field
from typing import Type
from template_system import TemplateRegistry

class GenerateImageToolInput(BaseModel):
    """Input for the Generate Image Tool."""
    product_image_url: str = Field(..., description="URL of the product image to be placed on the template.")
    product_name: str = Field(..., description="Name of the product.")
    original_price: str = Field(..., description="Original price of the product.")
    final_price: str = Field(..., description="Final price of the product.")
    coupon_code: str = Field(..., description="Coupon code to be displayed on the image.")
    template_name: str = Field(default="lidi_promo", description="Name of the template to use for generating the image.")

class GenerateImageTool(BaseTool):
    name: str = "Generate Image Tool"
    description: str = "Generates a promotional image for a product using a template and returns the file path."
    args_schema: Type[GenerateImageToolInput] = GenerateImageToolInput

    def _run(self, product_image_url: str, product_name: str, original_price: str, final_price: str, coupon_code: str, template_name: str = "lidi_promo") -> str:
        try:
            # Get the template instance
            template = TemplateRegistry.get_template(template_name)
            
            # Generate the image using the template
            return template.generate_image(
                product_image_url=product_image_url,
                product_name=product_name,
                original_price=original_price,
                final_price=final_price,
                coupon_code=coupon_code
            )
            
        except ValueError as e:
            return f"Error: {e}"
        except Exception as e:
            return f"An error occurred: {e}"