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

* update `tool_api`

* update `tool_api`

* update `get_plugin_prompt`

Author: braisedpork1964

lagent/actions/base_action.py

@@ -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
@@ -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``.
@@ -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'
@@ -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
@@ -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=[],
         )
@@ -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)
@@ -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')
@@ -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')):
@@ -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:
@@ -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:

lagent/agents/stream.py

@@ -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)
@@ -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)