From b1d11a314a7b09c818002170d41d21305b2bd948 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Sat, 27 Jul 2024 00:43:47 +0800 Subject: [PATCH] [dev] remove useless file InitializationHandler --- memoryscope/argument/init_handler.py | 28 -------- memoryscope/argument/memoryscope_arguments.py | 4 +- memoryscope/chat/api_memory_chat.py | 4 +- memoryscope/chat/base_memory_chat.py | 4 +- memoryscope/chat/cli_memory_chat.py | 8 ++- .../memory/operation/backend_operation.py | 11 ++- .../memory/operation/base_operation.py | 2 +- memoryscope/memory/operation/base_workflow.py | 14 ++-- .../memory/service/base_memory_service.py | 26 ++++--- .../memory/service/memory_scope_service.py | 72 ++++++++----------- memoryscope/memoryscope.py | 2 +- memoryscope/memoryscope_context.py | 2 + memoryscope/storage/dummy_memory_store.py | 22 +++--- 13 files changed, 86 insertions(+), 113 deletions(-) delete mode 100644 memoryscope/argument/init_handler.py diff --git a/memoryscope/argument/init_handler.py b/memoryscope/argument/init_handler.py deleted file mode 100644 index 86a2d840..00000000 --- a/memoryscope/argument/init_handler.py +++ /dev/null @@ -1,28 +0,0 @@ -class InitializationHandler(object): - - def __init__(self): - self.file_path: str = __file__ - - self.global_config_dict: dict = {} - - self.memory_chat_dict: dict = {} - - self.memory_service_dict: dict = {} - - self.worker_dict: dict = {} - - self.model_dict: dict = {} - - self.memory_store: dict = {} - - self.monitor: dict = {} - - def update_by_arguments(self): - pass - - def load_from_config(self): - pass - - - def load_from_file(self): - pass diff --git a/memoryscope/argument/memoryscope_arguments.py b/memoryscope/argument/memoryscope_arguments.py index fb10cd23..92572f25 100644 --- a/memoryscope/argument/memoryscope_arguments.py +++ b/memoryscope/argument/memoryscope_arguments.py @@ -55,7 +55,7 @@ class MemoryscopeArguments(object): es_url: str = field(default="http://localhost:9200") - # TODO at xianzhe - retrieve_type: str = field(default="dense", metadata={"help": "es_retrieve_type: dense, sparse, hybrid"}) + retrieve_mode: str = field(default="dense", metadata={ + "help": "retrieve_mode: dense, sparse(not implemented), hybrid(not implemented)"}) hybrid_alpha: float | None = field(default=1.0, metadata={"help": ""}) diff --git a/memoryscope/chat/api_memory_chat.py b/memoryscope/chat/api_memory_chat.py index da1d73de..10bc741d 100644 --- a/memoryscope/chat/api_memory_chat.py +++ b/memoryscope/chat/api_memory_chat.py @@ -31,10 +31,12 @@ class ApiMemoryChat(BaseMemoryChat): self.human_name: str = human_name if not self.human_name: self.human_name = DEFAULT_HUMAN_NAME[self.context.language] + self.context.meta_data["human_name"] = self.human_name self.assistant_name: str = assistant_name if not self.assistant_name: self.assistant_name = "AI" + self.context.meta_data["assistant_name"] = self.assistant_name self._prompt_handler: PromptHandler | None = None @@ -73,7 +75,7 @@ class ApiMemoryChat(BaseMemoryChat): self._memory_service: BaseMemoryService = self.context.memory_service_dict[self._memory_service] # init service & update kwargs - self._memory_service.init_service(human_name=self.human_name, assistant_name=self.assistant_name) + self._memory_service.init_service() return self._memory_service @property diff --git a/memoryscope/chat/base_memory_chat.py b/memoryscope/chat/base_memory_chat.py index dab56a0a..ea81c4fb 100644 --- a/memoryscope/chat/base_memory_chat.py +++ b/memoryscope/chat/base_memory_chat.py @@ -56,8 +56,8 @@ class BaseMemoryChat(metaclass=ABCMeta): def start_backend_service(self): self.memory_service.start_backend_service() - def do_memory_operation(self, op_name: str, **kwargs): - return self.memory_service.do_operation(op_name=op_name, **kwargs) + def do_memory_operation(self, operation_name: str, **kwargs): + return self.memory_service.do_operation(name=operation_name, **kwargs) def run(self): """ diff --git a/memoryscope/chat/cli_memory_chat.py b/memoryscope/chat/cli_memory_chat.py index 52c01d12..27a81485 100644 --- a/memoryscope/chat/cli_memory_chat.py +++ b/memoryscope/chat/cli_memory_chat.py @@ -46,10 +46,12 @@ class CliMemoryChat(BaseMemoryChat): self.human_name: str = human_name if not self.human_name: self.human_name = DEFAULT_HUMAN_NAME[self.context.language] + self.context.meta_data["human_name"] = self.human_name self.assistant_name: str = assistant_name if not self.assistant_name: self.assistant_name = "AI" + self.context.meta_data["assistant_name"] = self.assistant_name self._logo = char_logo("MemoryScope") self._prompt_handler: PromptHandler | None = None @@ -99,7 +101,7 @@ class CliMemoryChat(BaseMemoryChat): self._memory_service: BaseMemoryService = self.context.memory_service_dict[self._memory_service] # init service & update kwargs - self._memory_service.init_service(human_name=self.human_name, assistant_name=self.assistant_name) + self._memory_service.init_service() return self._memory_service @property @@ -257,7 +259,7 @@ class CliMemoryChat(BaseMemoryChat): refresh_time = int(refresh_time) self.memory_service.stop_backend_service() while True: - result = self.memory_service.do_operation(op_name=command, **kwargs) + result = self.memory_service.do_operation(name=command, **kwargs) os.system("clear") self.print_logo() if result: @@ -269,7 +271,7 @@ class CliMemoryChat(BaseMemoryChat): time.sleep(refresh_time) else: - result = self.memory_service.do_operation(op_name=command, **kwargs) + result = self.memory_service.do_operation(name=command, **kwargs) if result: if isinstance(result, list): result = "\n".join([str(x) for x in result]) diff --git a/memoryscope/memory/operation/backend_operation.py b/memoryscope/memory/operation/backend_operation.py index c52cf936..bd5d3ac9 100644 --- a/memoryscope/memory/operation/backend_operation.py +++ b/memoryscope/memory/operation/backend_operation.py @@ -115,11 +115,18 @@ class BackendOperation(BaseWorkflow, BaseOperation): if not self._loop_switch: self._loop_switch = True self._backend_task = G_CONTEXT.thread_pool.submit(self._loop_operation) + self.logger.info(f"start operation={operation.name}...") def stop_operation_backend(self, wait_task_end: bool = False): """ Stops the background operation loop by setting the _loop_switch to False. """ self._loop_switch = False - if wait_task_end and self._backend_task: - self._backend_task.result() + if self._backend_task: + if wait_task_end: + self._backend_task.result() + self.logger.info(f"stop operation={self.name}...") + else: + self.logger.info(f"send stop signal to operation={self.name}...") + + diff --git a/memoryscope/memory/operation/base_operation.py b/memoryscope/memory/operation/base_operation.py index 8904530c..df27ea2f 100644 --- a/memoryscope/memory/operation/base_operation.py +++ b/memoryscope/memory/operation/base_operation.py @@ -58,7 +58,7 @@ class BaseOperation(metaclass=ABCMeta): """ pass - def stop_operation_backend(self): + def stop_operation_backend(self, wait_task_end: bool = False): """ Placeholder method to stop any ongoing backend operations. Should be implemented in subclasses where backend operations are managed. diff --git a/memoryscope/memory/operation/base_workflow.py b/memoryscope/memory/operation/base_workflow.py index 926b59a1..1ec0a5f5 100644 --- a/memoryscope/memory/operation/base_workflow.py +++ b/memoryscope/memory/operation/base_workflow.py @@ -6,7 +6,7 @@ from typing import Dict, Any, List from memoryscope.constants.common_constants import WORKFLOW_NAME from memoryscope.memory.worker.base_worker import BaseWorker -from memoryscope.utils.global_context import G_CONTEXT +from memoryscope.memoryscope_context import MemoryscopeContext from memoryscope.utils.logger import Logger from memoryscope.utils.timer import Timer from memoryscope.utils.tool_functions import init_instance_by_config @@ -16,13 +16,13 @@ class BaseWorkflow(object): def __init__(self, name: str, + memoryscope_context: MemoryscopeContext, workflow: str = "", - thread_pool: ThreadPoolExecutor = G_CONTEXT.thread_pool, **kwargs): self.name: str = name + self.memoryscope_context: MemoryscopeContext = memoryscope_context self.workflow: str = workflow - self.thread_pool: ThreadPoolExecutor = thread_pool self.kwargs = kwargs self.workflow_worker_list: List[List[List[str]]] = [] @@ -128,17 +128,17 @@ class BaseWorkflow(object): This method modifies `self.worker_dict` in-place, replacing the keys with actual worker instances. """ for name in list(self.worker_dict.keys()): - if name not in G_CONTEXT.worker_config: - raise RuntimeError(f"worker={name} is not exists in worker_config!") + if name not in self.memoryscope_context.worker_conf_dict: + raise RuntimeError(f"worker={name} is not exists in worker config!") self.worker_dict[name] = init_instance_by_config( - config=G_CONTEXT.worker_config[name], + config=self.memoryscope_context.worker_conf_dict[name], suffix_name="worker", name=name, is_multi_thread=is_backend or self.worker_dict[name], context=self.context, context_lock=self.context_lock, - thread_pool=G_CONTEXT.thread_pool, + thread_pool=self.memoryscope_context.thread_pool, **kwargs) def _run_sub_workflow(self, worker_list: List[str]) -> bool: diff --git a/memoryscope/memory/service/base_memory_service.py b/memoryscope/memory/service/base_memory_service.py index c8d964c0..7b9fdee5 100644 --- a/memoryscope/memory/service/base_memory_service.py +++ b/memoryscope/memory/service/base_memory_service.py @@ -28,26 +28,25 @@ class BaseMemoryService(metaclass=ABCMeta): self.kwargs = kwargs self._operation_dict: Dict[str, BaseOperation] = {} - self._op_description_dict: Dict[str, str] = {} self.logger = Logger.get_logger() @property def op_description_dict(self) -> Dict[str, str]: """ Property to retrieve a dictionary mapping operation keys to their descriptions. - Lazily initializes the dictionary on first access. - Returns: Dict[str, str]: A dictionary where keys are operation identifiers and values are their descriptions. """ - if not self._op_description_dict: - self._op_description_dict = {k: v.description for k, v in self._operation_dict.items()} - return self._op_description_dict + return {k: v.description for k, v in self._operation_dict.items()} @abstractmethod def add_messages(self, messages: List[Message] | Message): raise NotImplementedError + @abstractmethod + def register_operation(self, name: str, operation_config: dict, **kwargs): + raise NotImplementedError + @abstractmethod def init_service(self, **kwargs): raise NotImplementedError @@ -58,12 +57,12 @@ class BaseMemoryService(metaclass=ABCMeta): def stop_backend_service(self): pass - def do_operation(self, op_name: str, **kwargs): + def do_operation(self, name: str, **kwargs): """ Executes a specific operation by its name with provided keyword arguments. Args: - op_name (str): The name of the operation to execute. + name (str): The name of the operation to execute. **kwargs: Keyword arguments for the operation's execution. Returns: @@ -72,12 +71,11 @@ class BaseMemoryService(metaclass=ABCMeta): Raises: Warning: If the operation name is not initialized in `_operation_dict`. """ - if op_name not in self._operation_dict: - self.logger.warning(f"op_name={op_name} is not inited!") + if name not in self._operation_dict: + self.logger.warning(f"operation={name} is not registered!") return - return self._operation_dict[op_name].run_operation(**kwargs) + return self._operation_dict[name].run_operation(**kwargs) def __getattr__(self, name: str): - return lambda **kwargs: self.do_operation(name, **kwargs) - - + assert name in self._operation_dict, f"operation={name} is not registered!" + return lambda **kwargs: self.do_operation(name=name, **kwargs) diff --git a/memoryscope/memory/service/memory_scope_service.py b/memoryscope/memory/service/memory_scope_service.py index fdcfbc30..eb9d14e8 100644 --- a/memoryscope/memory/service/memory_scope_service.py +++ b/memoryscope/memory/service/memory_scope_service.py @@ -1,3 +1,4 @@ +import threading from typing import List from memoryscope.memory.operation.base_operation import BaseOperation @@ -11,6 +12,8 @@ class MemoryScopeService(BaseMemoryService): history_msg_count: int = 100, contextual_msg_max_count: int = 20, contextual_msg_min_count: int = 0, + human_name: str = None, + assistant_name: str = None, **kwargs): """ init function. @@ -20,13 +23,19 @@ class MemoryScopeService(BaseMemoryService): it will not be included in the context to prevent token overflow. contextual_msg_min_count (int): The minimum context length in a conversation. If it is shorter than this length, no conversation summary will be made and no long-term memory will be generated. - kwargs (dict): other kwargs + human_name (str): human name. + assistant_name (str): assistant name. + kwargs (dict): other kwargs. """ super().__init__(**kwargs) self.history_msg_count: int = history_msg_count self.contextual_msg_max_count: int = contextual_msg_max_count self.contextual_msg_min_count: int = contextual_msg_min_count assert history_msg_count >= contextual_msg_max_count >= contextual_msg_min_count + if human_name: + self.context.meta_data["human_name"] = human_name + if assistant_name: + self.context.meta_data["assistant_name"] = assistant_name self.chat_messages: List[Message] = [] self.message_lock = threading.Lock() @@ -57,43 +66,28 @@ class MemoryScopeService(BaseMemoryService): for _ in range(gap_size): self.chat_messages.pop(0) - def do_operation(self, op_name: str, **kwargs): - """ - Executes a specific operation by its name with provided keyword arguments. - - Args: - op_name (str): The name of the operation to execute. - **kwargs: Keyword arguments for the operation's execution. - - Returns: - The result of the operation execution, if any. Otherwise, None. - - Raises: - Warning: If the operation name is not initialized in `_operation_dict`. - """ - if op_name not in self._operation_dict: - self.logger.warning(f"op_name={op_name} is not inited!") # Warn if operation not initialized + def register_operation(self, name: str, operation_config: dict, **kwargs): + if name in self._operation_dict: + self.logger.warning(f"op_name={name} is registered before!") return - return self._operation_dict[op_name].run_operation(**kwargs) # Execute the operation + + operation: BaseOperation = init_instance_by_config( + config=operation_config, + name=name, + chat_messages=self.chat_messages, + message_lock=self.message_lock, + context=self.context, + contextual_msg_max_count=self.contextual_msg_max_count, + contextual_msg_min_count=self.contextual_msg_min_count) + + # Initialize workflow for each operation + operation.init_workflow(**kwargs) + self._operation_dict[name] = operation + self.logger.info(f"service={self.__class__.__name__} init operation={name}") def init_service(self, **kwargs): - for name, operation_config in self.memory_operations.items(): - if name in self._operation_dict: - self.logger.warning(f"memory operation={name} is repeated!") - continue - - # ⭐ Initialize operation instance by its config - operation: BaseOperation = init_instance_by_config( - config=operation_config, - name=name, - chat_messages=self.chat_messages, - message_lock=self.message_lock, - contextual_msg_max_count=self.contextual_msg_max_count, - contextual_msg_min_count=self.contextual_msg_min_count) - operation.init_workflow(**kwargs) # Initialize workflow for each operation - - self._operation_dict[name] = operation - self.logger.info(f"service={self.__class__.__name__} init operation={name}") + for name, operation_config in self.memory_operations_conf.items(): + self.register_operation(name, operation_config, **kwargs) def start_backend_service(self): """ @@ -101,16 +95,12 @@ class MemoryScopeService(BaseMemoryService): """ for _, operation in self._operation_dict.items(): if operation.operation_type == "backend": - # Run backend operations operation.run_operation_backend() - self.logger.info(f"start operation={operation.name}...") - def stop_backend_service(self): + def stop_backend_service(self, wait_service_end: bool = False): """ Stops all backend operations that are currently running. """ for _, operation in self._operation_dict.items(): if operation.operation_type == "backend": - # Stop backend operations - operation.stop_operation_backend() - self.logger.info(f"stop operation={operation.name}...") + operation.stop_operation_backend(wait_task_end=wait_service_end) diff --git a/memoryscope/memoryscope.py b/memoryscope/memoryscope.py index 520bd127..0648871c 100644 --- a/memoryscope/memoryscope.py +++ b/memoryscope/memoryscope.py @@ -111,7 +111,7 @@ class MemoryScope(object): "embedding_model": "embedding_model", "index_name": arguments.es_index_name, "es_url": arguments.es_url, - "retrieve_type": arguments.retrieve_type, + "retrieve_mode": arguments.retrieve_mode, "hybrid_alpha": arguments.hybrid_alpha, } diff --git a/memoryscope/memoryscope_context.py b/memoryscope/memoryscope_context.py index 967bb77d..ac8d89fb 100644 --- a/memoryscope/memoryscope_context.py +++ b/memoryscope/memoryscope_context.py @@ -25,3 +25,5 @@ class MemoryscopeContext(object): model_dict: dict = field(default_factory=lambda: {}, metadata={"help": "name -> model"}) worker_conf_dict: dict = field(default_factory=lambda: {}, metadata={"help": "name -> worker_conf"}) + + meta_data: dict = field(default_factory=lambda: {}) diff --git a/memoryscope/storage/dummy_memory_store.py b/memoryscope/storage/dummy_memory_store.py index 304818bb..7c77f670 100644 --- a/memoryscope/storage/dummy_memory_store.py +++ b/memoryscope/storage/dummy_memory_store.py @@ -12,6 +12,17 @@ class DummyMemoryStore(BaseMemoryStore): semantic retrieval. Actual storage operations are not implemented. """ + def __init__(self, embedding_model: BaseModel, **kwargs): + """ + Initializes the DummyMemoryStore with an embedding model and additional keyword arguments. + + Args: + embedding_model (BaseModel): The model used to embed data for potential similarity-based retrieval. + **kwargs: Additional keyword arguments for configuration or future expansion. + """ + self.embedding_model: BaseModel = embedding_model + self.kwargs = kwargs + def retrieve_memories(self, query: str = "", top_k: int = 3, @@ -24,17 +35,6 @@ class DummyMemoryStore(BaseMemoryStore): filter_dict: Dict[str, List[str]] = None) -> List[MemoryNode]: pass - def __init__(self, embedding_model: BaseModel, **kwargs): - """ - Initializes the DummyMemoryStore with an embedding model and additional keyword arguments. - - Args: - embedding_model (BaseModel): The model used to embed data for potential similarity-based retrieval. - **kwargs: Additional keyword arguments for configuration or future expansion. - """ - self.embedding_model: BaseModel = embedding_model - self.kwargs = kwargs - def batch_insert(self, nodes: List[MemoryNode]): pass