From b271f8883e9eb426a4fd66a5df6d00e8464c3d60 Mon Sep 17 00:00:00 2001 From: Liu Zhicong Date: Tue, 26 Sep 2023 01:19:33 -0700 Subject: [PATCH] Add /enable /disable Feature --- src/aios_kernel/google_text_to_speech_node.py | 2 +- src/aios_kernel/knowledge_pipeline.py | 2 + src/aios_kernel/local_llama_compute_node.py | 3 + src/aios_kernel/storage.py | 23 +++- src/component/agent_manager/agent_manager.py | 4 +- .../workflow_manager/workflow_manager.py | 5 +- src/service/aios_shell/aios_shell.py | 119 +++++++++++++----- src/service/aios_shell/proxy.py | 8 +- 8 files changed, 124 insertions(+), 42 deletions(-) diff --git a/src/aios_kernel/google_text_to_speech_node.py b/src/aios_kernel/google_text_to_speech_node.py index b76d0fc..0fdba47 100644 --- a/src/aios_kernel/google_text_to_speech_node.py +++ b/src/aios_kernel/google_text_to_speech_node.py @@ -171,7 +171,7 @@ class GoogleTextToSpeechNode(ComputeNode): def is_local(self) -> bool: return False - def declare_user_config(self): + def declare_user_config(self,is_optional:bool = False): if os.getenv("GOOGLE_APPLICATION_CREDENTIALS") is None: user_config = AIStorage.get_instance().get_user_config() user_config.add_user_config("google_application_credentials", diff --git a/src/aios_kernel/knowledge_pipeline.py b/src/aios_kernel/knowledge_pipeline.py index d1ed109..7f90a08 100644 --- a/src/aios_kernel/knowledge_pipeline.py +++ b/src/aios_kernel/knowledge_pipeline.py @@ -322,6 +322,8 @@ class KnowledgePipline: if source_config['type'] == 'dir': self.add_dir_source(KnowledgeDirSource(source_config)) + return True + def __singleton_init__(self): self.knowledge_base = KnowledgeBase() self.email_sources = dict() diff --git a/src/aios_kernel/local_llama_compute_node.py b/src/aios_kernel/local_llama_compute_node.py index 7e50661..15eec79 100644 --- a/src/aios_kernel/local_llama_compute_node.py +++ b/src/aios_kernel/local_llama_compute_node.py @@ -78,6 +78,9 @@ class LocalLlama_ComputeNode(Queue_ComputeNode): } } + async def initial(self) -> bool: + return True + def display(self) -> str: return f"LocalLlama_ComputeNode: {self.node_id}" diff --git a/src/aios_kernel/storage.py b/src/aios_kernel/storage.py index 211129d..b0d34c3 100644 --- a/src/aios_kernel/storage.py +++ b/src/aios_kernel/storage.py @@ -140,6 +140,7 @@ class AIStorage: def __init__(self) -> None: self.is_dev_mode = False self.user_config = UserConfig() + self.feature_init_results = {} async def initial(self)->bool: self.user_config.user_config_path = str(self.get_myai_dir() / "etc/system.cfg.toml") @@ -147,16 +148,30 @@ class AIStorage: await self.user_config.load_value_from_file(self.user_config.user_config_path,True) async def enable_feature(self,feature_name:str) -> None: - pass + self.user_config.set_value(f"feature.{feature_name}","True") + await self.user_config.save_to_user_config() + async def disable_feature(self,feature_name:str) -> None: - pass + self.user_config.set_value(f"feature.{feature_name}","False") + await self.user_config.save_to_user_config() async def set_feature_init_result(self,feature_name:str,result:bool) -> None: - pass + self.feature_init_results[feature_name] = result async def is_feature_enable(self,feature_name:str) -> bool: - pass + is_enable = self.user_config.get_value(f"feature.{feature_name}") + if is_enable is None: + return False + + init_result = self.feature_init_results.get(feature_name) + if init_result: + if init_result is False: + return False + + if is_enable == "True": + return True + return False def get_user_config(self) -> UserConfig: return self.user_config diff --git a/src/component/agent_manager/agent_manager.py b/src/component/agent_manager/agent_manager.py index fecead2..083b680 100644 --- a/src/component/agent_manager/agent_manager.py +++ b/src/component/agent_manager/agent_manager.py @@ -41,8 +41,8 @@ class AgentManager: self.db_path = f"{user_data_dir}/messages.db" self.loaded_agent_instance = {} - if self.agent_templete_env is None: - raise Exception("agent_manager initial failed") + + return True async def scan_all_agent(self)->None: pass diff --git a/src/component/workflow_manager/workflow_manager.py b/src/component/workflow_manager/workflow_manager.py index 39be2e4..ee5ef9e 100644 --- a/src/component/workflow_manager/workflow_manager.py +++ b/src/component/workflow_manager/workflow_manager.py @@ -36,7 +36,10 @@ class WorkflowManager: self.db_file = os.path.abspath(f"{user_data_dir}/messages.db") if self.workflow_env is None: - raise Exception("WorkflowManager initial failed") + logger.error(f"load workflow env failed!") + return False + + return True async def get_agent_default_workflow(self,agent_id:str) -> Workflow: pass diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index ca43eca..3cde4b6 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/src/service/aios_shell/aios_shell.py @@ -6,6 +6,7 @@ import logging import re import toml import shlex +from logging.handlers import RotatingFileHandler from typing import Any, Optional, TypeVar, Tuple, Sequence import argparse @@ -52,8 +53,9 @@ class AIOS_Shell: self.is_working = True def declare_all_user_config(self): - cm_path = AIStorage.get_instance().get_myai_dir() / "contacts.toml" - cm = ContactManager.get_instance(cm_path) + user_data_dir = AIStorage.get_instance().get_myai_dir() + contact_config_path =os.path.abspath(f"{user_data_dir}/contacts.toml") + cm = ContactManager.get_instance(contact_config_path) cm.load_data() user_config = AIStorage.get_instance().get_user_config() @@ -61,6 +63,9 @@ class AIOS_Shell: user_config.add_user_config("telegram","Your telgram username",False,None) user_config.add_user_config("email","Your email",False,None) + user_config.add_user_config("feature.llama","enable Local-llama feature",True,"False") + user_config.add_user_config("feature.aigc","enable AIGC feature",True,"False") + openai_node = OpenAI_ComputeNode.get_instance() openai_node.declare_user_config() @@ -72,6 +77,9 @@ class AIOS_Shell: Local_Stability_ComputeNode.declare_user_config() + + + async def _handle_no_target_msg(self,bus:AIBus,msg:AgentMsg) -> bool: target_id = msg.target.split(".")[0] @@ -113,48 +121,55 @@ class AIOS_Shell: workspace_env = WorkspaceEnvironment("bash") Environment.set_env_by_id("bash",workspace_env) + if await AgentManager.get_instance().initial() is not True: + logger.error("agent manager initial failed!") + return False + if await WorkflowManager.get_instance().initial() is not True: + logger.error("workflow manager initial failed!") + return False - await AgentManager.get_instance().initial() - await WorkflowManager.get_instance().initial() - - open_ai_node = OpenAI_ComputeNode.get_instance() if await open_ai_node.initial() is not True: logger.error("openai node initial failed!") return False ComputeKernel.get_instance().add_compute_node(open_ai_node) - try: - google_text_to_speech_node = GoogleTextToSpeechNode.get_instance() - google_text_to_speech_node.init() - ComputeKernel.get_instance().add_compute_node(google_text_to_speech_node) - except Exception as e: - logger.error(f"google text to speech node initial failed! {e}") - return False + 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() + ComputeKernel.get_instance().add_compute_node(llama_ai_node) + else: + logger.error("llama node initial failed!") + await AIStorage.get_instance().set_feature_init_result("llama",False) - llama_ai_node = LocalLlama_ComputeNode() - await llama_ai_node.start() - # ComputeKernel.get_instance().add_compute_node(llama_ai_node) + if await AIStorage.get_instance().is_feature_enable("aigc"): + try: + google_text_to_speech_node = GoogleTextToSpeechNode.get_instance() + google_text_to_speech_node.init() + ComputeKernel.get_instance().add_compute_node(google_text_to_speech_node) + except Exception as e: + logger.error(f"google text to speech node initial failed! {e}") + await AIStorage.get_instance.set_feature_init_result("aigc",False) + + + local_sd_node = Local_Stability_ComputeNode.get_instance() + if await local_sd_node.initial() is True: + ComputeKernel.get_instance().add_compute_node(local_sd_node) + else: + logger.error("local stability node initial failed!") + await AIStorage.get_instance.set_feature_init_result("aigc",False) - local_sd_node = Local_Stability_ComputeNode.get_instance() - if await local_sd_node.initial() is not True: - logger.error("local stability node initial failed!") - - ComputeKernel.get_instance().add_compute_node(local_sd_node) await ComputeKernel.get_instance().start() AIBus().get_default_bus().register_unhandle_message_handler(self._handle_no_target_msg) AIBus().get_default_bus().register_message_handler(self.username,self._user_process_msg) + KnowledgePipline.get_instance().initial() TelegramTunnel.register_to_loader() EmailTunnel.register_to_loader() - - user_data_dir = AIStorage.get_instance().get_myai_dir() - contact_config_path =os.path.abspath(f"{user_data_dir}/contacts.toml") - cm = ContactManager.get_instance(contact_config_path) - cm.load_data() - + user_data_dir = str(AIStorage.get_instance().get_myai_dir()) tunnels_config_path = os.path.abspath(f"{user_data_dir}/etc/tunnels.cfg.toml") tunnel_config = None try: @@ -164,7 +179,7 @@ class AIOS_Shell: except Exception as e: logger.warning(f"load tunnels config from {tunnels_config_path} failed!") - KnowledgePipline.get_instance().initial() + return True @@ -392,6 +407,34 @@ class AIOS_Shell: AIStorage.get_instance().get_user_config().set_value("shell.current",f"{self.current_topic}@{self.current_target}") await AIStorage.get_instance().get_user_config().save_to_user_config() return show_text + case 'enable': + if len(args) >= 1: + feature = args[0] + else: + show_text = FormattedText([("class:error", "/enable Need Feature Name! like /enable llama")]) + return show_text + + if await AIStorage.get_instance().is_feature_enable(feature): + show_text = FormattedText([("class:title", f"Feature {feature} already enabled!")]) + return show_text + + AIStorage.get_instance().enable_feature(feature) + show_text = FormattedText([("class:title", f"Feature {feature} enabled!")]) + return show_text + case 'disable': + if len(args) >= 1: + feature = args[0] + else: + show_text = FormattedText([("class:error", "/disable Need Feature Name! like /disable llama")]) + return show_text + + if not await AIStorage.get_instance().is_feature_enable(feature): + show_text = FormattedText([("class:title", f"Feature {feature} already disabled!")]) + return show_text + + AIStorage.get_instance().disable_feature(feature) + show_text = FormattedText([("class:title", f"Feature {feature} disabled!")]) + return show_text case 'login': if len(args) >= 1: self.username = args[0] @@ -518,15 +561,25 @@ def print_welcome_screen(): async def main(): print_welcome_screen() print("Booting...") - logging.basicConfig(filename="aios_shell.log",filemode="w",encoding='utf-8',force=True, - level=logging.INFO, - format='[%(asctime)s]%(name)s[%(levelname)s]: %(message)s') if os.path.isdir(f"{directory}/../../../rootfs"): AIStorage.get_instance().is_dev_mode = True else: AIStorage.get_instance().is_dev_mode = False + + if AIStorage.get_instance().is_dev_mode: + logging.basicConfig(filename="aios_shell.log",filemode="w",encoding='utf-8',force=True, + level=logging.INFO, + format='[%(asctime)s]%(name)s[%(levelname)s]: %(message)s') + else: + log_file = f"{AIStorage.get_instance().get_myai_dir()}/logs/aios_shell.log" + handler = RotatingFileHandler(log_file, maxBytes=50*1024*1024, backupCount=100) + + logging.basicConfig(handlers=[handler], + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') + is_daemon = False if os.name != 'nt': if os.getppid() == 1: @@ -552,6 +605,7 @@ async def main(): return 1 else: print("aios shell initial failed!") + return 1 print(f"aios shell {shell.get_version()} ready.") if is_daemon: @@ -559,7 +613,6 @@ async def main(): proxy.apply_storage() - #TODO: read last input config completer = WordCompleter(['/send $target $msg $topic', '/open $target $topic', '/history $num $offset', @@ -570,6 +623,8 @@ async def main(): '/knowledge journal [$topn]', '/knowledge query $query' '/set_config $key', + '/enable $feature', + '/disable $feature', '/list_config', '/show', '/exit', diff --git a/src/service/aios_shell/proxy.py b/src/service/aios_shell/proxy.py index caf006d..f8b124c 100644 --- a/src/service/aios_shell/proxy.py +++ b/src/service/aios_shell/proxy.py @@ -4,12 +4,15 @@ import os import logging import socket import socks +import logging directory = os.path.dirname(__file__) sys.path.append(directory + '/../../') from aios_kernel import AIStorage +logger = logging.getLogger(__name__) + def apply_storage(): proxy_cfg = AIStorage.get_instance().get_user_config().get_config_item("proxy") if proxy_cfg is None: @@ -31,9 +34,10 @@ def apply_storage(): (host, port) = host.split(":") socks.set_default_proxy(socks.SOCKS5, host, int(port), username = username, password = password) socket.socket = socks.socksocket - print(f"proxy {host_url} will be used.") + logger.info(f"proxy {host_url} will be used.") case _: - print(f"the proxy type ({proxy_type}) has not support.") + logger.error(f"the proxy type ({proxy_type}) has not support. proxy will not be used.") + def declare_user_config(): user_config = AIStorage.get_instance().get_user_config()