| from typing import Optional | |
| import os | |
| from ...tool import Toolkit | |
| from ...storage_handler import FileStorageHandler | |
| from .image_generation_edit import FluxImageGenerationEditTool | |
| from ..openrouter_image_tools.image_analysis import ImageAnalysisTool | |
| class FluxImageGenerationToolkit(Toolkit): | |
| """ | |
| Flux toolkit combining generation tool and a common image analysis tool. | |
| """ | |
| def __init__( | |
| self, | |
| name: str = "FluxImageGenerationToolkit", | |
| api_key: Optional[str] = None, | |
| save_path: str = "./imgs", | |
| storage_handler: Optional[FileStorageHandler] = None, | |
| analysis_model: str = "openai/gpt-4o-mini", | |
| ): | |
| tools = [] | |
| gen_tool = FluxImageGenerationEditTool( | |
| api_key=api_key, | |
| storage_handler=storage_handler, | |
| base_path=save_path, | |
| ) | |
| tools.append(gen_tool) | |
| try: | |
| resolved_key = os.getenv("OPENROUTER_API_KEY") | |
| if resolved_key: | |
| analysis_tool = ImageAnalysisTool( | |
| api_key=resolved_key, | |
| model=analysis_model, | |
| ) | |
| tools.append(analysis_tool) | |
| except Exception: | |
| pass | |
| super().__init__(name=name, tools=tools) | |
| self.api_key = api_key | |
| self.save_path = save_path | |
| self.storage_handler = storage_handler | |