omagent_core/tool_system/base.py

import json
from abc import ABC
from distutils.util import strtobool
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union

import yaml
from omagent_core.base import BotBase
from omagent_core.models.od.schemas import Target
from omagent_core.services.handlers.sql_data_handler import SQLDataHandler
from omagent_core.utils.error import VQLError
from omagent_core.utils.logger import logging
from omagent_core.utils.plot import Annotator
from PIL import Image
from pydantic import BaseModel, model_validator


class ArgSchema(BaseModel):
    """ArgSchema defines the tool input schema. Only support one layer definition. Please prevent using complex structure."""

    class Config:
        """Configuration for this pydantic object."""

        extra = "allow"
        arbitrary_types_allowed = True

    class ArgInfo(BaseModel):
        description: Optional[str]
        type: str = "str"
        enum: Optional[List] = None
        required: Optional[bool] = True

    @model_validator(mode="before")
    @classmethod
    def validate_all(cls, values):
        for key, value in values.items():
            if type(value) is str:
                values[key] = cls.ArgInfo(name=value)
            elif type(value) is dict:
                values[key] = cls.ArgInfo(**value)
            elif type(value) is cls.ArgInfo:
                pass
            else:
                raise ValueError(
                    "The arg type must be one of string, dict or self.ArgInfo."
                )
        return values

    @classmethod
    def from_file(cls, schema_file: Union[str, Path]):
        if type(schema_file) is str:
            schema_file = Path(schema_file)
        if schema_file.suffix == ".json":
            with open(schema_file, "r") as f:
                schema = json.load(f)
        elif schema_file.suffix == ".yaml":
            with open(schema_file, "r") as f:
                schema = yaml.load(f, Loader=yaml.FullLoader)
        else:
            raise ValueError("Only support json and yaml file.")
        return cls(**schema)

    def generate_schema(self) -> Union[dict, list]:
        required_args = []
        parameters = {}
        for key, value in self.model_dump(exclude_none=True).items():
            parameters[key] = value
            if parameters[key].pop("required"):
                required_args.append(key)
        return parameters, required_args

    def validate_args(self, args: dict) -> dict:
        if type(args) is not dict:
            raise ValueError(
                "ArgSchema validate only support dict, not {}".format(type(args))
            )
        new_args = {}
        required_fields = set(
            [k for k, v in self.model_dump().items() if v["required"]]
        )
        name_mapping = {
            "str": "string",
            "int": "integer",
            "float": "number",
            "bool": "boolean",
        }
        for name, value in args.items():
            if name not in self.model_dump():
                logging.warning(
                    "The input args includes an unnecessary parameter {}. Removed from the args.".format(
                        name
                    )
                )
                continue
            if name_mapping[type(value).__name__] == self.model_dump()[name]["type"]:
                if (
                    self.model_dump()[name]["enum"]
                    and value not in self.model_dump()[name]["enum"]
                ):
                    raise ValueError(
                        "The value of {} should be one of {}, but got {}".format(
                            name, str(self.model_dump()[name]["enum"]), value
                        )
                    )
                new_args[name] = value
            elif self.model_dump()[name]["type"] == "string":
                try:
                    new_args[name] = str(value)
                except:
                    raise ValueError(
                        "Parameter {} type expect a str value, but got a {} {}".format(
                            name, type(value), value
                        )
                    )
            elif self.model_dump()[name]["type"] == "integer":
                try:
                    new_args[name] = int(value)
                except:
                    raise ValueError(
                        "Parameter {} type expect an int value, but got a {} {}".format(
                            name, type(value), value
                        )
                    )
            elif self.model_dump()[name]["type"] == "number":
                try:
                    new_args[name] = float(value)
                except:
                    raise ValueError(
                        "Parameter {} type expect a float value, but got a {} {}".format(
                            name, type(value), value
                        )
                    )
            elif self.model_dump()[name]["type"] == "boolean":
                if type(value) is bool:
                    new_args[name] = value
                else:
                    try:
                        new_args[name] = strtobool(str(value))
                    except:
                        raise ValueError(
                            "Parameter {} type expect a boolean value, but got a {} {}".format(
                                name, type(value), value
                            )
                        )
            else:
                raise ValueError(
                    "Parameter {} type expect one of string, integer, number and boolean, but got a {} {}".format(
                        name, self.model_dump()[name]["type"], type(value), value
                    )
                )

        if required_fields - set(new_args.keys()):
            raise VQLError(
                "The required fields {} are missing.".format(
                    required_fields - set(new_args.keys())
                )
            )
        return new_args


