Merge Llama support for release.

src/aios_kernel/__init__.py

@@ -24,6 +24,6 @@ from .workspace_env import WorkspaceEnvironment
 from .local_stability_node import Local_Stability_ComputeNode
 from .stability_node import Stability_ComputeNode
 from .local_st_compute_node import LocalSentenceTransformer_Text_ComputeNode,LocalSentenceTransformer_Image_ComputeNode
-
+from .compute_node_config import ComputeNodeConfig
 AIOS_Version = "0.5.1, build 2023-9-28"
 

src/aios_kernel/compute_kernel.py

@@ -52,7 +52,8 @@ class ComputeKernel:
                 task = await self.task_queue.get()
                 logger.info(f"compute_kernel get task: {task.display()}")
                 c_node: ComputeNode = self._schedule(task)
-                await c_node.push_task(task)
+                if c_node:
+                    await c_node.push_task(task)
 
             logger.warn("compute_kernel is stoped!")
 
@@ -62,6 +63,7 @@ class ComputeKernel:
         # find all the node which supports this task
         support_nodes = []
         total_weights = 0
+
         for node in self.compute_nodes.values():
             if node.is_support(task) is True:
                 support_nodes.append({
@@ -70,6 +72,10 @@ class ComputeKernel:
                 })
                 total_weights += node.weight()
 
+        if len(support_nodes) < 1:
+            logger.warning(f"task {task.display()} is not support by any compute node")
+            return None
+        
         # hit a random node with weight
         hit_pos = random.randint(0, total_weights - 1)
         for i in range(min(len(support_nodes) - 1, hit_pos), -1, -1):

src/service/aios_shell/aios_shell.py

@@ -139,6 +139,11 @@ class AIOS_Shell:
             return False
         ComputeKernel.get_instance().add_compute_node(open_ai_node)
 
+        llama_nodes = ComputeNodeConfig.get_instance().initial()
+        for llama_node in llama_nodes:
+            llama_node.start()
+            ComputeKernel.get_instance().add_compute_node(llama_node)
+
         if await AIStorage.get_instance().is_feature_enable("llama"):
             llama_ai_node = LocalLlama_ComputeNode()
             if await llama_ai_node.initial() is True:
@@ -379,6 +384,46 @@ class AIOS_Shell:
                 image = Image.open(io.BytesIO(image_data))
                 image.show()
 
+    async def handle_node_commands(self, args):
+        show_text = FormattedText([("class:title", "sub command not support!\n" 
+                              "/node add llama $model_name $url\n"
+                              "/node rm llama $model_name $url\n"
+                              "/node list\n")])
+        if len(args) < 1:
+            return show_text
+        sub_cmd = args[0]
+        if sub_cmd == "add":
+            if len(args) < 2:
+                return show_text
+            if args[1] == "llama":
+                if len(args) < 4:
+                    return show_text
+                
+                model_name = args[2]
+                url = args[3]
+                ComputeNodeConfig.get_instance().add_node("llama", url, model_name)
+                ComputeNodeConfig.get_instance().save()
+                node = LocalLlama_ComputeNode(url, model_name)
+                node.start()
+                ComputeKernel.get_instance().add_compute_node(node)
+            else:
+                return show_text
+        elif sub_cmd == "rm":
+            if len(args) < 2:
+                return show_text
+            if args[1] == "llama":
+                if len(args) < 4:
+                    return show_text
+                
+                model_name = args[3]
+                url = args[4]
+                ComputeNodeConfig.get_instance().remove_node("llama", url, model_name)
+                ComputeNodeConfig.get_instance().save()
+            else:
+                return show_text
+        elif sub_cmd == "list":
+            print_formatted_text(ComputeNodeConfig.get_instance().list())
+
     async def call_func(self,func_name, args):
         match func_name:
             case 'send':
@@ -504,10 +549,12 @@ class AIOS_Shell:
                         format_texts.append(("",f"\n-------------------\n"))
                     return FormattedText(format_texts)
                 return FormattedText([("class:title", f"chatsession not found")])
+            case 'node':
+                return await self.handle_node_commands(args)
             case 'exit':
                 os._exit(0)
             case 'help':
-                return FormattedText([("class:title", f"help~~~")])
+                return FormattedText([("class:title", f"GO to https://github.com/fiatrete/OpenDAN-Personal-AI-OS/issues ^_^")])
 
 
 ##########################################################################################################################
@@ -693,7 +740,8 @@ async def main():
                                '/set_config $key',
                                '/enable $feature',
                                '/disable $feature',
-                               '/list_config',
+                               '/node add llama $model_name $url',
+                               '/node rm llama $model_name $url',
                                '/show',
                                '/exit',
                                '/help'], ignore_case=True)