agents/all_agents.py

from smolagents import CodeAgent, LogLevel
from remote_tools.rag_tool import RemoteObjectDetectionModelRetrieverTool
from tools.bbox_drawing_tool import BBoxDrawingTool
from tools.cropping_tool import CroppingTool
from remote_tools.object_detection_tool import RemoteObjectDetectionTool
from remote_tools.upscaler import RemoteUpscalerTool


def get_master_agent(llm):
    description = """
    You are an agent that can perform tasks on an image.
    You can use the following tools to perform tasks on an image:
    - object_detection_tool: to detect objects in an image, you must provide the image to the agents.
    - object_detection_model_retriever: to retrieve object detection models, you must provide the type of class that a model can detect.

    If you don't know what model to use, you can use the object_detection_model_retriever tool to retrieve the model.
    Never assume an invented model name, always use the model name provided by the object_detection_model_retriever tool.
    Use batching to perform tasks on multiple images at once when a tool supports it.
    You have access to the variable "image" which is the image to perform tasks on, no need to load it, it is already loaded.
    You can also use opencv to draw the bounding boxes on the image.
    Always use the variable "image" to draw the bounding boxes on the image.
    """
    master_agent = CodeAgent(
        name="master_agent",
        description=description,
        model=llm,
        tools=[
            RemoteObjectDetectionTool(),
            BBoxDrawingTool(),
            CroppingTool(),
            RemoteUpscalerTool(),
            RemoteObjectDetectionModelRetrieverTool(),
        ],
        verbosity_level=LogLevel.DEBUG,
    )
    print("Loaded master agent")
    return master_agent