tool_specs¶

class ToolSpec(Serializable):
    """
    ToolSpec is an abstract class that represents a tool specification that describes all necessary information about a tool used by the LLM. 

    ToolSpec and its subclasses are responsible for providing four categories of interfaces:
    1. Transformations to LLM Tool: `to_tool`.
    2. Worker Creation: `create_worker`.
    3. Serialization and Deserialization.
    4. ToolSpec initialization from raw resources: `from_raw`.
    """
    _tool_id: Optional[Union[str, int]]
    """The unique ID of the tool, used to uniquely identify a tool across the entire system. This tool can be of various types."""
    _tool_name: Optional[str]
    """The name of the tool to be called"""
    _tool_description: Optional[str]
    """A description of what the tool does, used by the model to choose when and how to call the tool."""
    _tool_parameters: Optional[Dict[str, Any]]
    """The JSON schema of the tool's parameters"""

    def __init__(
        self,
        tool_name: Optional[str] = None,
        tool_description: Optional[str] = None,
        tool_parameters: Optional[Dict[str, Any]] = None,
    ):
        self._tool_id = None
        self._tool_name = tool_name
        self._tool_description = tool_description
        self._tool_parameters = tool_parameters

    @property
    def tool_name(self) -> Optional[str]:
        return self._tool_name

    @property
    def tool_description(self) -> Optional[str]:
        return self._tool_description

    @property
    def tool_parameters(self) -> Optional[Dict[str, Any]]:
        return self._tool_parameters

    def __str__(self) -> str:
        return f"{self.__class__.__name__}(tool_name={self._tool_name}, tool_description={self._tool_description}, tool_parameters={self._tool_parameters})"

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}(tool_name={self._tool_name}, tool_description={self._tool_description}, tool_parameters={self._tool_parameters})>"

    ###############################################################
    ######## Part One of interfaces: Transformations to Tool ######
    ###############################################################

    @abstractmethod
    def to_tool(self) -> Tool:
        """
        Transform this ToolSpec to a `Tool` object used by LLM.

        Returns
        -------
        Tool
            A `Tool` object that can be used by LLM.
        """
        ...

    ###############################################################
    ######## Part Two of interfaces: Worker Creation ##############
    ###############################################################

    @abstractmethod
    def create_worker(self) -> Worker:
        """
        Create a Worker from the information included in this ToolSpec.

        Returns
        -------
        Worker
            A new `Worker` object that can be added to an Automa to execute the tool.
        """
        ...

    ###############################################################
    ######## Part Three of interfaces: 
    ######## Serialization and Deserialization ####################
    ###############################################################

    @override
    def dump_to_dict(self) -> Dict[str, Any]:
        state_dict = {}
        if self._tool_id:
            state_dict["tool_id"] = self._tool_id
        if self._tool_name:
            state_dict["tool_name"] = self._tool_name
        if self._tool_description:
            state_dict["tool_description"] = self._tool_description
        if self._tool_parameters:
            state_dict["tool_parameters"] = self._tool_parameters
        return state_dict

    @override
    def load_from_dict(self, state_dict: Dict[str, Any]) -> None:
        self._tool_id = state_dict.get("tool_id")
        self._tool_name = state_dict.get("tool_name")
        self._tool_description = state_dict.get("tool_description")
        self._tool_parameters = state_dict.get("tool_parameters")

