rebase to main

rootfs/agents/math_teacher/agent.toml

@@ -1,6 +1,6 @@
 instance_id = "math_teacher"
 fullname = "the one"
-llm_model_name = "LLaMA2-70B"
+llm_model_name = "gpt-4-0613"
 [[prompt]]
 role = "system"
 content = "你是精通数学的老师"

src/aios_kernel/compute_node_config.py

@@ -0,0 +1,87 @@
+"""
+Configuration for nodes:
+
+```
+├── nodes
+│ └── llama
+|   └── 0
+|   |   └── url
+|   |   └── model_name
+|   └── 1
+|       └── url
+|       └── model_name
+```
+"""
+import logging
+from typing import List
+
+import os
+import toml
+
+from .local_llama_compute_node import LocalLlama_ComputeNode
+from .storage import AIStorage
+
+# define singleton class knowledge pipline
+class ComputeNodeConfig:
+    _instance = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = ComputeNodeConfig()
+            cls._instance.__singleton_init__()
+
+        return cls._instance
+    
+    def initial(self) -> List[LocalLlama_ComputeNode]:
+        config_path = self.__config_path()
+        logging.info(f"initial nodes from {config_path}")
+
+        if os.path.exists(config_path):
+            self.config = toml.load(self.__config_path())
+            if self.config is None:
+                return []
+            
+            nodes = []
+            llama_nodes_cfg = self.config["llama"]
+            if llama_nodes_cfg is not None:
+                for cfg in llama_nodes_cfg:
+                    node = LocalLlama_ComputeNode(url=cfg["url"], model_name=cfg["model_name"])
+                    nodes.append(node)
+
+            return nodes
+
+        return []
+        
+    def save(self):
+        with open(self.__config_path(), "w") as f:
+            toml.dump(self.config, f)
+        
+    def add_node(self, model_type: str, url: str, model_name: str):
+        if model_type == "llama":
+            llama_nodes_cfg = self.config.get("llama") or []
+            for cfg in llama_nodes_cfg:
+                if url == cfg["url"] and model_name == cfg["model_name"]:
+                    return
+            llama_nodes_cfg.append({"url": url, "model_name": model_name})
+            self.config["llama"] = llama_nodes_cfg
+    
+    
+    def remove_node(self, model_type: str, url: str, model_name: str):
+        if model_type == "llama":
+            llama_nodes_cfg = self.config.get("llama") or []
+            for i in range(0, len(llama_nodes_cfg)):
+                cfg = llama_nodes_cfg[i]
+                if url == cfg["url"] and model_name == cfg["model_name"]:
+                    llama_nodes_cfg.pop(i)
+
+    def list(self) -> str:
+        return toml.dumps(self.config)
+
+    def __singleton_init__(self):
+        self.config = {}
+
+    @classmethod
+    def __config_path(cls) -> str:
+        user_data_dir = AIStorage.get_instance().get_myai_dir()
+        return os.path.abspath(f"{user_data_dir}/etc/compute_nodes.cfg.toml")

src/aios_kernel/local_llama_compute_node.py

@@ -4,10 +4,10 @@ import logging
 import requests
 from typing import Optional, List
 from pydantic import BaseModel
-from llama_cpp import Llama
 
 from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType
 from .queue_compute_node import Queue_ComputeNode
+from .storage import AIStorage,UserConfig
 
 logger = logging.getLogger(__name__)
 
