611 lines
20 KiB
Python
611 lines
20 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Unit tests for rpc agent classes
|
|
"""
|
|
import unittest
|
|
import time
|
|
import os
|
|
import shutil
|
|
from typing import Optional, Union, Sequence
|
|
|
|
from loguru import logger
|
|
|
|
import agentscope
|
|
from agentscope.agents import AgentBase, DistConf
|
|
from agentscope.server import RpcAgentServerLauncher
|
|
from agentscope.message import Msg
|
|
from agentscope.message import PlaceholderMessage
|
|
from agentscope.message import deserialize
|
|
from agentscope.msghub import msghub
|
|
from agentscope.pipelines import sequentialpipeline
|
|
from agentscope.utils import MonitorFactory, QuotaExceededError
|
|
|
|
|
|
class DemoRpcAgent(AgentBase):
|
|
"""A demo Rpc agent for test usage."""
|
|
|
|
def __init__(self, **kwargs) -> None: # type: ignore[no-untyped-def]
|
|
super().__init__(**kwargs)
|
|
self.id = 0
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
"""Response after 2s"""
|
|
x.id = self.id
|
|
self.id += 1
|
|
time.sleep(2)
|
|
return x
|
|
|
|
|
|
class DemoRpcAgentAdd(AgentBase):
|
|
"""A demo Rpc agent for test usage"""
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
"""add the value, wait 1s"""
|
|
x.content["value"] += 1
|
|
time.sleep(1)
|
|
return x
|
|
|
|
|
|
class DemoLocalAgentAdd(AgentBase):
|
|
"""A demo local agent for test usage"""
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
"""add the value, wait 1s"""
|
|
x.content["value"] += 1
|
|
time.sleep(1)
|
|
return x
|
|
|
|
|
|
class DemoRpcAgentWithMemory(AgentBase):
|
|
"""A demo Rpc agent that count its memory"""
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
self.memory.add(x)
|
|
msg = Msg(
|
|
name=self.name,
|
|
content={"mem_size": self.memory.size()},
|
|
role="assistant",
|
|
)
|
|
self.memory.add(msg)
|
|
time.sleep(1)
|
|
return msg
|
|
|
|
|
|
class DemoRpcAgentWithMonitor(AgentBase):
|
|
"""A demo Rpc agent that use monitor"""
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
monitor = MonitorFactory.get_monitor()
|
|
try:
|
|
monitor.update({"msg_num": 1})
|
|
except QuotaExceededError:
|
|
x.content["quota_exceeded"] = True
|
|
logger.chat(
|
|
{
|
|
"name": self.name,
|
|
"content": "quota_exceeded",
|
|
},
|
|
)
|
|
return x
|
|
x.content["msg_num"] = monitor.get_value("msg_num")
|
|
logger.chat(
|
|
{
|
|
"name": self.name,
|
|
"content": f"msg_num {x.content['msg_num']}",
|
|
},
|
|
)
|
|
time.sleep(0.2)
|
|
return x
|
|
|
|
|
|
class DemoGeneratorAgent(AgentBase):
|
|
"""A demo agent to generate a number"""
|
|
|
|
def __init__(self, name: str, value: int) -> None:
|
|
super().__init__(name)
|
|
self.value = value
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
time.sleep(1)
|
|
return Msg(
|
|
name=self.name,
|
|
role="assistant",
|
|
content={
|
|
"value": self.value,
|
|
},
|
|
)
|
|
|
|
|
|
class DemoGatherAgent(AgentBase):
|
|
"""A demo agent to gather value"""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
agents: list[DemoGeneratorAgent],
|
|
to_dist: dict = None,
|
|
) -> None:
|
|
super().__init__(name, to_dist=to_dist)
|
|
self.agents = agents
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
result = []
|
|
stime = time.time()
|
|
for agent in self.agents:
|
|
result.append(agent())
|
|
value = 0
|
|
for r in result:
|
|
value += r.content["value"]
|
|
etime = time.time()
|
|
return Msg(
|
|
name=self.name,
|
|
role="assistant",
|
|
content={
|
|
"value": value,
|
|
"time": etime - stime,
|
|
},
|
|
)
|
|
|
|
|
|
class DemoErrorAgent(AgentBase):
|
|
"""A demo Rpc agent that raise Error"""
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
raise RuntimeError("Demo Error")
|
|
|
|
|
|
class BasicRpcAgentTest(unittest.TestCase):
|
|
"Test cases for Rpc Agent"
|
|
|
|
def setUp(self) -> None:
|
|
"""Init for Rpc Agent Test"""
|
|
agentscope.init(
|
|
project="test",
|
|
name="rpc_agent",
|
|
save_dir="./.unittest_runs",
|
|
save_log=True,
|
|
)
|
|
self.assertTrue(os.path.exists("./.unittest_runs"))
|
|
|
|
def tearDown(self) -> None:
|
|
MonitorFactory._instance = None # pylint: disable=W0212
|
|
logger.remove()
|
|
shutil.rmtree("./.unittest_runs")
|
|
|
|
def test_single_rpc_agent_server(self) -> None:
|
|
"""test setup a single rpc agent"""
|
|
agent_a = DemoRpcAgent(
|
|
name="a",
|
|
to_dist=True,
|
|
)
|
|
self.assertIsNotNone(agent_a)
|
|
msg = Msg(
|
|
name="System",
|
|
content={"text": "hello world"},
|
|
role="system",
|
|
)
|
|
result = agent_a(msg)
|
|
# get name without waiting for the server
|
|
self.assertEqual(result.name, "a")
|
|
self.assertEqual(result["name"], "a")
|
|
js_placeholder_result = result.serialize()
|
|
self.assertTrue(result._is_placeholder) # pylint: disable=W0212
|
|
placeholder_result = deserialize(js_placeholder_result)
|
|
self.assertTrue(isinstance(placeholder_result, PlaceholderMessage))
|
|
self.assertEqual(placeholder_result.name, "a")
|
|
self.assertEqual(
|
|
placeholder_result["name"], # type: ignore[call-overload]
|
|
"a",
|
|
)
|
|
self.assertTrue(
|
|
placeholder_result._is_placeholder, # pylint: disable=W0212
|
|
)
|
|
# wait to get content
|
|
self.assertEqual(result.content, msg.content)
|
|
self.assertFalse(result._is_placeholder) # pylint: disable=W0212
|
|
self.assertEqual(result.id, 0)
|
|
self.assertTrue(
|
|
placeholder_result._is_placeholder, # pylint: disable=W0212
|
|
)
|
|
self.assertEqual(placeholder_result.content, msg.content)
|
|
self.assertFalse(
|
|
placeholder_result._is_placeholder, # pylint: disable=W0212
|
|
)
|
|
self.assertEqual(placeholder_result.id, 0)
|
|
# check msg
|
|
js_msg_result = result.serialize()
|
|
msg_result = deserialize(js_msg_result)
|
|
self.assertTrue(isinstance(msg_result, Msg))
|
|
self.assertEqual(msg_result.content, msg.content)
|
|
self.assertEqual(msg_result.id, 0)
|
|
# check id increase
|
|
msg = agent_a(msg_result) # type: ignore[arg-type]
|
|
self.assertEqual(msg.id, 1)
|
|
|
|
def test_connect_to_an_existing_rpc_server(self) -> None:
|
|
"""test connecting to an existing server"""
|
|
launcher = RpcAgentServerLauncher(
|
|
# choose port automatically
|
|
host="127.0.0.1",
|
|
port=12010,
|
|
local_mode=False,
|
|
custom_agents=[DemoRpcAgent],
|
|
)
|
|
launcher.launch()
|
|
agent_a = DemoRpcAgent(
|
|
name="a",
|
|
).to_dist(
|
|
host="127.0.0.1",
|
|
port=launcher.port,
|
|
)
|
|
msg = Msg(
|
|
name="System",
|
|
content={"text": "hello world"},
|
|
role="system",
|
|
)
|
|
result = agent_a(msg)
|
|
# get name without waiting for the server
|
|
self.assertEqual(result.name, "a")
|
|
# waiting for server
|
|
self.assertEqual(result.content, msg.content)
|
|
# test dict usage
|
|
msg = Msg(
|
|
name="System",
|
|
content={"text": "hi world"},
|
|
role="system",
|
|
)
|
|
result = agent_a(msg)
|
|
# get name without waiting for the server
|
|
self.assertEqual(result["name"], "a")
|
|
# waiting for server
|
|
self.assertEqual(result["content"], msg.content)
|
|
# test to_str
|
|
msg = Msg(
|
|
name="System",
|
|
content={"text": "test"},
|
|
role="system",
|
|
)
|
|
result = agent_a(msg)
|
|
self.assertEqual(result.to_str(), "a: {'text': 'test'}")
|
|
launcher.shutdown()
|
|
|
|
def test_multi_rpc_agent(self) -> None:
|
|
"""test setup multi rpc agent"""
|
|
agent_a = DemoRpcAgentAdd(
|
|
name="a",
|
|
).to_dist(
|
|
lazy_launch=False,
|
|
)
|
|
agent_b = DemoRpcAgentAdd(
|
|
name="b",
|
|
).to_dist(
|
|
lazy_launch=False,
|
|
)
|
|
agent_c = DemoRpcAgentAdd(
|
|
name="c",
|
|
).to_dist(
|
|
lazy_launch=False,
|
|
)
|
|
|
|
# test sequential
|
|
msg = Msg(
|
|
name="System",
|
|
content={"value": 0},
|
|
role="system",
|
|
)
|
|
start_time = time.time()
|
|
msg = agent_a(msg)
|
|
self.assertTrue(isinstance(msg, PlaceholderMessage))
|
|
msg = agent_b(msg)
|
|
self.assertTrue(isinstance(msg, PlaceholderMessage))
|
|
msg = agent_c(msg)
|
|
self.assertTrue(isinstance(msg, PlaceholderMessage))
|
|
return_time = time.time()
|
|
# should return directly
|
|
self.assertTrue((return_time - start_time) < 1)
|
|
self.assertEqual(msg.content["value"], 3)
|
|
end_time = time.time()
|
|
# need at least 3s to finish
|
|
self.assertTrue((end_time - start_time) >= 3)
|
|
|
|
# test parallel
|
|
msg = Msg(
|
|
name="System",
|
|
content={"value": -1},
|
|
role="system",
|
|
)
|
|
start_time = time.time()
|
|
msg_a = agent_a(msg)
|
|
msg_b = agent_b(msg)
|
|
msg_c = agent_c(msg)
|
|
self.assertEqual(msg_a.content["value"], 0)
|
|
self.assertEqual(msg_b.content["value"], 0)
|
|
self.assertEqual(msg_c.content["value"], 0)
|
|
end_time = time.time()
|
|
# need 1s to finish
|
|
self.assertTrue((end_time - start_time) < 1.5)
|
|
|
|
def test_mix_rpc_agent_and_local_agent(self) -> None:
|
|
"""test to use local and rpc agent simultaneously"""
|
|
agent_a = DemoRpcAgentAdd(
|
|
name="a",
|
|
).to_dist(
|
|
lazy_launch=False,
|
|
)
|
|
# local agent b
|
|
agent_b = DemoLocalAgentAdd(
|
|
name="b",
|
|
)
|
|
# rpc agent c
|
|
agent_c = DemoRpcAgentAdd( # pylint: disable=E1123
|
|
name="c",
|
|
to_dist=DistConf(
|
|
lazy_launch=False,
|
|
),
|
|
)
|
|
msg = Msg(
|
|
name="System",
|
|
content={"value": 0},
|
|
role="system",
|
|
)
|
|
|
|
while msg.content["value"] < 4:
|
|
msg = agent_a(msg)
|
|
msg = agent_b(msg)
|
|
msg = agent_c(msg)
|
|
self.assertEqual(msg.content["value"], 6)
|
|
|
|
def test_msghub_compatibility(self) -> None:
|
|
"""test compatibility with msghub"""
|
|
agent_a = DemoRpcAgentWithMemory(
|
|
name="a",
|
|
).to_dist()
|
|
agent_b = DemoRpcAgentWithMemory(
|
|
name="b",
|
|
).to_dist()
|
|
agent_c = DemoRpcAgentWithMemory(
|
|
name="c",
|
|
to_dist=True,
|
|
)
|
|
participants = [agent_a, agent_b, agent_c]
|
|
annonuncement_msgs = [
|
|
Msg(name="System", content="Announcement 1", role="system"),
|
|
Msg(name="System", content="Announcement 2", role="system"),
|
|
]
|
|
with msghub(
|
|
participants=participants,
|
|
announcement=annonuncement_msgs,
|
|
):
|
|
x_a = agent_a()
|
|
x_b = agent_b(x_a)
|
|
x_c = agent_c(x_b)
|
|
self.assertEqual(x_a.content["mem_size"], 2)
|
|
self.assertEqual(x_b.content["mem_size"], 3)
|
|
self.assertEqual(x_c.content["mem_size"], 4)
|
|
x_a = agent_a(x_c)
|
|
self.assertEqual(x_a.content["mem_size"], 5)
|
|
x_b = agent_b(x_a)
|
|
self.assertEqual(x_b.content["mem_size"], 6)
|
|
x_c = agent_c(x_b)
|
|
self.assertEqual(x_c.content["mem_size"], 7)
|
|
x_c = sequentialpipeline(participants, x_c)
|
|
self.assertEqual(x_c.content["mem_size"], 10)
|
|
|
|
def test_standalone_multiprocess_init(self) -> None:
|
|
"""test compatibility with agentscope.init"""
|
|
monitor = MonitorFactory.get_monitor()
|
|
monitor.register("msg_num", quota=10)
|
|
# rpc agent a
|
|
agent_a = DemoRpcAgentWithMonitor(
|
|
name="a",
|
|
).to_dist(
|
|
lazy_launch=False,
|
|
)
|
|
# local agent b
|
|
agent_b = DemoRpcAgentWithMonitor(
|
|
name="b",
|
|
).to_dist(
|
|
lazy_launch=False,
|
|
)
|
|
msg = Msg(name="System", content={"msg_num": 0}, role="system")
|
|
j = 0
|
|
for _ in range(5):
|
|
msg = agent_a(msg)
|
|
self.assertEqual(msg["content"]["msg_num"], j + 1)
|
|
msg = agent_b(msg)
|
|
self.assertEqual(msg["content"]["msg_num"], j + 2)
|
|
j += 2
|
|
msg = agent_a(msg)
|
|
logger.chat(msg)
|
|
self.assertTrue(msg["content"]["quota_exceeded"])
|
|
msg = agent_b(msg)
|
|
logger.chat(msg)
|
|
self.assertTrue(msg["content"]["quota_exceeded"])
|
|
|
|
def test_multi_agent_in_same_server(self) -> None:
|
|
"""test agent server with multi agent"""
|
|
launcher = RpcAgentServerLauncher(
|
|
host="127.0.0.1",
|
|
port=12010,
|
|
local_mode=False,
|
|
custom_agents=[DemoRpcAgentWithMemory],
|
|
)
|
|
launcher.launch()
|
|
# although agent1 and agent2 connect to the same server
|
|
# they are different instances with different memories
|
|
agent1 = DemoRpcAgentWithMemory(
|
|
name="a",
|
|
)
|
|
oid = agent1.agent_id
|
|
agent1 = agent1.to_dist(
|
|
host="127.0.0.1",
|
|
port=launcher.port,
|
|
)
|
|
self.assertEqual(oid, agent1.agent_id)
|
|
self.assertEqual(oid, agent1.client.agent_id)
|
|
agent2 = DemoRpcAgentWithMemory( # pylint: disable=E1123
|
|
name="a",
|
|
to_dist={
|
|
"host": "127.0.0.1",
|
|
"port": launcher.port,
|
|
},
|
|
)
|
|
# agent3 has the same agent id as agent1
|
|
# so it share the same memory with agent1
|
|
agent3 = DemoRpcAgentWithMemory(
|
|
name="a",
|
|
).to_dist(
|
|
host="127.0.0.1",
|
|
port=launcher.port,
|
|
)
|
|
agent3._agent_id = agent1.agent_id # pylint: disable=W0212
|
|
agent3.client.agent_id = agent1.client.agent_id
|
|
msg1 = Msg(name="System", content="First Msg for agent1")
|
|
res1 = agent1(msg1)
|
|
self.assertEqual(res1.content["mem_size"], 1)
|
|
msg2 = Msg(name="System", content="First Msg for agent2")
|
|
res2 = agent2(msg2)
|
|
self.assertEqual(res2.content["mem_size"], 1)
|
|
msg3 = Msg(name="System", content="First Msg for agent3")
|
|
res3 = agent3(msg3)
|
|
self.assertEqual(res3.content["mem_size"], 3)
|
|
msg4 = Msg(name="System", content="Second Msg for agent2")
|
|
res4 = agent2(msg4)
|
|
self.assertEqual(res4.content["mem_size"], 3)
|
|
# delete existing agent
|
|
agent2.client.delete_agent()
|
|
msg2 = Msg(name="System", content="First Msg for agent2")
|
|
res2 = agent2(msg2)
|
|
self.assertRaises(ValueError, res2.__getattr__, "content")
|
|
|
|
# should override remote default parameter(e.g. name field)
|
|
agent4 = DemoRpcAgentWithMemory(
|
|
name="b",
|
|
).to_dist(
|
|
host="127.0.0.1",
|
|
port=launcher.port,
|
|
)
|
|
msg5 = Msg(name="System", content="Second Msg for agent4")
|
|
res5 = agent4(msg5)
|
|
self.assertEqual(res5.name, "b")
|
|
self.assertEqual(res5.content["mem_size"], 1)
|
|
launcher.shutdown()
|
|
|
|
def test_clone_instances(self) -> None:
|
|
"""Test the clone_instances method of RpcAgent"""
|
|
agent = DemoRpcAgentWithMemory(
|
|
name="a",
|
|
).to_dist()
|
|
# lazy launch will not init client
|
|
self.assertIsNone(agent.client)
|
|
# generate two agents (the first is it self)
|
|
agents = agent.clone_instances(2)
|
|
self.assertEqual(len(agents), 2)
|
|
agent1 = agents[0]
|
|
agent2 = agents[1]
|
|
self.assertTrue(agent1.agent_id.startswith("DemoRpcAgentWithMemory"))
|
|
self.assertTrue(agent2.agent_id.startswith("DemoRpcAgentWithMemory"))
|
|
self.assertTrue(
|
|
agent1.client.agent_id.startswith("DemoRpcAgentWithMemory"),
|
|
)
|
|
self.assertTrue(
|
|
agent2.client.agent_id.startswith("DemoRpcAgentWithMemory"),
|
|
)
|
|
self.assertNotEqual(agent1.agent_id, agent2.agent_id)
|
|
self.assertEqual(agent1.agent_id, agent1.client.agent_id)
|
|
self.assertEqual(agent2.agent_id, agent2.client.agent_id)
|
|
# clone instance will init client
|
|
self.assertIsNotNone(agent.client)
|
|
self.assertEqual(agent.agent_id, agent1.agent_id)
|
|
self.assertNotEqual(agent1.agent_id, agent2.agent_id)
|
|
self.assertIsNotNone(agent.server_launcher)
|
|
self.assertIsNotNone(agent1.server_launcher)
|
|
self.assertIsNone(agent2.server_launcher)
|
|
msg1 = Msg(name="System", content="First Msg for agent1")
|
|
res1 = agent1(msg1)
|
|
self.assertEqual(res1.content["mem_size"], 1)
|
|
msg2 = Msg(name="System", content="First Msg for agent2")
|
|
res2 = agent2(msg2)
|
|
self.assertEqual(res2.content["mem_size"], 1)
|
|
new_agents = agent.clone_instances(2, including_self=False)
|
|
agent3 = new_agents[0]
|
|
agent4 = new_agents[1]
|
|
self.assertEqual(len(new_agents), 2)
|
|
self.assertNotEqual(agent3.agent_id, agent.agent_id)
|
|
self.assertNotEqual(agent4.agent_id, agent.agent_id)
|
|
self.assertIsNone(agent3.server_launcher)
|
|
self.assertIsNone(agent4.server_launcher)
|
|
msg3 = Msg(name="System", content="First Msg for agent3")
|
|
res3 = agent3(msg3)
|
|
self.assertEqual(res1.content["mem_size"], 1)
|
|
msg4 = Msg(name="System", content="First Msg for agent4")
|
|
res4 = agent4(msg4)
|
|
self.assertEqual(res3.content["mem_size"], 1)
|
|
self.assertEqual(res4.content["mem_size"], 1)
|
|
|
|
def test_error_handling(self) -> None:
|
|
"""Test error handling"""
|
|
agent = DemoErrorAgent(name="a").to_dist()
|
|
x = agent()
|
|
self.assertRaises(RuntimeError, x.__getattr__, "content")
|
|
|
|
def test_agent_nesting(self) -> None:
|
|
"""Test agent nesting"""
|
|
host = "localhost"
|
|
launcher1 = RpcAgentServerLauncher(
|
|
# choose port automatically
|
|
host=host,
|
|
port=12010,
|
|
local_mode=False,
|
|
custom_agents=[DemoGatherAgent, DemoGeneratorAgent],
|
|
)
|
|
launcher2 = RpcAgentServerLauncher(
|
|
# choose port automatically
|
|
host=host,
|
|
port=12011,
|
|
local_mode=False,
|
|
custom_agents=[DemoGatherAgent, DemoGeneratorAgent],
|
|
)
|
|
launcher1.launch()
|
|
launcher2.launch()
|
|
agents = []
|
|
for i in range(8):
|
|
if i % 2:
|
|
agents.append(
|
|
DemoGeneratorAgent(name=f"a_{i}", value=i).to_dist(
|
|
host=host,
|
|
port=launcher1.port,
|
|
),
|
|
)
|
|
else:
|
|
agents.append(
|
|
DemoGeneratorAgent(name=f"a_{i}", value=i).to_dist(
|
|
host=host,
|
|
port=launcher2.port,
|
|
),
|
|
)
|
|
gather1 = DemoGatherAgent( # pylint: disable=E1123
|
|
name="g1",
|
|
agents=agents[:4],
|
|
to_dist=DistConf(
|
|
host=host,
|
|
port=launcher1.port,
|
|
),
|
|
)
|
|
gather2 = DemoGatherAgent( # pylint: disable=E1123
|
|
name="g2",
|
|
agents=agents[4:],
|
|
to_dist={
|
|
"host": host,
|
|
"port": launcher2.port,
|
|
},
|
|
)
|
|
r1 = gather1()
|
|
r2 = gather2()
|
|
self.assertEqual(r1.content["value"], 6)
|
|
self.assertEqual(r2.content["value"], 22)
|
|
self.assertTrue(0.5 < r1.content["time"] < 2)
|
|
self.assertTrue(0.5 < r2.content["time"] < 2)
|
|
launcher1.shutdown()
|
|
launcher2.shutdown()
|