File size: 1,388 Bytes
5374a2d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
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
|