Merge pull request #31 from dongyuanjushi/main

README.md

@@ -1,7 +1,7 @@
 # AIOS: LLM Agent Operating System
 
-<a href='https://arxiv.org/abs/2403.16971'><img src='https://img.shields.io/badge/Paper-PDF-red'></a> 
-<a href='https://arxiv.org/abs/2312.03815'><img src='https://img.shields.io/badge/Paper-PDF-blue'></a> 
+<a href='https://arxiv.org/abs/2403.16971'><img src='https://img.shields.io/badge/Paper-PDF-red'></a>
+<a href='https://arxiv.org/abs/2312.03815'><img src='https://img.shields.io/badge/Paper-PDF-blue'></a>
 [![Code License](https://img.shields.io/badge/Code%20License-MIT-green.svg)](https://github.com/agiresearch/AIOS/blob/main/LICENSE)
 
 AIOS, a Large Language Model (LLM) Agent operating system, embeds large language model into Operating Systems (OS) as the brain of the OS, enabling an operating system "with soul" -- an important step towards AGI. AIOS is designed to optimize resource allocation, facilitate context switch across agents, enable concurrent execution of agents, provide tool service for agents, maintain access control for agents, and provide a rich set of toolkits for LLM Agent developers.
@@ -23,39 +23,51 @@ AIOS, a Large Language Model (LLM) Agent operating system, embeds large language
 ```bash
 git clone https://github.com/agiresearch/AIOS.git
 ```
-**Make sure you have Python >= 3.9 and <= 3.11**  
-Install the required packages using pip  
+**Make sure you have Python >= 3.9 and <= 3.11**
+Install the required packages using pip
 ```bash
 pip install -r requirements.txt
 ```
 
 ### Usage
-
-### Run with locally-deployed LLM
-Set up [Hugging Face token](https://huggingface.co/settings/tokens) and cache directory
+If you use open-sourced models from huggingface, you need to setup your [Hugging Face token](https://huggingface.co/settings/tokens) and cache directory
 ```bash
 export HUGGING_FACE_HUB_TOKEN=<YOUR READ TOKEN>
 export HF_HOME=<YOUR CACHE DIRECTORY>
 ```
-Replace the max_gpu_memory and eval_device with your own and run
+If you use LLM APIs like Gemini-pro, you need to setup your [Gemini API Key](https://aistudio.google.com/app/apikey)
+```bash
+export GEMINI_API_KEY=<YOUR GEMINI API KEY>
+```
 
+Here we provide two modes to run the AIOS: interactive mode and deployment mode
+#### Interactive Mode
+In the interactive mode, you can interact with AIOS to see the output of each step in running multiple agents
 ```python
-# Use Gemma-2b-it
+# Use Gemma-2b-it, replace the max_gpu_memory and eval_device with your own and run
 python main.py --llm_name gemma-2b-it --max_gpu_memory '{"0": "24GB"}' --eval_device "cuda:0" --max_new_tokens 256
 ```
-
 ```python
-# Use Mixtral-8x7b-it
+# Use Mixtral-8x7b-it, replace the max_gpu_memory and eval_device with your own and run
 python main.py --llm_name mixtral-8x7b-it --max_gpu_memory '{"0": "48GB", "1": "48GB", "2": "48GB"}' --eval_device "cuda:0" --max_new_tokens 256
 ```
-### Run with LLM API
-Run with Gemini-pro, setup [Gemini API Key](https://aistudio.google.com/app/apikey)
-```bash
-export GEMINI_API_KEY=<YOUR GEMINI API KEY>
+```python
+# Use Gemini-pro, run with Gemini-pro
+python main.py --llm_name gemini-pro
+```
+#### Deployment Mode
+In the deployment mode, the outputs of running agents are stored in files. And in this mode, you are provided with multiple commands to run agents and see resource usage of agents (e.g., run \<xxxAgent\>: \<YOUR TASK\>, print agent)
+```python
+# Use Gemma-2b-it, replace the max_gpu_memory and eval_device with your own and run
+python simulator.py --llm_name gemma-2b-it --max_gpu_memory '{"0": "24GB"}' --eval_device "cuda:0" --max_new_tokens 256 --scheduler_log_mode file --agent_log_mode file
+```
+```python
+# Use Mixtral-8x7b-it
+python simulator.py --llm_name mixtral-8x7b-it --max_gpu_memory '{"0": "48GB", "1": "48GB", "2": "48GB"}' --eval_device "cuda:0" --max_new_tokens 256 --scheduler_log_mode file --agent_log_mode file
 ```
 ```python
 # Use Gemini-pro
-python main.py --llm_name gemini-pro
+python simulator.py --llm_name gemini-pro --scheduler_log_mode file --agent_log_mode file
 ```
 
 ## 🖋️ References
@@ -80,7 +92,7 @@ AIOS is dedicated to facilitating LLM agents' development and deployment in a sy
 For detailed information on how to contribute, see [CONTRIBUTE](https://github.com/agiresearch/AIOS/blob/main/CONTRIBUTE.md). If you would like to contribute to the codebase, [issues](https://github.com/agiresearch/AIOS/issues) or [pull requests](https://github.com/agiresearch/AIOS/pulls) are always welcome!
 
 ## 🌟 Discord Channel
-If you would like to join the community, ask questions, chat with fellows, learn about or propose new features, and participate in future developments, join our [Discord Community](https://discord.gg/aUg3b2Kd)! 
+If you would like to join the community, ask questions, chat with fellows, learn about or propose new features, and participate in future developments, join our [Discord Community](https://discord.gg/aUg3b2Kd)!
 
 ## 🌍 AIOS Contributors
 <a href="https://github.com/agiresearch/AIOS/graphs/contributors">

main.py

@@ -42,40 +42,43 @@ def main():
     llm_name = args.llm_name
     max_gpu_memory = args.max_gpu_memory
     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)
 
     # start the scheduler
-    scheduler = FIFOScheduler(llm)
+    scheduler = FIFOScheduler(
+        llm = llm,
+        log_mode = scheduler_log_mode
+    )
     scheduler.start()
 
     agent_factory = AgentFactory(
         llm = llm,
-        agent_process_queue = scheduler.agent_process_queue
+        agent_process_queue = scheduler.agent_process_queue,
+        agent_log_mode = agent_log_mode
     )
 
-    # assign maximum number of agents that can run in parallel
-    agent_thread_pool = ThreadPoolExecutor(max_workers=64)
-
     # construct agents
     math_agent = agent_factory.activate_agent(
-        agent_name = "MathAgent", 
+        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?",
     )
 
     narrative_agent = agent_factory.activate_agent(
-        agent_name = "NarrativeAgent", 
+        agent_name = "NarrativeAgent",
         task_input = "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", 
+        agent_name = "RecAgent",
         task_input = "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_thread_pool.submit(agent.run) for agent in agents]
+    tasks = [agent_factory.agent_thread_pool.submit(agent.run) for agent in agents]
 
     for r in as_completed(tasks):
         res = r.result()

simulator.py

@@ -0,0 +1,95 @@
+import os
+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.utils.utils import (
+    parse_global_args,
+    logger
+)
+
+from src.agents.agent_factory import AgentFactory
+
+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
+
+from src.utils.utils import logger
+
+def main():
+    warnings.filterwarnings("ignore")
+    parser = parse_global_args()
+    args = parser.parse_args()
+
+    llm_name = args.llm_name
+    max_gpu_memory = args.max_gpu_memory
+    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)
+
+    # start the scheduler
+    scheduler = FIFOScheduler(
+        llm = llm,
+        log_mode = scheduler_log_mode
+    )
+
+    agent_factory = AgentFactory(
+        llm = llm,
+        agent_process_queue = scheduler.agent_process_queue,
+        agent_log_mode = agent_log_mode
+    )
+
+    parser = PunctuationParser(
+        llm = llm
+    )
+
+    executor = Executor(agent_factory=agent_factory)
+
+    scheduler.start()
+
+    # agent_factory.start() # TODO add gabage recycle of agent ID
+
+    while True:
+        try:
+            # Read a command line input
+            command_line = input(f"[{llm_name}]>")
+            if command_line.strip().lower() == "exit":
+                print("Exiting...")
+                break
+
+            # Parse command
+            tokens = parser.parse(command_line)
+            # Execute the command
+            executor.execute(tokens)
+
+        except KeyboardInterrupt:
+            # Handle Ctrl+C gracefully
+            print("\nUse 'exit' to quit the shell.")
+        except EOFError:
+            pass
+    scheduler.stop()
+
+if __name__ == "__main__":
+    main()

src/agents/agent_factory.py

@@ -4,44 +4,92 @@
 
 from src.utils.global_param import (
     MAX_AID,
-    agent_pool,
     agent_table
 )
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+from threading import Thread
 class AgentFactory:
-    def __init__(self, llm, agent_process_queue):
+    def __init__(self, llm, agent_process_queue, agent_log_mode):
         self.MAX_AID = MAX_AID
         self.llm = llm
         self.aid_pool = [i for i in range(self.MAX_AID)]
         heapq.heapify(self.aid_pool)
         self.agent_process_queue = agent_process_queue
-        
+
         self.agent_table = agent_table
-        self.agent_pool = agent_pool
-
+        self.current_agents = {}
+
+        self.agent_thread_pool = ThreadPoolExecutor(max_workers=64)
+
+        self.thread = Thread(target=self.deactivate_agent)
+
+        self.agent_log_mode = agent_log_mode
+
     def activate_agent(self, agent_name, task_input):
         agent = self.agent_table[agent_name](
-            agent_name = agent_name, 
+            agent_name = agent_name,
             task_input = task_input,
             llm = self.llm,
-            agent_process_queue = self.agent_process_queue
+            agent_process_queue = self.agent_process_queue,
+            log_mode = self.agent_log_mode
         )
         aid = heapq.heappop(self.aid_pool)
-        
+
         agent.set_aid(aid)
         time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        
-        agent.set_status("Active")
+
+        agent.set_status("active")
         agent.set_created_time(time)
-        
-        self.agent_pool[aid] = agent
+
+        self.current_agents[aid] = agent
         return agent
-
-    def deactivate_agent(self, aid):
-        self.agent_pool.popitem(aid)
-        heapq.heappush(self.aid_pool, aid)
-        # self.aid_pool.heappush(aid)
-
-    def print(self):
-        for aid, agent in agent_pool:
-            print(f"| Agent ID: {aid} | Agent Name: {agent.agent_name} | Status: {agent.get_status()} | Activated Time: {agent.get_created_time()} |")
-
+
+    def print_agent(self):
+        headers = ["Agent ID", "Agent Name", "Created Time", "Status"]
+        data = []
+        for id, agent in self.current_agents.items():
+            data.append(
+                [id, agent.agent_name, agent.created_time, agent.status]
+            )
+        self.print(headers=headers, data=data)
+
+
+    def print(self, headers, data):
+        # align output
+        column_widths = [
+            max(len(str(row[i])) for row in [headers] + data) for i in range(len(headers))
+        ]
+        print("-" * (sum(column_widths) + len(headers) * 3 - 1))
+        print(self.format_row(headers, column_widths))
+        print("-" * (sum(column_widths) + len(headers) * 3 - 1))
+        for row in data:
+            print(self.format_row(row, column_widths))
+            print("-" * (sum(column_widths) + len(headers) * 3 - 1))
+
+
+    def format_row(self, row, widths, align="<"):
+        row_str = " | ".join(f"{str(item):{align}{widths[i]}}" for i, item in enumerate(row))
+        return row_str
+
+    def deactivate_agent(self):
+        import time
+        while True:
+            invalid_aids = []
+            items = self.current_agents.items()
+            for aid, agent in items:
+                if agent.get_status() == "Done":
+                    agent.set_status("Inactive")
+                    time.sleep(5)
+                    invalid_aids.append(aid)
+            for aid in invalid_aids:
+                self.current_agents.pop(aid)
+                heapq.heappush(self.aid_pool, aid)
+
+    def start(self):
+        """start the factory to check inactive agent"""
+        self.thread.start()
+
+    def stop(self):
+        self.thread.join()