Skip to content

[Feat] Allow custom argument scope in tool descriptions #292

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jan 16, 2025
+43 −63
Merged

[Feat] Allow custom argument scope in tool descriptions #292

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
88e4c82
update `tool_api`
braisedpork1964 Jan 15, 2025
3555ca1
update `tool_api`
braisedpork1964 Jan 15, 2025
5e07ee4
update `get_plugin_prompt`
braisedpork1964 Jan 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 42 additions & 55 deletions lagent/actions/base_action.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from abc import ABCMeta
from copy import deepcopy
from functools import wraps
from typing import Callable, Optional, Type, get_args, get_origin
from typing import Callable, Iterable, Optional, Type, get_args, get_origin

try:
from typing import Annotated
Expand All @@ -24,11 +24,15 @@
logging.getLogger('griffe').setLevel(logging.ERROR)


def tool_api(func: Optional[Callable] = None,
*,
explode_return: bool = False,
returns_named_value: bool = False,
**kwargs):
def tool_api(
func: Optional[Callable] = None,
*,
explode_return: bool = False,
returns_named_value: bool = False,
include_arguments: Optional[Iterable[str]] = None,
exclude_arguments: Optional[Iterable[str]] = None,
**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``.
Expand Down Expand Up @@ -90,6 +94,16 @@ def foo(a, b):
``return_data`` field will be added to ``api_description`` only
when ``explode_return`` or ``returns_named_value`` is enabled.
"""
if include_arguments is None:
exclude_arguments = exclude_arguments or set()
if isinstance(exclude_arguments, str):
exclude_arguments = {exclude_arguments}
elif not isinstance(exclude_arguments, set):
exclude_arguments = set(exclude_arguments)
if 'self' not in exclude_arguments:
exclude_arguments.add('self')
else:
include_arguments = {include_arguments} if isinstance(include_arguments, str) else set(include_arguments)

def _detect_type(string):
field_type = 'STRING'
Expand All @@ -106,10 +120,9 @@ def _detect_type(string):

def _explode(desc):
kvs = []
desc = '\nArgs:\n' + '\n'.join([
' ' + item.lstrip(' -+*#.')
for item in desc.split('\n')[1:] if item.strip()
])
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
Expand All @@ -125,13 +138,12 @@ def _explode(desc):

def _parse_tool(function):
# remove rst syntax
docs = Docstring(
re.sub(':(.+?):`(.+?)`', '\\2', function.__doc__ or '')).parse(
'google', returns_named_value=returns_named_value, **kwargs)
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 '',
description=docs[0].value if docs[0].kind is DocstringSectionKind.text else '',
parameters=[],
required=[],
)
Expand All @@ -155,17 +167,14 @@ def _parse_tool(function):

sig = inspect.signature(function)
for name, param in sig.parameters.items():
if name == 'self':
if name in exclude_arguments if include_arguments is None else name not in include_arguments:
continue
parameter = dict(
name=param.name,
type='STRING',
description=args_doc.get(param.name,
{}).get('description', ''))
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')
parameter['type'] = args_doc.get(param.name, {}).get('type', 'STRING')
else:
if get_origin(annotation) is Annotated:
annotation, info = get_args(annotation)
Expand Down Expand Up @@ -229,9 +238,8 @@ class ToolMeta(ABCMeta):

def __new__(mcs, name, base, attrs):
is_toolkit, tool_desc = True, dict(
name=name,
description=Docstring(attrs.get('__doc__',
'')).parse('google')[0].value)
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')
Expand All @@ -246,8 +254,7 @@ def __new__(mcs, name, base, attrs):
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')
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')):
Expand Down Expand Up @@ -346,26 +353,16 @@ 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)
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)
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)
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:
Expand Down Expand Up @@ -402,26 +399,16 @@ 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)
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)
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)
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:
Expand Down
9 changes: 1 addition & 8 deletions lagent/agents/stream.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
)


def get_plugin_prompt(actions, api_desc_template=API_PREFIX):
def get_plugin_prompt(actions, api_desc_template='{description}'):
plugin_descriptions = []
for action in actions if isinstance(actions, list) else [actions]:
action = create_object(action)
Expand All @@ -47,15 +47,8 @@ def get_plugin_prompt(actions, api_desc_template=API_PREFIX):
for api in action_desc['api_list']:
api['name'] = f"{action.name}.{api['name']}"
api['description'] = api_desc_template.format(tool_name=action.name, description=api['description'])
api['parameters'] = [param for param in api['parameters'] if param['name'] in api['required']]
plugin_descriptions.append(api)
else:
action_desc['description'] = api_desc_template.format(
tool_name=action.name, description=action_desc['description']
)
action_desc['parameters'] = [
param for param in action_desc['parameters'] if param['name'] in action_desc['required']
]
plugin_descriptions.append(action_desc)
return json.dumps(plugin_descriptions, ensure_ascii=False, indent=4)

Expand Down