Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add new features and update logging style #59

Merged
merged 7 commits into from
Apr 12, 2024
+243 −158
Merged

add new features and update logging style #59

Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
52de5cb
feat: add interrupt
dongyuanjushi Apr 12, 2024
af78a45
Merge branch 'dev' of https://github.com/dongyuanjushi/AIOS into dev
dongyuanjushi Apr 12, 2024
27c9ef6
feat: add round-robin scheduler with interrupt
dongyuanjushi Apr 12, 2024
15e9dd7
refactor: agent-process and agent construction
dongyuanjushi Apr 12, 2024
bfe98f5
style: update agent logging
dongyuanjushi Apr 12, 2024
975f8d0
style: update scheduler logging
dongyuanjushi Apr 12, 2024
0188b36
refactor: rebuild main entry
dongyuanjushi Apr 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
refactor: rebuild main entry
  • Loading branch information
@dongyuanjushi
dongyuanjushi committed Apr 12, 2024
commit 0188b36da829fd3f4d06e23b7a078217f62502ac
80 changes: 45 additions & 35 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,36 +2,22 @@
import sys
import json

# from src.command_parser import (
# PunctuationParser,
# ChatGPTParser
# )

# from src.command_executor import (
# Executor
# )

from src.scheduler.fifo_scheduler import FIFOScheduler

from src.scheduler.rr_scheduler import RRScheduler

from src.utils.utils import (
parse_global_args,
# logger
)

from src.agents.agent_factory import AgentFactory

from src.agents.agent_process import AgentProcessFactory

import warnings

from src.llms import llms

from src.agents.math_agent.math_agent import MathAgent

from src.agents.narrative_agent.narrative_agent import NarrativeAgent

from src.agents.rec_agent.rec_agent import RecAgent

from src.agents.travel_agent.travel_agent import TravelAgent

from concurrent.futures import ThreadPoolExecutor, as_completed

def main():
Expand All @@ -41,48 +27,72 @@ def main():

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 = llms.LLMKernel(llm_name, max_gpu_memory, max_new_tokens)
llm = llms.LLMKernel(
llm_name,
max_gpu_memory,
eval_device,
max_new_tokens
)

# start the scheduler
scheduler = FIFOScheduler(
# scheduler = FIFOScheduler(
# llm = llm,
# log_mode = scheduler_log_mode
# )

scheduler = RRScheduler(
llm = llm,
log_mode = scheduler_log_mode
)
scheduler.start()

agent_process_factory = AgentProcessFactory()

agent_factory = AgentFactory(
llm = llm,
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=64)

scheduler.start()

# construct agents
math_agent = agent_factory.activate_agent(
agent_name = "MathAgent",
task_input = "Solve the problem that Albert is wondering how much pizza he can eat in one day. He buys 2 large pizzas and 2 small pizzas. A large pizza has 16 slices and a small pizza has 8 slices. If he eats it all, how many pieces does he eat that day?",
math_agent = agent_thread_pool.submit(
lambda p: agent_factory.run_agent(*p),
[
"MathAgent",
f"Solve the problem that Albert is wondering how much pizza he can eat in one day. He buys 2 large pizzas and 2 small pizzas. A large pizza has 16 slices and a small pizza has 8 slices. If he eats it all, how many pieces does he eat that day?"
]
)

narrative_agent = agent_factory.activate_agent(
agent_name = "NarrativeAgent",
task_input = "Craft a tale about a valiant warrior on a quest to uncover priceless treasures hidden within a mystical island.",
narrative_agent = agent_thread_pool.submit(
lambda p: agent_factory.run_agent(*p),
[
"NarrativeAgent",
f"Craft a tale about a valiant warrior on a quest to uncover priceless treasures hidden within a mystical island."
]
)

rec_agent = agent_factory.activate_agent(
agent_name = "RecAgent",
task_input = "I want to take a tour to New York during the spring break, recommend some restaurants around for me.",
rec_agent = agent_thread_pool.submit(
lambda p: agent_factory.run_agent(*p),
[
"RecAgent",
f"I want to take a tour to New York during the spring break, recommend some restaurants around for me."
]
)
agents = [math_agent, narrative_agent, rec_agent]

# run agents concurrently
tasks = [agent_factory.agent_thread_pool.submit(agent.run) for agent in agents]
agent_tasks = [math_agent, narrative_agent, rec_agent]
# agent_tasks = [math_agent]

for r in as_completed(tasks):
for r in as_completed(agent_tasks):
res = r.result()
# logger.info(res)

scheduler.stop()

Expand Down