class FunctionToolSpec(ToolSpec):
    _func: Callable
    """The python function to be used as a tool"""

    def __init__(
        self,
        func: Callable,
        tool_name: Optional[str] = None,
        tool_description: Optional[str] = None,
        tool_parameters: Optional[Dict[str, Any]] = None,
    ):
        super().__init__(
            tool_name=tool_name,
            tool_description=tool_description,
            tool_parameters=tool_parameters
        )
        self._func = func

    @classmethod
    def from_raw(
        cls,
        func: Callable,
        tool_name: Optional[str] = None,
        tool_description: Optional[str] = None,
        tool_parameters: Optional[Dict[str, Any]] = None,
    ) -> "FunctionToolSpec":
        """
        Create a FunctionToolSpec from a python function. By default, the tool name, description and parameters' json schema will be extracted from the function's docstring and the parameters' type and description. However, these values can be customized by passing in the corresponding arguments.

        Parameters
        ----------
        func : Callable
            The python function to create a FunctionToolSpec from.
        tool_name : Optional[str]
            The name of the tool. If not provided, the function name will be used.
        tool_description : Optional[str]
            The description of the tool. If not provided, the function docstring will be used.
        tool_parameters : Optional[Dict[str, Any]]
            The JSON schema of the tool's parameters. If not provided, the JSON schema will be constructed properly from the parameters' annotations, the function's signature and/or docstring.

        Returns
        -------
        FunctionToolSpec
            A new `FunctionToolSpec` object.
        """
        if isinstance(func, MethodType):
            raise ValueError(f"`func` is not allowed to be a bound method: {func}.")

        if not tool_name:
            tool_name = func.__name__

        if not tool_description:
            tool_description = get_tool_description_from(func, tool_name)

        if not tool_parameters:
            tool_parameters = create_func_params_json_schema(func)
            # TODO: whether to remove the `title` field of the params_schema?

        return cls(
            func=func,
            tool_name=tool_name,
            tool_description=tool_description,
            tool_parameters=tool_parameters
        )

    @override
    def to_tool(self) -> Tool:
        """
        Transform this FunctionToolSpec to a `Tool` object used by LLM.

        Returns
        -------
        Tool
            A `Tool` object that can be used by LLM.
        """
        return Tool(
            name=self._tool_name,
            description=self._tool_description,
            parameters=self._tool_parameters
        )

    @override
    def create_worker(self) -> Worker:
        """
        Create a Worker from the information included in this FunctionToolSpec.

        Returns
        -------
        Worker
            A new `Worker` object that can be added to an Automa to execute the tool.
        """
        # TODO: some initialization arguments may be needed in future, e.g., `bound_needed`.
        return CallableWorker(self._func)

    @override
    def dump_to_dict(self) -> Dict[str, Any]:
        state_dict = super().dump_to_dict()
        state_dict["func"] = self._func.__module__ + "." + self._func.__qualname__
        return state_dict

    @override
    def load_from_dict(self, state_dict: Dict[str, Any]) -> None:
        super().load_from_dict(state_dict)
        self._func = load_qualified_class_or_func(state_dict["func"])

class AutomaToolSpec(ToolSpec):
    _automa_cls: Type[Automa]
    """The Automa class to be used as a tool"""
    _automa_init_kwargs: Dict[str, Any]
    """The initialization arguments for the Automa"""

    def __init__(
        self,
        automa_cls: Type[Automa],
        tool_name: Optional[str] = None,
        tool_description: Optional[str] = None,
        tool_parameters: Optional[Dict[str, Any]] = None,
        **automa_init_kwargs: Dict[str, Any],
    ):
        super().__init__(
            tool_name=tool_name,
            tool_description=tool_description,
            tool_parameters=tool_parameters
        )
        self._automa_cls = automa_cls
        self._automa_init_kwargs = automa_init_kwargs

    @classmethod
    def from_raw(
        cls,
        automa_cls: Type[Automa],
        tool_name: Optional[str] = None,
        tool_description: Optional[str] = None,
        tool_parameters: Optional[Dict[str, Any]] = None,
        **automa_init_kwargs: Dict[str, Any],
    ) -> "AutomaToolSpec":
        """
        Create an AutomaToolSpec from an Automa class.
        """

        def check_spec_func(automa_cls):
            if hasattr(automa_cls, "spec_func") and isinstance(automa_cls.spec_func, Callable):
                return
            raise ValueError(f"The Automa class {automa_cls} must be decorated with `@as_tool` in order to be used as a tool.")

        if (not tool_name) or (not tool_description) or (not tool_parameters):
            check_spec_func(automa_cls)

        if not tool_name:
            tool_name = automa_cls.spec_func.__name__

        if not tool_description:
            tool_description = get_tool_description_from(automa_cls.spec_func, tool_name)

        if not tool_parameters:
            tool_parameters = create_func_params_json_schema(automa_cls.spec_func)
            # TODO: whether to remove the `title` field of the params_schema?

        return cls(
            automa_cls=automa_cls,
            tool_name=tool_name,
            tool_description=tool_description,
            tool_parameters=tool_parameters,
            **automa_init_kwargs
        )

    @override
    def to_tool(self) -> Tool:
        """
        Transform this AutomaToolSpec to a `Tool` object used by LLM.

        Returns
        -------
        Tool
            A `Tool` object that can be used by LLM.
        """
        return Tool(
            name=self._tool_name,
            description=self._tool_description,
            parameters=self._tool_parameters
        )

    @override
    def create_worker(self) -> Worker:
        """
        Create a Worker from the information included in this AutomaToolSpec.

        Returns
        -------
        Worker
            A new `Worker` object that can be added to an Automa to execute the tool.
        """
        return self._automa_cls(**self._automa_init_kwargs)

    @override
    def dump_to_dict(self) -> Dict[str, Any]:
        state_dict = super().dump_to_dict()
        state_dict["automa_cls"] = self._automa_cls.__module__ + "." + self._automa_cls.__qualname__
        if self._automa_init_kwargs:
            state_dict["automa_init_kwargs"] = self._automa_init_kwargs
        return state_dict

    @override
    def load_from_dict(self, state_dict: Dict[str, Any]) -> None:
        super().load_from_dict(state_dict)
        self._automa_cls = load_qualified_class_or_func(state_dict["automa_cls"])
        self._automa_init_kwargs = state_dict.get("automa_init_kwargs") or {}