Feat: AgentLegoToolkit

examples/internlm2_agent_web_demo.py

@@ -6,6 +6,7 @@
 import streamlit as st
 
 from lagent.actions import ActionExecutor, ArxivSearch, IPythonInterpreter
+# from lagent.actions.agentlego_wrapper import AgentLegoToolkit
 from lagent.agents.internlm2_agent import INTERPRETER_CN, META_CN, PLUGIN_CN, Internlm2Agent, Internlm2Protocol
 from lagent.llms.lmdepoly_wrapper import LMDeployClient
 from lagent.llms.meta_template import INTERNLM2_META as META
@@ -21,9 +22,7 @@ def init_state(self):
         st.session_state['assistant'] = []
         st.session_state['user'] = []
 
-        action_list = [
-            ArxivSearch(),
-        ]
+        action_list = [ArxivSearch()]
         st.session_state['plugin_map'] = {
             action.name: action
             for action in action_list

lagent/actions/agentlego_wrapper.py

@@ -0,0 +1,48 @@
+from typing import Optional
+
+# from agentlego.parsers import DefaultParser
+from agentlego.tools.remote import RemoteTool
+
+from lagent import BaseAction
+from lagent.actions.parser import JsonParser
+
+
+class AgentLegoToolkit(BaseAction):
+
+    def __init__(self,
+                 name: str,
+                 url: Optional[str] = None,
+                 text: Optional[str] = None,
+                 spec_dict: Optional[dict] = None,
+                 parser=JsonParser,
+                 enable: bool = True):
+
+        if url is not None:
+            spec = dict(url=url)
+        elif text is not None:
+            spec = dict(text=text)
+        else:
+            assert spec_dict is not None
+            spec = dict(spec_dict=spec_dict)
+        if url is not None and not url.endswith('.json'):
+            api_list = [RemoteTool.from_url(url).to_lagent()]
+        else:
+            api_list = [
+                api.to_lagent() for api in RemoteTool.from_openapi(**spec)
+            ]
+        api_desc = []
+        for api in api_list:
+            api_desc.append(api.description)
+        if len(api_list) > 1:
+            tool_description = dict(name=name, api_list=api_desc)
+            for func in api_list:
+                setattr(self, func.name, func.run)
+        else:
+            tool_description = api_desc[0]
+            setattr(self, 'run', api_list[0].run)
+        super().__init__(
+            description=tool_description, parser=parser, enable=enable)
+
+    @property
+    def is_toolkit(self):
+        return 'api_list' in self.description

requirements/optional.txt

@@ -1,3 +1,4 @@
+agentlego
 google-search-results
 lmdeploy>=0.2.3
 pillow