from ai.tool import ToolSet

from .generate_image import GenerateImageTool
from .generate_image_anime import GenerateImageAnimeTool


class ImageGenerationToolSet(ToolSet):
    def __init__(self, fal_token: str, replicate_token: str):
        functions = [
            GenerateImageTool(fal_token),
            GenerateImageAnimeTool(replicate_token)
        ]
        with open("ai/tools/image_generation/prompt.md", "r") as f:
            system_prompt = f.read()
        super().__init__(functions=functions, system_prompt=system_prompt)


__all__ = ["GenerateImageTool", "GenerateImageAnimeTool", "ImageGenerationToolSet"]