import inspect
import logging
import re
from abc import ABCMeta
from copy import deepcopy
from functools import wraps
from typing import Callable, Optional, Type, get_args, get_origin
try:
from typing import Annotated
except ImportError:
from typing_extensions import Annotated
from griffe import Docstring
try:
from griffe import DocstringSectionKind
except ImportError:
from griffe.enumerations import DocstringSectionKind
from ..schema import ActionReturn, ActionStatusCode
from .parser import BaseParser, JsonParser, ParseError
logging.getLogger('griffe').setLevel(logging.ERROR)
def tool_api(func: Optional[Callable] = None,
*,
explode_return: bool = False,
returns_named_value: bool = False,
**kwargs):
"""Turn functions into tools. It will parse typehints as well as docstrings
to build the tool description and attach it to functions via an attribute
``api_description``.
Examples:
.. code-block:: python
# typehints has higher priority than docstrings
from typing import Annotated
@tool_api
def add(a: Annotated[int, 'augend'], b: Annotated[int, 'addend'] = 1):
'''Add operation
Args:
x (int): a
y (int): b
'''
return a + b
print(add.api_description)
Args:
func (Optional[Callable]): function to decorate. Defaults to ``None``.
explode_return (bool): whether to flatten the dictionary or tuple return
as the ``return_data`` field. When enabled, it is recommended to
annotate the member in docstrings. Defaults to ``False``.
.. code-block:: python
@tool_api(explode_return=True)
def foo(a, b):
'''A simple function
Args:
a (int): a
b (int): b
Returns:
dict: information of inputs
* x: value of a
* y: value of b
'''
return {'x': a, 'y': b}
print(foo.api_description)
returns_named_value (bool): whether to parse ``thing: Description`` in
returns sections as a name and description, rather than a type and
description. When true, type must be wrapped in parentheses:
``(int): Description``. When false, parentheses are optional but
the items cannot be named: ``int: Description``. Defaults to ``False``.
Returns:
Callable: wrapped function or partial decorator
Important:
``return_data`` field will be added to ``api_description`` only
when ``explode_return`` or ``returns_named_value`` is enabled.
"""
def _detect_type(string):
field_type = 'STRING'
if 'list' in string:
field_type = 'Array'
elif 'str' not in string:
if 'float' in string:
field_type = 'FLOAT'
elif 'int' in string:
field_type = 'NUMBER'
elif 'bool' in string:
field_type = 'BOOLEAN'
return field_type
def _explode(desc):
kvs = []
desc = '\nArgs:\n' + '\n'.join([
' ' + item.lstrip(' -+*#.')
for item in desc.split('\n')[1:] if item.strip()
])
docs = Docstring(desc).parse('google')
if not docs:
return kvs
if docs[0].kind is DocstringSectionKind.parameters:
for d in docs[0].value:
d = d.as_dict()
if not d['annotation']:
d.pop('annotation')
else:
d['type'] = _detect_type(d.pop('annotation').lower())
kvs.append(d)
return kvs
def _parse_tool(function):
# remove rst syntax
docs = Docstring(
re.sub(':(.+?):`(.+?)`', '\\2', function.__doc__ or '')).parse(
'google', returns_named_value=returns_named_value, **kwargs)
desc = dict(
name=function.__name__,
description=docs[0].value
if docs[0].kind is DocstringSectionKind.text else '',
parameters=[],
required=[],
)
args_doc, returns_doc = {}, []
for doc in docs:
if doc.kind is DocstringSectionKind.parameters:
for d in doc.value:
d = d.as_dict()
d['type'] = _detect_type(d.pop('annotation').lower())
args_doc[d['name']] = d
if doc.kind is DocstringSectionKind.returns:
for d in doc.value:
d = d.as_dict()
if not d['name']:
d.pop('name')
if not d['annotation']:
d.pop('annotation')
else:
d['type'] = _detect_type(d.pop('annotation').lower())
returns_doc.append(d)
sig = inspect.signature(function)
for name, param in sig.parameters.items():
if name == 'self':
continue
parameter = dict(
name=param.name,
type='STRING',
description=args_doc.get(param.name,
{}).get('description', ''))
annotation = param.annotation
if annotation is inspect.Signature.empty:
parameter['type'] = args_doc.get(param.name,
{}).get('type', 'STRING')
else:
if get_origin(annotation) is Annotated:
annotation, info = get_args(annotation)
if info:
parameter['description'] = info
while get_origin(annotation):
annotation = get_args(annotation)
parameter['type'] = _detect_type(str(annotation))
desc['parameters'].append(parameter)
if param.default is inspect.Signature.empty:
desc['required'].append(param.name)
return_data = []
if explode_return:
return_data = _explode(returns_doc[0]['description'])
elif returns_named_value:
return_data = returns_doc
if return_data:
desc['return_data'] = return_data
return desc
if callable(func):
if inspect.iscoroutinefunction(func):
@wraps(func)
async def wrapper(self, *args, **kwargs):
return await func(self, *args, **kwargs)
else:
@wraps(func)
def wrapper(self, *args, **kwargs):
return func(self, *args, **kwargs)
wrapper.api_description = _parse_tool(func)
return wrapper
def decorate(func):
if inspect.iscoroutinefunction(func):
@wraps(func)
async def wrapper(self, *args, **kwargs):
return await func(self, *args, **kwargs)
else:
@wraps(func)
def wrapper(self, *args, **kwargs):
return func(self, *args, **kwargs)
wrapper.api_description = _parse_tool(func)
return wrapper
return decorate
class ToolMeta(ABCMeta):
"""Metaclass of tools."""
def __new__(mcs, name, base, attrs):
is_toolkit, tool_desc = True, dict(
name=name,
description=Docstring(attrs.get('__doc__',
'')).parse('google')[0].value)
for key, value in attrs.items():
if callable(value) and hasattr(value, 'api_description'):
api_desc = getattr(value, 'api_description')
if key == 'run':
tool_desc['parameters'] = api_desc['parameters']
tool_desc['required'] = api_desc['required']
if api_desc['description']:
tool_desc['description'] = api_desc['description']
if api_desc.get('return_data'):
tool_desc['return_data'] = api_desc['return_data']
is_toolkit = False
else:
tool_desc.setdefault('api_list', []).append(api_desc)
if not is_toolkit and 'api_list' in tool_desc:
raise KeyError('`run` and other tool APIs can not be implemented '
'at the same time')
if is_toolkit and 'api_list' not in tool_desc:
is_toolkit = False
if callable(attrs.get('run')):
run_api = tool_api(attrs['run'])
api_desc = run_api.api_description
tool_desc['parameters'] = api_desc['parameters']
tool_desc['required'] = api_desc['required']
if api_desc['description']:
tool_desc['description'] = api_desc['description']
if api_desc.get('return_data'):
tool_desc['return_data'] = api_desc['return_data']
attrs['run'] = run_api
else:
tool_desc['parameters'], tool_desc['required'] = [], []
attrs['_is_toolkit'] = is_toolkit
attrs['__tool_description__'] = tool_desc
return super().__new__(mcs, name, base, attrs)
class BaseAction(metaclass=ToolMeta):
"""Base class for all actions.
Args:
description (:class:`Optional[dict]`): The description of the action.
Defaults to ``None``.
parser (:class:`Type[BaseParser]`): The parser class to process the
action's inputs and outputs. Defaults to :class:`JsonParser`.
Examples:
* simple tool
.. code-block:: python
class Bold(BaseAction):
'''Make text bold'''
def run(self, text: str):
'''
Args:
text (str): input text
Returns:
str: bold text
'''
return '**' + text + '**'
action = Bold()
* toolkit with multiple APIs
.. code-block:: python
class Calculator(BaseAction):
'''Calculator'''
@tool_api
def add(self, a, b):
'''Add operation
Args:
a (int): augend
b (int): addend
Returns:
int: sum
'''
return a + b
@tool_api
def sub(self, a, b):
'''Subtraction operation
Args:
a (int): minuend
b (int): subtrahend
Returns:
int: difference
'''
return a - b
action = Calculator()
"""
def __init__(
self,
description: Optional[dict] = None,
parser: Type[BaseParser] = JsonParser,
):
self._description = deepcopy(description or self.__tool_description__)
self._name = self._description['name']
self._parser = parser(self)
def __call__(self, inputs: str, name='run') -> ActionReturn:
fallback_args = {'inputs': inputs, 'name': name}
if not hasattr(self, name):
return ActionReturn(
fallback_args,
type=self.name,
errmsg=f'invalid API: {name}',
state=ActionStatusCode.API_ERROR)
try:
inputs = self._parser.parse_inputs(inputs, name)
except ParseError as exc:
return ActionReturn(
fallback_args,
type=self.name,
errmsg=exc.err_msg,
state=ActionStatusCode.ARGS_ERROR)
try:
outputs = getattr(self, name)(**inputs)
except Exception as exc:
return ActionReturn(
inputs,
type=self.name,
errmsg=str(exc),
state=ActionStatusCode.API_ERROR)
if isinstance(outputs, ActionReturn):
action_return = outputs
if not action_return.args:
action_return.args = inputs
if not action_return.type:
action_return.type = self.name
else:
result = self._parser.parse_outputs(outputs)
action_return = ActionReturn(inputs, type=self.name, result=result)
return action_return
@property
def name(self):
return self._name
@property
def is_toolkit(self):
return self._is_toolkit
@property
def description(self) -> dict:
"""Description of the tool."""
return self._description
def __repr__(self):
return f'{self.description}'
__str__ = __repr__
class AsyncActionMixin:
async def __call__(self, inputs: str, name='run') -> ActionReturn:
fallback_args = {'inputs': inputs, 'name': name}
if not hasattr(self, name):
return ActionReturn(
fallback_args,
type=self.name,
errmsg=f'invalid API: {name}',
state=ActionStatusCode.API_ERROR)
try:
inputs = self._parser.parse_inputs(inputs, name)
except ParseError as exc:
return ActionReturn(
fallback_args,
type=self.name,
errmsg=exc.err_msg,
state=ActionStatusCode.ARGS_ERROR)
try:
outputs = await getattr(self, name)(**inputs)
except Exception as exc:
return ActionReturn(
inputs,
type=self.name,
errmsg=str(exc),
state=ActionStatusCode.API_ERROR)
if isinstance(outputs, ActionReturn):
action_return = outputs
if not action_return.args:
action_return.args = inputs
if not action_return.type:
action_return.type = self.name
else:
result = self._parser.parse_outputs(outputs)
action_return = ActionReturn(inputs, type=self.name, result=result)
return action_return