class BaseTool(BotBase, ABC):
    description: str
    func: Optional[Callable] = None
    args_schema: Optional[ArgSchema]
    special_params: Dict = {}

    def model_post_init(self, __context: Any) -> None:
        for _, attr_value in self.__dict__.items():
            if isinstance(attr_value, BotBase):
                attr_value._parent = self

    @property
    def workflow_instance_id(self) -> str:
        if hasattr(self, "_parent"):
            return self._parent.workflow_instance_id
        return None

    @workflow_instance_id.setter
    def workflow_instance_id(self, value: str):
        if hasattr(self, "_parent"):
            self._parent.workflow_instance_id = value

    def _run(self, **input) -> str:
        """Implement this function or pass 'func' arg when initializing."""
        return self.func(**input)

    async def _arun(self, **input) -> str:
        """Implement this function or pass 'func' arg when initializing."""
        return await self.func(**input)

    def run(self, input: Any) -> str:
        if self.args_schema != None:
            if type(input) != dict:
                raise ValueError(
                    "The input type must be dict when args_schema is specified."
                )
            self.args_schema.validate_args(input)
        return self._run(**input, **self.special_params)

    async def arun(self, input: Any) -> str:
        if self.args_schema != None:
            if type(input) != dict:
                raise ValueError(
                    "The input type must be dict when args_schema is specified."
                )
            self.args_schema.validate_args(input)
        return await self._arun(**input, **self.special_params)

    def generate_schema(self):
        if not self.args_schema:
            return {
                "type": "function",
                "description": self.description,
                "function": {
                    "name": self.name,
                    "parameters": {
                        "type": "object",
                        "name": "input",
                        "required": ["input"],
                    },
                },
            }
        else:
            properties, required = self.args_schema.generate_schema()
            return {
                "type": "function",
                "function": {
                    "name": self.name,
                    "description": self.description,
                    "parameters": {
                        "type": "object",
                        "properties": properties,
                        "required": required,
                    },
                },
            }


class BaseModelTool(BaseTool, ABC):
    # data_handler: Optional[SQLDataHandler]

    def visual_prompting(
        self,
        image: Image.Image,
        annotation: List[Target],
        prompting_type: str = "label_on_img",
        include_labels: Union[List, set, tuple] = None,
        exclude_labels: Union[List, set, tuple] = None,
    ) -> List[Image.Image]:
        annotator = Annotator(image)
        for obj in annotation:
            if (exclude_labels is not None and obj.label in exclude_labels) or (
                include_labels is not None and obj.label not in include_labels
            ):
                continue
            if obj.bbox:
                annotator.box_label(obj.bbox, obj.label, color="red")
            # TODO: Add polygon support
        return annotator.result()

    def infer(self, images: List[Image.Image], kwargs) -> List[List[Target]]:
        """The model inference step. Only support OD type detection.

        Args:
            images (List[Image.Image]): The list of input images. Image should be PIL Image object.
            kwargs (dict): The additional arguments for the model.

        Returns:
            List[List[Target]]: The detection results.
        """

    def ainfer(self, images: List[Image.Image], kwargs) -> List[List[Target]]:
        """The async version of model inference step. Only support OD type detection.

        Args:
            images (List[Image.Image]): The list of input images. Image should be PIL Image object.
            kwargs (dict): The additional arguments for the model.

        Returns:
            List[List[Target]]: The detection results.
        """


class MemoryTool(BaseTool):
    memory_handler: Optional[SQLDataHandler]

    def generate_schema(self) -> dict:
        """Generate the data table schema in dict format.

        Returns:
            dict: The data table schema. Including the table name, and the name, data type and additional information of each column.
        """
        table = self.memory_handler.table
        schema = {"table_name": table.__tablename__, "columns": []}
        for column in table.__table__.columns:
            schema["columns"].append(
                {
                    "name": column.name,
                    "type": column.type.__visit_name__,
                    "info": column.info,
                }
            )
        return schema

    def generate_prompt(self):
        pass

    def _run(self):
        self.memory_handler.execute_sql()

    async def _arun(self):
        self.memory_handler.execute_sql()