From 193c627bdba3e9a05447f27a56b509431da7f774 Mon Sep 17 00:00:00 2001 From: streetycat <305190374@qq.com> Date: Thu, 28 Sep 2023 09:09:30 +0000 Subject: [PATCH] rebase to main --- rootfs/agents/math_teacher/agent.toml | 2 +- src/aios_kernel/compute_node_config.py | 87 +++++++++++ src/aios_kernel/local_llama_compute_node.py | 165 +++++++++++++------- src/service/aios_shell/aios_shell.py | 64 +++++++- 4 files changed, 255 insertions(+), 63 deletions(-) create mode 100644 src/aios_kernel/compute_node_config.py diff --git a/rootfs/agents/math_teacher/agent.toml b/rootfs/agents/math_teacher/agent.toml index f7a29a5..646c8e4 100644 --- a/rootfs/agents/math_teacher/agent.toml +++ b/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 = "你是精通数学的老师" diff --git a/src/aios_kernel/compute_node_config.py b/src/aios_kernel/compute_node_config.py new file mode 100644 index 0000000..8e25f50 --- /dev/null +++ b/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") diff --git a/src/aios_kernel/local_llama_compute_node.py b/src/aios_kernel/local_llama_compute_node.py index 7d8c87e..52e6af4 100644 --- a/src/aios_kernel/local_llama_compute_node.py +++ b/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: - embedding = self.llm.embed(input=input) - logger.info(f"local-llama({self.model_path}) response: {embedding}") - except Exception as e: - logger.error(f"call local-llama {model_name} run TEXT_EMBEDDING task error: {e}") - task.state = ComputeTaskState.ERROR - task.error_str = str(e) - result.error_str = str(e) - return result + self.embedding(input, result) - logger.info(f"local-llama({self.model_path}) response: {embedding}") - task.state = ComputeTaskState.DONE - result.result_code = ComputeTaskResultCode.OK - result.result = embedding + if result.result_code == ComputeTaskResultCode.OK: + task.state = ComputeTaskState.DONE + else: + task.state = ComputeTaskState.ERROR + task.error_str = result.error_str 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: - resp = self.llm.create_chat_completion(model=mode_name, - messages=prompts, - functions=llm_inner_functions, # function has not support? - max_tokens=max_token_size, - 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}") + self.completion(task, result) + + if result.result_code == ComputeTaskResultCode.OK: + task.state = ComputeTaskState.DONE + else: task.state = ComputeTaskState.ERROR - task.error_str = str(e) - result.error_str = str(e) - return result + task.error_str = result.error_str + + 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 - logger.info(f"local-llama({self.model_path}) response: {json.dumps(resp, indent=4)}") + 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}") - - return result - case _: - task.state = ComputeTaskState.ERROR + logger.info(f"local-llama({self.url}, {self.model_name}) success response: {result.result_str}") + elif response.status_code == 422: + resp = response.json() 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"LocalLlama_ComputeNode: {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 + 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 LLM_COMPLETION task error: {e}") + result.result_code = ComputeTaskResultCode.ERROR + result.error_str = str(e) + return result \ No newline at end of file diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index 1513317..ce447e5 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/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)