| from typing import Optional | |
| from ...tool import Toolkit | |
| from ...storage_handler import FileStorageHandler | |
| from .image_generation import OpenAIImageGenerationTool | |
| from .image_edit import OpenAIImageEditTool | |
| from .image_analysis_openai import OpenAIImageAnalysisTool | |
| class OpenAIImageToolkit(Toolkit): | |
| def __init__(self, name: str = "OpenAIImageToolkit", api_key: str = None, organization_id: str = None, | |
| generation_model: str = "dall-e-3", save_path: str = "./generated_images", | |
| storage_handler: Optional[FileStorageHandler] = None): | |
| gen_tool = OpenAIImageGenerationTool(api_key=api_key, organization_id=organization_id, | |
| model=generation_model, save_path=save_path, | |
| storage_handler=storage_handler) | |
| edit_tool = OpenAIImageEditTool(api_key=api_key, organization_id=organization_id, | |
| save_path=save_path, storage_handler=storage_handler) | |
| analysis_tool = OpenAIImageAnalysisTool(api_key=api_key, organization_id=organization_id, | |
| storage_handler=storage_handler) | |
| super().__init__(name=name, tools=[gen_tool, edit_tool, analysis_tool]) | |
| self.api_key = api_key | |
| self.organization_id = organization_id | |
| self.generation_model = generation_model | |
| self.save_path = save_path | |
| self.storage_handler = storage_handler | |