merge llm-refactor branch

.gitignore

@@ -1,4 +1,7 @@
 .ipynb_checkpoints
 __pycache__
 src/storage/*.json
-.idea
+/venv/
+.idea
+context_restoration/
+.pytest_cache

README.md

@@ -113,4 +113,4 @@ If you would like to join the community, ask questions, chat with fellows, learn
 For issues related to AIOS development, we encourage submitting [issues](https://github.com/agiresearch/AIOS/issues), [pull requests](https://github.com/agiresearch/AIOS/pulls), or initiating discussions in the AIOS [Discord Channel](https://discord.gg/m9Nw9fPR). For other issues please feel free to contact Kai Mei (marknju2018@gmail.com) and Yongfeng Zhang (yongfeng@email.com).
 
 ## 🌍 AIOS Contributors
-[![AIOS contributors](https://contrib.rocks/image?repo=agiresearch/AIOS&max=1000)](https://github.com/agiresearch/AIOS/graphs/contributors)
+[![AIOS contributors](https://contrib.rocks/image?repo=agiresearch/AIOS&max=300)](https://github.com/agiresearch/AIOS/graphs/contributors)

requirements.txt

@@ -1,18 +1,15 @@
-accelerate==0.28.0
-asposestorage==1.0.2
-django_cache_utils==0.7.2
 google_api_python_client==2.123.0
 googlemaps==4.10.0
-langchain_core==0.1.40
+langchain_community==0.0.33
+langchain_core==0.1.43
 numpy==1.26.4
-openai==1.16.2
-pandas==2.2.1
+openai==1.20.0
+pandas==2.2.2
 protobuf==5.26.1
-pydantic==2.6.4
+pydantic==2.7.0
 Pympler==1.0.1
 pytest==8.1.1
 Requests==2.31.0
-streamers==1.3.3
 torch==2.0.1
 transformers==4.38.1
 wolframalpha==5.0.0

src/agents/agent_config/NarrativeAgent.json

@@ -3,7 +3,7 @@
     "description": [
         "You are an expert who is good at writing novels. ",
         "Given a theme or background, you need to write a short story with a well-developed plot and characters, develop different sections of the story, such as introduction, rising action, climax, and conclusion."
-    ], 
+    ],
     "flow": [],
     "tool_info": []
-}
+}

src/agents/agent_config/RecAgent.json

@@ -1,7 +1,7 @@
 {
     "name": "RecAgent",
     "description": [
-        "You are an expert who is good at recommending restaurants or hotels for users.", 
+        "You are an expert who is good at recommending restaurants or hotels for users.",
         "Given a request, you need to first determine the right recommendation direction and then provide the recommendation lists."
     ],
     "flow": [],

src/agents/base.py

@@ -7,6 +7,8 @@
     # AgentProcessQueue
 )
 
+# from multiprocessing import Process
+
 import logging
 
 import time
@@ -32,6 +34,20 @@ def join(self):
         super().join()
         return self.result
 
+# class CustomizedProcess(Process):
+#     def __init__(self, target, args=()):
+#         super().__init__()
+#         self.target = target
+#         self.args = args
+#         self.result = None
+
+#     def run(self):
+#         self.result = self.target(*self.args)
+
+#     def join(self):
+#         super().join()
+#         return self.result
+
 class BaseAgent:
     def __init__(self,
                  agent_name,
@@ -74,6 +90,9 @@ def get_response(self, prompt, temperature=0.0):
         thread = CustomizedThread(target=self.query_loop, args=(prompt, ))
         thread.start()
         return thread.join()
+        # process = CustomizedProcess(target=self.query_loop, args=(prompt, ))
+        # process.start()
+        # return process.join()
         # return self.query_loop(prompt)
 
     def query_loop(self, prompt):
@@ -86,6 +105,7 @@ def query_loop(self, prompt):
         while agent_process.get_status() != "done":
             # print(agent_process.get_status())
             thread = Thread(target=self.listen, args=(agent_process, ))
+            # process = Process(target=self.listen, args=(agent_process, ))
             current_time = time.time()
 
             # reinitialize agent status
@@ -95,6 +115,8 @@ def query_loop(self, prompt):
 
             thread.start()
             thread.join()
+            # process.start()
+            # process.join()
 
             completed_response = agent_process.get_response()
             if agent_process.get_status() != "done":

src/command_parser.py

@@ -21,21 +21,21 @@ def parse(self, instruction):
         """
         if ": " in instruction:
             splited_command = instruction.split(": ")
-            
+
             command_head = splited_command[0].split(" ")
             command_body = splited_command[-1]
-            
+
             command_type = command_head[0]
             command_name = command_head[1]
             return {
                 "command_type": command_type,
                 "command_name": command_name,
                 "command_body": command_body
             }
-    
+
         elif " " in instruction:
             command_head = instruction.split(" ")
-            
+
             command_type = command_head[0]
             command_name = command_head[1]
             return {
@@ -58,4 +58,4 @@ def __init__(self, llm, parser_type = "gpt3.5"):
         pass
 
     def parse(self, instruction):
-        pass
+        pass

src/memory/base.py

@@ -10,8 +10,8 @@ def __init__(self, agent_id: int, round_id: int, operation_type: str, content: s
 class Memory:
     def __init__(self, size=1024):
         self.size = size
-        self.memory = (ctypes.c_ubyte * size)()  
-        self.free_blocks = [(0, size - 1)] 
+        self.memory = (ctypes.c_ubyte * size)()
+        self.free_blocks = [(0, size - 1)]
 
     def mem_alloc(self, size):
         for i, (start, end) in enumerate(self.free_blocks):
@@ -44,10 +44,10 @@ def mem_read(self, address, size):
 class BaseMemoryManager:
     def __init__(self, max_memory_block_size, memory_block_num):
         pass
-    
+
     def run(self):
         pass
-    
+
     def start(self):
         """start the scheduler"""
         self.active = True
@@ -68,4 +68,4 @@ def mem_alloc(self, agent_id):
         pass
 
     def mem_clear(self):
-        pass
+        pass

src/memory/shared_memory.py

@@ -11,4 +11,4 @@ def save(self):
         pass
 
     def load(self):
-        pass
+        pass

src/memory/single_memory.py

@@ -35,15 +35,15 @@ def __init__(self, max_memory_block_size, memory_block_num):
         self.compressor = ZLIBCompressor() # used for compressing data
         heapq.heapify(self.free_memory_blocks)
         self.memory_operation_queue = Queue() # TODO add lock to ensure parallel
-    
+
     def run(self):
         while self.active:
             try:
                 memory_request = self.memory_operation_queue.get(block=True, timeout=0.1)
                 self.execute_operation(memory_request)
             except Empty:
                 pass
-    
+
     def start(self):
         """start the scheduler"""
         self.active = True
@@ -71,7 +71,7 @@ def mem_write(self, agent_id, round_id: str, content: str):
         address = self.memory_blocks[
             self.aid_to_memory[agent_id][round_id]["memory_block_id"]
         ].mem_alloc(size)
-        
+
         self.memory_blocks[
             self.aid_to_memory[agent_id][round_id]["memory_block_id"]
         ].mem_write(address, compressed_content)
@@ -93,4 +93,3 @@ def mem_alloc(self, agent_id):
     def mem_clear(self, agent_id):
         self.aid_to_mid.pop(agent_id)
         heapq.heappush(agent_id)
-

src/scheduler/__init__.py

@@ -2,4 +2,4 @@
 
 # from .fifo_scheduling import FIFOScheduler
 
-# from .rr_scheduling import RRScheduler
+# from .rr_scheduling import RRScheduler

src/scheduler/rr_scheduler.py

@@ -13,13 +13,13 @@ class RRScheduler(BaseScheduler):
     def __init__(self, llm, log_mode):
         super().__init__(llm, log_mode)
         self.agent_process_queue = Queue()
-        self.time_limit = 15
+        self.time_limit = 5
         self.simple_context_manager = SimpleContextManager()
 
     def run(self):
         while self.active:
             try:
-                agent_process = self.agent_process_queue.get(block=True, timeout=1)
+                agent_process = self.agent_process_queue.get(block=True, timeout=0.2)
                 agent_process.set_time_limit(self.time_limit)
 
                 agent_process.set_status("executing")

src/storage/base.py

@@ -12,4 +12,4 @@ def sto_alloc(self, agent_id):
         pass
 
     def sto_clear(self, agent_id):
-        pass
+        pass

src/storage/file_storage.py

@@ -19,4 +19,4 @@ def sto_alloc(self, agent_id):
         return NotImplementedError
 
     def sto_clear(self, agent_id):
-        return NotImplementedError
+        return NotImplementedError

src/tools/online/arxiv.py

@@ -111,4 +111,4 @@ def run(self, query: str) -> str:
         if docs:
             return "\n\n".join(docs)[: self.doc_content_chars_max]
         else:
-            return "No good Arxiv Result was found"
+            return "No good Arxiv Result was found"

src/tools/online/bing_search.py

@@ -45,7 +45,7 @@ def run(self, query: str) -> str:
         response = self._bing_search_results(query, count=self.k)
         result = self.parse_result(response)
         return result
-    
+
     def parse_result(self, response):
         snippets = []
         if len(response) == 0:
@@ -54,5 +54,3 @@ def parse_result(self, response):
             snippets.append(result["snippet"])
 
         return " ".join(snippets)
-
-

src/tools/online/google_place.py

@@ -99,4 +99,4 @@ def format_place_details(self, place_details: Dict[str, Any]) -> Optional[str]:
             return formatted_details
         except Exception as e:
             logging.error(f"An error occurred while formatting place details: {e}")
-            return None
+            return None

src/tools/online/google_search.py

@@ -49,7 +49,7 @@ def __init__(self):
         self.google_cse_id = get_from_env("GOOGLE_CSE_ID")
         self.k: int = 10 # topk searched results
         self.search_engine = self.build_engine()
-        self.siterestrict: bool = False 
+        self.siterestrict: bool = False
 
     def build_engine(self):
         try:
@@ -77,7 +77,7 @@ def run(self, query: str) -> str:
         response = self._google_search_results(query, num=self.k)
         result = self.parse_result(response)
         return result
-    
+
     def parse_result(self, response):
         snippets = []
         if len(response) == 0:

src/tools/online/wikipedia.py

@@ -27,7 +27,7 @@ def build_client(self):
             import wikipedia
 
             wikipedia.set_lang(self.lang)
-            
+
         except ImportError:
             raise ImportError(
                 "Could not import wikipedia python package. "
@@ -115,4 +115,4 @@ def lazy_load(self, query: str) -> Iterator[Document]:
         for page_title in page_titles[: self.top_k_results]:
             if wiki_page := self._fetch_page(page_title):
                 if doc := self._page_to_document(page_title, wiki_page):
-                    yield doc
+                    yield doc

src/tools/online/wolfram_alpha.py

@@ -50,4 +50,4 @@ def run(self, query: str) -> str:
             # We don't want to return the assumption alone if answer is empty
             return "No good Wolfram Alpha Result was found"
         else:
-            return f"Assumption: {assumption} \nAnswer: {answer}"
+            return f"Assumption: {assumption} \nAnswer: {answer}"

src/utils/compressor.py

@@ -9,7 +9,7 @@ def compress(self, data):
 
     def decompress(self, compressed_data):
         pass
-    
+
 class ZLIBCompressor(Compressor):
     def __init__(self) -> None:
         pass
@@ -20,4 +20,4 @@ def compress(self, data):
 
     def decompress(self, compressed_data):
         decompressed_data = zlib.decompress(compressed_data)
-        return decompressed_data.decode('utf-8')
+        return decompressed_data.decode('utf-8')

tests/test_llms.py

@@ -2,31 +2,32 @@
 
 import os
 
-from src.llms.llms import LLMKernel
+from src.llm_kernel.llms import LLMKernel
 
 from src.agents.agent_process import AgentProcess
 
 from src.context.simple_context import SimpleContextManager
 
 def test_closed_llm():
-    if "GEMINI_API_KEY" not in os.environ or not os.environ["RAPID_API_KEY"]:
-        with pytest.raises(ValueError):
-            llm_type = "gemini-pro"
-            llm = LLMKernel(llm_type)
+    # if "GEMINI_API_KEY" not in os.environ or not os.environ["RAPID_API_KEY"]:
+    #     with pytest.raises(ValueError):
+    #         llm_type = "gemini-pro"
+    #         llm = LLMKernel(llm_type)
     llm_type = "gemini-pro"
     llm = LLMKernel(llm_type)
     agent_process = AgentProcess(
         agent_name = "Narrative Agent",
         prompt = "Craft a tale about a valiant warrior on a quest to uncover priceless treasures hidden within a mystical island."
     )
     llm.address_request(agent_process)
+    print(agent_process.get_response())
     assert isinstance(agent_process.get_response(), str)
     # print(response)
 
 def test_open_llm():
-    llm_type = "mistral-7b-it"
-    max_gpu_memory = {"2": "48GB"}
-    eval_device = "cuda:2"
+    llm_type = "gemma-2b-it"
+    max_gpu_memory = {"4": "48GB"}
+    eval_device = "cuda:4"
     max_new_tokens = 256
     llm = LLMKernel(
         llm_type,