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 15eec79..52e6af4 100644 --- a/src/aios_kernel/local_llama_compute_node.py +++ b/src/aios_kernel/local_llama_compute_node.py @@ -1,11 +1,13 @@ +import json import logging import requests from typing import Optional, List from pydantic import BaseModel -from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType +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__) @@ -14,81 +16,149 @@ This is a custom implementation, it should be redesigned. """ class LocalLlama_ComputeNode(Queue_ComputeNode): - async def execute_task(self, task: ComputeTask) -> { - "content": str, - "message": str, - "state": ComputeTaskState, - "error": { - "code": int, - "message": str, - } - }: - class GenerateResponse(BaseModel): - error: Optional[int] - msg: Optional[str] - results: Optional[List[str]] + def __init__(self, url: str, model_name: str): + super().__init__() + self.url = url + self.model_name = model_name - try: - prompt_msgs = [] - for prompt in task.params["prompts"]: - prompt_msgs.append(prompt["content"]) + 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 ({self.url}, {self.model_name}) {model_name} input: {input}") + + self.embedding(input, result) - body = { - "prompts": prompt_msgs - } - - response = requests.post("http://aigc:7880/generate", json = body, verify=False, headers={"Content-Type": "application/json"}) - response.close() - - logger.info(f"LocalLlama_ComputeNode task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}") - - if response.status_code != 200: - return { - "state": ComputeTaskState.ERROR, - "error": { - "code": response.status_code, - "message": "http request failed: " + str(response.status_code) - } - } - else: - resp = response.json() - if "error" in resp: - return { - "state": ComputeTaskState.ERROR, - "error": { - "code": resp["error"], - "message": "local llama failed:" + resp["msg"] - } - } + if result.result_code == ComputeTaskResultCode.OK: + task.state = ComputeTaskState.DONE else: - return { - "state": ComputeTaskState.DONE, - "content": str(resp["results"]), - "message": str(resp["results"]) - } - except Exception as err: - import traceback - logger.error(f"{traceback.format_exc()}, error: {err}") - - return { - "state": ComputeTaskState.ERROR, - "error": { - "code": -1, - "message": "unknown exception: " + str(err) - } - } + 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"] + + logger.info(f"local-llama({self.url}, {self.model_name}) prompts: {prompts}") + + self.completion(task, result) + + if result.result_code == ComputeTaskResultCode.OK: + task.state = ComputeTaskState.DONE + else: + task.state = ComputeTaskState.ERROR + 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"LocalLlama_ComputeNode: {self.node_id}" + return f"local-llama: {self.node_id}" def get_capacity(self): pass def is_support(self, task: ComputeTask) -> bool: - return task.task_type == ComputeTaskType.LLM_COMPLETION and (not task.params["model_name"] or task.params["model_name"] == "llama") + 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"] + + match status_code: + case "function_call": + task.state = ComputeTaskState.DONE + case "stop": + task.state = ComputeTaskState.DONE + case _: + task.state = ComputeTaskState.ERROR + task.error_str = f"The status code was {status_code}." + result.error_str = f"The status code was {status_code}." + result.result_code = ComputeTaskResultCode.ERROR + return None + + result.result_code = ComputeTaskResultCode.OK + result.result_str = resp["choices"][0]["message"]["content"] + result.result_message = resp["choices"][0]["message"] + if token_usage: + result.result_refers["token_usage"] = token_usage + + 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 + 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/aios_kernel/queue_compute_node.py b/src/aios_kernel/queue_compute_node.py index 6d97446..3c1bf3e 100644 --- a/src/aios_kernel/queue_compute_node.py +++ b/src/aios_kernel/queue_compute_node.py @@ -16,15 +16,7 @@ class Queue_ComputeNode(ComputeNode): self.is_start = False @abstractmethod - async def execute_task(self, task: ComputeTask) -> { - "content": str, - "message": str, - "state": ComputeTaskState, - "error": { - "code": int, - "message": str, - } - }: + async def execute_task(self, task: ComputeTask, result: ComputeTaskResult): pass async def push_task(self, task: ComputeTask, proiority: int = 0): @@ -36,24 +28,15 @@ class Queue_ComputeNode(ComputeNode): async def _run_task(self, task: ComputeTask): task.state = ComputeTaskState.RUNNING - - resp = await self.execute_task(task) + result = ComputeTaskResult() - - result.worker_id = self.node_id - task.state = resp["state"] - - if task.state == ComputeTaskState.ERROR: - result.result_code = ComputeTaskResultCode.ERROR - task.error_str = resp["error"]["message"] - else: - result.result_code = ComputeTaskResultCode.OK - result.result_str = resp["content"] - result.result_message = resp["message"] - + result.result_code = ComputeTaskResultCode.ERROR result.set_from_task(task) + result.worker_id = self.node_id + await self.execute_task(task, result) + return result def start(self): diff --git a/src/requirements.txt b/src/requirements.txt index c6128ee..b83e564 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -137,4 +137,4 @@ python-telegram-bot pydub stability_sdk sentence-transformers==2.2.2 -tiktoken +tiktoken \ 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 2f64d1c..2478e80 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,11 +141,15 @@ class AIOS_Shell: return False ComputeKernel.get_instance().add_compute_node(open_ai_node) + nodes = ComputeNodeConfig.get_instance().initial() + for node in nodes: + node.start() + ComputeKernel.get_instance().add_compute_node(node) if await AIStorage.get_instance().is_feature_enable("llama"): llama_ai_node = LocalLlama_ComputeNode() if await llama_ai_node.initial() is True: - await llama_ai_node.start() + llama_ai_node.start() ComputeKernel.get_instance().add_compute_node(llama_ai_node) else: logger.error("llama node initial failed!") @@ -164,9 +169,7 @@ class AIOS_Shell: # if await stability_api_node.initial() is not True: # logger.error("stability api node initial failed!") # ComputeKernel.get_instance().add_compute_node(stability_api_node) - - - + local_st_text_compute_node = LocalSentenceTransformer_Text_ComputeNode() if local_st_text_compute_node.initial() is not True: logger.error("local sentence transformer text embedding node initial failed!") @@ -179,7 +182,6 @@ class AIOS_Shell: else: ComputeKernel.get_instance().add_compute_node(local_st_image_compute_node) - await ComputeKernel.get_instance().start() AIBus().get_default_bus().register_unhandle_message_handler(self._handle_no_target_msg) @@ -358,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': @@ -483,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': @@ -671,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)