forked from agiresearch/AIOS
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
122 lines (87 loc) · 3.53 KB
/
main.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
# This is a main script that tests the functionality of specific agents.
# It requires no user input.
from aios.scheduler.fifo_scheduler import FIFOScheduler
from aios.utils.utils import (
parse_global_args,
)
from pyopenagi.agents.agent_factory import AgentFactory
from pyopenagi.agents.agent_process import AgentProcessFactory
import warnings
from aios.llm_core import llms
from concurrent.futures import ThreadPoolExecutor, as_completed
from aios.utils.utils import delete_directories
from dotenv import load_dotenv
def clean_cache(root_directory):
targets = {
".ipynb_checkpoints",
"__pycache__",
".pytest_cache",
"context_restoration",
}
delete_directories(root_directory, targets)
def main():
# parse arguments and set configuration for this run accordingly
warnings.filterwarnings("ignore")
parser = parse_global_args()
args = parser.parse_args()
llm_name = args.llm_name
max_gpu_memory = args.max_gpu_memory
eval_device = args.eval_device
max_new_tokens = args.max_new_tokens
scheduler_log_mode = args.scheduler_log_mode
agent_log_mode = args.agent_log_mode
llm_kernel_log_mode = args.llm_kernel_log_mode
use_backend = args.use_backend
load_dotenv()
llm = llms.LLM(
llm_name=llm_name,
max_gpu_memory=max_gpu_memory,
eval_device=eval_device,
max_new_tokens=max_new_tokens,
log_mode=llm_kernel_log_mode,
use_backend=use_backend
)
# run agents concurrently for maximum efficiency using a scheduler
scheduler = FIFOScheduler(llm=llm, log_mode=scheduler_log_mode)
agent_process_factory = AgentProcessFactory()
agent_factory = AgentFactory(
agent_process_queue=scheduler.agent_process_queue,
agent_process_factory=agent_process_factory,
agent_log_mode=agent_log_mode,
)
agent_thread_pool = ThreadPoolExecutor(max_workers=500)
scheduler.start()
# construct example agents
# travel_agent = agent_thread_pool.submit(
# agent_factory.run_agent,
# "example/travel_agent", "I want to take a trip to Paris, France from July 4th to July 10th 2024 and I am traveling from New York City. Help me plan this trip."
# )
# math_agent = agent_thread_pool.submit(
# agent_factory.run_agent,
# "example/math_agent",
# "Convert 15000 MXN to Canadian Dollars and find out how much it would be in USD if 1 CAD equals 0.79 USD."
# )
academic_agent = agent_thread_pool.submit(
agent_factory.run_agent,
"example/academic_agent",
"Summarize recent advancements in quantum computing from the past five years.",
)
# rec_agent = agent_thread_pool.submit(
# agent_factory.run_agent,
# "example/rec_agent", "Recommend two movies with groundbreaking visual effects released in the last fifteen years ranked between 1 and 20 with ratings above 8.0."
# )
creation_agent = agent_thread_pool.submit(
agent_factory.run_agent,
"example/creation_agent", "Create an image of a lush jungle with an ancient temple, evoking a sense of mystery and adventure."
)
# agent_tasks = [travel_agent, rec_agent, creation_agent, math_agent, academic_agent]
# agent_tasks = [rec_agent]
# agent_tasks = [creation_agent]
agent_tasks = [academic_agent, creation_agent]
# agent_tasks = [creation_agent]
for r in as_completed(agent_tasks):
_res = r.result()
scheduler.stop()
clean_cache(root_directory="./")
if __name__ == "__main__":
main()