Spaces:
Sleeping
Sleeping
Superkingjcj
commited on
Upload 111 files
Browse filesThis view is limited to 50 files because it contains too many changes.
See raw diff
- lagent/__init__.py +4 -0
- lagent/__pycache__/__init__.cpython-310.pyc +0 -0
- lagent/__pycache__/schema.cpython-310.pyc +0 -0
- lagent/__pycache__/version.cpython-310.pyc +0 -0
- lagent/actions/__init__.py +26 -0
- lagent/actions/__pycache__/__init__.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/action_executor.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/arxiv_search.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/base_action.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/bing_map.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/builtin_actions.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/google_scholar_search.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/google_search.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/ipython_interactive.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/ipython_interpreter.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/ipython_manager.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/parser.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/ppt.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/python_interpreter.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/weather_query.cpython-310.pyc +0 -0
- lagent/actions/__pycache__/web_browser.cpython-310.pyc +0 -0
- lagent/actions/action_executor.py +198 -0
- lagent/actions/arxiv_search.py +79 -0
- lagent/actions/base_action.py +434 -0
- lagent/actions/bing_map.py +268 -0
- lagent/actions/builtin_actions.py +109 -0
- lagent/actions/google_scholar_search.py +438 -0
- lagent/actions/google_search.py +244 -0
- lagent/actions/ipython_interactive.py +273 -0
- lagent/actions/ipython_interpreter.py +584 -0
- lagent/actions/ipython_manager.py +220 -0
- lagent/actions/parser.py +146 -0
- lagent/actions/ppt.py +233 -0
- lagent/actions/python_interpreter.py +176 -0
- lagent/actions/weather_query.py +71 -0
- lagent/actions/web_browser.py +908 -0
- lagent/agents/__init__.py +9 -0
- lagent/agents/__pycache__/__init__.cpython-310.pyc +0 -0
- lagent/agents/__pycache__/agent.cpython-310.pyc +0 -0
- lagent/agents/__pycache__/react.cpython-310.pyc +0 -0
- lagent/agents/__pycache__/stream.cpython-310.pyc +0 -0
- lagent/agents/agent.py +400 -0
- lagent/agents/aggregator/__init__.py +4 -0
- lagent/agents/aggregator/__pycache__/__init__.cpython-310.pyc +0 -0
- lagent/agents/aggregator/__pycache__/default_aggregator.cpython-310.pyc +0 -0
- lagent/agents/aggregator/__pycache__/tool_aggregator.cpython-310.pyc +0 -0
- lagent/agents/aggregator/default_aggregator.py +44 -0
- lagent/agents/aggregator/tool_aggregator.py +106 -0
- lagent/agents/react.py +161 -0
- lagent/agents/stream.py +316 -0
lagent/__init__.py
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) OpenMMLab. All rights reserved.
|
2 |
+
from .version import __version__, version_info
|
3 |
+
|
4 |
+
__all__ = ['__version__', 'version_info']
|
lagent/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (231 Bytes). View file
|
|
lagent/__pycache__/schema.cpython-310.pyc
ADDED
Binary file (3.46 kB). View file
|
|
lagent/__pycache__/version.cpython-310.pyc
ADDED
Binary file (744 Bytes). View file
|
|
lagent/actions/__init__.py
ADDED
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from .action_executor import ActionExecutor, AsyncActionExecutor
|
2 |
+
from .arxiv_search import ArxivSearch, AsyncArxivSearch
|
3 |
+
from .base_action import BaseAction, tool_api
|
4 |
+
from .bing_map import AsyncBINGMap, BINGMap
|
5 |
+
from .builtin_actions import FinishAction, InvalidAction, NoAction
|
6 |
+
from .google_scholar_search import AsyncGoogleScholar, GoogleScholar
|
7 |
+
from .google_search import AsyncGoogleSearch, GoogleSearch
|
8 |
+
from .ipython_interactive import AsyncIPythonInteractive, IPythonInteractive
|
9 |
+
from .ipython_interpreter import AsyncIPythonInterpreter, IPythonInterpreter
|
10 |
+
from .ipython_manager import IPythonInteractiveManager
|
11 |
+
from .parser import BaseParser, JsonParser, TupleParser
|
12 |
+
from .ppt import PPT, AsyncPPT
|
13 |
+
from .python_interpreter import AsyncPythonInterpreter, PythonInterpreter
|
14 |
+
from .web_browser import AsyncWebBrowser, WebBrowser
|
15 |
+
from .weather_query import WeatherQuery
|
16 |
+
|
17 |
+
__all__ = [
|
18 |
+
'BaseAction', 'ActionExecutor', 'AsyncActionExecutor', 'InvalidAction',
|
19 |
+
'FinishAction', 'NoAction', 'BINGMap', 'AsyncBINGMap', 'ArxivSearch',
|
20 |
+
'AsyncArxivSearch', 'GoogleSearch', 'AsyncGoogleSearch', 'GoogleScholar',
|
21 |
+
'AsyncGoogleScholar', 'IPythonInterpreter', 'AsyncIPythonInterpreter',
|
22 |
+
'IPythonInteractive', 'AsyncIPythonInteractive',
|
23 |
+
'IPythonInteractiveManager', 'PythonInterpreter', 'AsyncPythonInterpreter',
|
24 |
+
'PPT', 'AsyncPPT', 'WebBrowser', 'AsyncWebBrowser', 'BaseParser',
|
25 |
+
'JsonParser', 'TupleParser', 'tool_api', 'WeatherQuery'
|
26 |
+
]
|
lagent/actions/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (1.47 kB). View file
|
|
lagent/actions/__pycache__/action_executor.cpython-310.pyc
ADDED
Binary file (5.84 kB). View file
|
|
lagent/actions/__pycache__/arxiv_search.cpython-310.pyc
ADDED
Binary file (3.19 kB). View file
|
|
lagent/actions/__pycache__/base_action.cpython-310.pyc
ADDED
Binary file (11.6 kB). View file
|
|
lagent/actions/__pycache__/bing_map.cpython-310.pyc
ADDED
Binary file (7.79 kB). View file
|
|
lagent/actions/__pycache__/builtin_actions.cpython-310.pyc
ADDED
Binary file (3.89 kB). View file
|
|
lagent/actions/__pycache__/google_scholar_search.cpython-310.pyc
ADDED
Binary file (13 kB). View file
|
|
lagent/actions/__pycache__/google_search.cpython-310.pyc
ADDED
Binary file (6.93 kB). View file
|
|
lagent/actions/__pycache__/ipython_interactive.cpython-310.pyc
ADDED
Binary file (8.41 kB). View file
|
|
lagent/actions/__pycache__/ipython_interpreter.cpython-310.pyc
ADDED
Binary file (16.6 kB). View file
|
|
lagent/actions/__pycache__/ipython_manager.cpython-310.pyc
ADDED
Binary file (7.11 kB). View file
|
|
lagent/actions/__pycache__/parser.cpython-310.pyc
ADDED
Binary file (5.48 kB). View file
|
|
lagent/actions/__pycache__/ppt.cpython-310.pyc
ADDED
Binary file (6.81 kB). View file
|
|
lagent/actions/__pycache__/python_interpreter.cpython-310.pyc
ADDED
Binary file (5.38 kB). View file
|
|
lagent/actions/__pycache__/weather_query.cpython-310.pyc
ADDED
Binary file (2.66 kB). View file
|
|
lagent/actions/__pycache__/web_browser.cpython-310.pyc
ADDED
Binary file (28.8 kB). View file
|
|
lagent/actions/action_executor.py
ADDED
@@ -0,0 +1,198 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import inspect
|
2 |
+
from collections import OrderedDict
|
3 |
+
from typing import Callable, Dict, List, Union
|
4 |
+
|
5 |
+
from lagent.actions.base_action import BaseAction
|
6 |
+
from lagent.actions.builtin_actions import FinishAction, InvalidAction, NoAction
|
7 |
+
from lagent.hooks import Hook, RemovableHandle
|
8 |
+
from lagent.schema import ActionReturn, ActionValidCode, AgentMessage, FunctionCall
|
9 |
+
from lagent.utils import create_object
|
10 |
+
|
11 |
+
|
12 |
+
class ActionExecutor:
|
13 |
+
"""The action executor class.
|
14 |
+
|
15 |
+
Args:
|
16 |
+
actions (Union[BaseAction, List[BaseAction]]): The action or actions.
|
17 |
+
invalid_action (BaseAction, optional): The invalid action. Defaults to
|
18 |
+
InvalidAction().
|
19 |
+
no_action (BaseAction, optional): The no action.
|
20 |
+
Defaults to NoAction().
|
21 |
+
finish_action (BaseAction, optional): The finish action. Defaults to
|
22 |
+
FinishAction().
|
23 |
+
finish_in_action (bool, optional): Whether the finish action is in the
|
24 |
+
action list. Defaults to False.
|
25 |
+
"""
|
26 |
+
|
27 |
+
def __init__(
|
28 |
+
self,
|
29 |
+
actions: Union[BaseAction, List[BaseAction], Dict, List[Dict]],
|
30 |
+
invalid_action: BaseAction = dict(type=InvalidAction),
|
31 |
+
no_action: BaseAction = dict(type=NoAction),
|
32 |
+
finish_action: BaseAction = dict(type=FinishAction),
|
33 |
+
finish_in_action: bool = False,
|
34 |
+
hooks: List[Dict] = None,
|
35 |
+
):
|
36 |
+
|
37 |
+
if not isinstance(actions, list):
|
38 |
+
actions = [actions]
|
39 |
+
finish_action = create_object(finish_action)
|
40 |
+
if finish_in_action:
|
41 |
+
actions.append(finish_action)
|
42 |
+
for i, action in enumerate(actions):
|
43 |
+
actions[i] = create_object(action)
|
44 |
+
self.actions = {action.name: action for action in actions}
|
45 |
+
|
46 |
+
self.invalid_action = create_object(invalid_action)
|
47 |
+
self.no_action = create_object(no_action)
|
48 |
+
self.finish_action = finish_action
|
49 |
+
self._hooks: Dict[int, Hook] = OrderedDict()
|
50 |
+
if hooks:
|
51 |
+
for hook in hooks:
|
52 |
+
hook = create_object(hook)
|
53 |
+
self.register_hook(hook)
|
54 |
+
|
55 |
+
def description(self) -> List[Dict]:
|
56 |
+
actions = []
|
57 |
+
for action_name, action in self.actions.items():
|
58 |
+
if action.is_toolkit:
|
59 |
+
for api in action.description['api_list']:
|
60 |
+
api_desc = api.copy()
|
61 |
+
api_desc['name'] = f"{action_name}.{api_desc['name']}"
|
62 |
+
actions.append(api_desc)
|
63 |
+
else:
|
64 |
+
action_desc = action.description.copy()
|
65 |
+
actions.append(action_desc)
|
66 |
+
return actions
|
67 |
+
|
68 |
+
def __contains__(self, name: str):
|
69 |
+
return name in self.actions
|
70 |
+
|
71 |
+
def keys(self):
|
72 |
+
return list(self.actions.keys())
|
73 |
+
|
74 |
+
def __setitem__(self, name: str, action: Union[BaseAction, Dict]):
|
75 |
+
action = create_object(action)
|
76 |
+
self.actions[action.name] = action
|
77 |
+
|
78 |
+
def __delitem__(self, name: str):
|
79 |
+
del self.actions[name]
|
80 |
+
|
81 |
+
def forward(self, name, parameters, **kwargs) -> ActionReturn:
|
82 |
+
action_name, api_name = (
|
83 |
+
name.split('.') if '.' in name else (name, 'run'))
|
84 |
+
action_return: ActionReturn = ActionReturn()
|
85 |
+
if action_name not in self:
|
86 |
+
if name == self.no_action.name:
|
87 |
+
action_return = self.no_action(parameters)
|
88 |
+
elif name == self.finish_action.name:
|
89 |
+
action_return = self.finish_action(parameters)
|
90 |
+
else:
|
91 |
+
action_return = self.invalid_action(parameters)
|
92 |
+
else:
|
93 |
+
action_return = self.actions[action_name](parameters, api_name)
|
94 |
+
action_return.valid = ActionValidCode.OPEN
|
95 |
+
return action_return
|
96 |
+
|
97 |
+
def __call__(self,
|
98 |
+
message: AgentMessage,
|
99 |
+
session_id=0,
|
100 |
+
**kwargs) -> AgentMessage:
|
101 |
+
# message.receiver = self.name
|
102 |
+
for hook in self._hooks.values():
|
103 |
+
result = hook.before_action(self, message, session_id)
|
104 |
+
if result:
|
105 |
+
message = result
|
106 |
+
|
107 |
+
assert isinstance(message.content, FunctionCall) or (
|
108 |
+
isinstance(message.content, dict) and 'name' in message.content
|
109 |
+
and 'parameters' in message.content)
|
110 |
+
if isinstance(message.content, dict):
|
111 |
+
name = message.content.get('name')
|
112 |
+
parameters = message.content.get('parameters')
|
113 |
+
else:
|
114 |
+
name = message.content.name
|
115 |
+
parameters = message.content.parameters
|
116 |
+
|
117 |
+
response_message = self.forward(
|
118 |
+
name=name, parameters=parameters, **kwargs)
|
119 |
+
if not isinstance(response_message, AgentMessage):
|
120 |
+
response_message = AgentMessage(
|
121 |
+
sender=self.__class__.__name__,
|
122 |
+
content=response_message,
|
123 |
+
)
|
124 |
+
|
125 |
+
for hook in self._hooks.values():
|
126 |
+
result = hook.after_action(self, response_message, session_id)
|
127 |
+
if result:
|
128 |
+
response_message = result
|
129 |
+
return response_message
|
130 |
+
|
131 |
+
def register_hook(self, hook: Callable):
|
132 |
+
handle = RemovableHandle(self._hooks)
|
133 |
+
self._hooks[handle.id] = hook
|
134 |
+
return handle
|
135 |
+
|
136 |
+
|
137 |
+
class AsyncActionExecutor(ActionExecutor):
|
138 |
+
|
139 |
+
async def forward(self, name, parameters, **kwargs) -> ActionReturn:
|
140 |
+
action_name, api_name = (
|
141 |
+
name.split('.') if '.' in name else (name, 'run'))
|
142 |
+
action_return: ActionReturn = ActionReturn()
|
143 |
+
if action_name not in self:
|
144 |
+
if name == self.no_action.name:
|
145 |
+
action_return = self.no_action(parameters)
|
146 |
+
elif name == self.finish_action.name:
|
147 |
+
action_return = self.finish_action(parameters)
|
148 |
+
else:
|
149 |
+
action_return = self.invalid_action(parameters)
|
150 |
+
else:
|
151 |
+
action = self.actions[action_name]
|
152 |
+
if inspect.iscoroutinefunction(action.__call__):
|
153 |
+
action_return = await action(parameters, api_name)
|
154 |
+
else:
|
155 |
+
action_return = action(parameters, api_name)
|
156 |
+
action_return.valid = ActionValidCode.OPEN
|
157 |
+
return action_return
|
158 |
+
|
159 |
+
async def __call__(self,
|
160 |
+
message: AgentMessage,
|
161 |
+
session_id=0,
|
162 |
+
**kwargs) -> AgentMessage:
|
163 |
+
# message.receiver = self.name
|
164 |
+
for hook in self._hooks.values():
|
165 |
+
if inspect.iscoroutinefunction(hook.before_action):
|
166 |
+
result = await hook.before_action(self, message, session_id)
|
167 |
+
else:
|
168 |
+
result = hook.before_action(self, message, session_id)
|
169 |
+
if result:
|
170 |
+
message = result
|
171 |
+
|
172 |
+
assert isinstance(message.content, FunctionCall) or (
|
173 |
+
isinstance(message.content, dict) and 'name' in message.content
|
174 |
+
and 'parameters' in message.content)
|
175 |
+
if isinstance(message.content, dict):
|
176 |
+
name = message.content.get('name')
|
177 |
+
parameters = message.content.get('parameters')
|
178 |
+
else:
|
179 |
+
name = message.content.name
|
180 |
+
parameters = message.content.parameters
|
181 |
+
|
182 |
+
response_message = await self.forward(
|
183 |
+
name=name, parameters=parameters, **kwargs)
|
184 |
+
if not isinstance(response_message, AgentMessage):
|
185 |
+
response_message = AgentMessage(
|
186 |
+
sender=self.__class__.__name__,
|
187 |
+
content=response_message,
|
188 |
+
)
|
189 |
+
|
190 |
+
for hook in self._hooks.values():
|
191 |
+
if inspect.iscoroutinefunction(hook.after_action):
|
192 |
+
result = await hook.after_action(self, response_message,
|
193 |
+
session_id)
|
194 |
+
else:
|
195 |
+
result = hook.after_action(self, response_message, session_id)
|
196 |
+
if result:
|
197 |
+
response_message = result
|
198 |
+
return response_message
|
lagent/actions/arxiv_search.py
ADDED
@@ -0,0 +1,79 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Optional, Type
|
2 |
+
|
3 |
+
from asyncer import asyncify
|
4 |
+
|
5 |
+
from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
|
6 |
+
from lagent.actions.parser import BaseParser, JsonParser
|
7 |
+
from lagent.schema import ActionReturn, ActionStatusCode
|
8 |
+
|
9 |
+
|
10 |
+
class ArxivSearch(BaseAction):
|
11 |
+
"""Search information from Arxiv.org. \
|
12 |
+
Useful for when you need to answer questions about Physics, Mathematics, \
|
13 |
+
Computer Science, Quantitative Biology, Quantitative Finance, Statistics, \
|
14 |
+
Electrical Engineering, and Economics from scientific articles on arxiv.org.
|
15 |
+
"""
|
16 |
+
|
17 |
+
def __init__(
|
18 |
+
self,
|
19 |
+
top_k_results: int = 3,
|
20 |
+
max_query_len: int = 300,
|
21 |
+
doc_content_chars_max: int = 1500,
|
22 |
+
description: Optional[dict] = None,
|
23 |
+
parser: Type[BaseParser] = JsonParser,
|
24 |
+
):
|
25 |
+
super().__init__(description, parser)
|
26 |
+
self.top_k_results = top_k_results
|
27 |
+
self.max_query_len = max_query_len
|
28 |
+
self.doc_content_chars_max = doc_content_chars_max
|
29 |
+
|
30 |
+
@tool_api(explode_return=True)
|
31 |
+
def get_arxiv_article_information(self, query: str) -> dict:
|
32 |
+
"""Run Arxiv search and get the article meta information.
|
33 |
+
|
34 |
+
Args:
|
35 |
+
query (:class:`str`): the content of search query
|
36 |
+
|
37 |
+
Returns:
|
38 |
+
:class:`dict`: article information
|
39 |
+
* content (str): a list of 3 arxiv search papers
|
40 |
+
"""
|
41 |
+
import arxiv
|
42 |
+
|
43 |
+
try:
|
44 |
+
results = arxiv.Search( # type: ignore
|
45 |
+
query[: self.max_query_len], max_results=self.top_k_results
|
46 |
+
).results()
|
47 |
+
except Exception as exc:
|
48 |
+
return ActionReturn(errmsg=f'Arxiv exception: {exc}', state=ActionStatusCode.HTTP_ERROR)
|
49 |
+
docs = [
|
50 |
+
f'Published: {result.updated.date()}\nTitle: {result.title}\n'
|
51 |
+
f'Authors: {", ".join(a.name for a in result.authors)}\n'
|
52 |
+
f'Summary: {result.summary[:self.doc_content_chars_max]}'
|
53 |
+
for result in results
|
54 |
+
]
|
55 |
+
if docs:
|
56 |
+
return {'content': '\n\n'.join(docs)}
|
57 |
+
return {'content': 'No good Arxiv Result was found'}
|
58 |
+
|
59 |
+
|
60 |
+
class AsyncArxivSearch(AsyncActionMixin, ArxivSearch):
|
61 |
+
"""Search information from Arxiv.org. \
|
62 |
+
Useful for when you need to answer questions about Physics, Mathematics, \
|
63 |
+
Computer Science, Quantitative Biology, Quantitative Finance, Statistics, \
|
64 |
+
Electrical Engineering, and Economics from scientific articles on arxiv.org.
|
65 |
+
"""
|
66 |
+
|
67 |
+
@tool_api(explode_return=True)
|
68 |
+
@asyncify
|
69 |
+
def get_arxiv_article_information(self, query: str) -> dict:
|
70 |
+
"""Run Arxiv search and get the article meta information.
|
71 |
+
|
72 |
+
Args:
|
73 |
+
query (:class:`str`): the content of search query
|
74 |
+
|
75 |
+
Returns:
|
76 |
+
:class:`dict`: article information
|
77 |
+
* content (str): a list of 3 arxiv search papers
|
78 |
+
"""
|
79 |
+
return super().get_arxiv_article_information(query)
|
lagent/actions/base_action.py
ADDED
@@ -0,0 +1,434 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import inspect
|
2 |
+
import logging
|
3 |
+
import re
|
4 |
+
from abc import ABCMeta
|
5 |
+
from copy import deepcopy
|
6 |
+
from functools import wraps
|
7 |
+
from typing import Callable, Optional, Type, get_args, get_origin
|
8 |
+
|
9 |
+
try:
|
10 |
+
from typing import Annotated
|
11 |
+
except ImportError:
|
12 |
+
from typing_extensions import Annotated
|
13 |
+
|
14 |
+
from griffe import Docstring
|
15 |
+
|
16 |
+
try:
|
17 |
+
from griffe import DocstringSectionKind
|
18 |
+
except ImportError:
|
19 |
+
from griffe.enumerations import DocstringSectionKind
|
20 |
+
|
21 |
+
from ..schema import ActionReturn, ActionStatusCode
|
22 |
+
from .parser import BaseParser, JsonParser, ParseError
|
23 |
+
|
24 |
+
logging.getLogger('griffe').setLevel(logging.ERROR)
|
25 |
+
|
26 |
+
|
27 |
+
def tool_api(func: Optional[Callable] = None,
|
28 |
+
*,
|
29 |
+
explode_return: bool = False,
|
30 |
+
returns_named_value: bool = False,
|
31 |
+
**kwargs):
|
32 |
+
"""Turn functions into tools. It will parse typehints as well as docstrings
|
33 |
+
to build the tool description and attach it to functions via an attribute
|
34 |
+
``api_description``.
|
35 |
+
|
36 |
+
Examples:
|
37 |
+
|
38 |
+
.. code-block:: python
|
39 |
+
|
40 |
+
# typehints has higher priority than docstrings
|
41 |
+
from typing import Annotated
|
42 |
+
|
43 |
+
@tool_api
|
44 |
+
def add(a: Annotated[int, 'augend'], b: Annotated[int, 'addend'] = 1):
|
45 |
+
'''Add operation
|
46 |
+
|
47 |
+
Args:
|
48 |
+
x (int): a
|
49 |
+
y (int): b
|
50 |
+
'''
|
51 |
+
return a + b
|
52 |
+
|
53 |
+
print(add.api_description)
|
54 |
+
|
55 |
+
Args:
|
56 |
+
func (Optional[Callable]): function to decorate. Defaults to ``None``.
|
57 |
+
explode_return (bool): whether to flatten the dictionary or tuple return
|
58 |
+
as the ``return_data`` field. When enabled, it is recommended to
|
59 |
+
annotate the member in docstrings. Defaults to ``False``.
|
60 |
+
|
61 |
+
.. code-block:: python
|
62 |
+
|
63 |
+
@tool_api(explode_return=True)
|
64 |
+
def foo(a, b):
|
65 |
+
'''A simple function
|
66 |
+
|
67 |
+
Args:
|
68 |
+
a (int): a
|
69 |
+
b (int): b
|
70 |
+
|
71 |
+
Returns:
|
72 |
+
dict: information of inputs
|
73 |
+
* x: value of a
|
74 |
+
* y: value of b
|
75 |
+
'''
|
76 |
+
return {'x': a, 'y': b}
|
77 |
+
|
78 |
+
print(foo.api_description)
|
79 |
+
|
80 |
+
returns_named_value (bool): whether to parse ``thing: Description`` in
|
81 |
+
returns sections as a name and description, rather than a type and
|
82 |
+
description. When true, type must be wrapped in parentheses:
|
83 |
+
``(int): Description``. When false, parentheses are optional but
|
84 |
+
the items cannot be named: ``int: Description``. Defaults to ``False``.
|
85 |
+
|
86 |
+
Returns:
|
87 |
+
Callable: wrapped function or partial decorator
|
88 |
+
|
89 |
+
Important:
|
90 |
+
``return_data`` field will be added to ``api_description`` only
|
91 |
+
when ``explode_return`` or ``returns_named_value`` is enabled.
|
92 |
+
"""
|
93 |
+
|
94 |
+
def _detect_type(string):
|
95 |
+
field_type = 'STRING'
|
96 |
+
if 'list' in string:
|
97 |
+
field_type = 'Array'
|
98 |
+
elif 'str' not in string:
|
99 |
+
if 'float' in string:
|
100 |
+
field_type = 'FLOAT'
|
101 |
+
elif 'int' in string:
|
102 |
+
field_type = 'NUMBER'
|
103 |
+
elif 'bool' in string:
|
104 |
+
field_type = 'BOOLEAN'
|
105 |
+
return field_type
|
106 |
+
|
107 |
+
def _explode(desc):
|
108 |
+
kvs = []
|
109 |
+
desc = '\nArgs:\n' + '\n'.join([
|
110 |
+
' ' + item.lstrip(' -+*#.')
|
111 |
+
for item in desc.split('\n')[1:] if item.strip()
|
112 |
+
])
|
113 |
+
docs = Docstring(desc).parse('google')
|
114 |
+
if not docs:
|
115 |
+
return kvs
|
116 |
+
if docs[0].kind is DocstringSectionKind.parameters:
|
117 |
+
for d in docs[0].value:
|
118 |
+
d = d.as_dict()
|
119 |
+
if not d['annotation']:
|
120 |
+
d.pop('annotation')
|
121 |
+
else:
|
122 |
+
d['type'] = _detect_type(d.pop('annotation').lower())
|
123 |
+
kvs.append(d)
|
124 |
+
return kvs
|
125 |
+
|
126 |
+
def _parse_tool(function):
|
127 |
+
# remove rst syntax
|
128 |
+
docs = Docstring(
|
129 |
+
re.sub(':(.+?):`(.+?)`', '\\2', function.__doc__ or '')).parse(
|
130 |
+
'google', returns_named_value=returns_named_value, **kwargs)
|
131 |
+
desc = dict(
|
132 |
+
name=function.__name__,
|
133 |
+
description=docs[0].value
|
134 |
+
if docs[0].kind is DocstringSectionKind.text else '',
|
135 |
+
parameters=[],
|
136 |
+
required=[],
|
137 |
+
)
|
138 |
+
args_doc, returns_doc = {}, []
|
139 |
+
for doc in docs:
|
140 |
+
if doc.kind is DocstringSectionKind.parameters:
|
141 |
+
for d in doc.value:
|
142 |
+
d = d.as_dict()
|
143 |
+
d['type'] = _detect_type(d.pop('annotation').lower())
|
144 |
+
args_doc[d['name']] = d
|
145 |
+
if doc.kind is DocstringSectionKind.returns:
|
146 |
+
for d in doc.value:
|
147 |
+
d = d.as_dict()
|
148 |
+
if not d['name']:
|
149 |
+
d.pop('name')
|
150 |
+
if not d['annotation']:
|
151 |
+
d.pop('annotation')
|
152 |
+
else:
|
153 |
+
d['type'] = _detect_type(d.pop('annotation').lower())
|
154 |
+
returns_doc.append(d)
|
155 |
+
|
156 |
+
sig = inspect.signature(function)
|
157 |
+
for name, param in sig.parameters.items():
|
158 |
+
if name == 'self':
|
159 |
+
continue
|
160 |
+
parameter = dict(
|
161 |
+
name=param.name,
|
162 |
+
type='STRING',
|
163 |
+
description=args_doc.get(param.name,
|
164 |
+
{}).get('description', ''))
|
165 |
+
annotation = param.annotation
|
166 |
+
if annotation is inspect.Signature.empty:
|
167 |
+
parameter['type'] = args_doc.get(param.name,
|
168 |
+
{}).get('type', 'STRING')
|
169 |
+
else:
|
170 |
+
if get_origin(annotation) is Annotated:
|
171 |
+
annotation, info = get_args(annotation)
|
172 |
+
if info:
|
173 |
+
parameter['description'] = info
|
174 |
+
while get_origin(annotation):
|
175 |
+
annotation = get_args(annotation)
|
176 |
+
parameter['type'] = _detect_type(str(annotation))
|
177 |
+
desc['parameters'].append(parameter)
|
178 |
+
if param.default is inspect.Signature.empty:
|
179 |
+
desc['required'].append(param.name)
|
180 |
+
|
181 |
+
return_data = []
|
182 |
+
if explode_return:
|
183 |
+
return_data = _explode(returns_doc[0]['description'])
|
184 |
+
elif returns_named_value:
|
185 |
+
return_data = returns_doc
|
186 |
+
if return_data:
|
187 |
+
desc['return_data'] = return_data
|
188 |
+
return desc
|
189 |
+
|
190 |
+
if callable(func):
|
191 |
+
|
192 |
+
if inspect.iscoroutinefunction(func):
|
193 |
+
|
194 |
+
@wraps(func)
|
195 |
+
async def wrapper(self, *args, **kwargs):
|
196 |
+
return await func(self, *args, **kwargs)
|
197 |
+
|
198 |
+
else:
|
199 |
+
|
200 |
+
@wraps(func)
|
201 |
+
def wrapper(self, *args, **kwargs):
|
202 |
+
return func(self, *args, **kwargs)
|
203 |
+
|
204 |
+
wrapper.api_description = _parse_tool(func)
|
205 |
+
return wrapper
|
206 |
+
|
207 |
+
def decorate(func):
|
208 |
+
|
209 |
+
if inspect.iscoroutinefunction(func):
|
210 |
+
|
211 |
+
@wraps(func)
|
212 |
+
async def wrapper(self, *args, **kwargs):
|
213 |
+
return await func(self, *args, **kwargs)
|
214 |
+
|
215 |
+
else:
|
216 |
+
|
217 |
+
@wraps(func)
|
218 |
+
def wrapper(self, *args, **kwargs):
|
219 |
+
return func(self, *args, **kwargs)
|
220 |
+
|
221 |
+
wrapper.api_description = _parse_tool(func)
|
222 |
+
return wrapper
|
223 |
+
|
224 |
+
return decorate
|
225 |
+
|
226 |
+
|
227 |
+
class ToolMeta(ABCMeta):
|
228 |
+
"""Metaclass of tools."""
|
229 |
+
|
230 |
+
def __new__(mcs, name, base, attrs):
|
231 |
+
is_toolkit, tool_desc = True, dict(
|
232 |
+
name=name,
|
233 |
+
description=Docstring(attrs.get('__doc__',
|
234 |
+
'')).parse('google')[0].value)
|
235 |
+
for key, value in attrs.items():
|
236 |
+
if callable(value) and hasattr(value, 'api_description'):
|
237 |
+
api_desc = getattr(value, 'api_description')
|
238 |
+
if key == 'run':
|
239 |
+
tool_desc['parameters'] = api_desc['parameters']
|
240 |
+
tool_desc['required'] = api_desc['required']
|
241 |
+
if api_desc['description']:
|
242 |
+
tool_desc['description'] = api_desc['description']
|
243 |
+
if api_desc.get('return_data'):
|
244 |
+
tool_desc['return_data'] = api_desc['return_data']
|
245 |
+
is_toolkit = False
|
246 |
+
else:
|
247 |
+
tool_desc.setdefault('api_list', []).append(api_desc)
|
248 |
+
if not is_toolkit and 'api_list' in tool_desc:
|
249 |
+
raise KeyError('`run` and other tool APIs can not be implemented '
|
250 |
+
'at the same time')
|
251 |
+
if is_toolkit and 'api_list' not in tool_desc:
|
252 |
+
is_toolkit = False
|
253 |
+
if callable(attrs.get('run')):
|
254 |
+
run_api = tool_api(attrs['run'])
|
255 |
+
api_desc = run_api.api_description
|
256 |
+
tool_desc['parameters'] = api_desc['parameters']
|
257 |
+
tool_desc['required'] = api_desc['required']
|
258 |
+
if api_desc['description']:
|
259 |
+
tool_desc['description'] = api_desc['description']
|
260 |
+
if api_desc.get('return_data'):
|
261 |
+
tool_desc['return_data'] = api_desc['return_data']
|
262 |
+
attrs['run'] = run_api
|
263 |
+
else:
|
264 |
+
tool_desc['parameters'], tool_desc['required'] = [], []
|
265 |
+
attrs['_is_toolkit'] = is_toolkit
|
266 |
+
attrs['__tool_description__'] = tool_desc
|
267 |
+
return super().__new__(mcs, name, base, attrs)
|
268 |
+
|
269 |
+
|
270 |
+
class BaseAction(metaclass=ToolMeta):
|
271 |
+
"""Base class for all actions.
|
272 |
+
|
273 |
+
Args:
|
274 |
+
description (:class:`Optional[dict]`): The description of the action.
|
275 |
+
Defaults to ``None``.
|
276 |
+
parser (:class:`Type[BaseParser]`): The parser class to process the
|
277 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
278 |
+
|
279 |
+
Examples:
|
280 |
+
|
281 |
+
* simple tool
|
282 |
+
|
283 |
+
.. code-block:: python
|
284 |
+
|
285 |
+
class Bold(BaseAction):
|
286 |
+
'''Make text bold'''
|
287 |
+
|
288 |
+
def run(self, text: str):
|
289 |
+
'''
|
290 |
+
Args:
|
291 |
+
text (str): input text
|
292 |
+
|
293 |
+
Returns:
|
294 |
+
str: bold text
|
295 |
+
'''
|
296 |
+
return '**' + text + '**'
|
297 |
+
|
298 |
+
action = Bold()
|
299 |
+
|
300 |
+
* toolkit with multiple APIs
|
301 |
+
|
302 |
+
.. code-block:: python
|
303 |
+
|
304 |
+
class Calculator(BaseAction):
|
305 |
+
'''Calculator'''
|
306 |
+
|
307 |
+
@tool_api
|
308 |
+
def add(self, a, b):
|
309 |
+
'''Add operation
|
310 |
+
|
311 |
+
Args:
|
312 |
+
a (int): augend
|
313 |
+
b (int): addend
|
314 |
+
|
315 |
+
Returns:
|
316 |
+
int: sum
|
317 |
+
'''
|
318 |
+
return a + b
|
319 |
+
|
320 |
+
@tool_api
|
321 |
+
def sub(self, a, b):
|
322 |
+
'''Subtraction operation
|
323 |
+
|
324 |
+
Args:
|
325 |
+
a (int): minuend
|
326 |
+
b (int): subtrahend
|
327 |
+
|
328 |
+
Returns:
|
329 |
+
int: difference
|
330 |
+
'''
|
331 |
+
return a - b
|
332 |
+
|
333 |
+
action = Calculator()
|
334 |
+
"""
|
335 |
+
|
336 |
+
def __init__(
|
337 |
+
self,
|
338 |
+
description: Optional[dict] = None,
|
339 |
+
parser: Type[BaseParser] = JsonParser,
|
340 |
+
):
|
341 |
+
self._description = deepcopy(description or self.__tool_description__)
|
342 |
+
self._name = self._description['name']
|
343 |
+
self._parser = parser(self)
|
344 |
+
|
345 |
+
def __call__(self, inputs: str, name='run') -> ActionReturn:
|
346 |
+
fallback_args = {'inputs': inputs, 'name': name}
|
347 |
+
if not hasattr(self, name):
|
348 |
+
return ActionReturn(
|
349 |
+
fallback_args,
|
350 |
+
type=self.name,
|
351 |
+
errmsg=f'invalid API: {name}',
|
352 |
+
state=ActionStatusCode.API_ERROR)
|
353 |
+
try:
|
354 |
+
inputs = self._parser.parse_inputs(inputs, name)
|
355 |
+
except ParseError as exc:
|
356 |
+
return ActionReturn(
|
357 |
+
fallback_args,
|
358 |
+
type=self.name,
|
359 |
+
errmsg=exc.err_msg,
|
360 |
+
state=ActionStatusCode.ARGS_ERROR)
|
361 |
+
try:
|
362 |
+
outputs = getattr(self, name)(**inputs)
|
363 |
+
except Exception as exc:
|
364 |
+
return ActionReturn(
|
365 |
+
inputs,
|
366 |
+
type=self.name,
|
367 |
+
errmsg=str(exc),
|
368 |
+
state=ActionStatusCode.API_ERROR)
|
369 |
+
if isinstance(outputs, ActionReturn):
|
370 |
+
action_return = outputs
|
371 |
+
if not action_return.args:
|
372 |
+
action_return.args = inputs
|
373 |
+
if not action_return.type:
|
374 |
+
action_return.type = self.name
|
375 |
+
else:
|
376 |
+
result = self._parser.parse_outputs(outputs)
|
377 |
+
action_return = ActionReturn(inputs, type=self.name, result=result)
|
378 |
+
return action_return
|
379 |
+
|
380 |
+
@property
|
381 |
+
def name(self):
|
382 |
+
return self._name
|
383 |
+
|
384 |
+
@property
|
385 |
+
def is_toolkit(self):
|
386 |
+
return self._is_toolkit
|
387 |
+
|
388 |
+
@property
|
389 |
+
def description(self) -> dict:
|
390 |
+
"""Description of the tool."""
|
391 |
+
return self._description
|
392 |
+
|
393 |
+
def __repr__(self):
|
394 |
+
return f'{self.description}'
|
395 |
+
|
396 |
+
__str__ = __repr__
|
397 |
+
|
398 |
+
|
399 |
+
class AsyncActionMixin:
|
400 |
+
|
401 |
+
async def __call__(self, inputs: str, name='run') -> ActionReturn:
|
402 |
+
fallback_args = {'inputs': inputs, 'name': name}
|
403 |
+
if not hasattr(self, name):
|
404 |
+
return ActionReturn(
|
405 |
+
fallback_args,
|
406 |
+
type=self.name,
|
407 |
+
errmsg=f'invalid API: {name}',
|
408 |
+
state=ActionStatusCode.API_ERROR)
|
409 |
+
try:
|
410 |
+
inputs = self._parser.parse_inputs(inputs, name)
|
411 |
+
except ParseError as exc:
|
412 |
+
return ActionReturn(
|
413 |
+
fallback_args,
|
414 |
+
type=self.name,
|
415 |
+
errmsg=exc.err_msg,
|
416 |
+
state=ActionStatusCode.ARGS_ERROR)
|
417 |
+
try:
|
418 |
+
outputs = await getattr(self, name)(**inputs)
|
419 |
+
except Exception as exc:
|
420 |
+
return ActionReturn(
|
421 |
+
inputs,
|
422 |
+
type=self.name,
|
423 |
+
errmsg=str(exc),
|
424 |
+
state=ActionStatusCode.API_ERROR)
|
425 |
+
if isinstance(outputs, ActionReturn):
|
426 |
+
action_return = outputs
|
427 |
+
if not action_return.args:
|
428 |
+
action_return.args = inputs
|
429 |
+
if not action_return.type:
|
430 |
+
action_return.type = self.name
|
431 |
+
else:
|
432 |
+
result = self._parser.parse_outputs(outputs)
|
433 |
+
action_return = ActionReturn(inputs, type=self.name, result=result)
|
434 |
+
return action_return
|
lagent/actions/bing_map.py
ADDED
@@ -0,0 +1,268 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# flake8: noqa: E501
|
2 |
+
import json
|
3 |
+
import os
|
4 |
+
from typing import Optional, Type
|
5 |
+
|
6 |
+
import aiohttp
|
7 |
+
import requests
|
8 |
+
|
9 |
+
from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
|
10 |
+
from lagent.actions.parser import BaseParser, JsonParser
|
11 |
+
|
12 |
+
|
13 |
+
class BINGMap(BaseAction):
|
14 |
+
"""BING Map plugin for looking up map information."""
|
15 |
+
|
16 |
+
def __init__(
|
17 |
+
self,
|
18 |
+
key: Optional[str] = None,
|
19 |
+
description: Optional[dict] = None,
|
20 |
+
parser: Type[BaseParser] = JsonParser,
|
21 |
+
) -> None:
|
22 |
+
super().__init__(description, parser)
|
23 |
+
key = os.environ.get('BING_MAP_KEY', key)
|
24 |
+
if key is None:
|
25 |
+
raise ValueError(
|
26 |
+
'Please set BING Map API key either in the environment '
|
27 |
+
'as BING_MAP_KEY or pass it as `key` parameter.')
|
28 |
+
self.key = key
|
29 |
+
self.base_url = 'http://dev.virtualearth.net/REST/V1/'
|
30 |
+
|
31 |
+
@tool_api(explode_return=True)
|
32 |
+
def get_distance(self, start: str, end: str) -> dict:
|
33 |
+
"""Get the distance between two locations in km.
|
34 |
+
|
35 |
+
Args:
|
36 |
+
start (:class:`str`): The start location
|
37 |
+
end (:class:`str`): The end location
|
38 |
+
|
39 |
+
Returns:
|
40 |
+
:class:`dict`: distance information
|
41 |
+
* distance (str): the distance in km.
|
42 |
+
"""
|
43 |
+
# Request URL
|
44 |
+
url = self.base_url + 'Routes/Driving?o=json&wp.0=' + start + '&wp.1=' + end + '&key=' + self.key
|
45 |
+
# GET request
|
46 |
+
r = requests.get(url)
|
47 |
+
# TODO check request status?
|
48 |
+
data = json.loads(r.text)
|
49 |
+
# Extract route information
|
50 |
+
route = data['resourceSets'][0]['resources'][0]
|
51 |
+
# Extract distance in miles
|
52 |
+
distance = route['travelDistance']
|
53 |
+
return dict(distance=distance)
|
54 |
+
|
55 |
+
@tool_api(explode_return=True)
|
56 |
+
def get_route(self, start: str, end: str) -> dict:
|
57 |
+
"""Get the route between two locations in km.
|
58 |
+
|
59 |
+
Args:
|
60 |
+
start (:class:`str`): The start location
|
61 |
+
end (:class:`str`): The end location
|
62 |
+
|
63 |
+
Returns:
|
64 |
+
:class:`dict`: route information
|
65 |
+
* route (list): the route, a list of actions.
|
66 |
+
"""
|
67 |
+
# Request URL
|
68 |
+
url = self.base_url + 'Routes/Driving?o=json&wp.0=' + start + '&wp.1=' + end + '&key=' + self.key
|
69 |
+
# GET request
|
70 |
+
r = requests.get(url)
|
71 |
+
data = json.loads(r.text)
|
72 |
+
# Extract route information
|
73 |
+
route = data['resourceSets'][0]['resources'][0]
|
74 |
+
itinerary = route['routeLegs'][0]['itineraryItems']
|
75 |
+
# Extract route text information
|
76 |
+
route_text = []
|
77 |
+
for item in itinerary:
|
78 |
+
if 'instruction' in item:
|
79 |
+
route_text.append(item['instruction']['text'])
|
80 |
+
return dict(route=route_text)
|
81 |
+
|
82 |
+
@tool_api(explode_return=True)
|
83 |
+
def get_coordinates(self, location: str) -> dict:
|
84 |
+
"""Get the coordinates of a location.
|
85 |
+
|
86 |
+
Args:
|
87 |
+
location (:class:`str`): the location need to get coordinates.
|
88 |
+
|
89 |
+
Returns:
|
90 |
+
:class:`dict`: coordinates information
|
91 |
+
* latitude (float): the latitude of the location.
|
92 |
+
* longitude (float): the longitude of the location.
|
93 |
+
"""
|
94 |
+
url = self.base_url + 'Locations'
|
95 |
+
params = {'query': location, 'key': self.key}
|
96 |
+
response = requests.get(url, params=params)
|
97 |
+
json_data = response.json()
|
98 |
+
coordinates = json_data['resourceSets'][0]['resources'][0]['point'][
|
99 |
+
'coordinates']
|
100 |
+
return dict(latitude=coordinates[0], longitude=coordinates[1])
|
101 |
+
|
102 |
+
@tool_api(explode_return=True)
|
103 |
+
def search_nearby(self,
|
104 |
+
search_term: str,
|
105 |
+
places: str = 'unknown',
|
106 |
+
latitude: float = 0.0,
|
107 |
+
longitude: float = 0.0,
|
108 |
+
radius: int = 5000) -> dict:
|
109 |
+
"""Search for places nearby a location, within a given radius, and return the results into a list. You can use either the places name or the latitude and longitude.
|
110 |
+
|
111 |
+
Args:
|
112 |
+
search_term (:class:`str`): the place name.
|
113 |
+
places (:class:`str`): the name of the location. Defaults to ``'unknown'``.
|
114 |
+
latitude (:class:`float`): the latitude of the location. Defaults to ``0.0``.
|
115 |
+
longitude (:class:`float`): the longitude of the location. Defaults to ``0.0``.
|
116 |
+
radius (:class:`int`): radius in meters. Defaults to ``5000``.
|
117 |
+
|
118 |
+
Returns:
|
119 |
+
:class:`dict`: places information
|
120 |
+
* places (list): the list of places, each place is a dict with name and address, at most 5 places.
|
121 |
+
"""
|
122 |
+
url = self.base_url + 'LocalSearch'
|
123 |
+
if places != 'unknown':
|
124 |
+
pos = self.get_coordinates(**{'location': places})
|
125 |
+
latitude, longitude = pos[1]['latitude'], pos[1]['longitude']
|
126 |
+
# Build the request query string
|
127 |
+
params = {
|
128 |
+
'query': search_term,
|
129 |
+
'userLocation': f'{latitude},{longitude}',
|
130 |
+
'radius': radius,
|
131 |
+
'key': self.key
|
132 |
+
}
|
133 |
+
# Make the request
|
134 |
+
response = requests.get(url, params=params)
|
135 |
+
# Parse the response
|
136 |
+
response_data = json.loads(response.content)
|
137 |
+
# Get the results
|
138 |
+
results = response_data['resourceSets'][0]['resources']
|
139 |
+
addresses = []
|
140 |
+
for result in results:
|
141 |
+
name = result['name']
|
142 |
+
address = result['Address']['formattedAddress']
|
143 |
+
addresses.append(dict(name=name, address=address))
|
144 |
+
if len(addresses) == 5:
|
145 |
+
break
|
146 |
+
return dict(place=addresses)
|
147 |
+
|
148 |
+
|
149 |
+
class AsyncBINGMap(AsyncActionMixin, BINGMap):
|
150 |
+
"""BING Map plugin for looking up map information."""
|
151 |
+
|
152 |
+
@tool_api(explode_return=True)
|
153 |
+
async def get_distance(self, start: str, end: str) -> dict:
|
154 |
+
"""Get the distance between two locations in km.
|
155 |
+
|
156 |
+
Args:
|
157 |
+
start (:class:`str`): The start location
|
158 |
+
end (:class:`str`): The end location
|
159 |
+
|
160 |
+
Returns:
|
161 |
+
:class:`dict`: distance information
|
162 |
+
* distance (str): the distance in km.
|
163 |
+
"""
|
164 |
+
# Request URL
|
165 |
+
url = self.base_url + 'Routes/Driving?o=json&wp.0=' + start + '&wp.1=' + end + '&key=' + self.key
|
166 |
+
# GET request
|
167 |
+
async with aiohttp.ClientSession() as session:
|
168 |
+
async with session.get(url) as resp:
|
169 |
+
# TODO check request status?
|
170 |
+
data = await resp.json()
|
171 |
+
# Extract route information
|
172 |
+
route = data['resourceSets'][0]['resources'][0]
|
173 |
+
# Extract distance in miles
|
174 |
+
distance = route['travelDistance']
|
175 |
+
return dict(distance=distance)
|
176 |
+
|
177 |
+
@tool_api(explode_return=True)
|
178 |
+
async def get_route(self, start: str, end: str) -> dict:
|
179 |
+
"""Get the route between two locations in km.
|
180 |
+
|
181 |
+
Args:
|
182 |
+
start (:class:`str`): The start location
|
183 |
+
end (:class:`str`): The end location
|
184 |
+
|
185 |
+
Returns:
|
186 |
+
:class:`dict`: route information
|
187 |
+
* route (list): the route, a list of actions.
|
188 |
+
"""
|
189 |
+
# Request URL
|
190 |
+
url = self.base_url + 'Routes/Driving?o=json&wp.0=' + start + '&wp.1=' + end + '&key=' + self.key
|
191 |
+
# GET request
|
192 |
+
async with aiohttp.ClientSession() as session:
|
193 |
+
async with session.get(url) as resp:
|
194 |
+
data = await resp.json()
|
195 |
+
# Extract route information
|
196 |
+
route = data['resourceSets'][0]['resources'][0]
|
197 |
+
itinerary = route['routeLegs'][0]['itineraryItems']
|
198 |
+
# Extract route text information
|
199 |
+
route_text = []
|
200 |
+
for item in itinerary:
|
201 |
+
if 'instruction' in item:
|
202 |
+
route_text.append(item['instruction']['text'])
|
203 |
+
return dict(route=route_text)
|
204 |
+
|
205 |
+
@tool_api(explode_return=True)
|
206 |
+
async def get_coordinates(self, location: str) -> dict:
|
207 |
+
"""Get the coordinates of a location.
|
208 |
+
|
209 |
+
Args:
|
210 |
+
location (:class:`str`): the location need to get coordinates.
|
211 |
+
|
212 |
+
Returns:
|
213 |
+
:class:`dict`: coordinates information
|
214 |
+
* latitude (float): the latitude of the location.
|
215 |
+
* longitude (float): the longitude of the location.
|
216 |
+
"""
|
217 |
+
url = self.base_url + 'Locations'
|
218 |
+
params = {'query': location, 'key': self.key}
|
219 |
+
async with aiohttp.ClientSession() as session:
|
220 |
+
async with session.get(url, params=params) as resp:
|
221 |
+
data = await resp.json()
|
222 |
+
coordinates = data['resourceSets'][0]['resources'][0]['point'][
|
223 |
+
'coordinates']
|
224 |
+
return dict(latitude=coordinates[0], longitude=coordinates[1])
|
225 |
+
|
226 |
+
@tool_api(explode_return=True)
|
227 |
+
async def search_nearby(self,
|
228 |
+
search_term: str,
|
229 |
+
places: str = 'unknown',
|
230 |
+
latitude: float = 0.0,
|
231 |
+
longitude: float = 0.0,
|
232 |
+
radius: int = 5000) -> dict:
|
233 |
+
"""Search for places nearby a location, within a given radius, and return the results into a list. You can use either the places name or the latitude and longitude.
|
234 |
+
|
235 |
+
Args:
|
236 |
+
search_term (:class:`str`): the place name.
|
237 |
+
places (:class:`str`): the name of the location. Defaults to ``'unknown'``.
|
238 |
+
latitude (:class:`float`): the latitude of the location. Defaults to ``0.0``.
|
239 |
+
longitude (:class:`float`): the longitude of the location. Defaults to ``0.0``.
|
240 |
+
radius (:class:`int`): radius in meters. Defaults to ``5000``.
|
241 |
+
|
242 |
+
Returns:
|
243 |
+
:class:`dict`: places information
|
244 |
+
* places (list): the list of places, each place is a dict with name and address, at most 5 places.
|
245 |
+
"""
|
246 |
+
url = self.base_url + 'LocalSearch'
|
247 |
+
if places != 'unknown':
|
248 |
+
pos = self.get_coordinates(**{'location': places})
|
249 |
+
latitude, longitude = pos[1]['latitude'], pos[1]['longitude']
|
250 |
+
# Build the request query string
|
251 |
+
params = {
|
252 |
+
'query': search_term,
|
253 |
+
'userLocation': f'{latitude},{longitude}',
|
254 |
+
'radius': radius,
|
255 |
+
'key': self.key
|
256 |
+
}
|
257 |
+
async with aiohttp.ClientSession() as session:
|
258 |
+
async with session.get(url, params=params) as resp:
|
259 |
+
data = await resp.json()
|
260 |
+
results = data['resourceSets'][0]['resources']
|
261 |
+
addresses = []
|
262 |
+
for result in results:
|
263 |
+
name = result['name']
|
264 |
+
address = result['Address']['formattedAddress']
|
265 |
+
addresses.append(dict(name=name, address=address))
|
266 |
+
if len(addresses) == 5:
|
267 |
+
break
|
268 |
+
return dict(place=addresses)
|
lagent/actions/builtin_actions.py
ADDED
@@ -0,0 +1,109 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Optional
|
2 |
+
|
3 |
+
from lagent.actions.base_action import BaseAction, tool_api
|
4 |
+
from lagent.actions.parser import BaseParser
|
5 |
+
from lagent.schema import ActionReturn, ActionStatusCode, ActionValidCode
|
6 |
+
|
7 |
+
|
8 |
+
class InvalidAction(BaseAction):
|
9 |
+
"""This is a invalid action class, which is used to return error message
|
10 |
+
when the action is invalid.
|
11 |
+
|
12 |
+
Args:
|
13 |
+
err_msg (str): The error message. Defaults to 'The action is invalid,
|
14 |
+
please check the action name'.
|
15 |
+
|
16 |
+
Returns:
|
17 |
+
ActionReturn: The action return.
|
18 |
+
"""
|
19 |
+
|
20 |
+
def __init__(self,
|
21 |
+
err_msg:
|
22 |
+
str = 'The action is invalid, please check the action name.',
|
23 |
+
description: Optional[dict] = None,
|
24 |
+
parser=BaseParser) -> None:
|
25 |
+
super().__init__(description, parser)
|
26 |
+
self._err_msg = err_msg
|
27 |
+
|
28 |
+
@tool_api
|
29 |
+
def run(self, err_msg: Optional[str] = None) -> ActionReturn:
|
30 |
+
"""Return the error message.
|
31 |
+
|
32 |
+
Args:
|
33 |
+
err_msg (str, optional): The error message. If err_msg is not None,
|
34 |
+
it will be returned, otherwise the default error message will
|
35 |
+
be returned. Defaults to None.
|
36 |
+
"""
|
37 |
+
action_return = ActionReturn(
|
38 |
+
url=None,
|
39 |
+
args=dict(text=err_msg),
|
40 |
+
errmsg=err_msg or self._err_msg,
|
41 |
+
type=self.name,
|
42 |
+
valid=ActionValidCode.INVALID,
|
43 |
+
state=ActionStatusCode.API_ERROR)
|
44 |
+
return action_return
|
45 |
+
|
46 |
+
|
47 |
+
class NoAction(BaseAction):
|
48 |
+
"""This is a no action class, which is used to return error message when
|
49 |
+
the response does not follow the format.
|
50 |
+
|
51 |
+
Args:
|
52 |
+
err_msg (str): The error message. Defaults to
|
53 |
+
'Please follow the format'.
|
54 |
+
"""
|
55 |
+
|
56 |
+
def __init__(self,
|
57 |
+
err_msg: str = 'Please follow the format',
|
58 |
+
description: Optional[dict] = None,
|
59 |
+
parser=BaseParser):
|
60 |
+
super().__init__(description, parser)
|
61 |
+
self._err_msg = err_msg
|
62 |
+
|
63 |
+
@tool_api
|
64 |
+
def run(self, err_msg: Optional[str] = None) -> ActionReturn:
|
65 |
+
"""Return the error message.
|
66 |
+
|
67 |
+
Args:
|
68 |
+
err_msg (str, optional): The error message. If err_msg is not None,
|
69 |
+
it will be returned, otherwise the default error message will
|
70 |
+
be returned. Defaults to None.
|
71 |
+
|
72 |
+
Returns:
|
73 |
+
ActionReturn: The action return.
|
74 |
+
"""
|
75 |
+
action_return = ActionReturn(
|
76 |
+
url=None,
|
77 |
+
args=dict(text=err_msg),
|
78 |
+
type=self.name,
|
79 |
+
errmsg=err_msg or self._err_msg,
|
80 |
+
valid=ActionValidCode.INVALID,
|
81 |
+
state=ActionStatusCode.API_ERROR)
|
82 |
+
return action_return
|
83 |
+
|
84 |
+
|
85 |
+
class FinishAction(BaseAction):
|
86 |
+
"""This is a finish action class, which is used to return the final
|
87 |
+
result."""
|
88 |
+
|
89 |
+
def __init__(self, description: Optional[dict] = None, parser=BaseParser):
|
90 |
+
super().__init__(description, parser)
|
91 |
+
|
92 |
+
@tool_api
|
93 |
+
def run(self, response: str) -> ActionReturn:
|
94 |
+
"""Return the final result.
|
95 |
+
|
96 |
+
Args:
|
97 |
+
response (str): The final result.
|
98 |
+
|
99 |
+
Returns:
|
100 |
+
ActionReturn: The action return.
|
101 |
+
"""
|
102 |
+
action_return = ActionReturn(
|
103 |
+
url=None,
|
104 |
+
args=dict(text=response),
|
105 |
+
result=[dict(type='text', content=response)],
|
106 |
+
type=self.name,
|
107 |
+
valid=ActionValidCode.FINISH,
|
108 |
+
state=ActionStatusCode.SUCCESS)
|
109 |
+
return action_return
|
lagent/actions/google_scholar_search.py
ADDED
@@ -0,0 +1,438 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# flake8: noqa: E501
|
2 |
+
import os
|
3 |
+
from typing import Optional, Type
|
4 |
+
|
5 |
+
from asyncer import asyncify
|
6 |
+
|
7 |
+
from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
|
8 |
+
from lagent.schema import ActionReturn, ActionStatusCode
|
9 |
+
from .parser import BaseParser, JsonParser
|
10 |
+
|
11 |
+
|
12 |
+
class GoogleScholar(BaseAction):
|
13 |
+
"""Plugin for google scholar search.
|
14 |
+
|
15 |
+
Args:
|
16 |
+
api_key (str): API KEY to use serper google search API,
|
17 |
+
You can create a free API key at https://serper.dev.
|
18 |
+
description (dict): The description of the action. Defaults to ``None``.
|
19 |
+
parser (Type[BaseParser]): The parser class to process the
|
20 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
21 |
+
"""
|
22 |
+
|
23 |
+
def __init__(
|
24 |
+
self,
|
25 |
+
api_key: Optional[str] = None,
|
26 |
+
description: Optional[dict] = None,
|
27 |
+
parser: Type[BaseParser] = JsonParser,
|
28 |
+
):
|
29 |
+
super().__init__(description, parser)
|
30 |
+
api_key = os.environ.get('SERPER_API_KEY', api_key)
|
31 |
+
if api_key is None:
|
32 |
+
raise ValueError(
|
33 |
+
'Please set Serper API key either in the environment '
|
34 |
+
'as SERPER_API_KEY or pass it as `api_key` parameter.'
|
35 |
+
)
|
36 |
+
self.api_key = api_key
|
37 |
+
|
38 |
+
@tool_api(explode_return=True)
|
39 |
+
def search_google_scholar(
|
40 |
+
self,
|
41 |
+
query: str,
|
42 |
+
cites: Optional[str] = None,
|
43 |
+
as_ylo: Optional[int] = None,
|
44 |
+
as_yhi: Optional[int] = None,
|
45 |
+
scisbd: Optional[int] = None,
|
46 |
+
cluster: Optional[str] = None,
|
47 |
+
hl: Optional[str] = None,
|
48 |
+
lr: Optional[str] = None,
|
49 |
+
start: Optional[int] = None,
|
50 |
+
num: Optional[int] = None,
|
51 |
+
as_sdt: Optional[str] = None,
|
52 |
+
safe: Optional[str] = None,
|
53 |
+
filter: Optional[str] = None,
|
54 |
+
as_vis: Optional[str] = None,
|
55 |
+
) -> dict:
|
56 |
+
"""Search for scholarly articles based on a query according to the google scholar.
|
57 |
+
|
58 |
+
Args:
|
59 |
+
query (str): The query to search for.
|
60 |
+
cites (Optional[str]): The unique ID of an article for triggering "Cited By" searches.
|
61 |
+
as_ylo (Optional[int]): The starting year for results (e.g., if as_ylo=2018, results before this year will be omitted).
|
62 |
+
as_yhi (Optional[int]): The ending year for results (e.g., if as_yhi=2018, results after this year will be omitted).
|
63 |
+
scisbd (Optional[int]): Defines articles added in the last year, sorted by date. It can be set to 1 to include only abstracts, or 2 to include everything.
|
64 |
+
cluster (Optional[str]): The unique ID of an article for triggering "All Versions" searches.
|
65 |
+
hl (Optional[str]): The language to use for the Google Scholar search.
|
66 |
+
lr (Optional[str]): One or multiple languages to limit the search to.
|
67 |
+
start (Optional[int]): The result offset for pagination (0 is the first page of results, 10 is the 2nd page, etc.)
|
68 |
+
num (Optional[int]): The maximum number of results to return, limited to 20.
|
69 |
+
as_sdt (Optional[str]): Can be used either as a search type or a filter.
|
70 |
+
safe (Optional[str]): The level of filtering for adult content.
|
71 |
+
filter (Optional[str]): Defines if the filters for 'Similar Results' and 'Omitted Results' are on or off.
|
72 |
+
as_vis (Optional[str]): Defines whether to include citations or not.
|
73 |
+
|
74 |
+
Returns:
|
75 |
+
:class:`dict`: article information
|
76 |
+
- title: a list of the titles of the three selected papers
|
77 |
+
- cited_by: a list of the citation numbers of the three selected papers
|
78 |
+
- organic_id: a list of the organic results' ids of the three selected papers
|
79 |
+
- pub_info: publication information of selected papers
|
80 |
+
"""
|
81 |
+
from serpapi import GoogleSearch
|
82 |
+
|
83 |
+
params = {
|
84 |
+
'q': query,
|
85 |
+
'engine': 'google_scholar',
|
86 |
+
'api_key': self.api_key,
|
87 |
+
'cites': cites,
|
88 |
+
'as_ylo': as_ylo,
|
89 |
+
'as_yhi': as_yhi,
|
90 |
+
'scisbd': scisbd,
|
91 |
+
'cluster': cluster,
|
92 |
+
'hl': hl,
|
93 |
+
'lr': lr,
|
94 |
+
'start': start,
|
95 |
+
'num': num,
|
96 |
+
'as_sdt': as_sdt,
|
97 |
+
'safe': safe,
|
98 |
+
'filter': filter,
|
99 |
+
'as_vis': as_vis,
|
100 |
+
}
|
101 |
+
search = GoogleSearch(params)
|
102 |
+
try:
|
103 |
+
r = search.get_dict()
|
104 |
+
results = r['organic_results']
|
105 |
+
title = []
|
106 |
+
snippets = []
|
107 |
+
cited_by = []
|
108 |
+
organic_id = []
|
109 |
+
pub_info = []
|
110 |
+
for item in results[:3]:
|
111 |
+
title.append(item['title'])
|
112 |
+
pub_info.append(item['publication_info']['summary'])
|
113 |
+
citation = item['inline_links'].get('cited_by', {'total': ''})
|
114 |
+
cited_by.append(citation['total'])
|
115 |
+
snippets.append(item['snippet'])
|
116 |
+
organic_id.append(item['result_id'])
|
117 |
+
return dict(title=title, cited_by=cited_by, organic_id=organic_id, snippets=snippets)
|
118 |
+
except Exception as e:
|
119 |
+
return ActionReturn(errmsg=str(e), state=ActionStatusCode.HTTP_ERROR)
|
120 |
+
|
121 |
+
@tool_api(explode_return=True)
|
122 |
+
def get_author_information(
|
123 |
+
self,
|
124 |
+
author_id: str,
|
125 |
+
hl: Optional[str] = None,
|
126 |
+
view_op: Optional[str] = None,
|
127 |
+
sort: Optional[str] = None,
|
128 |
+
citation_id: Optional[str] = None,
|
129 |
+
start: Optional[int] = None,
|
130 |
+
num: Optional[int] = None,
|
131 |
+
no_cache: Optional[bool] = None,
|
132 |
+
async_req: Optional[bool] = None,
|
133 |
+
output: Optional[str] = None,
|
134 |
+
) -> dict:
|
135 |
+
"""Search for an author's information by author's id provided by get_author_id.
|
136 |
+
|
137 |
+
Args:
|
138 |
+
author_id (str): Required. The ID of an author.
|
139 |
+
hl (Optional[str]): The language to use for the Google Scholar Author search. Default is 'en'.
|
140 |
+
view_op (Optional[str]): Used for viewing specific parts of a page.
|
141 |
+
sort (Optional[str]): Used for sorting and refining articles.
|
142 |
+
citation_id (Optional[str]): Used for retrieving individual article citation.
|
143 |
+
start (Optional[int]): Defines the result offset. Default is 0.
|
144 |
+
num (Optional[int]): Defines the number of results to return. Default is 20.
|
145 |
+
no_cache (Optional[bool]): Forces SerpApi to fetch the results even if a cached version is already present. Default is False.
|
146 |
+
async_req (Optional[bool]): Defines the way you want to submit your search to SerpApi. Default is False.
|
147 |
+
output (Optional[str]): Defines the final output you want. Default is 'json'.
|
148 |
+
|
149 |
+
Returns:
|
150 |
+
:class:`dict`: author information
|
151 |
+
* name: author's name
|
152 |
+
* affliation: the affliation of the author
|
153 |
+
* articles: at most 3 articles by the author
|
154 |
+
* website: the author's homepage url
|
155 |
+
"""
|
156 |
+
from serpapi import GoogleSearch
|
157 |
+
|
158 |
+
params = {
|
159 |
+
'engine': 'google_scholar_author',
|
160 |
+
'author_id': author_id,
|
161 |
+
'api_key': self.api_key,
|
162 |
+
'hl': hl,
|
163 |
+
'view_op': view_op,
|
164 |
+
'sort': sort,
|
165 |
+
'citation_id': citation_id,
|
166 |
+
'start': start,
|
167 |
+
'num': num,
|
168 |
+
'no_cache': no_cache,
|
169 |
+
'async': async_req,
|
170 |
+
'output': output,
|
171 |
+
}
|
172 |
+
try:
|
173 |
+
search = GoogleSearch(params)
|
174 |
+
results = search.get_dict()
|
175 |
+
author = results['author']
|
176 |
+
articles = results.get('articles', [])
|
177 |
+
return dict(
|
178 |
+
name=author['name'],
|
179 |
+
affiliations=author.get('affiliations', ''),
|
180 |
+
website=author.get('website', ''),
|
181 |
+
articles=[dict(title=article['title'], authors=article['authors']) for article in articles[:3]],
|
182 |
+
)
|
183 |
+
except Exception as e:
|
184 |
+
return ActionReturn(errmsg=str(e), state=ActionStatusCode.HTTP_ERROR)
|
185 |
+
|
186 |
+
@tool_api(explode_return=True)
|
187 |
+
def get_citation_format(
|
188 |
+
self,
|
189 |
+
q: str,
|
190 |
+
no_cache: Optional[bool] = None,
|
191 |
+
async_: Optional[bool] = None,
|
192 |
+
output: Optional[str] = 'json',
|
193 |
+
) -> dict:
|
194 |
+
"""Function to get MLA citation format by an identification of organic_result's id provided by search_google_scholar.
|
195 |
+
|
196 |
+
Args:
|
197 |
+
q (str): ID of an individual Google Scholar organic search result.
|
198 |
+
no_cache (Optional[bool]): If set to True, will force SerpApi to fetch the Google Scholar Cite results even if a cached version is already present. Defaults to None.
|
199 |
+
async_ (Optional[bool]): If set to True, will submit search to SerpApi and retrieve results later. Defaults to None.
|
200 |
+
output (Optional[str]): Final output format. Set to 'json' to get a structured JSON of the results, or 'html' to get the raw html retrieved. Defaults to 'json'.
|
201 |
+
|
202 |
+
Returns:
|
203 |
+
:class:`dict`: citation format
|
204 |
+
* authors: the authors of the article
|
205 |
+
* citation: the citation format of the article
|
206 |
+
"""
|
207 |
+
from serpapi import GoogleSearch
|
208 |
+
|
209 |
+
params = {
|
210 |
+
'q': q,
|
211 |
+
'engine': 'google_scholar_cite',
|
212 |
+
'api_key': self.api_key,
|
213 |
+
'no_cache': no_cache,
|
214 |
+
'async': async_,
|
215 |
+
'output': output,
|
216 |
+
}
|
217 |
+
try:
|
218 |
+
search = GoogleSearch(params)
|
219 |
+
results = search.get_dict()
|
220 |
+
citation = results['citations']
|
221 |
+
citation_info = citation[0]['snippet']
|
222 |
+
return citation_info
|
223 |
+
except Exception as e:
|
224 |
+
return ActionReturn(errmsg=str(e), state=ActionStatusCode.HTTP_ERROR)
|
225 |
+
|
226 |
+
@tool_api(explode_return=True)
|
227 |
+
def get_author_id(
|
228 |
+
self,
|
229 |
+
mauthors: str,
|
230 |
+
hl: Optional[str] = 'en',
|
231 |
+
after_author: Optional[str] = None,
|
232 |
+
before_author: Optional[str] = None,
|
233 |
+
no_cache: Optional[bool] = False,
|
234 |
+
_async: Optional[bool] = False,
|
235 |
+
output: Optional[str] = 'json',
|
236 |
+
) -> dict:
|
237 |
+
"""The getAuthorId function is used to get the author's id by his or her name.
|
238 |
+
|
239 |
+
Args:
|
240 |
+
mauthors (str): Defines the author you want to search for.
|
241 |
+
hl (Optional[str]): Defines the language to use for the Google Scholar Profiles search. It's a two-letter language code. (e.g., 'en' for English, 'es' for Spanish, or 'fr' for French). Defaults to 'en'.
|
242 |
+
after_author (Optional[str]): Defines the next page token. It is used for retrieving the next page results. The parameter has the precedence over before_author parameter. Defaults to None.
|
243 |
+
before_author (Optional[str]): Defines the previous page token. It is used for retrieving the previous page results. Defaults to None.
|
244 |
+
no_cache (Optional[bool]): Will force SerpApi to fetch the Google Scholar Profiles results even if a cached version is already present. Defaults to False.
|
245 |
+
_async (Optional[bool]): Defines the way you want to submit your search to SerpApi. Defaults to False.
|
246 |
+
output (Optional[str]): Defines the final output you want. It can be set to 'json' (default) to get a structured JSON of the results, or 'html' to get the raw html retrieved. Defaults to 'json'.
|
247 |
+
|
248 |
+
Returns:
|
249 |
+
:class:`dict`: author id
|
250 |
+
* author_id: the author_id of the author
|
251 |
+
"""
|
252 |
+
from serpapi import GoogleSearch
|
253 |
+
|
254 |
+
params = {
|
255 |
+
'mauthors': mauthors,
|
256 |
+
'engine': 'google_scholar_profiles',
|
257 |
+
'api_key': self.api_key,
|
258 |
+
'hl': hl,
|
259 |
+
'after_author': after_author,
|
260 |
+
'before_author': before_author,
|
261 |
+
'no_cache': no_cache,
|
262 |
+
'async': _async,
|
263 |
+
'output': output,
|
264 |
+
}
|
265 |
+
try:
|
266 |
+
search = GoogleSearch(params)
|
267 |
+
results = search.get_dict()
|
268 |
+
profile = results['profiles']
|
269 |
+
author_info = dict(author_id=profile[0]['author_id'])
|
270 |
+
return author_info
|
271 |
+
except Exception as e:
|
272 |
+
return ActionReturn(errmsg=str(e), state=ActionStatusCode.HTTP_ERROR)
|
273 |
+
|
274 |
+
|
275 |
+
class AsyncGoogleScholar(AsyncActionMixin, GoogleScholar):
|
276 |
+
"""Plugin for google scholar search.
|
277 |
+
|
278 |
+
Args:
|
279 |
+
api_key (str): API KEY to use serper google search API,
|
280 |
+
You can create a free API key at https://serper.dev.
|
281 |
+
description (dict): The description of the action. Defaults to ``None``.
|
282 |
+
parser (Type[BaseParser]): The parser class to process the
|
283 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
284 |
+
"""
|
285 |
+
|
286 |
+
@tool_api(explode_return=True)
|
287 |
+
@asyncify
|
288 |
+
def search_google_scholar(
|
289 |
+
self,
|
290 |
+
query: str,
|
291 |
+
cites: Optional[str] = None,
|
292 |
+
as_ylo: Optional[int] = None,
|
293 |
+
as_yhi: Optional[int] = None,
|
294 |
+
scisbd: Optional[int] = None,
|
295 |
+
cluster: Optional[str] = None,
|
296 |
+
hl: Optional[str] = None,
|
297 |
+
lr: Optional[str] = None,
|
298 |
+
start: Optional[int] = None,
|
299 |
+
num: Optional[int] = None,
|
300 |
+
as_sdt: Optional[str] = None,
|
301 |
+
safe: Optional[str] = None,
|
302 |
+
filter: Optional[str] = None,
|
303 |
+
as_vis: Optional[str] = None,
|
304 |
+
) -> dict:
|
305 |
+
"""Search for scholarly articles based on a query according to the google scholar.
|
306 |
+
|
307 |
+
Args:
|
308 |
+
query (str): The query to search for.
|
309 |
+
cites (Optional[str]): The unique ID of an article for triggering "Cited By" searches.
|
310 |
+
as_ylo (Optional[int]): The starting year for results (e.g., if as_ylo=2018, results before this year will be omitted).
|
311 |
+
as_yhi (Optional[int]): The ending year for results (e.g., if as_yhi=2018, results after this year will be omitted).
|
312 |
+
scisbd (Optional[int]): Defines articles added in the last year, sorted by date. It can be set to 1 to include only abstracts, or 2 to include everything.
|
313 |
+
cluster (Optional[str]): The unique ID of an article for triggering "All Versions" searches.
|
314 |
+
hl (Optional[str]): The language to use for the Google Scholar search.
|
315 |
+
lr (Optional[str]): One or multiple languages to limit the search to.
|
316 |
+
start (Optional[int]): The result offset for pagination (0 is the first page of results, 10 is the 2nd page, etc.)
|
317 |
+
num (Optional[int]): The maximum number of results to return, limited to 20.
|
318 |
+
as_sdt (Optional[str]): Can be used either as a search type or a filter.
|
319 |
+
safe (Optional[str]): The level of filtering for adult content.
|
320 |
+
filter (Optional[str]): Defines if the filters for 'Similar Results' and 'Omitted Results' are on or off.
|
321 |
+
as_vis (Optional[str]): Defines whether to include citations or not.
|
322 |
+
|
323 |
+
Returns:
|
324 |
+
:class:`dict`: article information
|
325 |
+
- title: a list of the titles of the three selected papers
|
326 |
+
- cited_by: a list of the citation numbers of the three selected papers
|
327 |
+
- organic_id: a list of the organic results' ids of the three selected papers
|
328 |
+
- pub_info: publication information of selected papers
|
329 |
+
"""
|
330 |
+
return super().search_google_scholar(
|
331 |
+
query,
|
332 |
+
cites,
|
333 |
+
as_ylo,
|
334 |
+
as_yhi,
|
335 |
+
scisbd,
|
336 |
+
cluster,
|
337 |
+
hl,
|
338 |
+
lr,
|
339 |
+
start,
|
340 |
+
num,
|
341 |
+
as_sdt,
|
342 |
+
safe,
|
343 |
+
filter,
|
344 |
+
as_vis,
|
345 |
+
)
|
346 |
+
|
347 |
+
@tool_api(explode_return=True)
|
348 |
+
@asyncify
|
349 |
+
def get_author_information(
|
350 |
+
self,
|
351 |
+
author_id: str,
|
352 |
+
hl: Optional[str] = None,
|
353 |
+
view_op: Optional[str] = None,
|
354 |
+
sort: Optional[str] = None,
|
355 |
+
citation_id: Optional[str] = None,
|
356 |
+
start: Optional[int] = None,
|
357 |
+
num: Optional[int] = None,
|
358 |
+
no_cache: Optional[bool] = None,
|
359 |
+
async_req: Optional[bool] = None,
|
360 |
+
output: Optional[str] = None,
|
361 |
+
) -> dict:
|
362 |
+
"""Search for an author's information by author's id provided by get_author_id.
|
363 |
+
|
364 |
+
Args:
|
365 |
+
author_id (str): Required. The ID of an author.
|
366 |
+
hl (Optional[str]): The language to use for the Google Scholar Author search. Default is 'en'.
|
367 |
+
view_op (Optional[str]): Used for viewing specific parts of a page.
|
368 |
+
sort (Optional[str]): Used for sorting and refining articles.
|
369 |
+
citation_id (Optional[str]): Used for retrieving individual article citation.
|
370 |
+
start (Optional[int]): Defines the result offset. Default is 0.
|
371 |
+
num (Optional[int]): Defines the number of results to return. Default is 20.
|
372 |
+
no_cache (Optional[bool]): Forces SerpApi to fetch the results even if a cached version is already present. Default is False.
|
373 |
+
async_req (Optional[bool]): Defines the way you want to submit your search to SerpApi. Default is False.
|
374 |
+
output (Optional[str]): Defines the final output you want. Default is 'json'.
|
375 |
+
|
376 |
+
Returns:
|
377 |
+
:class:`dict`: author information
|
378 |
+
* name: author's name
|
379 |
+
* affliation: the affliation of the author
|
380 |
+
* articles: at most 3 articles by the author
|
381 |
+
* website: the author's homepage url
|
382 |
+
"""
|
383 |
+
return super().get_author_information(
|
384 |
+
author_id, hl, view_op, sort, citation_id, start, num, no_cache, async_req, output
|
385 |
+
)
|
386 |
+
|
387 |
+
@tool_api(explode_return=True)
|
388 |
+
@asyncify
|
389 |
+
def get_citation_format(
|
390 |
+
self,
|
391 |
+
q: str,
|
392 |
+
no_cache: Optional[bool] = None,
|
393 |
+
async_: Optional[bool] = None,
|
394 |
+
output: Optional[str] = 'json',
|
395 |
+
) -> dict:
|
396 |
+
"""Function to get MLA citation format by an identification of organic_result's id provided by search_google_scholar.
|
397 |
+
|
398 |
+
Args:
|
399 |
+
q (str): ID of an individual Google Scholar organic search result.
|
400 |
+
no_cache (Optional[bool]): If set to True, will force SerpApi to fetch the Google Scholar Cite results even if a cached version is already present. Defaults to None.
|
401 |
+
async_ (Optional[bool]): If set to True, will submit search to SerpApi and retrieve results later. Defaults to None.
|
402 |
+
output (Optional[str]): Final output format. Set to 'json' to get a structured JSON of the results, or 'html' to get the raw html retrieved. Defaults to 'json'.
|
403 |
+
|
404 |
+
Returns:
|
405 |
+
:class:`dict`: citation format
|
406 |
+
* authors: the authors of the article
|
407 |
+
* citation: the citation format of the article
|
408 |
+
"""
|
409 |
+
return super().get_citation_format(q, no_cache, async_, output)
|
410 |
+
|
411 |
+
@tool_api(explode_return=True)
|
412 |
+
@asyncify
|
413 |
+
def get_author_id(
|
414 |
+
self,
|
415 |
+
mauthors: str,
|
416 |
+
hl: Optional[str] = 'en',
|
417 |
+
after_author: Optional[str] = None,
|
418 |
+
before_author: Optional[str] = None,
|
419 |
+
no_cache: Optional[bool] = False,
|
420 |
+
_async: Optional[bool] = False,
|
421 |
+
output: Optional[str] = 'json',
|
422 |
+
) -> dict:
|
423 |
+
"""The getAuthorId function is used to get the author's id by his or her name.
|
424 |
+
|
425 |
+
Args:
|
426 |
+
mauthors (str): Defines the author you want to search for.
|
427 |
+
hl (Optional[str]): Defines the language to use for the Google Scholar Profiles search. It's a two-letter language code. (e.g., 'en' for English, 'es' for Spanish, or 'fr' for French). Defaults to 'en'.
|
428 |
+
after_author (Optional[str]): Defines the next page token. It is used for retrieving the next page results. The parameter has the precedence over before_author parameter. Defaults to None.
|
429 |
+
before_author (Optional[str]): Defines the previous page token. It is used for retrieving the previous page results. Defaults to None.
|
430 |
+
no_cache (Optional[bool]): Will force SerpApi to fetch the Google Scholar Profiles results even if a cached version is already present. Defaults to False.
|
431 |
+
_async (Optional[bool]): Defines the way you want to submit your search to SerpApi. Defaults to False.
|
432 |
+
output (Optional[str]): Defines the final output you want. It can be set to 'json' (default) to get a structured JSON of the results, or 'html' to get the raw html retrieved. Defaults to 'json'.
|
433 |
+
|
434 |
+
Returns:
|
435 |
+
:class:`dict`: author id
|
436 |
+
* author_id: the author_id of the author
|
437 |
+
"""
|
438 |
+
return super().get_author_id(mauthors, hl, after_author, before_author, no_cache, _async, output)
|
lagent/actions/google_search.py
ADDED
@@ -0,0 +1,244 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
from typing import List, Optional, Tuple, Type, Union
|
3 |
+
|
4 |
+
import aiohttp
|
5 |
+
import requests
|
6 |
+
|
7 |
+
from lagent.schema import ActionReturn, ActionStatusCode
|
8 |
+
from .base_action import AsyncActionMixin, BaseAction, tool_api
|
9 |
+
from .parser import BaseParser, JsonParser
|
10 |
+
|
11 |
+
|
12 |
+
class GoogleSearch(BaseAction):
|
13 |
+
"""Wrapper around the Serper.dev Google Search API.
|
14 |
+
|
15 |
+
To use, you should pass your serper API key to the constructor.
|
16 |
+
|
17 |
+
Code is modified from lang-chain GoogleSerperAPIWrapper
|
18 |
+
(https://github.com/langchain-ai/langchain/blob/ba5f
|
19 |
+
baba704a2d729a4b8f568ed70d7c53e799bb/libs/langchain/
|
20 |
+
langchain/utilities/google_serper.py)
|
21 |
+
|
22 |
+
Args:
|
23 |
+
api_key (str): API KEY to use serper google search API,
|
24 |
+
You can create a free API key at https://serper.dev.
|
25 |
+
timeout (int): Upper bound of waiting time for a serper request.
|
26 |
+
search_type (str): Serper API support ['search', 'images', 'news',
|
27 |
+
'places'] types of search, currently we only support 'search'.
|
28 |
+
description (dict): The description of the action. Defaults to ``None``.
|
29 |
+
parser (Type[BaseParser]): The parser class to process the
|
30 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
31 |
+
"""
|
32 |
+
result_key_for_type = {
|
33 |
+
'news': 'news',
|
34 |
+
'places': 'places',
|
35 |
+
'images': 'images',
|
36 |
+
'search': 'organic',
|
37 |
+
}
|
38 |
+
|
39 |
+
def __init__(
|
40 |
+
self,
|
41 |
+
api_key: Optional[str] = None,
|
42 |
+
timeout: int = 5,
|
43 |
+
search_type: str = 'search',
|
44 |
+
description: Optional[dict] = None,
|
45 |
+
parser: Type[BaseParser] = JsonParser,
|
46 |
+
):
|
47 |
+
super().__init__(description, parser)
|
48 |
+
api_key = os.environ.get('SERPER_API_KEY', api_key)
|
49 |
+
if api_key is None:
|
50 |
+
raise ValueError(
|
51 |
+
'Please set Serper API key either in the environment '
|
52 |
+
'as SERPER_API_KEY or pass it as `api_key` parameter.')
|
53 |
+
self.api_key = api_key
|
54 |
+
self.timeout = timeout
|
55 |
+
self.search_type = search_type
|
56 |
+
|
57 |
+
@tool_api
|
58 |
+
def run(self, query: str, k: int = 10) -> ActionReturn:
|
59 |
+
"""一个可以从谷歌搜索结果的API。当你需要对于一个特定问题找到简短明了的回答时,可以使用它。输入应该是一个搜索查询。
|
60 |
+
|
61 |
+
Args:
|
62 |
+
query (str): the search content
|
63 |
+
k (int): select first k results in the search results as response
|
64 |
+
"""
|
65 |
+
tool_return = ActionReturn(type=self.name)
|
66 |
+
status_code, response = self._search(query, k=k)
|
67 |
+
# convert search results to ToolReturn format
|
68 |
+
if status_code == -1:
|
69 |
+
tool_return.errmsg = response
|
70 |
+
tool_return.state = ActionStatusCode.HTTP_ERROR
|
71 |
+
elif status_code == 200:
|
72 |
+
parsed_res = self._parse_results(response, k)
|
73 |
+
tool_return.result = [dict(type='text', content=str(parsed_res))]
|
74 |
+
tool_return.state = ActionStatusCode.SUCCESS
|
75 |
+
else:
|
76 |
+
tool_return.errmsg = str(status_code)
|
77 |
+
tool_return.state = ActionStatusCode.API_ERROR
|
78 |
+
return tool_return
|
79 |
+
|
80 |
+
def _parse_results(self, results: dict, k: int) -> Union[str, List[str]]:
|
81 |
+
"""Parse the search results from Serper API.
|
82 |
+
|
83 |
+
Args:
|
84 |
+
results (dict): The search content from Serper API
|
85 |
+
in json format.
|
86 |
+
|
87 |
+
Returns:
|
88 |
+
List[str]: The parsed search results.
|
89 |
+
"""
|
90 |
+
|
91 |
+
snippets = []
|
92 |
+
|
93 |
+
if results.get('answerBox'):
|
94 |
+
answer_box = results.get('answerBox', {})
|
95 |
+
if answer_box.get('answer'):
|
96 |
+
return [answer_box.get('answer')]
|
97 |
+
elif answer_box.get('snippet'):
|
98 |
+
return [answer_box.get('snippet').replace('\n', ' ')]
|
99 |
+
elif answer_box.get('snippetHighlighted'):
|
100 |
+
return answer_box.get('snippetHighlighted')
|
101 |
+
|
102 |
+
if results.get('knowledgeGraph'):
|
103 |
+
kg = results.get('knowledgeGraph', {})
|
104 |
+
title = kg.get('title')
|
105 |
+
entity_type = kg.get('type')
|
106 |
+
if entity_type:
|
107 |
+
snippets.append(f'{title}: {entity_type}.')
|
108 |
+
description = kg.get('description')
|
109 |
+
if description:
|
110 |
+
snippets.append(description)
|
111 |
+
for attribute, value in kg.get('attributes', {}).items():
|
112 |
+
snippets.append(f'{title} {attribute}: {value}.')
|
113 |
+
|
114 |
+
for result in results[self.result_key_for_type[
|
115 |
+
self.search_type]][:k]:
|
116 |
+
if 'snippet' in result:
|
117 |
+
snippets.append(result['snippet'])
|
118 |
+
for attribute, value in result.get('attributes', {}).items():
|
119 |
+
snippets.append(f'{attribute}: {value}.')
|
120 |
+
|
121 |
+
if len(snippets) == 0:
|
122 |
+
return ['No good Google Search Result was found']
|
123 |
+
return snippets
|
124 |
+
|
125 |
+
def _search(self,
|
126 |
+
search_term: str,
|
127 |
+
search_type: Optional[str] = None,
|
128 |
+
**kwargs) -> Tuple[int, Union[dict, str]]:
|
129 |
+
"""HTTP requests to Serper API.
|
130 |
+
|
131 |
+
Args:
|
132 |
+
search_term (str): The search query.
|
133 |
+
search_type (str): search type supported by Serper API,
|
134 |
+
default to 'search'.
|
135 |
+
|
136 |
+
Returns:
|
137 |
+
tuple: the return value is a tuple contains:
|
138 |
+
- status_code (int): HTTP status code from Serper API.
|
139 |
+
- response (dict): response context with json format.
|
140 |
+
"""
|
141 |
+
headers = {
|
142 |
+
'X-API-KEY': self.api_key or '',
|
143 |
+
'Content-Type': 'application/json',
|
144 |
+
}
|
145 |
+
params = {
|
146 |
+
'q': search_term,
|
147 |
+
**{
|
148 |
+
key: value
|
149 |
+
for key, value in kwargs.items() if value is not None
|
150 |
+
},
|
151 |
+
}
|
152 |
+
try:
|
153 |
+
response = requests.post(
|
154 |
+
f'https://google.serper.dev/{search_type or self.search_type}',
|
155 |
+
headers=headers,
|
156 |
+
params=params,
|
157 |
+
timeout=self.timeout)
|
158 |
+
except Exception as e:
|
159 |
+
return -1, str(e)
|
160 |
+
return response.status_code, response.json()
|
161 |
+
|
162 |
+
|
163 |
+
class AsyncGoogleSearch(AsyncActionMixin, GoogleSearch):
|
164 |
+
"""Wrapper around the Serper.dev Google Search API.
|
165 |
+
|
166 |
+
To use, you should pass your serper API key to the constructor.
|
167 |
+
|
168 |
+
Code is modified from lang-chain GoogleSerperAPIWrapper
|
169 |
+
(https://github.com/langchain-ai/langchain/blob/ba5f
|
170 |
+
baba704a2d729a4b8f568ed70d7c53e799bb/libs/langchain/
|
171 |
+
langchain/utilities/google_serper.py)
|
172 |
+
|
173 |
+
Args:
|
174 |
+
api_key (str): API KEY to use serper google search API,
|
175 |
+
You can create a free API key at https://serper.dev.
|
176 |
+
timeout (int): Upper bound of waiting time for a serper request.
|
177 |
+
search_type (str): Serper API support ['search', 'images', 'news',
|
178 |
+
'places'] types of search, currently we only support 'search'.
|
179 |
+
description (dict): The description of the action. Defaults to ``None``.
|
180 |
+
parser (Type[BaseParser]): The parser class to process the
|
181 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
182 |
+
"""
|
183 |
+
|
184 |
+
@tool_api
|
185 |
+
async def run(self, query: str, k: int = 10) -> ActionReturn:
|
186 |
+
"""一个可以从谷歌搜索结果的API。当你需要对于一个特定问题找到简短明了的回答时,可以使用它。输入应该是一个搜索查询。
|
187 |
+
|
188 |
+
Args:
|
189 |
+
query (str): the search content
|
190 |
+
k (int): select first k results in the search results as response
|
191 |
+
"""
|
192 |
+
tool_return = ActionReturn(type=self.name)
|
193 |
+
status_code, response = await self._search(query, k=k)
|
194 |
+
# convert search results to ToolReturn format
|
195 |
+
if status_code == -1:
|
196 |
+
tool_return.errmsg = response
|
197 |
+
tool_return.state = ActionStatusCode.HTTP_ERROR
|
198 |
+
elif status_code == 200:
|
199 |
+
parsed_res = self._parse_results(response)
|
200 |
+
tool_return.result = [dict(type='text', content=str(parsed_res))]
|
201 |
+
tool_return.state = ActionStatusCode.SUCCESS
|
202 |
+
else:
|
203 |
+
tool_return.errmsg = str(status_code)
|
204 |
+
tool_return.state = ActionStatusCode.API_ERROR
|
205 |
+
return tool_return
|
206 |
+
|
207 |
+
async def _search(self,
|
208 |
+
search_term: str,
|
209 |
+
search_type: Optional[str] = None,
|
210 |
+
**kwargs) -> Tuple[int, Union[dict, str]]:
|
211 |
+
"""HTTP requests to Serper API.
|
212 |
+
|
213 |
+
Args:
|
214 |
+
search_term (str): The search query.
|
215 |
+
search_type (str): search type supported by Serper API,
|
216 |
+
default to 'search'.
|
217 |
+
|
218 |
+
Returns:
|
219 |
+
tuple: the return value is a tuple contains:
|
220 |
+
- status_code (int): HTTP status code from Serper API.
|
221 |
+
- response (dict): response context with json format.
|
222 |
+
"""
|
223 |
+
headers = {
|
224 |
+
'X-API-KEY': self.api_key or '',
|
225 |
+
'Content-Type': 'application/json',
|
226 |
+
}
|
227 |
+
params = {
|
228 |
+
'q': search_term,
|
229 |
+
**{
|
230 |
+
key: value
|
231 |
+
for key, value in kwargs.items() if value is not None
|
232 |
+
},
|
233 |
+
}
|
234 |
+
timeout = aiohttp.ClientTimeout(total=self.timeout)
|
235 |
+
async with aiohttp.ClientSession(timeout=timeout) as session:
|
236 |
+
try:
|
237 |
+
async with session.post(
|
238 |
+
f'https://google.serper.dev/{search_type or self.search_type}',
|
239 |
+
headers=headers,
|
240 |
+
params=params) as resp:
|
241 |
+
code, ret = resp.status, await resp.json()
|
242 |
+
except aiohttp.ClientError as e:
|
243 |
+
code, ret = -1, str(e)
|
244 |
+
return code, ret
|
lagent/actions/ipython_interactive.py
ADDED
@@ -0,0 +1,273 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import re
|
2 |
+
import signal
|
3 |
+
from contextlib import contextmanager, redirect_stdout
|
4 |
+
from dataclasses import dataclass
|
5 |
+
from enum import Enum
|
6 |
+
from io import StringIO
|
7 |
+
from typing import Optional, Type
|
8 |
+
|
9 |
+
from ..schema import ActionReturn, ActionStatusCode
|
10 |
+
from .base_action import AsyncActionMixin, BaseAction, tool_api
|
11 |
+
from .parser import BaseParser, JsonParser
|
12 |
+
|
13 |
+
|
14 |
+
class Status(str, Enum):
|
15 |
+
"""Execution status."""
|
16 |
+
SUCCESS = 'success'
|
17 |
+
FAILURE = 'failure'
|
18 |
+
|
19 |
+
|
20 |
+
@dataclass
|
21 |
+
class ExecutionResult:
|
22 |
+
"""Execution result."""
|
23 |
+
status: Status
|
24 |
+
value: Optional[str] = None
|
25 |
+
msg: Optional[str] = None
|
26 |
+
|
27 |
+
|
28 |
+
@contextmanager
|
29 |
+
def _raise_timeout(timeout):
|
30 |
+
|
31 |
+
def _handler(signum, frame):
|
32 |
+
raise TimeoutError()
|
33 |
+
|
34 |
+
signal.signal(signal.SIGALRM, _handler)
|
35 |
+
signal.alarm(timeout)
|
36 |
+
|
37 |
+
try:
|
38 |
+
yield
|
39 |
+
finally:
|
40 |
+
signal.alarm(0)
|
41 |
+
|
42 |
+
|
43 |
+
class IPythonInteractive(BaseAction):
|
44 |
+
"""An interactive IPython shell for code execution.
|
45 |
+
|
46 |
+
Args:
|
47 |
+
timeout (int): Upper bound of waiting time for Python script execution.
|
48 |
+
Defaults to ``20``.
|
49 |
+
max_out_len (int): maximum output length. No truncation occurs if negative.
|
50 |
+
Defaults to ``2048``.
|
51 |
+
use_signals (bool): whether signals should be used for timing function out
|
52 |
+
or the multiprocessing. Set to ``False`` when not running in the main
|
53 |
+
thread, e.g. web applications. Defaults to ``True``
|
54 |
+
description (dict): The description of the action. Defaults to ``None``.
|
55 |
+
parser (Type[BaseParser]): The parser class to process the
|
56 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
57 |
+
"""
|
58 |
+
|
59 |
+
def __init__(
|
60 |
+
self,
|
61 |
+
timeout: int = 30,
|
62 |
+
max_out_len: int = 8192,
|
63 |
+
use_signals: bool = True,
|
64 |
+
description: Optional[dict] = None,
|
65 |
+
parser: Type[BaseParser] = JsonParser,
|
66 |
+
):
|
67 |
+
super().__init__(description, parser)
|
68 |
+
self.timeout = timeout
|
69 |
+
self._executor = self.create_shell()
|
70 |
+
self._highlighting = re.compile(
|
71 |
+
r'(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]')
|
72 |
+
self._max_out_len = max_out_len if max_out_len >= 0 else None
|
73 |
+
self._use_signals = use_signals
|
74 |
+
|
75 |
+
def reset(self):
|
76 |
+
"""Clear the context."""
|
77 |
+
self._executor.reset()
|
78 |
+
|
79 |
+
@tool_api
|
80 |
+
def run(self, command: str, timeout: Optional[int] = None) -> ActionReturn:
|
81 |
+
"""Launch an IPython Interactive Shell to execute code.
|
82 |
+
|
83 |
+
Args:
|
84 |
+
command (:class:`str`): Python code snippet
|
85 |
+
timeout (:class:`Optional[int]`): timeout for execution.
|
86 |
+
This argument only works in the main thread. Defaults to ``None``.
|
87 |
+
"""
|
88 |
+
from timeout_decorator import timeout as timer
|
89 |
+
tool_return = ActionReturn(args={'text': command}, type=self.name)
|
90 |
+
ret = (
|
91 |
+
timer(timeout or self.timeout)(self.exec)(command)
|
92 |
+
if self._use_signals else self.exec(command))
|
93 |
+
if ret.status is Status.SUCCESS:
|
94 |
+
tool_return.result = [{'type': 'text', 'content': ret.value}]
|
95 |
+
tool_return.state = ActionStatusCode.SUCCESS
|
96 |
+
else:
|
97 |
+
tool_return.errmsg = ret.msg
|
98 |
+
tool_return.state = ActionStatusCode.API_ERROR
|
99 |
+
return tool_return
|
100 |
+
|
101 |
+
def exec(self, code: str) -> ExecutionResult:
|
102 |
+
"""Run Python scripts in IPython shell.
|
103 |
+
|
104 |
+
Args:
|
105 |
+
code (:class:`str`): code block
|
106 |
+
|
107 |
+
Returns:
|
108 |
+
:py:class:`ExecutionResult`: execution result
|
109 |
+
"""
|
110 |
+
with StringIO() as io:
|
111 |
+
with redirect_stdout(io):
|
112 |
+
ret = self._executor.run_cell(self.extract_code(code))
|
113 |
+
result = ret.result
|
114 |
+
if result is not None:
|
115 |
+
return ExecutionResult(Status.SUCCESS,
|
116 |
+
str(result)[:self._max_out_len])
|
117 |
+
outs = io.getvalue().strip().split('\n')
|
118 |
+
if not outs:
|
119 |
+
return ExecutionResult(Status.SUCCESS, '')
|
120 |
+
for i, out in enumerate(outs):
|
121 |
+
if re.search('Error|Traceback', out, re.S):
|
122 |
+
if 'TimeoutError' in out:
|
123 |
+
return ExecutionResult(
|
124 |
+
Status.FAILURE,
|
125 |
+
msg=('The code interpreter encountered '
|
126 |
+
'a timeout error.'))
|
127 |
+
err_idx = i
|
128 |
+
break
|
129 |
+
else:
|
130 |
+
return ExecutionResult(Status.SUCCESS,
|
131 |
+
outs[-1].strip()[:self._max_out_len])
|
132 |
+
return ExecutionResult(
|
133 |
+
Status.FAILURE,
|
134 |
+
msg=self._highlighting.sub(
|
135 |
+
'', '\n'.join(outs[err_idx:])[:self._max_out_len]),
|
136 |
+
)
|
137 |
+
|
138 |
+
@staticmethod
|
139 |
+
def create_shell():
|
140 |
+
from IPython import InteractiveShell
|
141 |
+
from traitlets.config import Config
|
142 |
+
|
143 |
+
c = Config()
|
144 |
+
c.HistoryManager.enabled = False
|
145 |
+
c.HistoryManager.hist_file = ':memory:'
|
146 |
+
return InteractiveShell(
|
147 |
+
user_ns={'_raise_timeout': _raise_timeout}, config=c)
|
148 |
+
|
149 |
+
@staticmethod
|
150 |
+
def extract_code(text: str) -> str:
|
151 |
+
"""Extract Python code from markup languages.
|
152 |
+
|
153 |
+
Args:
|
154 |
+
text (:class:`str`): Markdown-formatted text
|
155 |
+
|
156 |
+
Returns:
|
157 |
+
:class:`str`: Python code
|
158 |
+
"""
|
159 |
+
import json5
|
160 |
+
|
161 |
+
# Match triple backtick blocks first
|
162 |
+
triple_match = re.search(r'```[^\n]*\n(.+?)```', text, re.DOTALL)
|
163 |
+
# Match single backtick blocks second
|
164 |
+
single_match = re.search(r'`([^`]*)`', text, re.DOTALL)
|
165 |
+
if triple_match:
|
166 |
+
text = triple_match.group(1)
|
167 |
+
elif single_match:
|
168 |
+
text = single_match.group(1)
|
169 |
+
else:
|
170 |
+
try:
|
171 |
+
text = json5.loads(text)['code']
|
172 |
+
except Exception:
|
173 |
+
pass
|
174 |
+
# If no code blocks found, return original text
|
175 |
+
return text
|
176 |
+
|
177 |
+
@staticmethod
|
178 |
+
def wrap_code_with_timeout(code: str, timeout: int) -> str:
|
179 |
+
if not code.strip():
|
180 |
+
return code
|
181 |
+
code = code.strip('\n').rstrip()
|
182 |
+
indent = len(code) - len(code.lstrip())
|
183 |
+
handle = ' ' * indent + f'with _raise_timeout({timeout}):\n'
|
184 |
+
block = '\n'.join([' ' + line for line in code.split('\n')])
|
185 |
+
wrapped_code = handle + block
|
186 |
+
last_line = code.split('\n')[-1]
|
187 |
+
is_expression = True
|
188 |
+
try:
|
189 |
+
compile(last_line.lstrip(), '<stdin>', 'eval')
|
190 |
+
except SyntaxError:
|
191 |
+
is_expression = False
|
192 |
+
if is_expression:
|
193 |
+
wrapped_code += '\n' * 5 + last_line
|
194 |
+
return wrapped_code
|
195 |
+
|
196 |
+
|
197 |
+
class AsyncIPythonInteractive(AsyncActionMixin, IPythonInteractive):
|
198 |
+
"""An interactive IPython shell for code execution.
|
199 |
+
|
200 |
+
Args:
|
201 |
+
timeout (int): Upper bound of waiting time for Python script execution.
|
202 |
+
Defaults to ``20``.
|
203 |
+
max_out_len (int): maximum output length. No truncation occurs if negative.
|
204 |
+
Defaults to ``2048``.
|
205 |
+
use_signals (bool): whether signals should be used for timing function out
|
206 |
+
or the multiprocessing. Set to ``False`` when not running in the main
|
207 |
+
thread, e.g. web applications. Defaults to ``True``
|
208 |
+
description (dict): The description of the action. Defaults to ``None``.
|
209 |
+
parser (Type[BaseParser]): The parser class to process the
|
210 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
211 |
+
"""
|
212 |
+
|
213 |
+
@tool_api
|
214 |
+
async def run(self,
|
215 |
+
command: str,
|
216 |
+
timeout: Optional[int] = None) -> ActionReturn:
|
217 |
+
"""Launch an IPython Interactive Shell to execute code.
|
218 |
+
|
219 |
+
Args:
|
220 |
+
command (:class:`str`): Python code snippet
|
221 |
+
timeout (:class:`Optional[int]`): timeout for execution.
|
222 |
+
This argument only works in the main thread. Defaults to ``None``.
|
223 |
+
"""
|
224 |
+
tool_return = ActionReturn(args={'text': command}, type=self.name)
|
225 |
+
ret = await self.exec(command, timeout)
|
226 |
+
if ret.status is Status.SUCCESS:
|
227 |
+
tool_return.result = [{'type': 'text', 'content': ret.value}]
|
228 |
+
tool_return.state = ActionStatusCode.SUCCESS
|
229 |
+
else:
|
230 |
+
tool_return.errmsg = ret.msg
|
231 |
+
tool_return.state = ActionStatusCode.API_ERROR
|
232 |
+
return tool_return
|
233 |
+
|
234 |
+
async def exec(self, code: str, timeout: int = None) -> ExecutionResult:
|
235 |
+
"""Asynchronously run Python scripts in IPython shell.
|
236 |
+
|
237 |
+
Args:
|
238 |
+
code (:class:`str`): code block
|
239 |
+
timeout (:class:`int`): max waiting time for code execution
|
240 |
+
|
241 |
+
Returns:
|
242 |
+
:py:class:`ExecutionResult`: execution result
|
243 |
+
"""
|
244 |
+
with StringIO() as io:
|
245 |
+
with redirect_stdout(io):
|
246 |
+
ret = await self._executor.run_cell_async(
|
247 |
+
# ret = await self.create_shell().run_cell_async(
|
248 |
+
self.wrap_code_with_timeout(
|
249 |
+
self.extract_code(code), timeout or self.timeout))
|
250 |
+
result = ret.result
|
251 |
+
if result is not None:
|
252 |
+
return ExecutionResult(Status.SUCCESS,
|
253 |
+
str(result)[:self._max_out_len])
|
254 |
+
outs = io.getvalue().strip().split('\n')
|
255 |
+
if not outs:
|
256 |
+
return ExecutionResult(Status.SUCCESS, '')
|
257 |
+
for i, out in enumerate(outs):
|
258 |
+
if re.search('Error|Traceback', out, re.S):
|
259 |
+
if 'TimeoutError' in out:
|
260 |
+
return ExecutionResult(
|
261 |
+
Status.FAILURE,
|
262 |
+
msg=('The code interpreter encountered a '
|
263 |
+
'timeout error.'))
|
264 |
+
err_idx = i
|
265 |
+
break
|
266 |
+
else:
|
267 |
+
return ExecutionResult(Status.SUCCESS,
|
268 |
+
outs[-1].strip()[:self._max_out_len])
|
269 |
+
return ExecutionResult(
|
270 |
+
Status.FAILURE,
|
271 |
+
msg=self._highlighting.sub(
|
272 |
+
'', '\n'.join(outs[err_idx:])[:self._max_out_len]),
|
273 |
+
)
|
lagent/actions/ipython_interpreter.py
ADDED
@@ -0,0 +1,584 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# flake8: noqa: E501
|
2 |
+
import asyncio
|
3 |
+
import base64
|
4 |
+
import io
|
5 |
+
import json
|
6 |
+
import logging
|
7 |
+
import os
|
8 |
+
import queue
|
9 |
+
import re
|
10 |
+
import signal
|
11 |
+
import sys
|
12 |
+
import tempfile
|
13 |
+
import traceback
|
14 |
+
import uuid
|
15 |
+
from typing import Optional, Tuple, Type
|
16 |
+
|
17 |
+
from jupyter_client import AsyncKernelClient, AsyncKernelManager, AsyncMultiKernelManager
|
18 |
+
from tenacity import retry, retry_if_result, stop_after_attempt, wait_fixed
|
19 |
+
|
20 |
+
from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
|
21 |
+
from lagent.actions.parser import BaseParser, JsonParser
|
22 |
+
from lagent.schema import ActionReturn, ActionStatusCode
|
23 |
+
|
24 |
+
logger = logging.getLogger(__name__)
|
25 |
+
|
26 |
+
START_CODE = """
|
27 |
+
def input(*args, **kwargs):
|
28 |
+
raise NotImplementedError('Python input() function is disabled.')
|
29 |
+
|
30 |
+
get_ipython().system = lambda *args: print('Assume we have this package, ! is disabled!')
|
31 |
+
{}
|
32 |
+
""" # noqa
|
33 |
+
|
34 |
+
|
35 |
+
class TimeoutError(Exception):
|
36 |
+
pass
|
37 |
+
|
38 |
+
|
39 |
+
class KernelDeath(Exception):
|
40 |
+
pass
|
41 |
+
|
42 |
+
|
43 |
+
async def async_run_code(
|
44 |
+
km: AsyncKernelManager,
|
45 |
+
code,
|
46 |
+
*,
|
47 |
+
interrupt_after=30,
|
48 |
+
iopub_timeout=40,
|
49 |
+
wait_for_ready_timeout=60,
|
50 |
+
shutdown_kernel=True,
|
51 |
+
):
|
52 |
+
assert iopub_timeout > interrupt_after
|
53 |
+
try:
|
54 |
+
|
55 |
+
async def get_iopub_msg_with_death_detection(kc: AsyncKernelClient,
|
56 |
+
*,
|
57 |
+
timeout=None):
|
58 |
+
loop = asyncio.get_running_loop()
|
59 |
+
dead_fut = loop.create_future()
|
60 |
+
|
61 |
+
def restarting():
|
62 |
+
assert (
|
63 |
+
False
|
64 |
+
), "Restart shouldn't happen because config.KernelRestarter.restart_limit is expected to be set to 0"
|
65 |
+
|
66 |
+
def dead():
|
67 |
+
logger.info("Kernel has died, will NOT restart")
|
68 |
+
dead_fut.set_result(None)
|
69 |
+
|
70 |
+
msg_task = asyncio.create_task(kc.get_iopub_msg(timeout=timeout))
|
71 |
+
km.add_restart_callback(restarting, "restart")
|
72 |
+
km.add_restart_callback(dead, "dead")
|
73 |
+
try:
|
74 |
+
done, _ = await asyncio.wait(
|
75 |
+
[dead_fut, msg_task], return_when=asyncio.FIRST_COMPLETED)
|
76 |
+
if dead_fut in done:
|
77 |
+
raise KernelDeath()
|
78 |
+
assert msg_task in done
|
79 |
+
return await msg_task
|
80 |
+
finally:
|
81 |
+
msg_task.cancel()
|
82 |
+
km.remove_restart_callback(restarting, "restart")
|
83 |
+
km.remove_restart_callback(dead, "dead")
|
84 |
+
|
85 |
+
async def send_interrupt():
|
86 |
+
await asyncio.sleep(interrupt_after)
|
87 |
+
logger.info("Sending interrupt to kernel")
|
88 |
+
await km.interrupt_kernel()
|
89 |
+
|
90 |
+
@retry(
|
91 |
+
retry=retry_if_result(lambda ret: ret[-1].strip() in [
|
92 |
+
'KeyboardInterrupt',
|
93 |
+
f"Kernel didn't respond in {wait_for_ready_timeout} seconds",
|
94 |
+
] if isinstance(ret, tuple) else False),
|
95 |
+
stop=stop_after_attempt(3),
|
96 |
+
wait=wait_fixed(1),
|
97 |
+
retry_error_callback=lambda state: state.outcome.result())
|
98 |
+
async def run():
|
99 |
+
execute_result = None
|
100 |
+
error_traceback = None
|
101 |
+
stream_text_list = []
|
102 |
+
kc = km.client()
|
103 |
+
assert isinstance(kc, AsyncKernelClient)
|
104 |
+
kc.start_channels()
|
105 |
+
try:
|
106 |
+
await kc.wait_for_ready(timeout=wait_for_ready_timeout)
|
107 |
+
msg_id = kc.execute(code)
|
108 |
+
while True:
|
109 |
+
message = await get_iopub_msg_with_death_detection(
|
110 |
+
kc, timeout=iopub_timeout)
|
111 |
+
if logger.isEnabledFor(logging.DEBUG):
|
112 |
+
logger.debug(
|
113 |
+
json.dumps(message, indent=2, default=str))
|
114 |
+
assert message["parent_header"]["msg_id"] == msg_id
|
115 |
+
msg_type = message["msg_type"]
|
116 |
+
if msg_type == "status":
|
117 |
+
if message["content"]["execution_state"] == "idle":
|
118 |
+
break
|
119 |
+
elif msg_type == "stream":
|
120 |
+
stream_name = message["content"]["name"]
|
121 |
+
stream_text = message["content"]["text"]
|
122 |
+
stream_text_list.append(stream_text)
|
123 |
+
elif msg_type == "execute_result":
|
124 |
+
execute_result = message["content"]["data"]
|
125 |
+
elif msg_type == "error":
|
126 |
+
error_traceback_lines = message["content"]["traceback"]
|
127 |
+
error_traceback = "\n".join(error_traceback_lines)
|
128 |
+
elif msg_type == "execute_input":
|
129 |
+
pass
|
130 |
+
else:
|
131 |
+
assert False, f"Unknown message_type: {msg_type}"
|
132 |
+
finally:
|
133 |
+
kc.stop_channels()
|
134 |
+
return execute_result, error_traceback, "".join(stream_text_list)
|
135 |
+
|
136 |
+
if interrupt_after:
|
137 |
+
run_task = asyncio.create_task(run())
|
138 |
+
send_interrupt_task = asyncio.create_task(send_interrupt())
|
139 |
+
done, _ = await asyncio.wait([run_task, send_interrupt_task],
|
140 |
+
return_when=asyncio.FIRST_COMPLETED)
|
141 |
+
if run_task in done:
|
142 |
+
send_interrupt_task.cancel()
|
143 |
+
else:
|
144 |
+
assert send_interrupt_task in done
|
145 |
+
result = await run_task
|
146 |
+
else:
|
147 |
+
result = await run()
|
148 |
+
return result
|
149 |
+
finally:
|
150 |
+
if shutdown_kernel:
|
151 |
+
await km.shutdown_kernel()
|
152 |
+
|
153 |
+
|
154 |
+
class IPythonInterpreter(BaseAction):
|
155 |
+
"""A IPython executor that can execute Python scripts in a jupyter manner.
|
156 |
+
|
157 |
+
Args:
|
158 |
+
timeout (int): Upper bound of waiting time for Python script execution.
|
159 |
+
Defaults to 20.
|
160 |
+
user_data_dir (str, optional): Specified the user data directory for files
|
161 |
+
loading. If set to `ENV`, use `USER_DATA_DIR` environment variable.
|
162 |
+
Defaults to `ENV`.
|
163 |
+
work_dir (str, optional): Specify which directory to save output images to.
|
164 |
+
Defaults to ``'./work_dir/tmp_dir'``.
|
165 |
+
description (dict): The description of the action. Defaults to ``None``.
|
166 |
+
parser (Type[BaseParser]): The parser class to process the
|
167 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
168 |
+
"""
|
169 |
+
|
170 |
+
_KERNEL_CLIENTS = {}
|
171 |
+
|
172 |
+
def __init__(
|
173 |
+
self,
|
174 |
+
timeout: int = 20,
|
175 |
+
user_data_dir: str = 'ENV',
|
176 |
+
work_dir='./work_dir/tmp_dir',
|
177 |
+
description: Optional[dict] = None,
|
178 |
+
parser: Type[BaseParser] = JsonParser,
|
179 |
+
):
|
180 |
+
super().__init__(description, parser)
|
181 |
+
|
182 |
+
self.timeout = timeout
|
183 |
+
if user_data_dir == 'ENV':
|
184 |
+
user_data_dir = os.environ.get('USER_DATA_DIR', '')
|
185 |
+
|
186 |
+
if user_data_dir:
|
187 |
+
user_data_dir = os.path.dirname(user_data_dir)
|
188 |
+
user_data_dir = f"import os\nos.chdir('{user_data_dir}')"
|
189 |
+
self.user_data_dir = user_data_dir
|
190 |
+
self._initialized = False
|
191 |
+
self.work_dir = work_dir
|
192 |
+
if not os.path.exists(self.work_dir):
|
193 |
+
os.makedirs(self.work_dir, exist_ok=True)
|
194 |
+
|
195 |
+
@staticmethod
|
196 |
+
def start_kernel():
|
197 |
+
from jupyter_client import KernelManager
|
198 |
+
|
199 |
+
# start the kernel and manager
|
200 |
+
km = KernelManager()
|
201 |
+
km.start_kernel()
|
202 |
+
kc = km.client()
|
203 |
+
return km, kc
|
204 |
+
|
205 |
+
def initialize(self):
|
206 |
+
if self._initialized:
|
207 |
+
return
|
208 |
+
pid = os.getpid()
|
209 |
+
if pid not in self._KERNEL_CLIENTS:
|
210 |
+
self._KERNEL_CLIENTS[pid] = self.start_kernel()
|
211 |
+
self.kernel_manager, self.kernel_client = self._KERNEL_CLIENTS[pid]
|
212 |
+
self._initialized = True
|
213 |
+
self._call(START_CODE.format(self.user_data_dir), None)
|
214 |
+
|
215 |
+
def reset(self):
|
216 |
+
if not self._initialized:
|
217 |
+
self.initialize()
|
218 |
+
else:
|
219 |
+
code = "get_ipython().run_line_magic('reset', '-f')\n" + \
|
220 |
+
START_CODE.format(self.user_data_dir)
|
221 |
+
self._call(code, None)
|
222 |
+
|
223 |
+
def _call(self,
|
224 |
+
command: str,
|
225 |
+
timeout: Optional[int] = None) -> Tuple[str, bool]:
|
226 |
+
self.initialize()
|
227 |
+
command = extract_code(command)
|
228 |
+
|
229 |
+
# check previous remaining result
|
230 |
+
while True:
|
231 |
+
try:
|
232 |
+
msg = self.kernel_client.get_iopub_msg(timeout=5)
|
233 |
+
msg_type = msg['msg_type']
|
234 |
+
if msg_type == 'status':
|
235 |
+
if msg['content'].get('execution_state') == 'idle':
|
236 |
+
break
|
237 |
+
except queue.Empty:
|
238 |
+
# assume no result
|
239 |
+
break
|
240 |
+
|
241 |
+
self.kernel_client.execute(command)
|
242 |
+
|
243 |
+
def _inner_call():
|
244 |
+
result = ''
|
245 |
+
images = []
|
246 |
+
succeed = True
|
247 |
+
image_idx = 0
|
248 |
+
|
249 |
+
while True:
|
250 |
+
text = ''
|
251 |
+
image = ''
|
252 |
+
finished = False
|
253 |
+
msg_type = 'error'
|
254 |
+
try:
|
255 |
+
msg = self.kernel_client.get_iopub_msg(timeout=20)
|
256 |
+
msg_type = msg['msg_type']
|
257 |
+
if msg_type == 'status':
|
258 |
+
if msg['content'].get('execution_state') == 'idle':
|
259 |
+
finished = True
|
260 |
+
elif msg_type == 'execute_result':
|
261 |
+
text = msg['content']['data'].get('text/plain', '')
|
262 |
+
if 'image/png' in msg['content']['data']:
|
263 |
+
image_b64 = msg['content']['data']['image/png']
|
264 |
+
image_url = publish_image_to_local(
|
265 |
+
image_b64, self.work_dir)
|
266 |
+
image_idx += 1
|
267 |
+
image = '![fig-%03d](%s)' % (image_idx, image_url)
|
268 |
+
|
269 |
+
elif msg_type == 'display_data':
|
270 |
+
if 'image/png' in msg['content']['data']:
|
271 |
+
image_b64 = msg['content']['data']['image/png']
|
272 |
+
image_url = publish_image_to_local(
|
273 |
+
image_b64, self.work_dir)
|
274 |
+
image_idx += 1
|
275 |
+
image = '![fig-%03d](%s)' % (image_idx, image_url)
|
276 |
+
|
277 |
+
else:
|
278 |
+
text = msg['content']['data'].get('text/plain', '')
|
279 |
+
elif msg_type == 'stream':
|
280 |
+
msg_type = msg['content']['name'] # stdout, stderr
|
281 |
+
text = msg['content']['text']
|
282 |
+
elif msg_type == 'error':
|
283 |
+
succeed = False
|
284 |
+
text = escape_ansi('\n'.join(
|
285 |
+
msg['content']['traceback']))
|
286 |
+
if 'M6_CODE_INTERPRETER_TIMEOUT' in text:
|
287 |
+
text = f'Timeout. No response after {timeout} seconds.' # noqa
|
288 |
+
except queue.Empty:
|
289 |
+
# stop current task in case break next input.
|
290 |
+
self.kernel_manager.interrupt_kernel()
|
291 |
+
succeed = False
|
292 |
+
text = f'Timeout. No response after {timeout} seconds.'
|
293 |
+
finished = True
|
294 |
+
except Exception:
|
295 |
+
succeed = False
|
296 |
+
msg = ''.join(traceback.format_exception(*sys.exc_info()))
|
297 |
+
# text = 'The code interpreter encountered an unexpected error.' # noqa
|
298 |
+
text = msg
|
299 |
+
logging.warning(msg)
|
300 |
+
finished = True
|
301 |
+
if text:
|
302 |
+
# result += f'\n\n{msg_type}:\n\n```\n{text}\n```'
|
303 |
+
result += f'{text}'
|
304 |
+
|
305 |
+
if image:
|
306 |
+
images.append(image_url)
|
307 |
+
if finished:
|
308 |
+
return succeed, dict(text=result, image=images)
|
309 |
+
|
310 |
+
try:
|
311 |
+
if timeout:
|
312 |
+
|
313 |
+
def handler(signum, frame):
|
314 |
+
raise TimeoutError()
|
315 |
+
|
316 |
+
signal.signal(signal.SIGALRM, handler)
|
317 |
+
signal.alarm(timeout)
|
318 |
+
succeed, result = _inner_call()
|
319 |
+
except TimeoutError:
|
320 |
+
succeed = False
|
321 |
+
text = 'The code interpreter encountered an unexpected error.'
|
322 |
+
result = f'\n\nerror:\n\n```\n{text}\n```'
|
323 |
+
finally:
|
324 |
+
if timeout:
|
325 |
+
signal.alarm(0)
|
326 |
+
|
327 |
+
# result = result.strip('\n')
|
328 |
+
return succeed, result
|
329 |
+
|
330 |
+
@tool_api
|
331 |
+
def run(self, command: str, timeout: Optional[int] = None) -> ActionReturn:
|
332 |
+
r"""When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 60.0 seconds. The drive at '/mnt/data' can be used to save and persist user files. Internet access for this session is disabled. Do not make external web requests or API calls as they will fail.
|
333 |
+
|
334 |
+
Args:
|
335 |
+
command (:class:`str`): Python code
|
336 |
+
timeout (:class:`Optional[int]`): Upper bound of waiting time for Python script execution.
|
337 |
+
"""
|
338 |
+
tool_return = ActionReturn(url=None, args=None, type=self.name)
|
339 |
+
tool_return.args = dict(text=command)
|
340 |
+
succeed, result = self._call(command, timeout)
|
341 |
+
if succeed:
|
342 |
+
text = result['text']
|
343 |
+
image = result.get('image', [])
|
344 |
+
resp = [dict(type='text', content=text)]
|
345 |
+
if image:
|
346 |
+
resp.extend([dict(type='image', content=im) for im in image])
|
347 |
+
tool_return.result = resp
|
348 |
+
# tool_return.result = dict(
|
349 |
+
# text=result['text'], image=result.get('image', [])[0])
|
350 |
+
tool_return.state = ActionStatusCode.SUCCESS
|
351 |
+
else:
|
352 |
+
tool_return.errmsg = result.get('text', '') if isinstance(
|
353 |
+
result, dict) else result
|
354 |
+
tool_return.state = ActionStatusCode.API_ERROR
|
355 |
+
return tool_return
|
356 |
+
|
357 |
+
|
358 |
+
class AsyncIPythonInterpreter(AsyncActionMixin, IPythonInterpreter):
|
359 |
+
"""A IPython executor that can execute Python scripts in a jupyter manner.
|
360 |
+
|
361 |
+
Args:
|
362 |
+
timeout (int): Upper bound of waiting time for Python script execution.
|
363 |
+
Defaults to 20.
|
364 |
+
user_data_dir (str, optional): Specified the user data directory for files
|
365 |
+
loading. If set to `ENV`, use `USER_DATA_DIR` environment variable.
|
366 |
+
Defaults to `ENV`.
|
367 |
+
work_dir (str, optional): Specify which directory to save output images to.
|
368 |
+
Defaults to ``'./work_dir/tmp_dir'``.
|
369 |
+
description (dict): The description of the action. Defaults to ``None``.
|
370 |
+
parser (Type[BaseParser]): The parser class to process the
|
371 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
372 |
+
"""
|
373 |
+
|
374 |
+
_UNBOUND_KERNEL_CLIENTS = asyncio.Queue()
|
375 |
+
|
376 |
+
def __init__(
|
377 |
+
self,
|
378 |
+
timeout: int = 20,
|
379 |
+
user_data_dir: str = 'ENV',
|
380 |
+
work_dir=os.path.join(tempfile.gettempdir(), 'tmp_dir'),
|
381 |
+
max_kernels: Optional[int] = None,
|
382 |
+
reuse_kernel: bool = True,
|
383 |
+
startup_rate: bool = 32,
|
384 |
+
connection_dir: str = tempfile.gettempdir(),
|
385 |
+
description: Optional[dict] = None,
|
386 |
+
parser: Type[BaseParser] = JsonParser,
|
387 |
+
):
|
388 |
+
super().__init__(timeout, user_data_dir, work_dir, description, parser)
|
389 |
+
from traitlets.config import Config
|
390 |
+
|
391 |
+
c = Config()
|
392 |
+
c.KernelManager.transport = 'ipc'
|
393 |
+
self._amkm = AsyncMultiKernelManager(
|
394 |
+
config=c, connection_dir=connection_dir)
|
395 |
+
self._max_kernels = max_kernels
|
396 |
+
self._reuse_kernel = reuse_kernel
|
397 |
+
self._sem = asyncio.Semaphore(startup_rate)
|
398 |
+
self._lock = asyncio.Lock()
|
399 |
+
|
400 |
+
async def initialize(self, session_id: str):
|
401 |
+
session_id = str(session_id)
|
402 |
+
while True:
|
403 |
+
if session_id in self._KERNEL_CLIENTS:
|
404 |
+
return self._KERNEL_CLIENTS[session_id]
|
405 |
+
if self._reuse_kernel and not self._UNBOUND_KERNEL_CLIENTS.empty():
|
406 |
+
self._KERNEL_CLIENTS[
|
407 |
+
session_id] = await self._UNBOUND_KERNEL_CLIENTS.get()
|
408 |
+
return self._KERNEL_CLIENTS[session_id]
|
409 |
+
async with self._sem:
|
410 |
+
if self._max_kernels is None or len(
|
411 |
+
self._KERNEL_CLIENTS
|
412 |
+
) + self._UNBOUND_KERNEL_CLIENTS.qsize() < self._max_kernels:
|
413 |
+
kernel_id = None
|
414 |
+
try:
|
415 |
+
kernel_id = await self._amkm.start_kernel()
|
416 |
+
kernel = self._amkm.get_kernel(kernel_id)
|
417 |
+
client = kernel.client()
|
418 |
+
_, error_stacktrace, stream_text = await async_run_code(
|
419 |
+
kernel,
|
420 |
+
START_CODE.format(self.user_data_dir),
|
421 |
+
shutdown_kernel=False)
|
422 |
+
# check if the output of START_CODE meets expectations
|
423 |
+
if not (error_stacktrace is None
|
424 |
+
and stream_text == ''):
|
425 |
+
raise RuntimeError
|
426 |
+
except Exception as e:
|
427 |
+
print(f'Starting kernel error: {e}')
|
428 |
+
if kernel_id:
|
429 |
+
await self._amkm.shutdown_kernel(kernel_id)
|
430 |
+
self._amkm.remove_kernel(kernel_id)
|
431 |
+
await asyncio.sleep(1)
|
432 |
+
continue
|
433 |
+
if self._max_kernels is None:
|
434 |
+
self._KERNEL_CLIENTS[session_id] = (kernel_id, kernel,
|
435 |
+
client)
|
436 |
+
return kernel_id, kernel, client
|
437 |
+
async with self._lock:
|
438 |
+
if len(self._KERNEL_CLIENTS
|
439 |
+
) + self._UNBOUND_KERNEL_CLIENTS.qsize(
|
440 |
+
) < self._max_kernels:
|
441 |
+
self._KERNEL_CLIENTS[session_id] = (kernel_id,
|
442 |
+
kernel, client)
|
443 |
+
return kernel_id, kernel, client
|
444 |
+
await self._amkm.shutdown_kernel(kernel_id)
|
445 |
+
self._amkm.remove_kernel(kernel_id)
|
446 |
+
await asyncio.sleep(1)
|
447 |
+
|
448 |
+
async def reset(self, session_id: str):
|
449 |
+
session_id = str(session_id)
|
450 |
+
if session_id not in self._KERNEL_CLIENTS:
|
451 |
+
return
|
452 |
+
_, kernel, _ = self._KERNEL_CLIENTS[session_id]
|
453 |
+
code = "get_ipython().run_line_magic('reset', '-f')\n" + \
|
454 |
+
START_CODE.format(self.user_data_dir)
|
455 |
+
await async_run_code(kernel, code, shutdown_kernel=False)
|
456 |
+
|
457 |
+
async def shutdown(self, session_id: str):
|
458 |
+
session_id = str(session_id)
|
459 |
+
if session_id in self._KERNEL_CLIENTS:
|
460 |
+
kernel_id, _, _ = self._KERNEL_CLIENTS.get(session_id)
|
461 |
+
await self._amkm.shutdown_kernel(kernel_id)
|
462 |
+
self._amkm.remove_kernel(kernel_id)
|
463 |
+
del self._KERNEL_CLIENTS[session_id]
|
464 |
+
|
465 |
+
async def close_session(self, session_id: str):
|
466 |
+
session_id = str(session_id)
|
467 |
+
if self._reuse_kernel:
|
468 |
+
if session_id in self._KERNEL_CLIENTS:
|
469 |
+
await self.reset(session_id)
|
470 |
+
await self._UNBOUND_KERNEL_CLIENTS.put(
|
471 |
+
self._KERNEL_CLIENTS.pop(session_id))
|
472 |
+
else:
|
473 |
+
await self.shutdown(session_id)
|
474 |
+
|
475 |
+
async def _call(self, command, timeout=None, session_id=None):
|
476 |
+
_, kernel, _ = await self.initialize(str(session_id))
|
477 |
+
result = await async_run_code(
|
478 |
+
kernel,
|
479 |
+
extract_code(command),
|
480 |
+
interrupt_after=timeout or self.timeout,
|
481 |
+
shutdown_kernel=False)
|
482 |
+
execute_result, error_stacktrace, stream_text = result
|
483 |
+
if error_stacktrace is not None:
|
484 |
+
ret = re.sub('^-*\n', '', escape_ansi(error_stacktrace))
|
485 |
+
if ret.endswith('KeyboardInterrupt: '):
|
486 |
+
ret = 'The code interpreter encountered a timeout error.'
|
487 |
+
status, ret = False, ret.strip()
|
488 |
+
elif execute_result is not None:
|
489 |
+
status, ret = True, dict(text=execute_result.get('text/plain', ''))
|
490 |
+
else:
|
491 |
+
status, ret = True, dict(text=stream_text.strip())
|
492 |
+
return status, ret
|
493 |
+
|
494 |
+
@tool_api
|
495 |
+
async def run(self,
|
496 |
+
command: str,
|
497 |
+
timeout: Optional[int] = None,
|
498 |
+
session_id: Optional[str] = None) -> ActionReturn:
|
499 |
+
r"""When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 60.0 seconds. The drive at '/mnt/data' can be used to save and persist user files. Internet access for this session is disabled. Do not make external web requests or API calls as they will fail.
|
500 |
+
|
501 |
+
Args:
|
502 |
+
command (:class:`str`): Python code
|
503 |
+
timeout (:class:`Optional[int]`): Upper bound of waiting time for Python script execution.
|
504 |
+
"""
|
505 |
+
tool_return = ActionReturn(url=None, args=None, type=self.name)
|
506 |
+
tool_return.args = dict(text=command)
|
507 |
+
succeed, result = await self._call(command, timeout, session_id)
|
508 |
+
if succeed:
|
509 |
+
text = result['text']
|
510 |
+
image = result.get('image', [])
|
511 |
+
resp = [dict(type='text', content=text)]
|
512 |
+
if image:
|
513 |
+
resp.extend([dict(type='image', content=im) for im in image])
|
514 |
+
tool_return.result = resp
|
515 |
+
# tool_return.result = dict(
|
516 |
+
# text=result['text'], image=result.get('image', [])[0])
|
517 |
+
tool_return.state = ActionStatusCode.SUCCESS
|
518 |
+
else:
|
519 |
+
tool_return.errmsg = result.get('text', '') if isinstance(
|
520 |
+
result, dict) else result
|
521 |
+
tool_return.state = ActionStatusCode.API_ERROR
|
522 |
+
return tool_return
|
523 |
+
|
524 |
+
|
525 |
+
def extract_code(text):
|
526 |
+
import json5
|
527 |
+
|
528 |
+
# Match triple backtick blocks first
|
529 |
+
triple_match = re.search(r'```[^\n]*\n(.+?)```', text, re.DOTALL)
|
530 |
+
# Match single backtick blocks second
|
531 |
+
single_match = re.search(r'`([^`]*)`', text, re.DOTALL)
|
532 |
+
if triple_match:
|
533 |
+
text = triple_match.group(1)
|
534 |
+
elif single_match:
|
535 |
+
text = single_match.group(1)
|
536 |
+
else:
|
537 |
+
try:
|
538 |
+
text = json5.loads(text)['code']
|
539 |
+
except Exception:
|
540 |
+
pass
|
541 |
+
# If no code blocks found, return original text
|
542 |
+
return text
|
543 |
+
|
544 |
+
|
545 |
+
def escape_ansi(line):
|
546 |
+
ansi_escape = re.compile(r'(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]')
|
547 |
+
return ansi_escape.sub('', line)
|
548 |
+
|
549 |
+
|
550 |
+
def publish_image_to_local(image_base64: str, work_dir='./work_dir/tmp_dir'):
|
551 |
+
import PIL.Image
|
552 |
+
image_file = str(uuid.uuid4()) + '.png'
|
553 |
+
local_image_file = os.path.join(work_dir, image_file)
|
554 |
+
|
555 |
+
png_bytes = base64.b64decode(image_base64)
|
556 |
+
assert isinstance(png_bytes, bytes)
|
557 |
+
bytes_io = io.BytesIO(png_bytes)
|
558 |
+
PIL.Image.open(bytes_io).save(local_image_file, 'png')
|
559 |
+
|
560 |
+
return local_image_file
|
561 |
+
|
562 |
+
|
563 |
+
# local test for code interpreter
|
564 |
+
def get_multiline_input(hint):
|
565 |
+
print(hint)
|
566 |
+
print('// Press ENTER to make a new line. Press CTRL-D to end input.')
|
567 |
+
lines = []
|
568 |
+
while True:
|
569 |
+
try:
|
570 |
+
line = input()
|
571 |
+
except EOFError: # CTRL-D
|
572 |
+
break
|
573 |
+
lines.append(line)
|
574 |
+
print('// Input received.')
|
575 |
+
if lines:
|
576 |
+
return '\n'.join(lines)
|
577 |
+
else:
|
578 |
+
return ''
|
579 |
+
|
580 |
+
|
581 |
+
if __name__ == '__main__':
|
582 |
+
code_interpreter = IPythonInterpreter()
|
583 |
+
while True:
|
584 |
+
print(code_interpreter(get_multiline_input('Enter python code:')))
|
lagent/actions/ipython_manager.py
ADDED
@@ -0,0 +1,220 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import re
|
2 |
+
import sys
|
3 |
+
from collections import defaultdict
|
4 |
+
from contextlib import nullcontext
|
5 |
+
from io import StringIO
|
6 |
+
from multiprocessing import Process, Queue
|
7 |
+
from typing import List, Optional, Type, Union
|
8 |
+
|
9 |
+
from filelock import FileLock
|
10 |
+
from timeout_decorator import timeout as tm
|
11 |
+
|
12 |
+
from ..schema import ActionReturn, ActionStatusCode
|
13 |
+
from .base_action import BaseAction
|
14 |
+
from .parser import BaseParser, JsonParser
|
15 |
+
|
16 |
+
|
17 |
+
class IPythonProcess(Process):
|
18 |
+
|
19 |
+
def __init__(self,
|
20 |
+
in_q: Queue,
|
21 |
+
out_q: Queue,
|
22 |
+
timeout: int = 20,
|
23 |
+
ci_lock: str = None,
|
24 |
+
daemon: bool = True):
|
25 |
+
super().__init__(daemon=daemon)
|
26 |
+
self.in_q = in_q
|
27 |
+
self.out_q = out_q
|
28 |
+
self.timeout = timeout
|
29 |
+
self.session_id2shell = defaultdict(self.create_shell)
|
30 |
+
self.ci_lock = FileLock(
|
31 |
+
ci_lock) if ci_lock else nullcontext() # avoid core corruption
|
32 |
+
self._highlighting = re.compile(r'\x1b\[\d{,3}(;\d{,3}){,3}m')
|
33 |
+
|
34 |
+
def run(self):
|
35 |
+
while True:
|
36 |
+
msg = self.in_q.get()
|
37 |
+
if msg == 'reset':
|
38 |
+
for session_id, shell in self.session_id2shell.items():
|
39 |
+
with self.ci_lock:
|
40 |
+
try:
|
41 |
+
shell.reset(new_session=False)
|
42 |
+
# shell.run_line_magic('reset', '-sf')
|
43 |
+
except Exception:
|
44 |
+
self.session_id2shell[
|
45 |
+
session_id] = self.create_shell()
|
46 |
+
self.out_q.put('ok')
|
47 |
+
elif isinstance(msg, tuple) and len(msg) == 3:
|
48 |
+
i, session_id, code = msg
|
49 |
+
res = self.exec(session_id, code)
|
50 |
+
self.out_q.put((i, session_id, res))
|
51 |
+
|
52 |
+
def exec(self, session_id, code):
|
53 |
+
try:
|
54 |
+
shell = self.session_id2shell[session_id]
|
55 |
+
with StringIO() as io:
|
56 |
+
old_stdout = sys.stdout
|
57 |
+
sys.stdout = io
|
58 |
+
if self.timeout is False or self.timeout < 0:
|
59 |
+
shell.run_cell(self.extract_code(code))
|
60 |
+
else:
|
61 |
+
tm(self.timeout)(shell.run_cell)(self.extract_code(code))
|
62 |
+
sys.stdout = old_stdout
|
63 |
+
output = self._highlighting.sub('', io.getvalue().strip())
|
64 |
+
output = re.sub(r'^Out\[\d+\]: ', '', output)
|
65 |
+
if 'Error' in output or 'Traceback' in output:
|
66 |
+
output = output.lstrip('-').strip()
|
67 |
+
if output.startswith('TimeoutError'):
|
68 |
+
output = 'The code interpreter encountered a timeout error.'
|
69 |
+
return {'status': 'FAILURE', 'msg': output, 'code': code}
|
70 |
+
return {'status': 'SUCCESS', 'value': output, 'code': code}
|
71 |
+
except Exception as e:
|
72 |
+
return {'status': 'FAILURE', 'msg': str(e), 'code': code}
|
73 |
+
|
74 |
+
@staticmethod
|
75 |
+
def create_shell(enable_history: bool = False, in_memory: bool = True):
|
76 |
+
from IPython import InteractiveShell
|
77 |
+
from traitlets.config import Config
|
78 |
+
|
79 |
+
c = Config()
|
80 |
+
c.HistoryManager.enabled = enable_history
|
81 |
+
if in_memory:
|
82 |
+
c.HistoryManager.hist_file = ':memory:'
|
83 |
+
shell = InteractiveShell(config=c)
|
84 |
+
return shell
|
85 |
+
|
86 |
+
@staticmethod
|
87 |
+
def extract_code(text: str) -> str:
|
88 |
+
"""Extract Python code from markup languages.
|
89 |
+
|
90 |
+
Args:
|
91 |
+
text (:class:`str`): Markdown-formatted text
|
92 |
+
|
93 |
+
Returns:
|
94 |
+
:class:`str`: Python code
|
95 |
+
"""
|
96 |
+
import json5
|
97 |
+
|
98 |
+
# Match triple backtick blocks first
|
99 |
+
triple_match = re.search(r'```[^\n]*\n(.+?)```', text, re.DOTALL)
|
100 |
+
# Match single backtick blocks second
|
101 |
+
single_match = re.search(r'`([^`]*)`', text, re.DOTALL)
|
102 |
+
if triple_match:
|
103 |
+
text = triple_match.group(1)
|
104 |
+
elif single_match:
|
105 |
+
text = single_match.group(1)
|
106 |
+
else:
|
107 |
+
try:
|
108 |
+
text = json5.loads(text)['code']
|
109 |
+
except Exception:
|
110 |
+
pass
|
111 |
+
# If no code blocks found, return original text
|
112 |
+
return text
|
113 |
+
|
114 |
+
|
115 |
+
class IPythonInteractiveManager(BaseAction):
|
116 |
+
"""An interactive IPython shell manager for code execution"""
|
117 |
+
|
118 |
+
def __init__(
|
119 |
+
self,
|
120 |
+
max_workers: int = 50,
|
121 |
+
timeout: int = 20,
|
122 |
+
ci_lock: str = None,
|
123 |
+
description: Optional[dict] = None,
|
124 |
+
parser: Type[BaseParser] = JsonParser,
|
125 |
+
):
|
126 |
+
super().__init__(description, parser)
|
127 |
+
self.max_workers = max_workers
|
128 |
+
self.timeout = timeout
|
129 |
+
self.ci_lock = ci_lock
|
130 |
+
self.id2queue = defaultdict(Queue)
|
131 |
+
self.id2process = {}
|
132 |
+
self.out_queue = Queue()
|
133 |
+
|
134 |
+
def __call__(self,
|
135 |
+
commands: Union[str, List[str]],
|
136 |
+
session_ids: Union[int, List[int]] = None):
|
137 |
+
if isinstance(commands, list):
|
138 |
+
batch_size = len(commands)
|
139 |
+
is_batch = True
|
140 |
+
else:
|
141 |
+
batch_size = 1
|
142 |
+
commands = [commands]
|
143 |
+
is_batch = False
|
144 |
+
if session_ids is None:
|
145 |
+
session_ids = range(batch_size)
|
146 |
+
elif isinstance(session_ids, int):
|
147 |
+
session_ids = [session_ids]
|
148 |
+
if len(session_ids) != batch_size or len(session_ids) != len(
|
149 |
+
set(session_ids)):
|
150 |
+
raise ValueError(
|
151 |
+
'the size of `session_ids` must equal that of `commands`')
|
152 |
+
try:
|
153 |
+
exec_results = self.run_code_blocks([
|
154 |
+
(session_id, command)
|
155 |
+
for session_id, command in zip(session_ids, commands)
|
156 |
+
])
|
157 |
+
except KeyboardInterrupt:
|
158 |
+
self.clear()
|
159 |
+
exit(1)
|
160 |
+
action_returns = []
|
161 |
+
for result, code in zip(exec_results, commands):
|
162 |
+
action_return = ActionReturn({'command': code}, type=self.name)
|
163 |
+
if result['status'] == 'SUCCESS':
|
164 |
+
action_return.result = [
|
165 |
+
dict(type='text', content=result['value'])
|
166 |
+
]
|
167 |
+
action_return.state = ActionStatusCode.SUCCESS
|
168 |
+
else:
|
169 |
+
action_return.errmsg = result['msg']
|
170 |
+
action_return.state = ActionStatusCode.API_ERROR
|
171 |
+
action_returns.append(action_return)
|
172 |
+
if not is_batch:
|
173 |
+
return action_returns[0]
|
174 |
+
return action_returns
|
175 |
+
|
176 |
+
def process_code(self, index, session_id, code):
|
177 |
+
ipy_id = session_id % self.max_workers
|
178 |
+
input_queue = self.id2queue[ipy_id]
|
179 |
+
proc = self.id2process.setdefault(
|
180 |
+
ipy_id,
|
181 |
+
IPythonProcess(
|
182 |
+
input_queue,
|
183 |
+
self.out_queue,
|
184 |
+
self.timeout,
|
185 |
+
self.ci_lock,
|
186 |
+
daemon=True))
|
187 |
+
if not proc.is_alive():
|
188 |
+
proc.start()
|
189 |
+
input_queue.put((index, session_id, code))
|
190 |
+
|
191 |
+
def run_code_blocks(self, session_code_pairs):
|
192 |
+
size = len(session_code_pairs)
|
193 |
+
for index, (session_id, code) in enumerate(session_code_pairs):
|
194 |
+
self.process_code(index, session_id, code)
|
195 |
+
results = []
|
196 |
+
while len(results) < size:
|
197 |
+
msg = self.out_queue.get()
|
198 |
+
if isinstance(msg, tuple) and len(msg) == 3:
|
199 |
+
index, _, result = msg
|
200 |
+
results.append((index, result))
|
201 |
+
results.sort()
|
202 |
+
return [item[1] for item in results]
|
203 |
+
|
204 |
+
def clear(self):
|
205 |
+
self.id2queue.clear()
|
206 |
+
for proc in self.id2process.values():
|
207 |
+
proc.terminate()
|
208 |
+
self.id2process.clear()
|
209 |
+
while not self.out_queue.empty():
|
210 |
+
self.out_queue.get()
|
211 |
+
|
212 |
+
def reset(self):
|
213 |
+
cnt = 0
|
214 |
+
for q in self.id2queue.values():
|
215 |
+
q.put('reset')
|
216 |
+
cnt += 1
|
217 |
+
while cnt > 0:
|
218 |
+
msg = self.out_queue.get()
|
219 |
+
if msg == 'ok':
|
220 |
+
cnt -= 1
|
lagent/actions/parser.py
ADDED
@@ -0,0 +1,146 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
import re
|
3 |
+
from ast import literal_eval
|
4 |
+
from typing import Any, List, Union
|
5 |
+
|
6 |
+
|
7 |
+
class ParseError(Exception):
|
8 |
+
"""Parsing exception class."""
|
9 |
+
|
10 |
+
def __init__(self, err_msg: str):
|
11 |
+
self.err_msg = err_msg
|
12 |
+
|
13 |
+
|
14 |
+
class BaseParser:
|
15 |
+
"""Base parser to process inputs and outputs of actions.
|
16 |
+
|
17 |
+
Args:
|
18 |
+
action (:class:`BaseAction`): action to validate
|
19 |
+
|
20 |
+
Attributes:
|
21 |
+
PARAMETER_DESCRIPTION (:class:`str`): declare the input format which
|
22 |
+
LLMs should follow when generating arguments for decided tools.
|
23 |
+
"""
|
24 |
+
|
25 |
+
PARAMETER_DESCRIPTION: str = ''
|
26 |
+
|
27 |
+
def __init__(self, action):
|
28 |
+
self.action = action
|
29 |
+
self._api2param = {}
|
30 |
+
self._api2required = {}
|
31 |
+
# perform basic argument validation
|
32 |
+
if action.description:
|
33 |
+
for api in action.description.get('api_list',
|
34 |
+
[action.description]):
|
35 |
+
name = (f'{action.name}.{api["name"]}'
|
36 |
+
if self.action.is_toolkit else api['name'])
|
37 |
+
required_parameters = set(api['required'])
|
38 |
+
all_parameters = {j['name'] for j in api['parameters']}
|
39 |
+
if not required_parameters.issubset(all_parameters):
|
40 |
+
raise ValueError(
|
41 |
+
f'unknown parameters for function "{name}": '
|
42 |
+
f'{required_parameters - all_parameters}')
|
43 |
+
if self.PARAMETER_DESCRIPTION:
|
44 |
+
api['parameter_description'] = self.PARAMETER_DESCRIPTION
|
45 |
+
api_name = api['name'] if self.action.is_toolkit else 'run'
|
46 |
+
self._api2param[api_name] = api['parameters']
|
47 |
+
self._api2required[api_name] = api['required']
|
48 |
+
|
49 |
+
def parse_inputs(self, inputs: str, name: str = 'run') -> dict:
|
50 |
+
"""Parse inputs LLMs generate for the action.
|
51 |
+
|
52 |
+
Args:
|
53 |
+
inputs (:class:`str`): input string extracted from responses
|
54 |
+
|
55 |
+
Returns:
|
56 |
+
:class:`dict`: processed input
|
57 |
+
"""
|
58 |
+
inputs = {self._api2param[name][0]['name']: inputs}
|
59 |
+
return inputs
|
60 |
+
|
61 |
+
def parse_outputs(self, outputs: Any) -> List[dict]:
|
62 |
+
"""Parser outputs returned by the action.
|
63 |
+
|
64 |
+
Args:
|
65 |
+
outputs (:class:`Any`): raw output of the action
|
66 |
+
|
67 |
+
Returns:
|
68 |
+
:class:`List[dict]`: processed output of which each member is a
|
69 |
+
dictionary with two keys - 'type' and 'content'.
|
70 |
+
"""
|
71 |
+
if isinstance(outputs, dict):
|
72 |
+
outputs = json.dumps(outputs, ensure_ascii=False)
|
73 |
+
elif not isinstance(outputs, str):
|
74 |
+
outputs = str(outputs)
|
75 |
+
return [{
|
76 |
+
'type': 'text',
|
77 |
+
'content': outputs.encode('gbk', 'ignore').decode('gbk')
|
78 |
+
}]
|
79 |
+
|
80 |
+
|
81 |
+
class JsonParser(BaseParser):
|
82 |
+
"""Json parser to convert input string into a dictionary.
|
83 |
+
|
84 |
+
Args:
|
85 |
+
action (:class:`BaseAction`): action to validate
|
86 |
+
"""
|
87 |
+
|
88 |
+
PARAMETER_DESCRIPTION = (
|
89 |
+
'If you call this tool, you must pass arguments in '
|
90 |
+
'the JSON format {key: value}, where the key is the parameter name.')
|
91 |
+
|
92 |
+
def parse_inputs(self,
|
93 |
+
inputs: Union[str, dict],
|
94 |
+
name: str = 'run') -> dict:
|
95 |
+
if not isinstance(inputs, dict):
|
96 |
+
try:
|
97 |
+
match = re.search(r'^\s*(```json\n)?(.*)\n```\s*$', inputs,
|
98 |
+
re.S)
|
99 |
+
if match:
|
100 |
+
inputs = match.group(2).strip()
|
101 |
+
inputs = json.loads(inputs)
|
102 |
+
except json.JSONDecodeError as exc:
|
103 |
+
raise ParseError(f'invalid json format: {inputs}') from exc
|
104 |
+
input_keys = set(inputs)
|
105 |
+
all_keys = {param['name'] for param in self._api2param[name]}
|
106 |
+
if not input_keys.issubset(all_keys):
|
107 |
+
raise ParseError(f'unknown arguments: {input_keys - all_keys}')
|
108 |
+
required_keys = set(self._api2required[name])
|
109 |
+
if not input_keys.issuperset(required_keys):
|
110 |
+
raise ParseError(
|
111 |
+
f'missing required arguments: {required_keys - input_keys}')
|
112 |
+
return inputs
|
113 |
+
|
114 |
+
|
115 |
+
class TupleParser(BaseParser):
|
116 |
+
"""Tuple parser to convert input string into a tuple.
|
117 |
+
|
118 |
+
Args:
|
119 |
+
action (:class:`BaseAction`): action to validate
|
120 |
+
"""
|
121 |
+
|
122 |
+
PARAMETER_DESCRIPTION = (
|
123 |
+
'If you call this tool, you must pass arguments in the tuple format '
|
124 |
+
'like (arg1, arg2, arg3), and the arguments are ordered.')
|
125 |
+
|
126 |
+
def parse_inputs(self,
|
127 |
+
inputs: Union[str, tuple],
|
128 |
+
name: str = 'run') -> dict:
|
129 |
+
if not isinstance(inputs, tuple):
|
130 |
+
try:
|
131 |
+
inputs = literal_eval(inputs)
|
132 |
+
except Exception as exc:
|
133 |
+
raise ParseError(f'invalid tuple format: {inputs}') from exc
|
134 |
+
if len(inputs) < len(self._api2required[name]):
|
135 |
+
raise ParseError(
|
136 |
+
f'API takes {len(self._api2required[name])} required positional '
|
137 |
+
f'arguments but {len(inputs)} were given')
|
138 |
+
if len(inputs) > len(self._api2param[name]):
|
139 |
+
raise ParseError(
|
140 |
+
f'API takes {len(self._api2param[name])} positional arguments '
|
141 |
+
f'but {len(inputs)} were given')
|
142 |
+
inputs = {
|
143 |
+
self._api2param[name][i]['name']: item
|
144 |
+
for i, item in enumerate(inputs)
|
145 |
+
}
|
146 |
+
return inputs
|
lagent/actions/ppt.py
ADDED
@@ -0,0 +1,233 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Dict, Optional, Type
|
2 |
+
|
3 |
+
from asyncer import asyncify
|
4 |
+
|
5 |
+
from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
|
6 |
+
from lagent.actions.parser import BaseParser, JsonParser
|
7 |
+
|
8 |
+
THEME_MAPPING = {
|
9 |
+
'Default': {
|
10 |
+
'template': None,
|
11 |
+
'title': 'Title Slide',
|
12 |
+
'single': 'Title and Content',
|
13 |
+
'two': 'Two Content',
|
14 |
+
}
|
15 |
+
}
|
16 |
+
|
17 |
+
|
18 |
+
class PPT(BaseAction):
|
19 |
+
"""Plugin to create ppt slides with text, paragraph, images in good looking styles."""
|
20 |
+
|
21 |
+
def __init__(
|
22 |
+
self,
|
23 |
+
theme_mapping: Optional[Dict[str, dict]] = None,
|
24 |
+
description: Optional[dict] = None,
|
25 |
+
parser: Type[BaseParser] = JsonParser,
|
26 |
+
):
|
27 |
+
super().__init__(description, parser)
|
28 |
+
self.theme_mapping = theme_mapping or THEME_MAPPING
|
29 |
+
self.pointer = None
|
30 |
+
self.location = None
|
31 |
+
|
32 |
+
@tool_api(explode_return=True)
|
33 |
+
def create_file(self, theme: str, abs_location: str) -> dict:
|
34 |
+
"""Create a pptx file with specific themes.
|
35 |
+
|
36 |
+
Args:
|
37 |
+
theme (:class:`str`): the theme used. The value should be one of ['Default'].
|
38 |
+
abs_location (:class:`str`): the ppt file's absolute location
|
39 |
+
|
40 |
+
Returns:
|
41 |
+
:class:`dict`: operation status
|
42 |
+
* status: the result of the execution
|
43 |
+
"""
|
44 |
+
from pptx import Presentation
|
45 |
+
|
46 |
+
self.location = abs_location
|
47 |
+
try:
|
48 |
+
self.pointer = Presentation(self.theme_mapping[theme]['template'])
|
49 |
+
self.pointer.slide_master.name = theme
|
50 |
+
# print('created')
|
51 |
+
except Exception as e:
|
52 |
+
print(e)
|
53 |
+
return dict(status='created a ppt file.')
|
54 |
+
|
55 |
+
@tool_api(explode_return=True)
|
56 |
+
def add_first_page(self, title: str, subtitle: str) -> dict:
|
57 |
+
"""Add the first page of ppt.
|
58 |
+
|
59 |
+
Args:
|
60 |
+
title (:class:`str`): the title of ppt
|
61 |
+
subtitle (:class:`str`): the subtitle of ppt
|
62 |
+
|
63 |
+
Returns:
|
64 |
+
:class:`dict`: operation status
|
65 |
+
* status: the result of the execution
|
66 |
+
"""
|
67 |
+
layout_name = self.theme_mapping[self.pointer.slide_master.name]['title']
|
68 |
+
layout = next(i for i in self.pointer.slide_master.slide_layouts if i.name == layout_name)
|
69 |
+
slide = self.pointer.slides.add_slide(layout)
|
70 |
+
ph_title, ph_subtitle = slide.placeholders
|
71 |
+
ph_title.text = title
|
72 |
+
if subtitle:
|
73 |
+
ph_subtitle.text = subtitle
|
74 |
+
return dict(status='added page')
|
75 |
+
|
76 |
+
@tool_api(explode_return=True)
|
77 |
+
def add_text_page(self, title: str, bullet_items: str) -> dict:
|
78 |
+
"""Add text page of ppt.
|
79 |
+
|
80 |
+
Args:
|
81 |
+
title (:class:`str`): the title of the page
|
82 |
+
bullet_items (:class:`str`): bullet_items should be string, for multiple bullet items, please use [SPAN] to separate them.
|
83 |
+
|
84 |
+
Returns:
|
85 |
+
:class:`dict`: operation status
|
86 |
+
* status: the result of the execution
|
87 |
+
""" # noqa: E501
|
88 |
+
layout_name = self.theme_mapping[self.pointer.slide_master.name]['single']
|
89 |
+
layout = next(i for i in self.pointer.slide_master.slide_layouts if i.name == layout_name)
|
90 |
+
slide = self.pointer.slides.add_slide(layout)
|
91 |
+
ph_title, ph_body = slide.placeholders
|
92 |
+
ph_title.text = title
|
93 |
+
ph = ph_body
|
94 |
+
tf = ph.text_frame
|
95 |
+
for i, item in enumerate(bullet_items.split('[SPAN]')):
|
96 |
+
if i == 0:
|
97 |
+
p = tf.paragraphs[0]
|
98 |
+
else:
|
99 |
+
p = tf.add_paragraph()
|
100 |
+
p.text = item.strip()
|
101 |
+
p.level = 0
|
102 |
+
return dict(status='added page')
|
103 |
+
|
104 |
+
@tool_api(explode_return=True)
|
105 |
+
def add_text_image_page(self, title: str, bullet_items: str, image: str) -> dict:
|
106 |
+
"""Add a text page with one image. Image should be a path.
|
107 |
+
|
108 |
+
Args:
|
109 |
+
title (:class:`str`): the title of the page
|
110 |
+
bullet_items (:class:`str`): bullet_items should be string, for multiple bullet items, please use [SPAN] to separate them.
|
111 |
+
image (:class:`str`): the path of the image
|
112 |
+
|
113 |
+
Returns:
|
114 |
+
:class:`dict`: operation status
|
115 |
+
* status: the result of the execution
|
116 |
+
""" # noqa: E501
|
117 |
+
from PIL import Image
|
118 |
+
|
119 |
+
layout_name = self.theme_mapping[self.pointer.slide_master.name]['two']
|
120 |
+
layout = next(i for i in self.pointer.slide_master.slide_layouts if i.name == layout_name)
|
121 |
+
slide = self.pointer.slides.add_slide(layout)
|
122 |
+
ph_title, ph_body1, ph_body2 = slide.placeholders
|
123 |
+
ph_title.text = title
|
124 |
+
ph = ph_body2
|
125 |
+
image = Image.open(image)
|
126 |
+
image_pil = image.to_pil()
|
127 |
+
left = ph.left
|
128 |
+
width = ph.width
|
129 |
+
height = int(width / image_pil.width * image_pil.height)
|
130 |
+
top = (ph.top + (ph.top + ph.height)) // 2 - height // 2
|
131 |
+
slide.shapes.add_picture(image.to_path(), left, top, width, height)
|
132 |
+
|
133 |
+
ph = ph_body1
|
134 |
+
tf = ph.text_frame
|
135 |
+
for i, item in enumerate(bullet_items.split('[SPAN]')):
|
136 |
+
if i == 0:
|
137 |
+
p = tf.paragraphs[0]
|
138 |
+
else:
|
139 |
+
p = tf.add_paragraph()
|
140 |
+
p.text = item.strip()
|
141 |
+
p.level = 0
|
142 |
+
|
143 |
+
return dict(status='added page')
|
144 |
+
|
145 |
+
@tool_api(explode_return=True)
|
146 |
+
def submit_file(self) -> dict:
|
147 |
+
"""When all steps done, YOU MUST use submit_file() to submit your work.
|
148 |
+
|
149 |
+
Returns:
|
150 |
+
:class:`dict`: operation status
|
151 |
+
* status: the result of the execution
|
152 |
+
"""
|
153 |
+
# file_path = os.path.join(self.CACHE_DIR, f'{self._return_timestamp()}.pptx')
|
154 |
+
# self.pointer.save(file_path)
|
155 |
+
# retreival_url = upload_file(file_path)
|
156 |
+
self.pointer.save(self.location)
|
157 |
+
return dict(status=f'submitted. view ppt at {self.location}')
|
158 |
+
|
159 |
+
|
160 |
+
class AsyncPPT(AsyncActionMixin, PPT):
|
161 |
+
"""Plugin to create ppt slides with text, paragraph, images in good looking styles."""
|
162 |
+
|
163 |
+
@tool_api(explode_return=True)
|
164 |
+
@asyncify
|
165 |
+
def create_file(self, theme: str, abs_location: str) -> dict:
|
166 |
+
"""Create a pptx file with specific themes.
|
167 |
+
|
168 |
+
Args:
|
169 |
+
theme (:class:`str`): the theme used. The value should be one of ['Default'].
|
170 |
+
abs_location (:class:`str`): the ppt file's absolute location
|
171 |
+
|
172 |
+
Returns:
|
173 |
+
:class:`dict`: operation status
|
174 |
+
* status: the result of the execution
|
175 |
+
"""
|
176 |
+
return super().create_file(theme, abs_location)
|
177 |
+
|
178 |
+
@tool_api(explode_return=True)
|
179 |
+
@asyncify
|
180 |
+
def add_first_page(self, title: str, subtitle: str) -> dict:
|
181 |
+
"""Add the first page of ppt.
|
182 |
+
|
183 |
+
Args:
|
184 |
+
title (:class:`str`): the title of ppt
|
185 |
+
subtitle (:class:`str`): the subtitle of ppt
|
186 |
+
|
187 |
+
Returns:
|
188 |
+
:class:`dict`: operation status
|
189 |
+
* status: the result of the execution
|
190 |
+
"""
|
191 |
+
return super().add_first_page(title, subtitle)
|
192 |
+
|
193 |
+
@tool_api(explode_return=True)
|
194 |
+
@asyncify
|
195 |
+
def add_text_page(self, title: str, bullet_items: str) -> dict:
|
196 |
+
"""Add text page of ppt.
|
197 |
+
|
198 |
+
Args:
|
199 |
+
title (:class:`str`): the title of the page
|
200 |
+
bullet_items (:class:`str`): bullet_items should be string, for multiple bullet items, please use [SPAN] to separate them.
|
201 |
+
|
202 |
+
Returns:
|
203 |
+
:class:`dict`: operation status
|
204 |
+
* status: the result of the execution
|
205 |
+
""" # noqa: E501
|
206 |
+
return super().add_text_page(title, bullet_items)
|
207 |
+
|
208 |
+
@tool_api(explode_return=True)
|
209 |
+
@asyncify
|
210 |
+
def add_text_image_page(self, title: str, bullet_items: str, image: str) -> dict:
|
211 |
+
"""Add a text page with one image. Image should be a path.
|
212 |
+
|
213 |
+
Args:
|
214 |
+
title (:class:`str`): the title of the page
|
215 |
+
bullet_items (:class:`str`): bullet_items should be string, for multiple bullet items, please use [SPAN] to separate them.
|
216 |
+
image (:class:`str`): the path of the image
|
217 |
+
|
218 |
+
Returns:
|
219 |
+
:class:`dict`: operation status
|
220 |
+
* status: the result of the execution
|
221 |
+
""" # noqa: E501
|
222 |
+
return super().add_text_image_page(title, bullet_items, image)
|
223 |
+
|
224 |
+
@tool_api(explode_return=True)
|
225 |
+
@asyncify
|
226 |
+
def submit_file(self) -> dict:
|
227 |
+
"""When all steps done, YOU MUST use submit_file() to submit your work.
|
228 |
+
|
229 |
+
Returns:
|
230 |
+
:class:`dict`: operation status
|
231 |
+
* status: the result of the execution
|
232 |
+
"""
|
233 |
+
return super().submit_file()
|
lagent/actions/python_interpreter.py
ADDED
@@ -0,0 +1,176 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# flake8: noqa: E501
|
2 |
+
import copy
|
3 |
+
import io
|
4 |
+
from contextlib import redirect_stdout
|
5 |
+
from typing import Any, Optional, Type
|
6 |
+
|
7 |
+
from asyncer import asyncify
|
8 |
+
|
9 |
+
from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
|
10 |
+
from lagent.actions.parser import BaseParser, JsonParser
|
11 |
+
from lagent.schema import ActionReturn, ActionStatusCode
|
12 |
+
|
13 |
+
|
14 |
+
class GenericRuntime:
|
15 |
+
GLOBAL_DICT = {}
|
16 |
+
LOCAL_DICT = None
|
17 |
+
HEADERS = []
|
18 |
+
|
19 |
+
def __init__(self):
|
20 |
+
self._global_vars = copy.copy(self.GLOBAL_DICT)
|
21 |
+
self._local_vars = copy.copy(self.LOCAL_DICT) if self.LOCAL_DICT else None
|
22 |
+
|
23 |
+
for c in self.HEADERS:
|
24 |
+
self.exec_code(c)
|
25 |
+
|
26 |
+
def exec_code(self, code_piece: str) -> None:
|
27 |
+
exec(code_piece, self._global_vars)
|
28 |
+
|
29 |
+
def eval_code(self, expr: str) -> Any:
|
30 |
+
return eval(expr, self._global_vars)
|
31 |
+
|
32 |
+
|
33 |
+
class PythonInterpreter(BaseAction):
|
34 |
+
"""A Python executor that can execute Python scripts.
|
35 |
+
|
36 |
+
Args:
|
37 |
+
answer_symbol (str, Optional): the answer symbol from LLM. Defaults to ``None``.
|
38 |
+
answer_expr (str, Optional): the answer function name of the Python
|
39 |
+
script. Defaults to ``'solution()'``.
|
40 |
+
answer_from_stdout (boolean, Optional): whether the execution results is from
|
41 |
+
stdout. Defaults to ``False``.
|
42 |
+
timeout (int, Optional): Upper bound of waiting time for Python script execution.
|
43 |
+
Defaults to ``20``.
|
44 |
+
description (dict, Optional): The description of the action. Defaults to ``None``.
|
45 |
+
parser (Type[BaseParser]): The parser class to process the
|
46 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
47 |
+
"""
|
48 |
+
|
49 |
+
def __init__(
|
50 |
+
self,
|
51 |
+
answer_symbol: Optional[str] = None,
|
52 |
+
answer_expr: Optional[str] = 'solution()',
|
53 |
+
answer_from_stdout: bool = False,
|
54 |
+
timeout: int = 20,
|
55 |
+
description: Optional[dict] = None,
|
56 |
+
parser: Type[BaseParser] = JsonParser,
|
57 |
+
) -> None:
|
58 |
+
super().__init__(description, parser)
|
59 |
+
self.answer_symbol = answer_symbol
|
60 |
+
self.answer_expr = answer_expr
|
61 |
+
self.answer_from_stdout = answer_from_stdout
|
62 |
+
self.timeout = timeout
|
63 |
+
|
64 |
+
@tool_api
|
65 |
+
def run(self, command: str) -> ActionReturn:
|
66 |
+
"""用来执行Python代码。代码必须是一个函数,函数名必须得是 'solution',代码对应你的思考过程。代码实例格式如下:
|
67 |
+
|
68 |
+
```python
|
69 |
+
# import 依赖包
|
70 |
+
import xxx
|
71 |
+
def solution():
|
72 |
+
# 初始化一些变量
|
73 |
+
variable_names_with_real_meaning = xxx
|
74 |
+
# 步骤一
|
75 |
+
mid_variable = func(variable_names_with_real_meaning)
|
76 |
+
# 步骤 x
|
77 |
+
mid_variable = func(mid_variable)
|
78 |
+
# 最后结果
|
79 |
+
final_answer = func(mid_variable)
|
80 |
+
return final_answer
|
81 |
+
```
|
82 |
+
|
83 |
+
Args:
|
84 |
+
command (:class:`str`): Python code snippet
|
85 |
+
"""
|
86 |
+
from func_timeout import FunctionTimedOut, func_set_timeout
|
87 |
+
|
88 |
+
self.runtime = GenericRuntime()
|
89 |
+
try:
|
90 |
+
tool_return = func_set_timeout(self.timeout)(self._call)(command)
|
91 |
+
except FunctionTimedOut as e:
|
92 |
+
tool_return = ActionReturn(type=self.name)
|
93 |
+
tool_return.errmsg = repr(e)
|
94 |
+
tool_return.state = ActionStatusCode.API_ERROR
|
95 |
+
return tool_return
|
96 |
+
|
97 |
+
def _call(self, command: str) -> ActionReturn:
|
98 |
+
tool_return = ActionReturn(type=self.name)
|
99 |
+
try:
|
100 |
+
if '```python' in command:
|
101 |
+
command = command.split('```python')[1].split('```')[0]
|
102 |
+
elif '```' in command:
|
103 |
+
command = command.split('```')[1].split('```')[0]
|
104 |
+
tool_return.args = dict(text='```python\n' + command + '\n```')
|
105 |
+
command = command.split('\n')
|
106 |
+
|
107 |
+
if self.answer_from_stdout:
|
108 |
+
program_io = io.StringIO()
|
109 |
+
with redirect_stdout(program_io):
|
110 |
+
self.runtime.exec_code('\n'.join(command))
|
111 |
+
program_io.seek(0)
|
112 |
+
res = program_io.readlines()[-1]
|
113 |
+
elif self.answer_symbol:
|
114 |
+
self.runtime.exec_code('\n'.join(command))
|
115 |
+
res = self.runtime._global_vars[self.answer_symbol]
|
116 |
+
elif self.answer_expr:
|
117 |
+
self.runtime.exec_code('\n'.join(command))
|
118 |
+
res = self.runtime.eval_code(self.answer_expr)
|
119 |
+
else:
|
120 |
+
self.runtime.exec_code('\n'.join(command[:-1]))
|
121 |
+
res = self.runtime.eval_code(command[-1])
|
122 |
+
except Exception as e:
|
123 |
+
tool_return.errmsg = repr(e)
|
124 |
+
tool_return.type = self.name
|
125 |
+
tool_return.state = ActionStatusCode.API_ERROR
|
126 |
+
return tool_return
|
127 |
+
try:
|
128 |
+
tool_return.result = [dict(type='text', content=str(res))]
|
129 |
+
tool_return.state = ActionStatusCode.SUCCESS
|
130 |
+
except Exception as e:
|
131 |
+
tool_return.errmsg = repr(e)
|
132 |
+
tool_return.type = self.name
|
133 |
+
tool_return.state = ActionStatusCode.API_ERROR
|
134 |
+
return tool_return
|
135 |
+
|
136 |
+
|
137 |
+
class AsyncPythonInterpreter(AsyncActionMixin, PythonInterpreter):
|
138 |
+
"""A Python executor that can execute Python scripts.
|
139 |
+
|
140 |
+
Args:
|
141 |
+
answer_symbol (str, Optional): the answer symbol from LLM. Defaults to ``None``.
|
142 |
+
answer_expr (str, Optional): the answer function name of the Python
|
143 |
+
script. Defaults to ``'solution()'``.
|
144 |
+
answer_from_stdout (boolean, Optional): whether the execution results is from
|
145 |
+
stdout. Defaults to ``False``.
|
146 |
+
timeout (int, Optional): Upper bound of waiting time for Python script execution.
|
147 |
+
Defaults to ``20``.
|
148 |
+
description (dict, Optional): The description of the action. Defaults to ``None``.
|
149 |
+
parser (Type[BaseParser]): The parser class to process the
|
150 |
+
action's inputs and outputs. Defaults to :class:`JsonParser`.
|
151 |
+
"""
|
152 |
+
|
153 |
+
@tool_api
|
154 |
+
@asyncify
|
155 |
+
def run(self, command: str) -> ActionReturn:
|
156 |
+
"""用来执行Python代码。代码必须是一个函数,函数名必须得是 'solution',代码对应你的思考过程。代码实例格式如下:
|
157 |
+
|
158 |
+
```python
|
159 |
+
# import 依赖包
|
160 |
+
import xxx
|
161 |
+
def solution():
|
162 |
+
# 初始化一些变量
|
163 |
+
variable_names_with_real_meaning = xxx
|
164 |
+
# 步骤一
|
165 |
+
mid_variable = func(variable_names_with_real_meaning)
|
166 |
+
# 步骤 x
|
167 |
+
mid_variable = func(mid_variable)
|
168 |
+
# 最后结果
|
169 |
+
final_answer = func(mid_variable)
|
170 |
+
return final_answer
|
171 |
+
```
|
172 |
+
|
173 |
+
Args:
|
174 |
+
command (:class:`str`): Python code snippet
|
175 |
+
"""
|
176 |
+
return super().run(command)
|
lagent/actions/weather_query.py
ADDED
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import requests
|
3 |
+
from lagent.actions.base_action import BaseAction, tool_api
|
4 |
+
from lagent.schema import ActionReturn, ActionStatusCode
|
5 |
+
|
6 |
+
class WeatherQuery(BaseAction):
|
7 |
+
def __init__(self):
|
8 |
+
super().__init__()
|
9 |
+
self.api_key = os.getenv("weather_token")
|
10 |
+
print(self.api_key)
|
11 |
+
if not self.api_key:
|
12 |
+
raise EnvironmentError("未找到环境变量 'token'。请设置你的和风天气 API Key 到 'weather_token' 环境变量中,比如export weather_token='xxx' ")
|
13 |
+
|
14 |
+
@tool_api
|
15 |
+
def run(self, location: str) -> dict:
|
16 |
+
"""
|
17 |
+
查询实时天气信息。
|
18 |
+
|
19 |
+
Args:
|
20 |
+
location (str): 要查询的地点名称、LocationID 或经纬度坐标(如 "101010100" 或 "116.41,39.92")。
|
21 |
+
|
22 |
+
Returns:
|
23 |
+
dict: 包含天气信息的字典
|
24 |
+
* location: 地点名称
|
25 |
+
* weather: 天气状况
|
26 |
+
* temperature: 当前温度
|
27 |
+
* wind_direction: 风向
|
28 |
+
* wind_speed: 风速(公里/小时)
|
29 |
+
* humidity: 相对湿度(%)
|
30 |
+
* report_time: 数据报告时间
|
31 |
+
"""
|
32 |
+
try:
|
33 |
+
# 如果 location 不是坐标格式(例如 "116.41,39.92"),则调用 GeoAPI 获取 LocationID
|
34 |
+
if not ("," in location and location.replace(",", "").replace(".", "").isdigit()):
|
35 |
+
# 使用 GeoAPI 获取 LocationID
|
36 |
+
geo_url = f"https://geoapi.qweather.com/v2/city/lookup?location={location}&key={self.api_key}"
|
37 |
+
geo_response = requests.get(geo_url)
|
38 |
+
geo_data = geo_response.json()
|
39 |
+
|
40 |
+
if geo_data.get("code") != "200" or not geo_data.get("location"):
|
41 |
+
raise Exception(f"GeoAPI 返回错误码:{geo_data.get('code')} 或未找到位置")
|
42 |
+
|
43 |
+
location = geo_data["location"][0]["id"]
|
44 |
+
|
45 |
+
# 构建天气查询的 API 请求 URL
|
46 |
+
weather_url = f"https://devapi.qweather.com/v7/weather/now?location={location}&key={self.api_key}"
|
47 |
+
response = requests.get(weather_url)
|
48 |
+
data = response.json()
|
49 |
+
|
50 |
+
# 检查 API 响应码
|
51 |
+
if data.get("code") != "200":
|
52 |
+
raise Exception(f"Weather API 返回错误码:{data.get('code')}")
|
53 |
+
|
54 |
+
# 解析和组织天气信息
|
55 |
+
weather_info = {
|
56 |
+
"location": location,
|
57 |
+
"weather": data["now"]["text"],
|
58 |
+
"temperature": data["now"]["temp"] + "°C",
|
59 |
+
"wind_direction": data["now"]["windDir"],
|
60 |
+
"wind_speed": data["now"]["windSpeed"] + " km/h",
|
61 |
+
"humidity": data["now"]["humidity"] + "%",
|
62 |
+
"report_time": data["updateTime"]
|
63 |
+
}
|
64 |
+
|
65 |
+
return {"result": weather_info}
|
66 |
+
|
67 |
+
except Exception as exc:
|
68 |
+
return ActionReturn(
|
69 |
+
errmsg=f"WeatherQuery 异常:{exc}",
|
70 |
+
state=ActionStatusCode.HTTP_ERROR
|
71 |
+
)
|
lagent/actions/web_browser.py
ADDED
@@ -0,0 +1,908 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import asyncio
|
2 |
+
import hashlib
|
3 |
+
import hmac
|
4 |
+
import json
|
5 |
+
import logging
|
6 |
+
import random
|
7 |
+
import re
|
8 |
+
import time
|
9 |
+
import warnings
|
10 |
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
11 |
+
from datetime import datetime
|
12 |
+
from http.client import HTTPSConnection
|
13 |
+
from typing import List, Optional, Tuple, Type, Union
|
14 |
+
|
15 |
+
import aiohttp
|
16 |
+
import aiohttp.client_exceptions
|
17 |
+
import requests
|
18 |
+
from asyncache import cached as acached
|
19 |
+
from bs4 import BeautifulSoup
|
20 |
+
from cachetools import TTLCache, cached
|
21 |
+
from duckduckgo_search import DDGS, AsyncDDGS
|
22 |
+
|
23 |
+
from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
|
24 |
+
from lagent.actions.parser import BaseParser, JsonParser
|
25 |
+
from lagent.utils import async_as_completed
|
26 |
+
|
27 |
+
|
28 |
+
class BaseSearch:
|
29 |
+
|
30 |
+
def __init__(self, topk: int = 3, black_list: List[str] = None):
|
31 |
+
self.topk = topk
|
32 |
+
self.black_list = black_list
|
33 |
+
|
34 |
+
def _filter_results(self, results: List[tuple]) -> dict:
|
35 |
+
filtered_results = {}
|
36 |
+
count = 0
|
37 |
+
for url, snippet, title in results:
|
38 |
+
if all(domain not in url
|
39 |
+
for domain in self.black_list) and not url.endswith('.pdf'):
|
40 |
+
filtered_results[count] = {
|
41 |
+
'url': url,
|
42 |
+
'summ': json.dumps(snippet, ensure_ascii=False)[1:-1],
|
43 |
+
'title': title
|
44 |
+
}
|
45 |
+
count += 1
|
46 |
+
if count >= self.topk:
|
47 |
+
break
|
48 |
+
return filtered_results
|
49 |
+
|
50 |
+
|
51 |
+
class DuckDuckGoSearch(BaseSearch):
|
52 |
+
|
53 |
+
def __init__(self,
|
54 |
+
topk: int = 3,
|
55 |
+
black_list: List[str] = [
|
56 |
+
'enoN',
|
57 |
+
'youtube.com',
|
58 |
+
'bilibili.com',
|
59 |
+
'researchgate.net',
|
60 |
+
],
|
61 |
+
**kwargs):
|
62 |
+
self.proxy = kwargs.get('proxy')
|
63 |
+
self.timeout = kwargs.get('timeout', 30)
|
64 |
+
super().__init__(topk, black_list)
|
65 |
+
|
66 |
+
@cached(cache=TTLCache(maxsize=100, ttl=600))
|
67 |
+
def search(self, query: str, max_retry: int = 3) -> dict:
|
68 |
+
for attempt in range(max_retry):
|
69 |
+
try:
|
70 |
+
response = self._call_ddgs(
|
71 |
+
query, timeout=self.timeout, proxy=self.proxy)
|
72 |
+
return self._parse_response(response)
|
73 |
+
except Exception as e:
|
74 |
+
logging.exception(str(e))
|
75 |
+
warnings.warn(
|
76 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
77 |
+
time.sleep(random.randint(2, 5))
|
78 |
+
raise Exception(
|
79 |
+
'Failed to get search results from DuckDuckGo after retries.')
|
80 |
+
|
81 |
+
@acached(cache=TTLCache(maxsize=100, ttl=600))
|
82 |
+
async def asearch(self, query: str, max_retry: int = 3) -> dict:
|
83 |
+
for attempt in range(max_retry):
|
84 |
+
try:
|
85 |
+
ddgs = AsyncDDGS(timeout=self.timeout, proxy=self.proxy)
|
86 |
+
response = await ddgs.atext(query.strip("'"), max_results=10)
|
87 |
+
return self._parse_response(response)
|
88 |
+
except Exception as e:
|
89 |
+
if isinstance(e, asyncio.TimeoutError):
|
90 |
+
logging.exception('Request to DDGS timed out.')
|
91 |
+
logging.exception(str(e))
|
92 |
+
warnings.warn(
|
93 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
94 |
+
await asyncio.sleep(random.randint(2, 5))
|
95 |
+
raise Exception(
|
96 |
+
'Failed to get search results from DuckDuckGo after retries.')
|
97 |
+
|
98 |
+
async def _async_call_ddgs(self, query: str, **kwargs) -> dict:
|
99 |
+
ddgs = DDGS(**kwargs)
|
100 |
+
try:
|
101 |
+
response = await asyncio.wait_for(
|
102 |
+
asyncio.to_thread(ddgs.text, query.strip("'"), max_results=10),
|
103 |
+
timeout=self.timeout)
|
104 |
+
return response
|
105 |
+
except asyncio.TimeoutError:
|
106 |
+
logging.exception('Request to DDGS timed out.')
|
107 |
+
raise
|
108 |
+
|
109 |
+
def _call_ddgs(self, query: str, **kwargs) -> dict:
|
110 |
+
loop = asyncio.new_event_loop()
|
111 |
+
asyncio.set_event_loop(loop)
|
112 |
+
try:
|
113 |
+
response = loop.run_until_complete(
|
114 |
+
self._async_call_ddgs(query, **kwargs))
|
115 |
+
return response
|
116 |
+
finally:
|
117 |
+
loop.close()
|
118 |
+
|
119 |
+
def _parse_response(self, response: dict) -> dict:
|
120 |
+
raw_results = []
|
121 |
+
for item in response:
|
122 |
+
raw_results.append(
|
123 |
+
(item['href'], item['description']
|
124 |
+
if 'description' in item else item['body'], item['title']))
|
125 |
+
return self._filter_results(raw_results)
|
126 |
+
|
127 |
+
|
128 |
+
class BingSearch(BaseSearch):
|
129 |
+
|
130 |
+
def __init__(self,
|
131 |
+
api_key: str,
|
132 |
+
region: str = 'zh-CN',
|
133 |
+
topk: int = 3,
|
134 |
+
black_list: List[str] = [
|
135 |
+
'enoN',
|
136 |
+
'youtube.com',
|
137 |
+
'bilibili.com',
|
138 |
+
'researchgate.net',
|
139 |
+
],
|
140 |
+
**kwargs):
|
141 |
+
self.api_key = api_key
|
142 |
+
self.market = region
|
143 |
+
self.proxy = kwargs.get('proxy')
|
144 |
+
super().__init__(topk, black_list)
|
145 |
+
|
146 |
+
@cached(cache=TTLCache(maxsize=100, ttl=600))
|
147 |
+
def search(self, query: str, max_retry: int = 3) -> dict:
|
148 |
+
for attempt in range(max_retry):
|
149 |
+
try:
|
150 |
+
response = self._call_bing_api(query)
|
151 |
+
return self._parse_response(response)
|
152 |
+
except Exception as e:
|
153 |
+
logging.exception(str(e))
|
154 |
+
warnings.warn(
|
155 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
156 |
+
time.sleep(random.randint(2, 5))
|
157 |
+
raise Exception(
|
158 |
+
'Failed to get search results from Bing Search after retries.')
|
159 |
+
|
160 |
+
@acached(cache=TTLCache(maxsize=100, ttl=600))
|
161 |
+
async def asearch(self, query: str, max_retry: int = 3) -> dict:
|
162 |
+
for attempt in range(max_retry):
|
163 |
+
try:
|
164 |
+
response = await self._async_call_bing_api(query)
|
165 |
+
return self._parse_response(response)
|
166 |
+
except Exception as e:
|
167 |
+
logging.exception(str(e))
|
168 |
+
warnings.warn(
|
169 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
170 |
+
await asyncio.sleep(random.randint(2, 5))
|
171 |
+
raise Exception(
|
172 |
+
'Failed to get search results from Bing Search after retries.')
|
173 |
+
|
174 |
+
def _call_bing_api(self, query: str) -> dict:
|
175 |
+
endpoint = 'https://api.bing.microsoft.com/v7.0/search'
|
176 |
+
params = {'q': query, 'mkt': self.market, 'count': f'{self.topk * 2}'}
|
177 |
+
headers = {'Ocp-Apim-Subscription-Key': self.api_key}
|
178 |
+
response = requests.get(
|
179 |
+
endpoint, headers=headers, params=params, proxies=self.proxy)
|
180 |
+
response.raise_for_status()
|
181 |
+
return response.json()
|
182 |
+
|
183 |
+
async def _async_call_bing_api(self, query: str) -> dict:
|
184 |
+
endpoint = 'https://api.bing.microsoft.com/v7.0/search'
|
185 |
+
params = {'q': query, 'mkt': self.market, 'count': f'{self.topk * 2}'}
|
186 |
+
headers = {'Ocp-Apim-Subscription-Key': self.api_key}
|
187 |
+
async with aiohttp.ClientSession(raise_for_status=True) as session:
|
188 |
+
async with session.get(
|
189 |
+
endpoint,
|
190 |
+
headers=headers,
|
191 |
+
params=params,
|
192 |
+
proxy=self.proxy and
|
193 |
+
(self.proxy.get('http') or self.proxy.get('https'))) as resp:
|
194 |
+
return await resp.json()
|
195 |
+
|
196 |
+
def _parse_response(self, response: dict) -> dict:
|
197 |
+
webpages = {
|
198 |
+
w['id']: w
|
199 |
+
for w in response.get('webPages', {}).get('value', [])
|
200 |
+
}
|
201 |
+
raw_results = []
|
202 |
+
|
203 |
+
for item in response.get('rankingResponse',
|
204 |
+
{}).get('mainline', {}).get('items', []):
|
205 |
+
if item['answerType'] == 'WebPages':
|
206 |
+
webpage = webpages.get(item['value']['id'])
|
207 |
+
if webpage:
|
208 |
+
raw_results.append(
|
209 |
+
(webpage['url'], webpage['snippet'], webpage['name']))
|
210 |
+
elif item['answerType'] == 'News' and item['value'][
|
211 |
+
'id'] == response.get('news', {}).get('id'):
|
212 |
+
for news in response.get('news', {}).get('value', []):
|
213 |
+
raw_results.append(
|
214 |
+
(news['url'], news['description'], news['name']))
|
215 |
+
|
216 |
+
return self._filter_results(raw_results)
|
217 |
+
|
218 |
+
|
219 |
+
class BraveSearch(BaseSearch):
|
220 |
+
"""
|
221 |
+
Wrapper around the Brave Search API.
|
222 |
+
|
223 |
+
To use, you should pass your Brave Search API key to the constructor.
|
224 |
+
|
225 |
+
Args:
|
226 |
+
api_key (str): API KEY to use Brave Search API.
|
227 |
+
You can create a free API key at https://api.search.brave.com/app/keys.
|
228 |
+
search_type (str): Brave Search API supports ['web', 'news', 'images', 'videos'],
|
229 |
+
currently only supports 'news' and 'web'.
|
230 |
+
topk (int): The number of search results returned in response from API search results.
|
231 |
+
region (str): The country code string. Specifies the country where the search results come from.
|
232 |
+
language (str): The language code string. Specifies the preferred language for the search results.
|
233 |
+
extra_snippets (bool): Allows retrieving up to 5 additional snippets, which are alternative excerpts from the search results.
|
234 |
+
**kwargs: Any other parameters related to the Brave Search API. Find more details at
|
235 |
+
https://api.search.brave.com/app/documentation/web-search/get-started.
|
236 |
+
"""
|
237 |
+
|
238 |
+
def __init__(self,
|
239 |
+
api_key: str,
|
240 |
+
region: str = 'ALL',
|
241 |
+
language: str = 'zh-hans',
|
242 |
+
extra_snippests: bool = True,
|
243 |
+
topk: int = 3,
|
244 |
+
black_list: List[str] = [
|
245 |
+
'enoN',
|
246 |
+
'youtube.com',
|
247 |
+
'bilibili.com',
|
248 |
+
'researchgate.net',
|
249 |
+
],
|
250 |
+
**kwargs):
|
251 |
+
self.api_key = api_key
|
252 |
+
self.market = region
|
253 |
+
self.proxy = kwargs.get('proxy')
|
254 |
+
self.language = language
|
255 |
+
self.extra_snippests = extra_snippests
|
256 |
+
self.search_type = kwargs.get('search_type', 'web')
|
257 |
+
self.kwargs = kwargs
|
258 |
+
super().__init__(topk, black_list)
|
259 |
+
|
260 |
+
@cached(cache=TTLCache(maxsize=100, ttl=600))
|
261 |
+
def search(self, query: str, max_retry: int = 3) -> dict:
|
262 |
+
for attempt in range(max_retry):
|
263 |
+
try:
|
264 |
+
response = self._call_brave_api(query)
|
265 |
+
return self._parse_response(response)
|
266 |
+
except Exception as e:
|
267 |
+
logging.exception(str(e))
|
268 |
+
warnings.warn(
|
269 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
270 |
+
time.sleep(random.randint(2, 5))
|
271 |
+
raise Exception(
|
272 |
+
'Failed to get search results from Brave Search after retries.')
|
273 |
+
|
274 |
+
@acached(cache=TTLCache(maxsize=100, ttl=600))
|
275 |
+
async def asearch(self, query: str, max_retry: int = 3) -> dict:
|
276 |
+
for attempt in range(max_retry):
|
277 |
+
try:
|
278 |
+
response = await self._async_call_brave_api(query)
|
279 |
+
return self._parse_response(response)
|
280 |
+
except Exception as e:
|
281 |
+
logging.exception(str(e))
|
282 |
+
warnings.warn(
|
283 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
284 |
+
await asyncio.sleep(random.randint(2, 5))
|
285 |
+
raise Exception(
|
286 |
+
'Failed to get search results from Brave Search after retries.')
|
287 |
+
|
288 |
+
def _call_brave_api(self, query: str) -> dict:
|
289 |
+
endpoint = f'https://api.search.brave.com/res/v1/{self.search_type}/search'
|
290 |
+
params = {
|
291 |
+
'q': query,
|
292 |
+
'country': self.market,
|
293 |
+
'search_lang': self.language,
|
294 |
+
'extra_snippets': self.extra_snippests,
|
295 |
+
'count': self.topk,
|
296 |
+
**{
|
297 |
+
key: value
|
298 |
+
for key, value in self.kwargs.items() if value is not None
|
299 |
+
},
|
300 |
+
}
|
301 |
+
headers = {
|
302 |
+
'X-Subscription-Token': self.api_key or '',
|
303 |
+
'Accept': 'application/json'
|
304 |
+
}
|
305 |
+
response = requests.get(
|
306 |
+
endpoint, headers=headers, params=params, proxies=self.proxy)
|
307 |
+
response.raise_for_status()
|
308 |
+
return response.json()
|
309 |
+
|
310 |
+
async def _async_call_brave_api(self, query: str) -> dict:
|
311 |
+
endpoint = f'https://api.search.brave.com/res/v1/{self.search_type}/search'
|
312 |
+
params = {
|
313 |
+
'q': query,
|
314 |
+
'country': self.market,
|
315 |
+
'search_lang': self.language,
|
316 |
+
'extra_snippets': self.extra_snippests,
|
317 |
+
'count': self.topk,
|
318 |
+
**{
|
319 |
+
key: value
|
320 |
+
for key, value in self.kwargs.items() if value is not None
|
321 |
+
},
|
322 |
+
}
|
323 |
+
headers = {
|
324 |
+
'X-Subscription-Token': self.api_key or '',
|
325 |
+
'Accept': 'application/json'
|
326 |
+
}
|
327 |
+
async with aiohttp.ClientSession(raise_for_status=True) as session:
|
328 |
+
async with session.get(
|
329 |
+
endpoint,
|
330 |
+
headers=headers,
|
331 |
+
params=params,
|
332 |
+
proxy=self.proxy and
|
333 |
+
(self.proxy.get('http') or self.proxy.get('https'))) as resp:
|
334 |
+
return await resp.json()
|
335 |
+
|
336 |
+
def _parse_response(self, response: dict) -> dict:
|
337 |
+
if self.search_type == 'web':
|
338 |
+
filtered_result = response.get('web', {}).get('results', [])
|
339 |
+
else:
|
340 |
+
filtered_result = response.get('results', {})
|
341 |
+
raw_results = []
|
342 |
+
|
343 |
+
for item in filtered_result:
|
344 |
+
raw_results.append((
|
345 |
+
item.get('url', ''),
|
346 |
+
' '.join(
|
347 |
+
filter(None, [
|
348 |
+
item.get('description'),
|
349 |
+
*item.get('extra_snippets', [])
|
350 |
+
])),
|
351 |
+
item.get('title', ''),
|
352 |
+
))
|
353 |
+
return self._filter_results(raw_results)
|
354 |
+
|
355 |
+
|
356 |
+
class GoogleSearch(BaseSearch):
|
357 |
+
"""
|
358 |
+
Wrapper around the Serper.dev Google Search API.
|
359 |
+
|
360 |
+
To use, you should pass your serper API key to the constructor.
|
361 |
+
|
362 |
+
Args:
|
363 |
+
api_key (str): API KEY to use serper google search API.
|
364 |
+
You can create a free API key at https://serper.dev.
|
365 |
+
search_type (str): Serper API supports ['search', 'images', 'news',
|
366 |
+
'places'] types of search, currently we only support 'search' and 'news'.
|
367 |
+
topk (int): The number of search results returned in response from api search results.
|
368 |
+
**kwargs: Any other parameters related to the Serper API. Find more details at
|
369 |
+
https://serper.dev/playground
|
370 |
+
"""
|
371 |
+
|
372 |
+
result_key_for_type = {
|
373 |
+
'news': 'news',
|
374 |
+
'places': 'places',
|
375 |
+
'images': 'images',
|
376 |
+
'search': 'organic',
|
377 |
+
}
|
378 |
+
|
379 |
+
def __init__(self,
|
380 |
+
api_key: str,
|
381 |
+
topk: int = 3,
|
382 |
+
black_list: List[str] = [
|
383 |
+
'enoN',
|
384 |
+
'youtube.com',
|
385 |
+
'bilibili.com',
|
386 |
+
'researchgate.net',
|
387 |
+
],
|
388 |
+
**kwargs):
|
389 |
+
self.api_key = api_key
|
390 |
+
self.proxy = kwargs.get('proxy')
|
391 |
+
self.search_type = kwargs.get('search_type', 'search')
|
392 |
+
self.kwargs = kwargs
|
393 |
+
super().__init__(topk, black_list)
|
394 |
+
|
395 |
+
@cached(cache=TTLCache(maxsize=100, ttl=600))
|
396 |
+
def search(self, query: str, max_retry: int = 3) -> dict:
|
397 |
+
for attempt in range(max_retry):
|
398 |
+
try:
|
399 |
+
response = self._call_serper_api(query)
|
400 |
+
return self._parse_response(response)
|
401 |
+
except Exception as e:
|
402 |
+
logging.exception(str(e))
|
403 |
+
warnings.warn(
|
404 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
405 |
+
time.sleep(random.randint(2, 5))
|
406 |
+
raise Exception(
|
407 |
+
'Failed to get search results from Google Serper Search after retries.'
|
408 |
+
)
|
409 |
+
|
410 |
+
@acached(cache=TTLCache(maxsize=100, ttl=600))
|
411 |
+
async def asearch(self, query: str, max_retry: int = 3) -> dict:
|
412 |
+
for attempt in range(max_retry):
|
413 |
+
try:
|
414 |
+
response = await self._async_call_serper_api(query)
|
415 |
+
return self._parse_response(response)
|
416 |
+
except Exception as e:
|
417 |
+
logging.exception(str(e))
|
418 |
+
warnings.warn(
|
419 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
420 |
+
await asyncio.sleep(random.randint(2, 5))
|
421 |
+
raise Exception(
|
422 |
+
'Failed to get search results from Google Serper Search after retries.'
|
423 |
+
)
|
424 |
+
|
425 |
+
def _call_serper_api(self, query: str) -> dict:
|
426 |
+
endpoint = f'https://google.serper.dev/{self.search_type}'
|
427 |
+
params = {
|
428 |
+
'q': query,
|
429 |
+
'num': self.topk,
|
430 |
+
**{
|
431 |
+
key: value
|
432 |
+
for key, value in self.kwargs.items() if value is not None
|
433 |
+
},
|
434 |
+
}
|
435 |
+
headers = {
|
436 |
+
'X-API-KEY': self.api_key or '',
|
437 |
+
'Content-Type': 'application/json'
|
438 |
+
}
|
439 |
+
response = requests.get(
|
440 |
+
endpoint, headers=headers, params=params, proxies=self.proxy)
|
441 |
+
response.raise_for_status()
|
442 |
+
return response.json()
|
443 |
+
|
444 |
+
async def _async_call_serper_api(self, query: str) -> dict:
|
445 |
+
endpoint = f'https://google.serper.dev/{self.search_type}'
|
446 |
+
params = {
|
447 |
+
'q': query,
|
448 |
+
'num': self.topk,
|
449 |
+
**{
|
450 |
+
key: value
|
451 |
+
for key, value in self.kwargs.items() if value is not None
|
452 |
+
},
|
453 |
+
}
|
454 |
+
headers = {
|
455 |
+
'X-API-KEY': self.api_key or '',
|
456 |
+
'Content-Type': 'application/json'
|
457 |
+
}
|
458 |
+
async with aiohttp.ClientSession(raise_for_status=True) as session:
|
459 |
+
async with session.get(
|
460 |
+
endpoint,
|
461 |
+
headers=headers,
|
462 |
+
params=params,
|
463 |
+
proxy=self.proxy and
|
464 |
+
(self.proxy.get('http') or self.proxy.get('https'))) as resp:
|
465 |
+
return await resp.json()
|
466 |
+
|
467 |
+
def _parse_response(self, response: dict) -> dict:
|
468 |
+
raw_results = []
|
469 |
+
|
470 |
+
if response.get('answerBox'):
|
471 |
+
answer_box = response.get('answerBox', {})
|
472 |
+
if answer_box.get('answer'):
|
473 |
+
raw_results.append(('', answer_box.get('answer'), ''))
|
474 |
+
elif answer_box.get('snippet'):
|
475 |
+
raw_results.append(
|
476 |
+
('', answer_box.get('snippet').replace('\n', ' '), ''))
|
477 |
+
elif answer_box.get('snippetHighlighted'):
|
478 |
+
raw_results.append(
|
479 |
+
('', answer_box.get('snippetHighlighted'), ''))
|
480 |
+
|
481 |
+
if response.get('knowledgeGraph'):
|
482 |
+
kg = response.get('knowledgeGraph', {})
|
483 |
+
description = kg.get('description', '')
|
484 |
+
attributes = '. '.join(
|
485 |
+
f'{attribute}: {value}'
|
486 |
+
for attribute, value in kg.get('attributes', {}).items())
|
487 |
+
raw_results.append(
|
488 |
+
(kg.get('descriptionLink', ''),
|
489 |
+
f'{description}. {attributes}' if attributes else description,
|
490 |
+
f"{kg.get('title', '')}: {kg.get('type', '')}."))
|
491 |
+
|
492 |
+
for result in response[self.result_key_for_type[
|
493 |
+
self.search_type]][:self.topk]:
|
494 |
+
description = result.get('snippet', '')
|
495 |
+
attributes = '. '.join(
|
496 |
+
f'{attribute}: {value}'
|
497 |
+
for attribute, value in result.get('attributes', {}).items())
|
498 |
+
raw_results.append(
|
499 |
+
(result.get('link', ''),
|
500 |
+
f'{description}. {attributes}' if attributes else description,
|
501 |
+
result.get('title', '')))
|
502 |
+
|
503 |
+
return self._filter_results(raw_results)
|
504 |
+
|
505 |
+
|
506 |
+
class TencentSearch(BaseSearch):
|
507 |
+
"""Wrapper around the tencentclound Search API.
|
508 |
+
|
509 |
+
To use, you should pass your secret_id and secret_key to the constructor.
|
510 |
+
|
511 |
+
Args:
|
512 |
+
secret_id (str): Your Tencent Cloud secret ID for accessing the API.
|
513 |
+
For more details, refer to the documentation: https://cloud.tencent.com/document/product/598/40488.
|
514 |
+
secret_key (str): Your Tencent Cloud secret key for accessing the API.
|
515 |
+
api_key (str, optional): Additional API key, if required.
|
516 |
+
action (str): The action for this interface, use `SearchCommon`.
|
517 |
+
version (str): The API version, use `2020-12-29`.
|
518 |
+
service (str): The service name, use `tms`.
|
519 |
+
host (str): The API host, use `tms.tencentcloudapi.com`.
|
520 |
+
topk (int): The maximum number of search results to return.
|
521 |
+
tsn (int): Time filter for search results. Valid values:
|
522 |
+
1 (within 1 day), 2 (within 1 week), 3 (within 1 month),
|
523 |
+
4 (within 1 year), 5 (within 6 months), 6 (within 3 years).
|
524 |
+
insite (str): Specify a site to search within (supports only a single site).
|
525 |
+
If not specified, the entire web is searched. Example: `zhihu.com`.
|
526 |
+
category (str): Vertical category for filtering results. Optional values include:
|
527 |
+
`baike` (encyclopedia), `weather`, `calendar`, `medical`, `news`, `train`, `star` (horoscope).
|
528 |
+
vrid (str): Result card type(s). Different `vrid` values represent different types of result cards.
|
529 |
+
Supports multiple values separated by commas. Example: `30010255`.
|
530 |
+
"""
|
531 |
+
|
532 |
+
def __init__(self,
|
533 |
+
secret_id: str = 'Your SecretId',
|
534 |
+
secret_key: str = 'Your SecretKey',
|
535 |
+
api_key: str = '',
|
536 |
+
action: str = 'SearchCommon',
|
537 |
+
version: str = '2020-12-29',
|
538 |
+
service: str = 'tms',
|
539 |
+
host: str = 'tms.tencentcloudapi.com',
|
540 |
+
topk: int = 3,
|
541 |
+
tsn: int = None,
|
542 |
+
insite: str = None,
|
543 |
+
category: str = None,
|
544 |
+
vrid: str = None,
|
545 |
+
black_list: List[str] = [
|
546 |
+
'enoN',
|
547 |
+
'youtube.com',
|
548 |
+
'bilibili.com',
|
549 |
+
'researchgate.net',
|
550 |
+
]):
|
551 |
+
self.secret_id = secret_id
|
552 |
+
self.secret_key = secret_key
|
553 |
+
self.api_key = api_key
|
554 |
+
self.action = action
|
555 |
+
self.version = version
|
556 |
+
self.service = service
|
557 |
+
self.host = host
|
558 |
+
self.tsn = tsn
|
559 |
+
self.insite = insite
|
560 |
+
self.category = category
|
561 |
+
self.vrid = vrid
|
562 |
+
super().__init__(topk, black_list=black_list)
|
563 |
+
|
564 |
+
@cached(cache=TTLCache(maxsize=100, ttl=600))
|
565 |
+
def search(self, query: str, max_retry: int = 3) -> dict:
|
566 |
+
for attempt in range(max_retry):
|
567 |
+
try:
|
568 |
+
response = self._call_tencent_api(query)
|
569 |
+
return self._parse_response(response)
|
570 |
+
except Exception as e:
|
571 |
+
logging.exception(str(e))
|
572 |
+
warnings.warn(
|
573 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
574 |
+
time.sleep(random.randint(2, 5))
|
575 |
+
raise Exception(
|
576 |
+
'Failed to get search results from Bing Search after retries.')
|
577 |
+
|
578 |
+
@acached(cache=TTLCache(maxsize=100, ttl=600))
|
579 |
+
async def asearch(self, query: str, max_retry: int = 3) -> dict:
|
580 |
+
for attempt in range(max_retry):
|
581 |
+
try:
|
582 |
+
response = await self._async_call_tencent_api(query)
|
583 |
+
return self._parse_response(response)
|
584 |
+
except Exception as e:
|
585 |
+
logging.exception(str(e))
|
586 |
+
warnings.warn(
|
587 |
+
f'Retry {attempt + 1}/{max_retry} due to error: {e}')
|
588 |
+
await asyncio.sleep(random.randint(2, 5))
|
589 |
+
raise Exception(
|
590 |
+
'Failed to get search results from Bing Search after retries.')
|
591 |
+
|
592 |
+
def _get_headers_and_payload(self, query: str) -> tuple:
|
593 |
+
|
594 |
+
def sign(key, msg):
|
595 |
+
return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()
|
596 |
+
|
597 |
+
params = dict(Query=query)
|
598 |
+
# if self.topk:
|
599 |
+
# params['Cnt'] = self.topk
|
600 |
+
if self.tsn:
|
601 |
+
params['Tsn'] = self.tsn
|
602 |
+
if self.insite:
|
603 |
+
params['Insite'] = self.insite
|
604 |
+
if self.category:
|
605 |
+
params['Category'] = self.category
|
606 |
+
if self.vrid:
|
607 |
+
params['Vrid'] = self.vrid
|
608 |
+
payload = json.dumps(params)
|
609 |
+
algorithm = 'TC3-HMAC-SHA256'
|
610 |
+
timestamp = int(time.time())
|
611 |
+
date = datetime.utcfromtimestamp(timestamp).strftime('%Y-%m-%d')
|
612 |
+
|
613 |
+
# ************* 步骤 1:拼接规范请求串 *************
|
614 |
+
http_request_method = 'POST'
|
615 |
+
canonical_uri = '/'
|
616 |
+
canonical_querystring = ''
|
617 |
+
ct = 'application/json; charset=utf-8'
|
618 |
+
canonical_headers = f'content-type:{ct}\nhost:{self.host}\nx-tc-action:{self.action.lower()}\n'
|
619 |
+
signed_headers = 'content-type;host;x-tc-action'
|
620 |
+
hashed_request_payload = hashlib.sha256(
|
621 |
+
payload.encode('utf-8')).hexdigest()
|
622 |
+
canonical_request = (
|
623 |
+
http_request_method + '\n' + canonical_uri + '\n' +
|
624 |
+
canonical_querystring + '\n' + canonical_headers + '\n' +
|
625 |
+
signed_headers + '\n' + hashed_request_payload)
|
626 |
+
|
627 |
+
# ************* 步骤 2:拼接待签名字符串 *************
|
628 |
+
credential_scope = date + '/' + self.service + '/' + 'tc3_request'
|
629 |
+
hashed_canonical_request = hashlib.sha256(
|
630 |
+
canonical_request.encode('utf-8')).hexdigest()
|
631 |
+
string_to_sign = (
|
632 |
+
algorithm + '\n' + str(timestamp) + '\n' + credential_scope +
|
633 |
+
'\n' + hashed_canonical_request)
|
634 |
+
|
635 |
+
# ************* 步骤 3:计算签名 *************
|
636 |
+
secret_date = sign(('TC3' + self.secret_key).encode('utf-8'), date)
|
637 |
+
secret_service = sign(secret_date, self.service)
|
638 |
+
secret_signing = sign(secret_service, 'tc3_request')
|
639 |
+
signature = hmac.new(secret_signing, string_to_sign.encode('utf-8'),
|
640 |
+
hashlib.sha256).hexdigest()
|
641 |
+
|
642 |
+
# ************* 步骤 4:拼接 Authorization *************
|
643 |
+
authorization = (
|
644 |
+
algorithm + ' ' + 'Credential=' + self.secret_id + '/' +
|
645 |
+
credential_scope + ', ' + 'SignedHeaders=' + signed_headers +
|
646 |
+
', ' + 'Signature=' + signature)
|
647 |
+
|
648 |
+
# ************* 步骤 5:构造并发起请求 *************
|
649 |
+
headers = {
|
650 |
+
'Authorization': authorization,
|
651 |
+
'Content-Type': 'application/json; charset=utf-8',
|
652 |
+
'Host': self.host,
|
653 |
+
'X-TC-Action': self.action,
|
654 |
+
'X-TC-Timestamp': str(timestamp),
|
655 |
+
'X-TC-Version': self.version
|
656 |
+
}
|
657 |
+
# if self.region:
|
658 |
+
# headers["X-TC-Region"] = self.region
|
659 |
+
if self.api_key:
|
660 |
+
headers['X-TC-Token'] = self.api_key
|
661 |
+
return headers, payload
|
662 |
+
|
663 |
+
def _call_tencent_api(self, query: str) -> dict:
|
664 |
+
headers, payload = self._get_headers_and_payload(query)
|
665 |
+
req = HTTPSConnection(self.host)
|
666 |
+
req.request('POST', '/', headers=headers, body=payload.encode('utf-8'))
|
667 |
+
resp = req.getresponse()
|
668 |
+
try:
|
669 |
+
resp = json.loads(resp.read().decode('utf-8'))
|
670 |
+
except Exception as e:
|
671 |
+
logging.warning(str(e))
|
672 |
+
import ast
|
673 |
+
resp = ast.literal_eval(resp)
|
674 |
+
return resp.get('Response', dict())
|
675 |
+
|
676 |
+
async def _async_call_tencent_api(self, query: str):
|
677 |
+
headers, payload = self._get_headers_and_payload(query)
|
678 |
+
async with aiohttp.ClientSession(raise_for_status=True) as session:
|
679 |
+
async with session.post(
|
680 |
+
'https://' + self.host.lstrip('/'),
|
681 |
+
headers=headers,
|
682 |
+
data=payload) as resp:
|
683 |
+
return (await resp.json()).get('Response', {})
|
684 |
+
|
685 |
+
def _parse_response(self, response: dict) -> dict:
|
686 |
+
raw_results = []
|
687 |
+
for item in response.get('Pages', []):
|
688 |
+
display = json.loads(item['Display'])
|
689 |
+
if not display['url']:
|
690 |
+
continue
|
691 |
+
raw_results.append((display['url'], display['content']
|
692 |
+
or display['abstract_info'], display['title']))
|
693 |
+
return self._filter_results(raw_results)
|
694 |
+
|
695 |
+
|
696 |
+
class ContentFetcher:
|
697 |
+
|
698 |
+
def __init__(self, timeout: int = 5):
|
699 |
+
self.timeout = timeout
|
700 |
+
|
701 |
+
@cached(cache=TTLCache(maxsize=100, ttl=600))
|
702 |
+
def fetch(self, url: str) -> Tuple[bool, str]:
|
703 |
+
try:
|
704 |
+
response = requests.get(url, timeout=self.timeout)
|
705 |
+
response.raise_for_status()
|
706 |
+
html = response.content
|
707 |
+
except requests.RequestException as e:
|
708 |
+
return False, str(e)
|
709 |
+
|
710 |
+
text = BeautifulSoup(html, 'html.parser').get_text()
|
711 |
+
cleaned_text = re.sub(r'\n+', '\n', text)
|
712 |
+
return True, cleaned_text
|
713 |
+
|
714 |
+
@acached(cache=TTLCache(maxsize=100, ttl=600))
|
715 |
+
async def afetch(self, url: str) -> Tuple[bool, str]:
|
716 |
+
try:
|
717 |
+
async with aiohttp.ClientSession(
|
718 |
+
raise_for_status=True,
|
719 |
+
timeout=aiohttp.ClientTimeout(self.timeout)) as session:
|
720 |
+
async with session.get(url) as resp:
|
721 |
+
html = await resp.text(errors='ignore')
|
722 |
+
text = BeautifulSoup(html, 'html.parser').get_text()
|
723 |
+
cleaned_text = re.sub(r'\n+', '\n', text)
|
724 |
+
return True, cleaned_text
|
725 |
+
except Exception as e:
|
726 |
+
return False, str(e)
|
727 |
+
|
728 |
+
|
729 |
+
class WebBrowser(BaseAction):
|
730 |
+
"""Wrapper around the Web Browser Tool.
|
731 |
+
"""
|
732 |
+
|
733 |
+
def __init__(self,
|
734 |
+
searcher_type: str = 'DuckDuckGoSearch',
|
735 |
+
timeout: int = 5,
|
736 |
+
black_list: Optional[List[str]] = [
|
737 |
+
'enoN',
|
738 |
+
'youtube.com',
|
739 |
+
'bilibili.com',
|
740 |
+
'researchgate.net',
|
741 |
+
],
|
742 |
+
topk: int = 20,
|
743 |
+
description: Optional[dict] = None,
|
744 |
+
parser: Type[BaseParser] = JsonParser,
|
745 |
+
**kwargs):
|
746 |
+
self.searcher = eval(searcher_type)(
|
747 |
+
black_list=black_list, topk=topk, **kwargs)
|
748 |
+
self.fetcher = ContentFetcher(timeout=timeout)
|
749 |
+
self.search_results = None
|
750 |
+
super().__init__(description, parser)
|
751 |
+
|
752 |
+
@tool_api
|
753 |
+
def search(self, query: Union[str, List[str]]) -> dict:
|
754 |
+
"""BING search API
|
755 |
+
Args:
|
756 |
+
query (List[str]): list of search query strings
|
757 |
+
"""
|
758 |
+
queries = query if isinstance(query, list) else [query]
|
759 |
+
search_results = {}
|
760 |
+
|
761 |
+
with ThreadPoolExecutor() as executor:
|
762 |
+
future_to_query = {
|
763 |
+
executor.submit(self.searcher.search, q): q
|
764 |
+
for q in queries
|
765 |
+
}
|
766 |
+
|
767 |
+
for future in as_completed(future_to_query):
|
768 |
+
query = future_to_query[future]
|
769 |
+
try:
|
770 |
+
results = future.result()
|
771 |
+
except Exception as exc:
|
772 |
+
warnings.warn(f'{query} generated an exception: {exc}')
|
773 |
+
else:
|
774 |
+
for result in results.values():
|
775 |
+
if result['url'] not in search_results:
|
776 |
+
search_results[result['url']] = result
|
777 |
+
else:
|
778 |
+
search_results[
|
779 |
+
result['url']]['summ'] += f"\n{result['summ']}"
|
780 |
+
|
781 |
+
self.search_results = {
|
782 |
+
idx: result
|
783 |
+
for idx, result in enumerate(search_results.values())
|
784 |
+
}
|
785 |
+
return self.search_results
|
786 |
+
|
787 |
+
@tool_api
|
788 |
+
def select(self, select_ids: List[int]) -> dict:
|
789 |
+
"""get the detailed content on the selected pages.
|
790 |
+
|
791 |
+
Args:
|
792 |
+
select_ids (List[int]): list of index to select. Max number of index to be selected is no more than 4.
|
793 |
+
"""
|
794 |
+
if not self.search_results:
|
795 |
+
raise ValueError('No search results to select from.')
|
796 |
+
|
797 |
+
new_search_results = {}
|
798 |
+
with ThreadPoolExecutor() as executor:
|
799 |
+
future_to_id = {
|
800 |
+
executor.submit(self.fetcher.fetch, self.search_results[select_id]['url']): select_id
|
801 |
+
for select_id in select_ids if select_id in self.search_results
|
802 |
+
}
|
803 |
+
for future in as_completed(future_to_id):
|
804 |
+
select_id = future_to_id[future]
|
805 |
+
try:
|
806 |
+
web_success, web_content = future.result()
|
807 |
+
except Exception as exc:
|
808 |
+
warnings.warn(f'{select_id} generated an exception: {exc}')
|
809 |
+
else:
|
810 |
+
if web_success:
|
811 |
+
self.search_results[select_id][
|
812 |
+
'content'] = web_content[:8192]
|
813 |
+
new_search_results[select_id] = self.search_results[
|
814 |
+
select_id].copy()
|
815 |
+
new_search_results[select_id].pop('summ')
|
816 |
+
|
817 |
+
return new_search_results
|
818 |
+
|
819 |
+
@tool_api
|
820 |
+
def open_url(self, url: str) -> dict:
|
821 |
+
print(f'Start Browsing: {url}')
|
822 |
+
web_success, web_content = self.fetcher.fetch(url)
|
823 |
+
if web_success:
|
824 |
+
return {'type': 'text', 'content': web_content}
|
825 |
+
else:
|
826 |
+
return {'error': web_content}
|
827 |
+
|
828 |
+
|
829 |
+
class AsyncWebBrowser(AsyncActionMixin, WebBrowser):
|
830 |
+
"""Wrapper around the Web Browser Tool.
|
831 |
+
"""
|
832 |
+
|
833 |
+
@tool_api
|
834 |
+
async def search(self, query: Union[str, List[str]]) -> dict:
|
835 |
+
"""BING search API
|
836 |
+
|
837 |
+
Args:
|
838 |
+
query (List[str]): list of search query strings
|
839 |
+
"""
|
840 |
+
queries = query if isinstance(query, list) else [query]
|
841 |
+
search_results = {}
|
842 |
+
|
843 |
+
tasks = []
|
844 |
+
for q in queries:
|
845 |
+
task = asyncio.create_task(self.searcher.asearch(q))
|
846 |
+
task.query = q
|
847 |
+
tasks.append(task)
|
848 |
+
async for future in async_as_completed(tasks):
|
849 |
+
query = future.query
|
850 |
+
try:
|
851 |
+
results = await future
|
852 |
+
except Exception as exc:
|
853 |
+
warnings.warn(f'{query} generated an exception: {exc}')
|
854 |
+
else:
|
855 |
+
for result in results.values():
|
856 |
+
if result['url'] not in search_results:
|
857 |
+
search_results[result['url']] = result
|
858 |
+
else:
|
859 |
+
search_results[
|
860 |
+
result['url']]['summ'] += f"\n{result['summ']}"
|
861 |
+
|
862 |
+
self.search_results = {
|
863 |
+
idx: result
|
864 |
+
for idx, result in enumerate(search_results.values())
|
865 |
+
}
|
866 |
+
return self.search_results
|
867 |
+
|
868 |
+
@tool_api
|
869 |
+
async def select(self, select_ids: List[int]) -> dict:
|
870 |
+
"""get the detailed content on the selected pages.
|
871 |
+
|
872 |
+
Args:
|
873 |
+
select_ids (List[int]): list of index to select. Max number of index to be selected is no more than 4.
|
874 |
+
"""
|
875 |
+
if not self.search_results:
|
876 |
+
raise ValueError('No search results to select from.')
|
877 |
+
|
878 |
+
new_search_results = {}
|
879 |
+
tasks = []
|
880 |
+
for select_id in select_ids:
|
881 |
+
if select_id in self.search_results:
|
882 |
+
task = asyncio.create_task(
|
883 |
+
self.fetcher.afetch(self.search_results[select_id]['url']))
|
884 |
+
task.select_id = select_id
|
885 |
+
tasks.append(task)
|
886 |
+
async for future in async_as_completed(tasks):
|
887 |
+
select_id = future.select_id
|
888 |
+
try:
|
889 |
+
web_success, web_content = await future
|
890 |
+
except Exception as exc:
|
891 |
+
warnings.warn(f'{select_id} generated an exception: {exc}')
|
892 |
+
else:
|
893 |
+
if web_success:
|
894 |
+
self.search_results[select_id][
|
895 |
+
'content'] = web_content[:8192]
|
896 |
+
new_search_results[select_id] = self.search_results[
|
897 |
+
select_id].copy()
|
898 |
+
new_search_results[select_id].pop('summ')
|
899 |
+
return new_search_results
|
900 |
+
|
901 |
+
@tool_api
|
902 |
+
async def open_url(self, url: str) -> dict:
|
903 |
+
print(f'Start Browsing: {url}')
|
904 |
+
web_success, web_content = await self.fetcher.afetch(url)
|
905 |
+
if web_success:
|
906 |
+
return {'type': 'text', 'content': web_content}
|
907 |
+
else:
|
908 |
+
return {'error': web_content}
|
lagent/agents/__init__.py
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from .agent import Agent, AgentDict, AgentList, AsyncAgent, AsyncSequential, Sequential
|
2 |
+
from .react import AsyncReAct, ReAct
|
3 |
+
from .stream import AgentForInternLM, AsyncAgentForInternLM, AsyncMathCoder, MathCoder
|
4 |
+
|
5 |
+
__all__ = [
|
6 |
+
'Agent', 'AgentDict', 'AgentList', 'AsyncAgent', 'AgentForInternLM',
|
7 |
+
'AsyncAgentForInternLM', 'MathCoder', 'AsyncMathCoder', 'ReAct',
|
8 |
+
'AsyncReAct', 'Sequential', 'AsyncSequential'
|
9 |
+
]
|
lagent/agents/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (547 Bytes). View file
|
|
lagent/agents/__pycache__/agent.cpython-310.pyc
ADDED
Binary file (12.9 kB). View file
|
|
lagent/agents/__pycache__/react.cpython-310.pyc
ADDED
Binary file (4.85 kB). View file
|
|
lagent/agents/__pycache__/stream.cpython-310.pyc
ADDED
Binary file (8.95 kB). View file
|
|
lagent/agents/agent.py
ADDED
@@ -0,0 +1,400 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import copy
|
2 |
+
import warnings
|
3 |
+
from collections import OrderedDict, UserDict, UserList, abc
|
4 |
+
from functools import wraps
|
5 |
+
from itertools import chain, repeat
|
6 |
+
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Union
|
7 |
+
|
8 |
+
from lagent.agents.aggregator import DefaultAggregator
|
9 |
+
from lagent.hooks import Hook, RemovableHandle
|
10 |
+
from lagent.llms import BaseLLM
|
11 |
+
from lagent.memory import Memory, MemoryManager
|
12 |
+
from lagent.prompts.parsers import StrParser
|
13 |
+
from lagent.prompts.prompt_template import PromptTemplate
|
14 |
+
from lagent.schema import AgentMessage
|
15 |
+
from lagent.utils import create_object
|
16 |
+
|
17 |
+
|
18 |
+
class Agent:
|
19 |
+
"""Agent is the basic unit of the system. It is responsible for
|
20 |
+
communicating with the LLM, managing the memory, and handling the
|
21 |
+
message aggregation and parsing. It can also be extended with hooks
|
22 |
+
|
23 |
+
Args:
|
24 |
+
llm (Union[BaseLLM, Dict]): The language model used by the agent.
|
25 |
+
template (Union[PromptTemplate, str]): The template used to format the
|
26 |
+
messages.
|
27 |
+
memory (Dict): The memory used by the agent.
|
28 |
+
output_format (Dict): The output format used by the agent.
|
29 |
+
aggregator (Dict): The aggregator used by the agent.
|
30 |
+
name (Optional[str]): The name of the agent.
|
31 |
+
description (Optional[str]): The description of the agent.
|
32 |
+
hooks (Optional[Union[List[Dict], Dict]]): The hooks used by the agent.
|
33 |
+
|
34 |
+
Returns:
|
35 |
+
AgentMessage: The response message.
|
36 |
+
"""
|
37 |
+
|
38 |
+
def __init__(
|
39 |
+
self,
|
40 |
+
llm: Union[BaseLLM, Dict] = None,
|
41 |
+
template: Union[PromptTemplate, str, dict, List[dict]] = None,
|
42 |
+
memory: Dict = dict(type=Memory),
|
43 |
+
output_format: Optional[Dict] = None,
|
44 |
+
aggregator: Dict = dict(type=DefaultAggregator),
|
45 |
+
name: Optional[str] = None,
|
46 |
+
description: Optional[str] = None,
|
47 |
+
hooks: Optional[Union[List[Dict], Dict]] = None,
|
48 |
+
):
|
49 |
+
self.name = name or self.__class__.__name__
|
50 |
+
self.llm: BaseLLM = create_object(llm)
|
51 |
+
self.memory: MemoryManager = MemoryManager(memory) if memory else None
|
52 |
+
self.output_format: StrParser = create_object(output_format)
|
53 |
+
self.template = template
|
54 |
+
self.description = description
|
55 |
+
self.aggregator: DefaultAggregator = create_object(aggregator)
|
56 |
+
self._hooks: Dict[int, Hook] = OrderedDict()
|
57 |
+
if hooks:
|
58 |
+
for hook in hooks:
|
59 |
+
hook = create_object(hook)
|
60 |
+
self.register_hook(hook)
|
61 |
+
|
62 |
+
def update_memory(self, message, session_id=0):
|
63 |
+
if self.memory:
|
64 |
+
self.memory.add(message, session_id=session_id)
|
65 |
+
|
66 |
+
def __call__(
|
67 |
+
self,
|
68 |
+
*message: Union[str, AgentMessage, List[AgentMessage]],
|
69 |
+
session_id=0,
|
70 |
+
**kwargs,
|
71 |
+
) -> AgentMessage:
|
72 |
+
# message.receiver = self.name
|
73 |
+
message = [
|
74 |
+
AgentMessage(sender='user', content=m)
|
75 |
+
if isinstance(m, str) else copy.deepcopy(m) for m in message
|
76 |
+
]
|
77 |
+
for hook in self._hooks.values():
|
78 |
+
result = hook.before_agent(self, message, session_id)
|
79 |
+
if result:
|
80 |
+
message = result
|
81 |
+
self.update_memory(message, session_id=session_id)
|
82 |
+
response_message = self.forward(
|
83 |
+
*message, session_id=session_id, **kwargs)
|
84 |
+
if not isinstance(response_message, AgentMessage):
|
85 |
+
response_message = AgentMessage(
|
86 |
+
sender=self.name,
|
87 |
+
content=response_message,
|
88 |
+
)
|
89 |
+
self.update_memory(response_message, session_id=session_id)
|
90 |
+
response_message = copy.deepcopy(response_message)
|
91 |
+
for hook in self._hooks.values():
|
92 |
+
result = hook.after_agent(self, response_message, session_id)
|
93 |
+
if result:
|
94 |
+
response_message = result
|
95 |
+
return response_message
|
96 |
+
|
97 |
+
def forward(self,
|
98 |
+
*message: AgentMessage,
|
99 |
+
session_id=0,
|
100 |
+
**kwargs) -> Union[AgentMessage, str]:
|
101 |
+
formatted_messages = self.aggregator.aggregate(
|
102 |
+
self.memory.get(session_id),
|
103 |
+
self.name,
|
104 |
+
self.output_format,
|
105 |
+
self.template,
|
106 |
+
)
|
107 |
+
llm_response = self.llm.chat(formatted_messages, **kwargs)
|
108 |
+
if self.output_format:
|
109 |
+
formatted_messages = self.output_format.parse_response(
|
110 |
+
llm_response)
|
111 |
+
return AgentMessage(
|
112 |
+
sender=self.name,
|
113 |
+
content=llm_response,
|
114 |
+
formatted=formatted_messages,
|
115 |
+
)
|
116 |
+
return llm_response
|
117 |
+
|
118 |
+
def __setattr__(self, __name: str, __value: Any) -> None:
|
119 |
+
if isinstance(__value, Agent):
|
120 |
+
_agents = getattr(self, '_agents', OrderedDict())
|
121 |
+
_agents[__name] = __value
|
122 |
+
super().__setattr__('_agents', _agents)
|
123 |
+
super().__setattr__(__name, __value)
|
124 |
+
|
125 |
+
def state_dict(self, session_id=0):
|
126 |
+
state_dict, stack = {}, [('', self)]
|
127 |
+
while stack:
|
128 |
+
prefix, node = stack.pop()
|
129 |
+
key = prefix + 'memory'
|
130 |
+
if node.memory is not None:
|
131 |
+
if session_id not in node.memory.memory_map:
|
132 |
+
warnings.warn(f'No session id {session_id} in {key}')
|
133 |
+
memory = node.memory.get(session_id)
|
134 |
+
state_dict[key] = memory and memory.save() or []
|
135 |
+
if hasattr(node, '_agents'):
|
136 |
+
for name, value in reversed(node._agents.items()):
|
137 |
+
stack.append((prefix + name + '.', value))
|
138 |
+
return state_dict
|
139 |
+
|
140 |
+
def load_state_dict(self, state_dict: Dict, session_id=0):
|
141 |
+
_state_dict = self.state_dict()
|
142 |
+
missing_keys = set(_state_dict) - set(state_dict)
|
143 |
+
if missing_keys:
|
144 |
+
raise KeyError(f'Missing keys: {missing_keys}')
|
145 |
+
extra_keys = set(state_dict) - set(_state_dict)
|
146 |
+
if extra_keys:
|
147 |
+
warnings.warn(f'Mismatch keys which are not used: {extra_keys}')
|
148 |
+
for key in _state_dict:
|
149 |
+
obj = self
|
150 |
+
for attr in key.split('.')[:-1]:
|
151 |
+
if isinstance(obj, AgentList):
|
152 |
+
assert attr.isdigit()
|
153 |
+
obj = obj[int(attr)]
|
154 |
+
elif isinstance(obj, AgentDict):
|
155 |
+
obj = obj[attr]
|
156 |
+
else:
|
157 |
+
obj = getattr(obj, attr)
|
158 |
+
if obj.memory is not None:
|
159 |
+
if session_id not in obj.memory.memory_map:
|
160 |
+
obj.memory.create_instance(session_id)
|
161 |
+
obj.memory.memory_map[session_id].load(state_dict[key] or [])
|
162 |
+
|
163 |
+
def register_hook(self, hook: Callable):
|
164 |
+
handle = RemovableHandle(self._hooks)
|
165 |
+
self._hooks[handle.id] = hook
|
166 |
+
return handle
|
167 |
+
|
168 |
+
def reset(self,
|
169 |
+
session_id=0,
|
170 |
+
keypath: Optional[str] = None,
|
171 |
+
recursive: bool = False):
|
172 |
+
assert not (keypath and
|
173 |
+
recursive), 'keypath and recursive can\'t be used together'
|
174 |
+
if keypath:
|
175 |
+
keys, agent = keypath.split('.'), self
|
176 |
+
for key in keys:
|
177 |
+
agents = getattr(agent, '_agents', {})
|
178 |
+
if key not in agents:
|
179 |
+
raise KeyError(f'No sub-agent named {key} in {agent}')
|
180 |
+
agent = agents[key]
|
181 |
+
agent.reset(session_id, recursive=False)
|
182 |
+
else:
|
183 |
+
if self.memory:
|
184 |
+
self.memory.reset(session_id=session_id)
|
185 |
+
if recursive:
|
186 |
+
for agent in getattr(self, '_agents', {}).values():
|
187 |
+
agent.reset(session_id, recursive=True)
|
188 |
+
|
189 |
+
def __repr__(self):
|
190 |
+
|
191 |
+
def _rcsv_repr(agent, n_indent=1):
|
192 |
+
res = agent.__class__.__name__ + (f"(name='{agent.name}')"
|
193 |
+
if agent.name else '')
|
194 |
+
modules = [
|
195 |
+
f"{n_indent * ' '}({name}): {_rcsv_repr(agent, n_indent + 1)}"
|
196 |
+
for name, agent in getattr(agent, '_agents', {}).items()
|
197 |
+
]
|
198 |
+
if modules:
|
199 |
+
res += '(\n' + '\n'.join(
|
200 |
+
modules) + f'\n{(n_indent - 1) * " "})'
|
201 |
+
elif not res.endswith(')'):
|
202 |
+
res += '()'
|
203 |
+
return res
|
204 |
+
|
205 |
+
return _rcsv_repr(self)
|
206 |
+
|
207 |
+
|
208 |
+
class AsyncAgent(Agent):
|
209 |
+
|
210 |
+
async def __call__(self,
|
211 |
+
*message: AgentMessage | List[AgentMessage],
|
212 |
+
session_id=0,
|
213 |
+
**kwargs) -> AgentMessage:
|
214 |
+
message = [
|
215 |
+
AgentMessage(sender='user', content=m)
|
216 |
+
if isinstance(m, str) else copy.deepcopy(m) for m in message
|
217 |
+
]
|
218 |
+
for hook in self._hooks.values():
|
219 |
+
result = hook.before_agent(self, message, session_id)
|
220 |
+
if result:
|
221 |
+
message = result
|
222 |
+
self.update_memory(message, session_id=session_id)
|
223 |
+
response_message = await self.forward(
|
224 |
+
*message, session_id=session_id, **kwargs)
|
225 |
+
if not isinstance(response_message, AgentMessage):
|
226 |
+
response_message = AgentMessage(
|
227 |
+
sender=self.name,
|
228 |
+
content=response_message,
|
229 |
+
)
|
230 |
+
self.update_memory(response_message, session_id=session_id)
|
231 |
+
response_message = copy.deepcopy(response_message)
|
232 |
+
for hook in self._hooks.values():
|
233 |
+
result = hook.after_agent(self, response_message, session_id)
|
234 |
+
if result:
|
235 |
+
response_message = result
|
236 |
+
return response_message
|
237 |
+
|
238 |
+
async def forward(self,
|
239 |
+
*message: AgentMessage,
|
240 |
+
session_id=0,
|
241 |
+
**kwargs) -> Union[AgentMessage, str]:
|
242 |
+
formatted_messages = self.aggregator.aggregate(
|
243 |
+
self.memory.get(session_id),
|
244 |
+
self.name,
|
245 |
+
self.output_format,
|
246 |
+
self.template,
|
247 |
+
)
|
248 |
+
llm_response = await self.llm.chat(formatted_messages, session_id,
|
249 |
+
**kwargs)
|
250 |
+
if self.output_format:
|
251 |
+
formatted_messages = self.output_format.parse_response(
|
252 |
+
llm_response)
|
253 |
+
return AgentMessage(
|
254 |
+
sender=self.name,
|
255 |
+
content=llm_response,
|
256 |
+
formatted=formatted_messages,
|
257 |
+
)
|
258 |
+
return llm_response
|
259 |
+
|
260 |
+
|
261 |
+
class Sequential(Agent):
|
262 |
+
"""Sequential is an agent container that forwards messages to each agent
|
263 |
+
in the order they are added."""
|
264 |
+
|
265 |
+
def __init__(self, *agents: Union[Agent, AsyncAgent, Iterable], **kwargs):
|
266 |
+
super().__init__(**kwargs)
|
267 |
+
self._agents = OrderedDict()
|
268 |
+
if not agents:
|
269 |
+
raise ValueError('At least one agent should be provided')
|
270 |
+
if isinstance(agents[0],
|
271 |
+
Iterable) and not isinstance(agents[0], Agent):
|
272 |
+
if not agents[0]:
|
273 |
+
raise ValueError('At least one agent should be provided')
|
274 |
+
agents = agents[0]
|
275 |
+
for key, agent in enumerate(agents):
|
276 |
+
if isinstance(agents, Mapping):
|
277 |
+
key, agent = agent, agents[agent]
|
278 |
+
elif isinstance(agent, tuple):
|
279 |
+
key, agent = agent
|
280 |
+
self.add_agent(key, agent)
|
281 |
+
|
282 |
+
def add_agent(self, name: str, agent: Union[Agent, AsyncAgent]):
|
283 |
+
assert isinstance(
|
284 |
+
agent, (Agent, AsyncAgent
|
285 |
+
)), f'{type(agent)} is not an Agent or AsyncAgent subclass'
|
286 |
+
self._agents[str(name)] = agent
|
287 |
+
|
288 |
+
def forward(self,
|
289 |
+
*message: AgentMessage,
|
290 |
+
session_id=0,
|
291 |
+
exit_at: Optional[int] = None,
|
292 |
+
**kwargs) -> AgentMessage:
|
293 |
+
assert exit_at is None or exit_at >= 0, 'exit_at should be greater than or equal to 0'
|
294 |
+
if exit_at is None:
|
295 |
+
exit_at = len(self) - 1
|
296 |
+
iterator = chain.from_iterable(repeat(self._agents.values()))
|
297 |
+
for _ in range(exit_at + 1):
|
298 |
+
agent = next(iterator)
|
299 |
+
if isinstance(message, AgentMessage):
|
300 |
+
message = (message, )
|
301 |
+
message = agent(*message, session_id=session_id, **kwargs)
|
302 |
+
return message
|
303 |
+
|
304 |
+
def __getitem__(self, key):
|
305 |
+
if isinstance(key, int) and key < 0:
|
306 |
+
assert key >= -len(self), 'index out of range'
|
307 |
+
key = len(self) + key
|
308 |
+
return self._agents[str(key)]
|
309 |
+
|
310 |
+
def __len__(self):
|
311 |
+
return len(self._agents)
|
312 |
+
|
313 |
+
|
314 |
+
class AsyncSequential(Sequential, AsyncAgent):
|
315 |
+
|
316 |
+
async def forward(self,
|
317 |
+
*message: AgentMessage,
|
318 |
+
session_id=0,
|
319 |
+
exit_at: Optional[int] = None,
|
320 |
+
**kwargs) -> AgentMessage:
|
321 |
+
assert exit_at is None or exit_at >= 0, 'exit_at should be greater than or equal to 0'
|
322 |
+
if exit_at is None:
|
323 |
+
exit_at = len(self) - 1
|
324 |
+
iterator = chain.from_iterable(repeat(self._agents.values()))
|
325 |
+
for _ in range(exit_at + 1):
|
326 |
+
agent = next(iterator)
|
327 |
+
if isinstance(message, AgentMessage):
|
328 |
+
message = (message, )
|
329 |
+
message = await agent(*message, session_id=session_id, **kwargs)
|
330 |
+
return message
|
331 |
+
|
332 |
+
|
333 |
+
class AgentContainerMixin:
|
334 |
+
|
335 |
+
def __init_subclass__(cls):
|
336 |
+
super().__init_subclass__()
|
337 |
+
|
338 |
+
def wrap_api(func):
|
339 |
+
|
340 |
+
@wraps(func)
|
341 |
+
def wrapped_func(self, *args, **kwargs):
|
342 |
+
data = self.data.copy() if hasattr(self, 'data') else None
|
343 |
+
|
344 |
+
def _backup(d):
|
345 |
+
if d is None:
|
346 |
+
self.data.clear()
|
347 |
+
else:
|
348 |
+
self.data = d
|
349 |
+
|
350 |
+
ret = func(self, *args, **kwargs)
|
351 |
+
agents = OrderedDict()
|
352 |
+
for k, item in (self.data.items() if isinstance(
|
353 |
+
self.data, abc.Mapping) else enumerate(self.data)):
|
354 |
+
if isinstance(self.data,
|
355 |
+
abc.Mapping) and not isinstance(k, str):
|
356 |
+
_backup(data)
|
357 |
+
raise KeyError(
|
358 |
+
f'agent name should be a string, got {type(k)}')
|
359 |
+
if isinstance(k, str) and '.' in k:
|
360 |
+
_backup(data)
|
361 |
+
raise KeyError(
|
362 |
+
f'agent name can\'t contain ".", got {k}')
|
363 |
+
if not isinstance(item, (Agent, AsyncAgent)):
|
364 |
+
_backup(data)
|
365 |
+
raise TypeError(
|
366 |
+
f'{type(item)} is not an Agent or AsyncAgent subclass'
|
367 |
+
)
|
368 |
+
agents[str(k)] = item
|
369 |
+
self._agents = agents
|
370 |
+
return ret
|
371 |
+
|
372 |
+
return wrapped_func
|
373 |
+
|
374 |
+
for method in [
|
375 |
+
'append', 'sort', 'reverse', 'pop', 'clear', 'update',
|
376 |
+
'insert', 'extend', 'remove', '__init__', '__setitem__',
|
377 |
+
'__delitem__', '__add__', '__iadd__', '__radd__', '__mul__',
|
378 |
+
'__imul__', '__rmul__'
|
379 |
+
]:
|
380 |
+
if hasattr(cls, method):
|
381 |
+
setattr(cls, method, wrap_api(getattr(cls, method)))
|
382 |
+
|
383 |
+
|
384 |
+
class AgentList(Agent, UserList, AgentContainerMixin):
|
385 |
+
|
386 |
+
def __init__(self,
|
387 |
+
agents: Optional[Iterable[Union[Agent, AsyncAgent]]] = None):
|
388 |
+
Agent.__init__(self, memory=None)
|
389 |
+
UserList.__init__(self, agents)
|
390 |
+
self.name = None
|
391 |
+
|
392 |
+
|
393 |
+
class AgentDict(Agent, UserDict, AgentContainerMixin):
|
394 |
+
|
395 |
+
def __init__(self,
|
396 |
+
agents: Optional[Mapping[str, Union[Agent,
|
397 |
+
AsyncAgent]]] = None):
|
398 |
+
Agent.__init__(self, memory=None)
|
399 |
+
UserDict.__init__(self, agents)
|
400 |
+
self.name = None
|
lagent/agents/aggregator/__init__.py
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from .default_aggregator import DefaultAggregator
|
2 |
+
from .tool_aggregator import InternLMToolAggregator
|
3 |
+
|
4 |
+
__all__ = ['DefaultAggregator', 'InternLMToolAggregator']
|
lagent/agents/aggregator/__pycache__/__init__.cpython-310.pyc
ADDED
Binary file (305 Bytes). View file
|
|
lagent/agents/aggregator/__pycache__/default_aggregator.cpython-310.pyc
ADDED
Binary file (1.6 kB). View file
|
|
lagent/agents/aggregator/__pycache__/tool_aggregator.cpython-310.pyc
ADDED
Binary file (2.71 kB). View file
|
|
lagent/agents/aggregator/default_aggregator.py
ADDED
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Dict, List
|
2 |
+
|
3 |
+
from lagent.memory import Memory
|
4 |
+
from lagent.prompts import StrParser
|
5 |
+
|
6 |
+
|
7 |
+
class DefaultAggregator:
|
8 |
+
|
9 |
+
def aggregate(self,
|
10 |
+
messages: Memory,
|
11 |
+
name: str,
|
12 |
+
parser: StrParser = None,
|
13 |
+
system_instruction: str = None) -> List[Dict[str, str]]:
|
14 |
+
_message = []
|
15 |
+
messages = messages.get_memory()
|
16 |
+
if system_instruction:
|
17 |
+
_message.extend(
|
18 |
+
self.aggregate_system_intruction(system_instruction))
|
19 |
+
for message in messages:
|
20 |
+
if message.sender == name:
|
21 |
+
_message.append(
|
22 |
+
dict(role='assistant', content=str(message.content)))
|
23 |
+
else:
|
24 |
+
user_message = message.content
|
25 |
+
if len(_message) > 0 and _message[-1]['role'] == 'user':
|
26 |
+
_message[-1]['content'] += user_message
|
27 |
+
else:
|
28 |
+
_message.append(dict(role='user', content=user_message))
|
29 |
+
return _message
|
30 |
+
|
31 |
+
@staticmethod
|
32 |
+
def aggregate_system_intruction(system_intruction) -> List[dict]:
|
33 |
+
if isinstance(system_intruction, str):
|
34 |
+
system_intruction = dict(role='system', content=system_intruction)
|
35 |
+
if isinstance(system_intruction, dict):
|
36 |
+
system_intruction = [system_intruction]
|
37 |
+
if isinstance(system_intruction, list):
|
38 |
+
for msg in system_intruction:
|
39 |
+
if not isinstance(msg, dict):
|
40 |
+
raise TypeError(f'Unsupported message type: {type(msg)}')
|
41 |
+
if not ('role' in msg and 'content' in msg):
|
42 |
+
raise KeyError(
|
43 |
+
f"Missing required key 'role' or 'content': {msg}")
|
44 |
+
return system_intruction
|
lagent/agents/aggregator/tool_aggregator.py
ADDED
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Dict, List, Optional, Union
|
2 |
+
|
3 |
+
from lagent.agents.aggregator.default_aggregator import DefaultAggregator
|
4 |
+
from lagent.memory.base_memory import Memory
|
5 |
+
from lagent.prompts.parsers.tool_parser import MixedToolParser, ToolParser, ToolStatusCode
|
6 |
+
|
7 |
+
|
8 |
+
class InternLMToolAggregator(DefaultAggregator):
|
9 |
+
|
10 |
+
def __init__(self,
|
11 |
+
environment_role='environment',
|
12 |
+
environment_begin='',
|
13 |
+
environment_end='',
|
14 |
+
user_names: Optional[List[str]] = None,
|
15 |
+
few_shot: Optional[List[List[dict]]] = None):
|
16 |
+
self.environment_role = environment_role
|
17 |
+
self.environment_begin = environment_begin
|
18 |
+
self.environment_end = environment_end
|
19 |
+
self.user_names = user_names or ['user']
|
20 |
+
self.few_shot = few_shot or []
|
21 |
+
|
22 |
+
def aggregate(self,
|
23 |
+
messages: Memory,
|
24 |
+
name: str,
|
25 |
+
parser: Union[ToolParser, MixedToolParser],
|
26 |
+
system_instruction: str = None) -> List[Dict[str, str]]:
|
27 |
+
_message = []
|
28 |
+
messages = messages.get_memory()
|
29 |
+
if system_instruction:
|
30 |
+
_message.extend(
|
31 |
+
self.aggregate_system_intruction(system_instruction))
|
32 |
+
tool_instruction = parser.format_instruction()
|
33 |
+
if tool_instruction:
|
34 |
+
if isinstance(tool_instruction, str):
|
35 |
+
tool_instruction = dict(
|
36 |
+
role='system', content=tool_instruction)
|
37 |
+
if parser.tool_type:
|
38 |
+
tool_instruction['name'] = parser.tool_type
|
39 |
+
if isinstance(tool_instruction, dict):
|
40 |
+
tool_instruction = [tool_instruction]
|
41 |
+
_message.extend(tool_instruction)
|
42 |
+
|
43 |
+
for shot in self.few_shot:
|
44 |
+
i = 0
|
45 |
+
while i < len(shot):
|
46 |
+
msg = shot[i]
|
47 |
+
if msg['role'] in ['assistant', 'user', 'system']:
|
48 |
+
_message.append(msg)
|
49 |
+
elif msg['role'] == self.environment_role:
|
50 |
+
if not msg['content'].startswith(self.environment_begin):
|
51 |
+
msg['content'] = self.environment_begin + msg['content']
|
52 |
+
if not msg['content'].endswith(self.environment_end):
|
53 |
+
msg['content'] += self.environment_end
|
54 |
+
_message.append(msg)
|
55 |
+
elif msg['role'] in ['thought', 'language']:
|
56 |
+
if i < len(shot) - 1 and shot[i + 1]['role'] == 'tool':
|
57 |
+
_message.append(
|
58 |
+
dict(
|
59 |
+
role='assistant',
|
60 |
+
content=parser.format_response(
|
61 |
+
dict(
|
62 |
+
tool_type=shot[i + 1]['name'],
|
63 |
+
thought=msg['content'],
|
64 |
+
action=shot[i + 1]['content'],
|
65 |
+
status=None))))
|
66 |
+
i += 1
|
67 |
+
else:
|
68 |
+
_message.append(
|
69 |
+
dict(
|
70 |
+
role='assistant',
|
71 |
+
content=parser.format_response(
|
72 |
+
dict(
|
73 |
+
tool_type=None,
|
74 |
+
thought=msg['content'],
|
75 |
+
action=None,
|
76 |
+
status=None))))
|
77 |
+
else:
|
78 |
+
raise KeyError(f'Unkown role: {msg["role"]}')
|
79 |
+
i += 1
|
80 |
+
|
81 |
+
tool_type = None
|
82 |
+
for message in messages:
|
83 |
+
if message.sender == name:
|
84 |
+
if isinstance(message.formatted, dict):
|
85 |
+
parsed = message.formatted
|
86 |
+
if parsed['status'] == ToolStatusCode.PARSING_ERROR:
|
87 |
+
continue
|
88 |
+
_message.append(
|
89 |
+
dict(
|
90 |
+
role='assistant',
|
91 |
+
content=parser.format_response(parsed)))
|
92 |
+
tool_type = parsed['tool_type']
|
93 |
+
else:
|
94 |
+
_message.append(
|
95 |
+
dict(role='assistant', content=str(message.content)))
|
96 |
+
elif message.sender in self.user_names:
|
97 |
+
_message.append(dict(role='user', content=message.content))
|
98 |
+
else:
|
99 |
+
msg = dict(
|
100 |
+
role=self.environment_role,
|
101 |
+
content=self.environment_begin + str(message.content) +
|
102 |
+
self.environment_end)
|
103 |
+
if tool_type:
|
104 |
+
msg['name'] = tool_type
|
105 |
+
_message.append(msg)
|
106 |
+
return _message
|
lagent/agents/react.py
ADDED
@@ -0,0 +1,161 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
from typing import Callable, Dict, List, Union
|
3 |
+
|
4 |
+
from pydantic import BaseModel, Field
|
5 |
+
|
6 |
+
from lagent.actions import ActionExecutor, AsyncActionExecutor, BaseAction
|
7 |
+
from lagent.agents.agent import Agent, AsyncAgent
|
8 |
+
from lagent.agents.aggregator import DefaultAggregator
|
9 |
+
from lagent.hooks import ActionPreprocessor
|
10 |
+
from lagent.llms import BaseLLM
|
11 |
+
from lagent.memory import Memory
|
12 |
+
from lagent.prompts.parsers.json_parser import JSONParser
|
13 |
+
from lagent.prompts.prompt_template import PromptTemplate
|
14 |
+
from lagent.schema import AgentMessage
|
15 |
+
from lagent.utils import create_object
|
16 |
+
|
17 |
+
select_action_template = """你是一个可以调用外部工具的助手,可以使用的工具包括:
|
18 |
+
{action_info}
|
19 |
+
{output_format}
|
20 |
+
开始!"""
|
21 |
+
|
22 |
+
output_format_template = """如果使用工具请遵循以下格式回复:
|
23 |
+
{function_format}
|
24 |
+
|
25 |
+
如果你已经知道了答案,或者你不需要工具,请遵循以下格式回复
|
26 |
+
{finish_format}"""
|
27 |
+
|
28 |
+
|
29 |
+
class ReAct(Agent):
|
30 |
+
|
31 |
+
def __init__(self,
|
32 |
+
llm: Union[BaseLLM, Dict],
|
33 |
+
actions: Union[BaseAction, List[BaseAction]],
|
34 |
+
template: Union[PromptTemplate, str] = None,
|
35 |
+
memory: Dict = dict(type=Memory),
|
36 |
+
output_format: Dict = dict(type=JSONParser),
|
37 |
+
aggregator: Dict = dict(type=DefaultAggregator),
|
38 |
+
hooks: List = [dict(type=ActionPreprocessor)],
|
39 |
+
finish_condition: Callable[[AgentMessage], bool] = lambda m:
|
40 |
+
'conclusion' in m.content or 'conclusion' in m.formatted,
|
41 |
+
max_turn: int = 5,
|
42 |
+
**kwargs):
|
43 |
+
self.max_turn = max_turn
|
44 |
+
self.finish_condition = finish_condition
|
45 |
+
actions = dict(
|
46 |
+
type=ActionExecutor,
|
47 |
+
actions=actions,
|
48 |
+
hooks=hooks,
|
49 |
+
)
|
50 |
+
self.actions: ActionExecutor = create_object(actions)
|
51 |
+
select_agent = dict(
|
52 |
+
type=Agent,
|
53 |
+
llm=llm,
|
54 |
+
template=template.format(
|
55 |
+
action_info=json.dumps(self.actions.description()),
|
56 |
+
output_format=output_format.format_instruction()),
|
57 |
+
output_format=output_format,
|
58 |
+
memory=memory,
|
59 |
+
aggregator=aggregator,
|
60 |
+
hooks=hooks,
|
61 |
+
)
|
62 |
+
self.select_agent = create_object(select_agent)
|
63 |
+
super().__init__(**kwargs)
|
64 |
+
|
65 |
+
def forward(self, message: AgentMessage, **kwargs) -> AgentMessage:
|
66 |
+
for _ in range(self.max_turn):
|
67 |
+
message = self.select_agent(message)
|
68 |
+
if self.finish_condition(message):
|
69 |
+
return message
|
70 |
+
message = self.actions(message)
|
71 |
+
return message
|
72 |
+
|
73 |
+
|
74 |
+
class AsyncReAct(AsyncAgent):
|
75 |
+
|
76 |
+
def __init__(self,
|
77 |
+
llm: Union[BaseLLM, Dict],
|
78 |
+
actions: Union[BaseAction, List[BaseAction]],
|
79 |
+
template: Union[PromptTemplate, str] = None,
|
80 |
+
memory: Dict = dict(type=Memory),
|
81 |
+
output_format: Dict = dict(type=JSONParser),
|
82 |
+
aggregator: Dict = dict(type=DefaultAggregator),
|
83 |
+
hooks: List = [dict(type=ActionPreprocessor)],
|
84 |
+
finish_condition: Callable[[AgentMessage], bool] = lambda m:
|
85 |
+
'conclusion' in m.content or 'conclusion' in m.formatted,
|
86 |
+
max_turn: int = 5,
|
87 |
+
**kwargs):
|
88 |
+
self.max_turn = max_turn
|
89 |
+
self.finish_condition = finish_condition
|
90 |
+
actions = dict(
|
91 |
+
type=AsyncActionExecutor,
|
92 |
+
actions=actions,
|
93 |
+
hooks=hooks,
|
94 |
+
)
|
95 |
+
self.actions: AsyncActionExecutor = create_object(actions)
|
96 |
+
select_agent = dict(
|
97 |
+
type=AsyncAgent,
|
98 |
+
llm=llm,
|
99 |
+
template=template.format(
|
100 |
+
action_info=json.dumps(self.actions.description()),
|
101 |
+
output_format=output_format.format_instruction()),
|
102 |
+
output_format=output_format,
|
103 |
+
memory=memory,
|
104 |
+
aggregator=aggregator,
|
105 |
+
hooks=hooks,
|
106 |
+
)
|
107 |
+
self.select_agent = create_object(select_agent)
|
108 |
+
super().__init__(**kwargs)
|
109 |
+
|
110 |
+
async def forward(self, message: AgentMessage, **kwargs) -> AgentMessage:
|
111 |
+
for _ in range(self.max_turn):
|
112 |
+
message = await self.select_agent(message)
|
113 |
+
if self.finish_condition(message):
|
114 |
+
return message
|
115 |
+
message = await self.actions(message)
|
116 |
+
return message
|
117 |
+
|
118 |
+
|
119 |
+
if __name__ == '__main__':
|
120 |
+
from lagent.llms import GPTAPI
|
121 |
+
|
122 |
+
class ActionCall(BaseModel):
|
123 |
+
name: str = Field(description='调用的函数名称')
|
124 |
+
parameters: Dict = Field(description='调用函数的参数')
|
125 |
+
|
126 |
+
class ActionFormat(BaseModel):
|
127 |
+
thought_process: str = Field(
|
128 |
+
description='描述当前所处的状态和已知信息。这有助于明确目前所掌握的信息和接下来的搜索方向。')
|
129 |
+
action: ActionCall = Field(description='当前步骤需要执行的操作,包括函数名��和参数。')
|
130 |
+
|
131 |
+
class FinishFormat(BaseModel):
|
132 |
+
thought_process: str = Field(
|
133 |
+
description='描述当前所处的状态和已知信息。这有助于明确目前所掌握的信息和接下来的搜索方向。')
|
134 |
+
conclusion: str = Field(description='总结当前的搜索结果,回答问题。')
|
135 |
+
|
136 |
+
prompt_template = PromptTemplate(select_action_template)
|
137 |
+
output_format = JSONParser(
|
138 |
+
output_format_template,
|
139 |
+
function_format=ActionFormat,
|
140 |
+
finish_format=FinishFormat)
|
141 |
+
|
142 |
+
llm = dict(
|
143 |
+
type=GPTAPI,
|
144 |
+
model_type='gpt-4o-2024-05-13',
|
145 |
+
key=None,
|
146 |
+
max_new_tokens=4096,
|
147 |
+
proxies=dict(),
|
148 |
+
retry=1000)
|
149 |
+
|
150 |
+
agent = ReAct(
|
151 |
+
llm=llm,
|
152 |
+
template=prompt_template,
|
153 |
+
output_format=output_format,
|
154 |
+
aggregator=dict(type='DefaultAggregator'),
|
155 |
+
actions=[dict(type='PythonInterpreter')],
|
156 |
+
)
|
157 |
+
response = agent(
|
158 |
+
AgentMessage(sender='user', content='用 Python 计算一下 3 ** 5'))
|
159 |
+
print(response)
|
160 |
+
response = agent(AgentMessage(sender='user', content=' 2 ** 5 呢'))
|
161 |
+
print(response)
|
lagent/agents/stream.py
ADDED
@@ -0,0 +1,316 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
import warnings
|
3 |
+
from copy import deepcopy
|
4 |
+
from typing import Callable, Dict, List, Union
|
5 |
+
|
6 |
+
from lagent.actions import ActionExecutor, AsyncActionExecutor, AsyncIPythonInterpreter, IPythonInteractive
|
7 |
+
from lagent.agents.agent import Agent, AsyncAgent
|
8 |
+
from lagent.agents.aggregator import InternLMToolAggregator
|
9 |
+
from lagent.hooks import InternLMActionProcessor
|
10 |
+
from lagent.llms import BaseLLM
|
11 |
+
from lagent.memory import Memory
|
12 |
+
from lagent.prompts.parsers import InterpreterParser, MixedToolParser, PluginParser, ToolStatusCode
|
13 |
+
from lagent.schema import AgentMessage
|
14 |
+
from lagent.utils import create_object
|
15 |
+
|
16 |
+
API_PREFIX = (
|
17 |
+
"This is the subfunction for tool '{tool_name}', you can use this tool. "
|
18 |
+
'The description of this function is: \n{description}')
|
19 |
+
|
20 |
+
META_CN = ('当开启工具以及代码时,根据需求选择合适的工具进行调用')
|
21 |
+
|
22 |
+
INTERPRETER_CN = ('你现在已经能够在一个有状态的 Jupyter 笔记本环境中运行 Python 代码。'
|
23 |
+
'当你向 python 发送含有 Python 代码的消息时,它将在该环境中执行。'
|
24 |
+
'这个工具适用于多种场景,如数据分析或处理(包括数据操作、统计分析、图表绘制),'
|
25 |
+
'复杂的计算问题(解决数学和物理难题),编程示例(理解编程概念或特性),'
|
26 |
+
'文本处理和分析(比如文本解析和自然语言处理),'
|
27 |
+
'机器学习和数据科学(用于展示模型训练和数据可视化),'
|
28 |
+
'以及文件操作和数据导入(处理CSV、JSON等格式的文件)。')
|
29 |
+
|
30 |
+
PLUGIN_CN = ('你可以使用如下工具:'
|
31 |
+
'\n{prompt}\n'
|
32 |
+
'如果你已经获得足够信息,请直接给出答案. 避免不必要的工具调用! '
|
33 |
+
'同时注意你可以使用的工具,不要随意捏造!')
|
34 |
+
|
35 |
+
|
36 |
+
def get_plugin_prompt(actions, api_desc_template=API_PREFIX):
|
37 |
+
plugin_descriptions = []
|
38 |
+
for action in actions if isinstance(actions, list) else [actions]:
|
39 |
+
action = create_object(action)
|
40 |
+
action_desc = deepcopy(action.description)
|
41 |
+
if action.is_toolkit:
|
42 |
+
for api in action_desc['api_list']:
|
43 |
+
api['name'] = f"{action.name}.{api['name']}"
|
44 |
+
api['description'] = api_desc_template.format(
|
45 |
+
tool_name=action.name, description=api['description'])
|
46 |
+
api['parameters'] = [
|
47 |
+
param for param in api['parameters']
|
48 |
+
if param['name'] in api['required']
|
49 |
+
]
|
50 |
+
plugin_descriptions.append(api)
|
51 |
+
else:
|
52 |
+
action_desc['description'] = api_desc_template.format(
|
53 |
+
tool_name=action.name, description=action_desc['description'])
|
54 |
+
action_desc['parameters'] = [
|
55 |
+
param for param in action_desc['parameters']
|
56 |
+
if param['name'] in action_desc['required']
|
57 |
+
]
|
58 |
+
plugin_descriptions.append(action_desc)
|
59 |
+
return json.dumps(plugin_descriptions, ensure_ascii=False, indent=4)
|
60 |
+
|
61 |
+
|
62 |
+
class AgentForInternLM(Agent):
|
63 |
+
|
64 |
+
_INTERNAL_AGENT_CLS = Agent
|
65 |
+
|
66 |
+
def __init__(
|
67 |
+
self,
|
68 |
+
llm: Union[BaseLLM, Dict],
|
69 |
+
plugins: Union[dict, List[dict]] = None,
|
70 |
+
interpreter: dict = None,
|
71 |
+
template: Union[str, dict, List[dict]] = None,
|
72 |
+
memory: Dict = dict(type=Memory),
|
73 |
+
output_format: Dict = dict(
|
74 |
+
type=MixedToolParser,
|
75 |
+
template=META_CN,
|
76 |
+
parsers=[
|
77 |
+
dict(type=PluginParser, template=PLUGIN_CN),
|
78 |
+
dict(type=InterpreterParser, template=INTERPRETER_CN),
|
79 |
+
]),
|
80 |
+
aggregator: Dict = dict(type=InternLMToolAggregator),
|
81 |
+
action_hooks: List = [dict(type=InternLMActionProcessor)],
|
82 |
+
finish_condition: Callable[
|
83 |
+
[AgentMessage],
|
84 |
+
bool] = lambda m: m.formatted['status'] == ToolStatusCode.NO_TOOL,
|
85 |
+
max_turn: int = 4,
|
86 |
+
**kwargs,
|
87 |
+
):
|
88 |
+
agent = dict(
|
89 |
+
type=self._INTERNAL_AGENT_CLS,
|
90 |
+
llm=llm,
|
91 |
+
template=template,
|
92 |
+
output_format=output_format,
|
93 |
+
memory=memory,
|
94 |
+
aggregator=aggregator,
|
95 |
+
hooks=kwargs.pop('hooks', None),
|
96 |
+
)
|
97 |
+
self.agent = create_object(agent)
|
98 |
+
self.plugin_executor = plugins and ActionExecutor(
|
99 |
+
plugins, hooks=action_hooks)
|
100 |
+
self.interpreter_executor = interpreter and ActionExecutor(
|
101 |
+
interpreter, hooks=action_hooks)
|
102 |
+
if not (self.plugin_executor or self.interpreter_executor):
|
103 |
+
warnings.warn(
|
104 |
+
'Neither plugin nor interpreter executor is initialized. '
|
105 |
+
'An exception will be thrown when the agent call a tool.')
|
106 |
+
self.finish_condition = finish_condition
|
107 |
+
self.max_turn = max_turn
|
108 |
+
super().__init__(**kwargs)
|
109 |
+
|
110 |
+
def forward(self, message: AgentMessage, session_id=0, **kwargs):
|
111 |
+
if isinstance(message, str):
|
112 |
+
message = AgentMessage(sender='user', content=message)
|
113 |
+
for _ in range(self.max_turn):
|
114 |
+
message = self.agent(message, session_id=session_id, **kwargs)
|
115 |
+
assert isinstance(message.formatted, dict)
|
116 |
+
if self.finish_condition(message):
|
117 |
+
return message
|
118 |
+
if message.formatted['tool_type']:
|
119 |
+
tool_type = message.formatted["tool_type"]
|
120 |
+
executor = getattr(self, f'{tool_type}_executor', None)
|
121 |
+
if not executor:
|
122 |
+
raise RuntimeError(f'No available {tool_type} executor')
|
123 |
+
message = executor(message, session_id=session_id)
|
124 |
+
return message
|
125 |
+
|
126 |
+
def get_steps(self, session_id=0):
|
127 |
+
steps, tool_type = [], None
|
128 |
+
for msg in self.agent.memory.get_memory(session_id):
|
129 |
+
if msg.sender == self.agent.name:
|
130 |
+
steps.append(
|
131 |
+
dict(role='thought', content=msg.formatted['thought']))
|
132 |
+
if msg.formatted['tool_type']:
|
133 |
+
tool_type = msg.formatted['tool_type']
|
134 |
+
steps.append(
|
135 |
+
dict(
|
136 |
+
role='tool',
|
137 |
+
content=msg.formatted['action'],
|
138 |
+
name=tool_type))
|
139 |
+
elif msg.sender != 'user':
|
140 |
+
feedback = dict(role='environment', content=msg.content)
|
141 |
+
if tool_type:
|
142 |
+
feedback['name'] = tool_type
|
143 |
+
steps.append(feedback)
|
144 |
+
return steps
|
145 |
+
|
146 |
+
|
147 |
+
class MathCoder(AgentForInternLM):
|
148 |
+
|
149 |
+
def __init__(
|
150 |
+
self,
|
151 |
+
llm: Union[BaseLLM, Dict],
|
152 |
+
interpreter: dict = dict(
|
153 |
+
type=IPythonInteractive, timeout=20, max_out_len=8192),
|
154 |
+
template: Union[str, dict, List[dict]] = None,
|
155 |
+
memory: Dict = dict(type=Memory),
|
156 |
+
output_format: Dict = dict(
|
157 |
+
type=InterpreterParser,
|
158 |
+
template=
|
159 |
+
('Integrate step-by-step reasoning and Python code to solve math problems '
|
160 |
+
'using the following guidelines:\n'
|
161 |
+
'- Analyze the question and write jupyter code to solve the problem;\n'
|
162 |
+
r"- Present the final result in LaTeX using a '\boxed{{}}' without any "
|
163 |
+
'units. \n')),
|
164 |
+
aggregator: Dict = dict(type=InternLMToolAggregator),
|
165 |
+
action_hooks: List = [dict(type=InternLMActionProcessor)],
|
166 |
+
finish_condition: Callable[
|
167 |
+
[AgentMessage],
|
168 |
+
bool] = lambda m: m.formatted['status'] == ToolStatusCode.NO_TOOL,
|
169 |
+
max_turn: int = 6,
|
170 |
+
**kwargs,
|
171 |
+
):
|
172 |
+
kwargs.pop('plugins', None)
|
173 |
+
super().__init__(
|
174 |
+
llm=llm,
|
175 |
+
interpreter=interpreter,
|
176 |
+
template=template,
|
177 |
+
memory=memory,
|
178 |
+
output_format=output_format,
|
179 |
+
aggregator=aggregator,
|
180 |
+
action_hooks=action_hooks,
|
181 |
+
finish_condition=finish_condition,
|
182 |
+
max_turn=max_turn,
|
183 |
+
**kwargs)
|
184 |
+
|
185 |
+
|
186 |
+
class AsyncAgentForInternLM(AsyncAgent):
|
187 |
+
|
188 |
+
_INTERNAL_AGENT_CLS = AsyncAgent
|
189 |
+
|
190 |
+
def __init__(
|
191 |
+
self,
|
192 |
+
llm: Union[BaseLLM, Dict],
|
193 |
+
plugins: Union[dict, List[dict]] = None,
|
194 |
+
interpreter: dict = None,
|
195 |
+
template: Union[str, dict, List[dict]] = None,
|
196 |
+
memory: Dict = dict(type=Memory),
|
197 |
+
output_format: Dict = dict(
|
198 |
+
type=MixedToolParser,
|
199 |
+
template=META_CN,
|
200 |
+
parsers=[
|
201 |
+
dict(type=PluginParser, template=PLUGIN_CN),
|
202 |
+
dict(type=InterpreterParser, template=INTERPRETER_CN),
|
203 |
+
]),
|
204 |
+
aggregator: Dict = dict(type=InternLMToolAggregator),
|
205 |
+
action_hooks: List = [dict(type=InternLMActionProcessor)],
|
206 |
+
finish_condition: Callable[
|
207 |
+
[AgentMessage],
|
208 |
+
bool] = lambda m: m.formatted['status'] == ToolStatusCode.NO_TOOL,
|
209 |
+
max_turn: int = 4,
|
210 |
+
**kwargs,
|
211 |
+
):
|
212 |
+
agent = dict(
|
213 |
+
type=self._INTERNAL_AGENT_CLS,
|
214 |
+
llm=llm,
|
215 |
+
template=template,
|
216 |
+
output_format=output_format,
|
217 |
+
memory=memory,
|
218 |
+
aggregator=aggregator,
|
219 |
+
hooks=kwargs.pop('hooks', None),
|
220 |
+
)
|
221 |
+
self.agent = create_object(agent)
|
222 |
+
self.plugin_executor = plugins and AsyncActionExecutor(
|
223 |
+
plugins, hooks=action_hooks)
|
224 |
+
self.interpreter_executor = interpreter and AsyncActionExecutor(
|
225 |
+
interpreter, hooks=action_hooks)
|
226 |
+
if not (self.plugin_executor or self.interpreter_executor):
|
227 |
+
warnings.warn(
|
228 |
+
'Neither plugin nor interpreter executor is initialized. '
|
229 |
+
'An exception will be thrown when the agent call a tool.')
|
230 |
+
self.finish_condition = finish_condition
|
231 |
+
self.max_turn = max_turn
|
232 |
+
super().__init__(**kwargs)
|
233 |
+
|
234 |
+
async def forward(self, message: AgentMessage, session_id=0, **kwargs):
|
235 |
+
if isinstance(message, str):
|
236 |
+
message = AgentMessage(sender='user', content=message)
|
237 |
+
for _ in range(self.max_turn):
|
238 |
+
message = await self.agent(
|
239 |
+
message, session_id=session_id, **kwargs)
|
240 |
+
assert isinstance(message.formatted, dict)
|
241 |
+
if self.finish_condition(message):
|
242 |
+
return message
|
243 |
+
if message.formatted['tool_type']:
|
244 |
+
tool_type = message.formatted["tool_type"]
|
245 |
+
executor = getattr(self, f'{tool_type}_executor', None)
|
246 |
+
if not executor:
|
247 |
+
raise RuntimeError(f'No available {tool_type} executor')
|
248 |
+
message = await executor(message, session_id=session_id)
|
249 |
+
return message
|
250 |
+
|
251 |
+
def get_steps(self, session_id=0):
|
252 |
+
steps, tool_type = [], None
|
253 |
+
for msg in self.agent.memory.get_memory(session_id):
|
254 |
+
if msg.sender == self.agent.name:
|
255 |
+
steps.append(
|
256 |
+
dict(role='thought', content=msg.formatted['thought']))
|
257 |
+
if msg.formatted['tool_type']:
|
258 |
+
tool_type = msg.formatted['tool_type']
|
259 |
+
steps.append(
|
260 |
+
dict(
|
261 |
+
role='tool',
|
262 |
+
content=msg.formatted['action'],
|
263 |
+
name=tool_type))
|
264 |
+
elif msg.sender != 'user':
|
265 |
+
feedback = dict(role='environment', content=msg.content)
|
266 |
+
if tool_type:
|
267 |
+
feedback['name'] = tool_type
|
268 |
+
steps.append(feedback)
|
269 |
+
return steps
|
270 |
+
|
271 |
+
|
272 |
+
class AsyncMathCoder(AsyncAgentForInternLM):
|
273 |
+
|
274 |
+
def __init__(
|
275 |
+
self,
|
276 |
+
llm: Union[BaseLLM, Dict],
|
277 |
+
interpreter: dict = dict(type=AsyncIPythonInterpreter),
|
278 |
+
template: Union[str, dict, List[dict]] = None,
|
279 |
+
memory: Dict = dict(type=Memory),
|
280 |
+
output_format: Dict = dict(
|
281 |
+
type=InterpreterParser,
|
282 |
+
template=
|
283 |
+
('Integrate step-by-step reasoning and Python code to solve math problems '
|
284 |
+
'using the following guidelines:\n'
|
285 |
+
'- Analyze the question and write jupyter code to solve the problem;\n'
|
286 |
+
r"- Present the final result in LaTeX using a '\boxed{{}}' without any "
|
287 |
+
'units. \n')),
|
288 |
+
aggregator: Dict = dict(type=InternLMToolAggregator),
|
289 |
+
action_hooks: List = [dict(type=InternLMActionProcessor)],
|
290 |
+
finish_condition: Callable[
|
291 |
+
[AgentMessage],
|
292 |
+
bool] = lambda m: m.formatted['status'] == ToolStatusCode.NO_TOOL,
|
293 |
+
max_turn: int = 6,
|
294 |
+
**kwargs,
|
295 |
+
):
|
296 |
+
kwargs.pop('plugins', None)
|
297 |
+
super().__init__(
|
298 |
+
llm=llm,
|
299 |
+
interpreter=interpreter,
|
300 |
+
template=template,
|
301 |
+
memory=memory,
|
302 |
+
output_format=output_format,
|
303 |
+
aggregator=aggregator,
|
304 |
+
action_hooks=action_hooks,
|
305 |
+
finish_condition=finish_condition,
|
306 |
+
max_turn=max_turn,
|
307 |
+
**kwargs)
|
308 |
+
|
309 |
+
async def forward(self, message: AgentMessage, session_id=0, **kwargs):
|
310 |
+
try:
|
311 |
+
return await super().forward(message, session_id, **kwargs)
|
312 |
+
finally:
|
313 |
+
interpreter = next(
|
314 |
+
iter(self.interpreter_executor.actions.values()))
|
315 |
+
if interpreter.name == 'AsyncIPythonInterpreter':
|
316 |
+
await interpreter.close_session(session_id)
|