-
Notifications
You must be signed in to change notification settings - Fork 314
/
conversation_moa.py
139 lines (119 loc) · 4.35 KB
/
conversation_moa.py
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
# -*- coding: utf-8 -*-
# pylint: disable=line-too-long
"""
Implementation of paper [MoA](https://github.com/togethercomputer/MoA).
Here is a simple example for conversation with MoA in Agentscope.
"""
from typing import Optional, Union, Sequence
import agentscope
from agentscope.agents import AgentBase, UserAgent
from agentscope.strategy import MixtureOfAgents
from agentscope.message import Msg
class DialogAgentWithMoA(AgentBase):
"""A simple agent used to perform a dialogue.
We will show all the modifcation need to use MoA as the main model"""
def __init__(
self,
name: str,
moa_module: MixtureOfAgents, # changed to passing moa_module here
use_memory: bool = True,
) -> None:
"""Initialize the dialog agent.
Arguments:
name (`str`):
The name of the agent.
sys_prompt (`Optional[str]`):
The system prompt of the agent, which can be passed by args
or hard-coded in the agent.
moa_module (`MixutreOfAgents`):
The inited MoA module you want to use as the main module.
use_memory (`bool`, defaults to `True`):
Whether the agent has memory.
"""
super().__init__(
name=name,
sys_prompt="",
use_memory=use_memory,
)
self.moa_module = moa_module # change model init to moa_module
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
"""Reply function of the agent. Processes the input data,
generates a prompt using the current dialogue memory and system
prompt, and invokes the language model to produce a response. The
response is then formatted and added to the dialogue memory.
Args:
x (`Optional[Union[Msg, Sequence[Msg]]]`, defaults to `None`):
The input message(s) to the agent, which also can be omitted if
the agent doesn't need any input.
Returns:
`Msg`: The output message generated by the agent.
"""
# record the input if needed
if self.memory:
self.memory.add(x)
# use the module as below:
response = self.moa_module(
Msg("system", self.sys_prompt, role="system"),
self.memory
and self.memory.get_memory()
or x, # type: ignore[arg-type]
)
msg = Msg(self.name, response, role="assistant")
# Print/speak the message in this agent's voice
self.speak(msg)
# Record the message in memory
if self.memory:
self.memory.add(msg)
return msg
if __name__ == "__main__":
# fill you api keys, or host local models using vllm or ollama.
model_configs = [
{
"config_name": "qwen-max",
"model_type": "dashscope_chat",
"model_name": "qwen-max",
"api_key": "{your_api_key}",
"generate_args": {
"temperature": 0.7,
},
},
{
"config_name": "gemini-pro",
"model_type": "gemini_chat",
"model_name": "gemini-pro",
"api_key": "{your_api_key}",
},
{
"config_name": "gpt-4",
"model_type": "openai_chat",
"model_name": "gpt-4",
"api_key": "{your_api_key}",
"client_args": {
"max_retries": 3,
},
"generate_args": {
"temperature": 0.7,
},
},
]
agentscope.init(model_configs=model_configs, project="Mixture of Agents")
user_agent = UserAgent()
your_moa_module = MixtureOfAgents(
main_model="qwen-max", # the models you use
reference_models=["gpt-4", "qwen-max", "gemini-pro"],
show_internal=False, # set to True to see the internal of MoA modules
rounds=1, # can range from 0 to inf
)
# Init two agents
dialog_agent = DialogAgentWithMoA(
name="Assistant",
moa_module=your_moa_module,
use_memory=True, # whether to use memory for this agent
)
user_agent = UserAgent()
# start the conversation between user and assistant
while True:
q = user_agent(None)
if q.content == "exit": # type "exit" to break the loop
break
q = dialog_agent(q)