File size: 12,813 Bytes
01523b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
import json
import ast
import openai
from string import Template
from colorama import Fore
from aiohttp import ClientSession
from copy import deepcopy
from typing import TYPE_CHECKING, Any, List, Tuple

from agentverse.agents import ExecutorAgent
from agentverse.message import Message, ExecutorMessage, SolverMessage
from agentverse.logging import logger

from . import BaseExecutor, executor_registry
import asyncio


url = "http://127.0.0.1:8080"
# url = "http://8.217.97.110:8080"

SUMMARIZE_PROMPT = """Here is the text gathered from a webpage, and a question you need to answer from the webpage. 
-- Webpage -- 
${webpage}
-- Question --
${question}

Now summarize the webpage to answer the question. If the question cannot be answer from the webpage, return the summarization of the webpage."""


@executor_registry.register("tool-using")
class ToolUsingExecutor(BaseExecutor):
    num_agents: int = 3
    max_tool_call_times: int = 10
    tools: List[dict] = []
    tool_names: List[str] = []
    tool_config: str = None
    cookies: dict = {}
    tool_retrieval: bool = False
    real_execution_agents: dict = {}
    agent_names: List[str] = []
    # tool_description: str

    def __init__(self, *args, **kwargs):
        assert kwargs.get("tool_config", None) is not None
        with open(kwargs.get("tool_config"), "r") as f:
            tools_dict = json.load(f)
        tools = tools_dict["tools_json"]
        tool_names = [t["name"] for t in tools]

        # For each tool, we manually add a "thought" argument to achieve
        # chain-of-thought in OpenAI's function call.
        for t in tools:
            properties = t["parameters"]["properties"]
            thought = {
                "thought": {
                    "type": "string",
                    "description": "Your internal reasoning and thoughts on the task, and how you plan to solve it based on the current attempts.",
                }
            }
            thought.update(properties)
            t["parameters"]["properties"] = thought
            t["parameters"]["required"].insert(0, "thought")
        super().__init__(
            tools=tools,
            tool_names=tool_names,
            # tool_description=tool_description,
            *args,
            **kwargs,
        )

    async def astep(
        self,
        agent: ExecutorAgent,
        task_description: str,
        plans: List[SolverMessage],
        *args,
        **kwargs,
    ):
        plan_this_turn = {}
        agent_name_this_turn = []
        for i in range(len(plans)):
            name = plans[i].content.split("-")[0].strip()
            if name not in self.real_execution_agents:
                self.real_execution_agents[name] = deepcopy(agent)
                self.real_execution_agents[name].name = name
                self.agent_names.append(name)
            plan_this_turn[name] = plans[i].content.split("-")[1].strip()
            agent_name_this_turn.append(name)
        # agents = [deepcopy(agent) for _ in range(len(plans))]

        if self.tool_retrieval:
            # We retrieve 5 related tools for each agent
            tools_and_cookies = await asyncio.gather(
                *[
                    self.retrieve_tools(plan_this_turn[name], self.tools)
                    for name in agent_name_this_turn
                ]
            )
            tools = {
                name: t[0] for name, t in zip(agent_name_this_turn, tools_and_cookies)
            }
            cookies = {
                name: t[1] for name, t in zip(agent_name_this_turn, tools_and_cookies)
            }
            self.update_cookies(cookies)
        else:
            # We just use the tools that are provided in the config file
            tools = {name: self.tools for name in agent_name_this_turn}

        # Record the indices of agents that have finished their tasks
        # so that they will not be called again
        finished_agent_names = set()
        # result = ["" for _ in range(len(plan_this_turn))]
        result = {name: "" for name in agent_name_this_turn}
        for current_turn in range(self.max_tool_call_times):
            if len(finished_agent_names) == len(agent_name_this_turn):
                # All agents have finished their tasks. Break the loop.
                break

            # Filter out agents that have finished and gather tool actions for the rest
            tool_calls = []
            active_agents_names = [
                name
                for name in agent_name_this_turn
                if name not in finished_agent_names
            ]
            for name in active_agents_names:
                if current_turn == self.max_tool_call_times - 1:
                    tool = [t for t in tools[name] if t["name"] == "submit_task"]
                else:
                    tool = tools[name]
                tool_calls.append(
                    self.real_execution_agents[name].astep(
                        task_description,
                        plan_this_turn[name],
                        tool,
                        current_turn=current_turn + 1,
                    )
                )
            # Use asyncio.gather to run astep concurrently
            tool_call_decisions = await asyncio.gather(*tool_calls)
            for name, tool_call_result in zip(active_agents_names, tool_call_decisions):
                self.real_execution_agents[name].add_message_to_memory(
                    [tool_call_result]
                )

            # Actually call the tool and get the observation
            tool_responses = await asyncio.gather(
                *[
                    ToolUsingExecutor.call_tool(
                        tool.tool_name,
                        tool.tool_input,
                        self.cookies.get(name, None),
                    )
                    for name, tool in zip(active_agents_names, tool_call_decisions)
                ]
            )
            # Update each agent's memory and check if they have finished
            cookies = {}
            for name, response in zip(active_agents_names, tool_responses):
                observation = response["observation"]
                is_finish = response["is_finish"]
                cookies[name] = response["cookies"]
                self.real_execution_agents[name].add_message_to_memory([observation])
                logger.info(
                    f"\nTool: {observation.tool_name}\nTool Input: {observation.tool_input}\nObservation: {observation.content}",
                    name,
                    Fore.YELLOW,
                )
                if is_finish:
                    finished_agent_names.add(name)
                    result[name] = observation.content
            self.update_cookies(cookies)

        message_result = []
        for name, conclusion in result.items():
            if conclusion != "":
                message_result.append(
                    ExecutorMessage(
                        content=f"[{name}]: My execution result:\n{conclusion}",
                        sender=name,
                    )
                )
        return message_result

    def update_cookies(self, cookies: dict):
        for name, cookie in cookies.items():
            self.cookies[name] = cookie

    @classmethod
    async def retrieve_tools(
        cls, plan: SolverMessage, curr_tools: List = [], cookies=None
    ):
        async with ClientSession(cookies=cookies) as session:
            if cookies is None:
                async with session.post(f"{url}/get_cookie", timeout=30) as response:
                    cookies = response.cookies
                    session.cookie_jar.update_cookies(cookies)
                    await response.text()
                    # Sometimes the toolserver's docker container is not ready yet
                    # So we need to wait for a while
                    await asyncio.sleep(10)
            async with session.post(
                f"{url}/retrieving_tools", json={"question": plan.content, "top_k": 5}
            ) as response:
                retrieved_tools = await response.json()
                retrieved_tools = ast.literal_eval(retrieved_tools)
        tools = deepcopy(curr_tools)
        existed_tool_names = set([t["name"] for t in tools])
        # Add the retrieved tools into the final tools
        for tool in retrieved_tools["tools_json"]:
            if tool["name"] not in existed_tool_names:
                existed_tool_names.add(tool["name"])
                tools.append(tool)
        return tools, cookies

    @classmethod
    async def call_tool(cls, command: str, arguments: dict, cookies=None):
        async def _summarize_webpage(webpage, question):
            summarize_prompt = Template(SUMMARIZE_PROMPT).safe_substitute(
                webpage=webpage, question=question
            )
            for _ in range(3):
                try:
                    response = await openai.ChatCompletion.acreate(
                        messages=[{"role": "user", "content": summarize_prompt}],
                        model="gpt-3.5-turbo-16k",
                    )
                except:
                    continue
            return response["choices"][0]["message"]["content"]

        if command == "submit_task":
            return {
                "observation": ExecutorMessage(
                    content=f"Task Status: {arguments['status']}\nConclusion: {arguments['conclusion']}",
                    sender="function",
                    tool_name=command,
                    tool_input=arguments,
                ),
                "is_finish": True,
                "cookies": cookies,
            }
        if command == "":
            return {
                "observation": ExecutorMessage(
                    content=f"The function calling format is incorrect.",
                    sender="function",
                    tool_name=command,
                    tool_input=arguments,
                ),
                "is_finish": False,
                "cookies": cookies,
            }

        for i in range(3):
            try:
                async with ClientSession(cookies=cookies) as session:
                    if cookies is None:
                        async with session.post(
                            f"{url}/get_cookie", timeout=30
                        ) as response:
                            cookies = response.cookies
                            session.cookie_jar.update_cookies(cookies)
                            await response.text()
                            # Sometimes the toolserver's docker container is not ready yet
                            # So we need to wait for a while
                            await asyncio.sleep(10)

                    payload_arguments = deepcopy(arguments)
                    if "thought" in payload_arguments:
                        del payload_arguments["thought"]
                    payload = {
                        "tool_name": command,
                        "arguments": payload_arguments,
                    }
                    # async with ClientSession() as session:
                    async with session.post(
                        f"{url}/execute_tool",
                        json=payload,
                        headers={
                            "toolbench_key": "p5ZASSLBO0EknAQLE5ecNZ7kq5i1YfY9eoWUXNxL3TM6lXwdXs"
                        },
                        timeout=30,
                    ) as response:
                        content = await response.text()
                        if command == "WebEnv_browse_website":
                            content = await _summarize_webpage(
                                content, arguments["question"]
                            )

                        message = ExecutorMessage(
                            content=content,
                            sender="function",
                            tool_name=command,
                            tool_input=arguments,
                        )
                    # async with session.post(
                    #     f"{url}/release_session", timeout=30
                    # ) as response:
                    #     await response.text()
                break
            except Exception as e:
                message = ExecutorMessage(
                    content="Failed to call the tool. Exception: " + str(e),
                    sender="function",
                    tool_name=command,
                    tool_input=arguments,
                )
                continue
        return {"observation": message, "is_finish": False, "cookies": cookies}

    def broadcast_messages(self, agents, messages) -> None:
        for agent in agents:
            agent.add_message_to_memory(messages)