-
Notifications
You must be signed in to change notification settings - Fork 9
/
agent_setup.py
95 lines (74 loc) · 2.72 KB
/
agent_setup.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
from typing import List
from autogen import AssistantAgent, UserProxyAgent, GroupChat, GroupChatManager
from cobol_converter.autogen_config import llm_config
from cobol_converter.autogen_config import config_list
from cobol_converter.config import cfg
from cobol_converter.toml_support import prompts
from cobol_converter.service.terminators import terminate_lambda
class AgentType:
AGENT_PYTHON_CODER = "Python_Coder"
AGENT_UNIT_TESTER = "Unit_Tester"
AGENT_CODE_CRITIC = "Code_Critic"
AGENT_REST_INTERFACE_GENERATOR = "Code_REST_Generator"
class UserProxyType:
MAIN_PROXY = "user_proxy"
REST_PROXY = "rest_user_proxy"
llm_config = {
"cache_seed": cfg.seed, # seed for caching and reproducibility
"config_list": config_list, # a list of OpenAI API configurations
"temperature": cfg.temperature, # temperature for sampling
}
def cobol_convert_agent_factory() -> AssistantAgent:
assistant = AssistantAgent(
name=AgentType.AGENT_PYTHON_CODER,
system_message=prompts["agents"]["python_coder"]["system_message"],
llm_config=llm_config,
)
return assistant
def python_test_agent_factory() -> AssistantAgent:
assistant = AssistantAgent(
name=AgentType.AGENT_UNIT_TESTER,
system_message=prompts["agents"]["python_unit_tester"]["system_message"],
llm_config=llm_config,
)
return assistant
def python_code_critic_factory() -> AssistantAgent:
assistant = AssistantAgent(
name=AgentType.AGENT_CODE_CRITIC,
system_message=prompts["agents"]["code_critic"]["system_message"],
llm_config=llm_config,
is_termination_msg=terminate_lambda,
)
return assistant
def user_proxy_factory() -> UserProxyAgent:
user_proxy_agent = UserProxyAgent(
name=UserProxyType.MAIN_PROXY,
human_input_mode="NEVER",
max_consecutive_auto_reply=10,
is_termination_msg=terminate_lambda,
code_execution_config={"work_dir": cfg.code_execution_dir, "use_docker": False},
system_message=prompts["agents"]["userproxy"]["system_message"],
)
return user_proxy_agent
def create_group_chat_manager(
user_proxy: UserProxyAgent,
agents: List[AssistantAgent],
) -> GroupChatManager:
groupchat = GroupChat(
agents=[user_proxy, *agents],
messages=[],
max_round=cfg.max_tries,
)
manager = GroupChatManager(groupchat=groupchat, llm_config=llm_config)
return manager
user_proxy = user_proxy_factory()
conversion_manager = create_group_chat_manager(
user_proxy,
[
cobol_convert_agent_factory(),
python_test_agent_factory(),
python_code_critic_factory(),
],
)
if __name__ == "__main__":
assert conversion_manager is not None