@@ -16,59 +16,117 @@ This is a custom implementation, it should be redesigned.
 """
 
 class LocalLlama_ComputeNode(Queue_ComputeNode):
-    def __init__(self, model_path: str, model_name: str):
+    def __init__(self, url: str, model_name: str):
         super().__init__()
-        self.model_path = model_path
+        self.url = url
         self.model_name = model_name
-        self.llm = Llama(model_path=model_path)
 
-    async def execute_task(self, task: ComputeTask, result: ComputeTaskResult) -> ComputeTaskResult:
+    async def execute_task(self, task: ComputeTask, result: ComputeTaskResult):
         match task.task_type:
             case ComputeTaskType.TEXT_EMBEDDING:
                 model_name = task.params["model_name"]
                 input = task.params["input"]
-                logger.info(f"call local-llama {model_name} input: {input}")
+                logger.info(f"call local-llama ({self.url}, {self.model_name}) {model_name} input: {input}")
 
-                try:
+                self.embedding(input, result)
-                    embedding = self.llm.embed(input=input)
+                
-                    logger.info(f"local-llama({self.model_path}) response: {embedding}")
+                if result.result_code == ComputeTaskResultCode.OK:
-                except Exception as e:
+                    task.state = ComputeTaskState.DONE
-                    logger.error(f"call local-llama {model_name} run TEXT_EMBEDDING task error: {e}")
+                else:
                     task.state = ComputeTaskState.ERROR
-                    task.error_str = str(e)
+                    task.error_str = result.error_str
-                    result.error_str = str(e)
-                    return result
-                
-                logger.info(f"local-llama({self.model_path}) response: {embedding}")
-                task.state = ComputeTaskState.DONE
-                result.result_code = ComputeTaskResultCode.OK
-                result.result = embedding
 
                 return result
             case ComputeTaskType.LLM_COMPLETION:
                 mode_name = task.params["model_name"]
                 prompts = task.params["prompts"]
-                max_token_size = task.params.get("max_token_size")
-                llm_inner_functions = task.params.get("inner_functions")
-                if max_token_size is None:
-                    max_token_size = 4000
                     
-                logger.info(f"local-llama({self.model_path}) prompts: {prompts}")
+                logger.info(f"local-llama({self.url}, {self.model_name}) prompts: {prompts}")
 
-                try:
+                self.completion(task, result)
-                    resp = self.llm.create_chat_completion(model=mode_name,
+
-                                                    messages=prompts,
+                if result.result_code == ComputeTaskResultCode.OK:
-                                                    functions=llm_inner_functions, # function has not support?
+                    task.state = ComputeTaskState.DONE
-                                                    max_tokens=max_token_size,
+                else:
-                                                    temperature=0.7) # TODO: add temperature to task params?
-                except Exception as e:
-                    logger.error(f"local-llama({self.model_path}) run LLM_COMPLETION task error: {e}")
                     task.state = ComputeTaskState.ERROR
-                    task.error_str = str(e)
+                    task.error_str = result.error_str
-                    result.error_str = str(e)
-                    return result
                 
-                logger.info(f"local-llama({self.model_path}) response: {json.dumps(resp, indent=4)}")
+            case _:
+                task.state = ComputeTaskState.ERROR
+                result.result_code = ComputeTaskResultCode.ERROR
+                task.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
+                result.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
+                return None
+
+    async def initial(self) -> bool:
+        return True
+
+    def display(self) -> str:
+        return f"local-llama: {self.node_id}"
+
+    def get_capacity(self):
+        pass
+
+    def is_support(self, task: ComputeTask) -> bool:
+        return (task.task_type == ComputeTaskType.TEXT_EMBEDDING or task.task_type == ComputeTaskType.LLM_COMPLETION) and (not task.params["model_name"] or task.params["model_name"] == self.model_name)
+
+    def is_local(self) -> bool:
+        return True
+
+    def embedding(self, input: str, result: ComputeTaskResult):
+        body = {
+            "input": input
+        }
+        
+        try:
+            response = requests.post(self.url + "/v1/embeddings", json = body, verify=False, headers={"Content-Type": "application/json"})
+            response.close()
+
+            logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
+
+            if response.status_code == 200:
+                resp = response.json()
+                result.result = resp["data"][0]["embedding"]
+            elif response.status_code == 422:
+                resp = response.json()
+                result.result_code = ComputeTaskResultCode.ERROR
+                result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
+            else:
+                result.result_code = ComputeTaskResultCode.ERROR
+                result.error_str = "http request failed: " + str(response.status_code)
+        except Exception as e:
+            logger.error(f"call local-llama({self.url}, {self.model_name}) run TEXT_EMBEDDING task error: {e}")
+            result.result_code = ComputeTaskResultCode.ERROR
+            result.error_str = str(e)
+            return result
+        
+    def completion(self, task: ComputeTask, result: ComputeTaskResult):
+        mode_name = task.params["model_name"]
+        prompts = task.params["prompts"]
+        max_token_size = task.params.get("max_token_size")
+        llm_inner_functions = task.params.get("inner_functions")
+        if max_token_size is None:
+            max_token_size = max_token_size
+            
+        body = {
+            "messages": [],
+            "max_tokens": 4000
+        }
+
+        for prompt in prompts:
+            body["messages"].append({
+                "role": prompt["role"],
+                "content": prompt["content"]
+            })
+        
+        try:
+            response = requests.post(self.url + "/v1/chat/completions", json = body, verify=False, headers={"Content-Type": "application/json"})
+            response.close()
+
+            logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
+
+            if response.status_code == 200:
+                resp = response.json()
 
                 status_code = resp["choices"][0]["finish_reason"]
                 token_usage = resp["usage"]
@@ -91,27 +149,16 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
                 if token_usage:
                     result.result_refers["token_usage"] = token_usage
 
-                logger.info(f"local-llama({self.model_path}) success response: {result.result_str}")
+                logger.info(f"local-llama({self.url}, {self.model_name}) success response: {result.result_str}")
-
+            elif response.status_code == 422:
-                return result
+                resp = response.json()
-            case _:
-                task.state = ComputeTaskState.ERROR
                 result.result_code = ComputeTaskResultCode.ERROR
-                task.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
+                result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
-                result.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
+            else:
-                return None
+                result.result_code = ComputeTaskResultCode.ERROR
-
+                result.error_str = "http request failed: " + str(response.status_code)
-    async def initial(self) -> bool:
+        except Exception as e:
-        return True
+            logger.error(f"call local-llama({self.url}, {self.model_name}) run LLM_COMPLETION task error: {e}")
-
+            result.result_code = ComputeTaskResultCode.ERROR
-    def display(self) -> str:
+            result.error_str = str(e)
-        return f"LocalLlama_ComputeNode: {self.node_id}"
+            return result
-
-    def get_capacity(self):
-        pass
-
-    def is_support(self, task: ComputeTask) -> bool:
-        return (task.task_type == ComputeTaskType.TEXT_EMBEDDING or task.task_type == ComputeTaskType.LLM_COMPLETION) and (not task.params["model_name"] or task.params["model_name"] == self.model_name)
-
-    def is_local(self) -> bool:
-        return True

src/service/aios_shell/aios_shell.py

@@ -20,14 +20,15 @@ from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
 from prompt_toolkit.completion import WordCompleter
 from prompt_toolkit.styles import Style
 
+
 directory = os.path.dirname(__file__)
 sys.path.append(directory + '/../../')
 
-
+from aios_kernel import AIOS_Version,AgentMsgType,UserConfigItem,AIStorage,Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment,EmailTunnel,LocalLlama_ComputeNode,Local_Stability_ComputeNode,Stability_ComputeNode,PaintEnvironment
-
+from aios_kernel import ContactManager,Contact
 import proxy
 from aios_kernel import *
-
+from aios_kernel.compute_node_config import ComputeNodeConfig
 
 sys.path.append(directory + '/../../component/')
 from agent_manager import AgentManager
@@ -140,6 +141,10 @@ class AIOS_Shell:
             return False
         ComputeKernel.get_instance().add_compute_node(open_ai_node)
 
+        nodes = ComputeNodeConfig.get_instance().initial()
+        for node in nodes:
+            await node.start()
+            ComputeKernel.get_instance().add_compute_node(node)
 
         if await AIStorage.get_instance().is_feature_enable("llama"):
             llama_ai_node = LocalLlama_ComputeNode()
@@ -355,6 +360,54 @@ class AIOS_Shell:
             journals = [str(journal) for journal in KnowledgePipline.get_instance().get_latest_journals(topn)]
             print_formatted_text("\r\n".join(journals))
 
+        if sub_cmd == "query":
+            if len(args) < 2:
+                return show_text
+            prompt = AgentPrompt()
+            prompt.messages.append({"role": "user", "content":" ".join(args[1:])})
+            result = await KnowledgeBase().query_prompt(prompt)
+            print_formatted_text(result.as_str())
+    
+    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':
@@ -480,6 +533,8 @@ 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':
@@ -668,6 +723,9 @@ async def main():
                                '/enable $feature',
                                '/disable $feature',
                                '/list_config',
+                               '/node add llama $model_name $url',
+                               '/node rm llama $model_name $url',
+                               '/node list',
                                '/show',
                                '/exit',
                                '/help'], ignore_case=True)