From d815ba988ab0132b16c4d9c2efd39e3a3394980b Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 26 Jun 2024 14:01:48 +0800 Subject: [PATCH] [dev] add operation --- .../{workflow => operation}/__init__.py | 0 .../memory/operation/base_operation.py | 15 ++++ .../{workflow => operation}/base_workflow.py | 8 +-- .../memory/operation/read_operation.py | 23 ++++++ .../memory/operation/summary_operation.py | 37 ++++++++++ .../memory/operation/write_operation.py | 66 +++++++++++++++++ memory_scope/memory/worker/base_worker.py | 66 +++++++---------- memory_scope/memory/worker/dummy_worker.py | 6 +- .../memory/worker/memory_base_worker.py | 70 ------------------- .../memory/workflow/backend_v1_workflow.py | 43 ------------ .../memory/workflow/frontend_workflow.py | 12 ---- 11 files changed, 171 insertions(+), 175 deletions(-) rename memory_scope/memory/{workflow => operation}/__init__.py (100%) create mode 100644 memory_scope/memory/operation/base_operation.py rename memory_scope/memory/{workflow => operation}/base_workflow.py (93%) create mode 100644 memory_scope/memory/operation/read_operation.py create mode 100644 memory_scope/memory/operation/summary_operation.py create mode 100644 memory_scope/memory/operation/write_operation.py delete mode 100644 memory_scope/memory/worker/memory_base_worker.py delete mode 100644 memory_scope/memory/workflow/backend_v1_workflow.py delete mode 100644 memory_scope/memory/workflow/frontend_workflow.py diff --git a/memory_scope/memory/workflow/__init__.py b/memory_scope/memory/operation/__init__.py similarity index 100% rename from memory_scope/memory/workflow/__init__.py rename to memory_scope/memory/operation/__init__.py diff --git a/memory_scope/memory/operation/base_operation.py b/memory_scope/memory/operation/base_operation.py new file mode 100644 index 00000000..26a9b9ca --- /dev/null +++ b/memory_scope/memory/operation/base_operation.py @@ -0,0 +1,15 @@ +from abc import ABCMeta, abstractmethod +from typing import Literal + +OPERATION_TYPE = Literal["frontend", "backend"] + + +class BaseOperation(metaclass=ABCMeta): + operation_type: OPERATION_TYPE = "frontend" + + @abstractmethod + def run_operation(self): + raise NotImplementedError + + def run_operation_backend(self): + pass diff --git a/memory_scope/memory/workflow/base_workflow.py b/memory_scope/memory/operation/base_workflow.py similarity index 93% rename from memory_scope/memory/workflow/base_workflow.py rename to memory_scope/memory/operation/base_workflow.py index de79c569..8b5ac67c 100644 --- a/memory_scope/memory/workflow/base_workflow.py +++ b/memory_scope/memory/operation/base_workflow.py @@ -6,7 +6,6 @@ from typing import Dict, Any, List from memory_scope.chat_v2.global_context import G_CONTEXT from memory_scope.memory.worker.base_worker import BaseWorker -from memory_scope.scheme.message import Message from memory_scope.utils.logger import Logger from memory_scope.utils.timer import Timer from memory_scope.utils.tool_functions import init_instance_by_config @@ -18,15 +17,11 @@ class BaseWorkflow(object): name: str, workflow: str, thread_pool: ThreadPoolExecutor, - chat_messages: List[Message], - max_history_message_count: int, **kwargs): self.name: str = name self.workflow: str = workflow self.thread_pool: ThreadPoolExecutor = thread_pool - self.chat_messages: List[Message] = chat_messages - self.max_history_message_count: int = max_history_message_count self.kwargs = kwargs self.workflow_worker_list: List[List[List[str]]] = [] @@ -115,7 +110,8 @@ class BaseWorkflow(object): else: t_list = [] for sub_workflow in workflow_part: - t_list.append(G_CONTEXT.thread_pool.submit(self._run_sub_workflow, sub_workflow)) + t_list.append(G_CONTEXT.thread_pool.submit( + self._run_sub_workflow, sub_workflow)) flag = True for future in as_completed(t_list): diff --git a/memory_scope/memory/operation/read_operation.py b/memory_scope/memory/operation/read_operation.py new file mode 100644 index 00000000..9f8f294f --- /dev/null +++ b/memory_scope/memory/operation/read_operation.py @@ -0,0 +1,23 @@ +from typing import List + +from memory_scope.constants.common_constants import RESULT, CHAT_MESSAGES +from memory_scope.memory.operation.base_operation import BaseOperation, OPERATION_TYPE +from memory_scope.memory.operation.base_workflow import BaseWorkflow +from memory_scope.scheme.message import Message + + +class ReadOperation(BaseWorkflow, BaseOperation): + operation_type: OPERATION_TYPE = "frontend" + + def __init__(self, chat_messages: List[Message], max_his_msg_count: int = 0, **kwargs): + super().__init__(**kwargs) + self.chat_messages: List[Message] = chat_messages + self.max_his_msg_count: int = max_his_msg_count + + def run_operation(self): + max_count = 1 + self.max_his_msg_count + self.context[CHAT_MESSAGES] = [x.copy() for x in self.chat_messages[-max_count:]] + self.run_workflow() + result = self.context.get(RESULT) + self.context.clear() + return result diff --git a/memory_scope/memory/operation/summary_operation.py b/memory_scope/memory/operation/summary_operation.py new file mode 100644 index 00000000..68d98c41 --- /dev/null +++ b/memory_scope/memory/operation/summary_operation.py @@ -0,0 +1,37 @@ +import time + +from memory_scope.memory.base_workflow import BaseWorkflow + +from memory_scope.chat_v2.global_context import G_CONTEXT +from memory_scope.memory.operation.base_operation import BaseOperation, OPERATION_TYPE + + +class SummaryOperation(BaseWorkflow, BaseOperation): + operation_type: OPERATION_TYPE = "backend" + + def __init__(self, interval_time: int = 300, **kwargs): + super().__init__(**kwargs) + + self.interval_time: int = interval_time + + self._operation_status_run: bool = False + self._loop_switch: bool = False + + def run_operation(self): + if self._operation_status_run: + return + + self._operation_status_run = True + self.run_workflow() + self.context.clear() + self._operation_status_run = False + + def _loop_operation(self): + while self._loop_switch: + time.sleep(self.interval_time) + self.run_operation() + + def run_operation_backend(self): + if not self._loop_switch: + self._loop_switch = True + return G_CONTEXT.thread_pool.submit(self._loop_operation) diff --git a/memory_scope/memory/operation/write_operation.py b/memory_scope/memory/operation/write_operation.py new file mode 100644 index 00000000..efd5d4d1 --- /dev/null +++ b/memory_scope/memory/operation/write_operation.py @@ -0,0 +1,66 @@ +import time +from typing import List + +from memory_scope.chat_v2.global_context import G_CONTEXT +from memory_scope.constants.common_constants import CHAT_MESSAGES +from memory_scope.memory.operation.base_operation import BaseOperation, OPERATION_TYPE +from memory_scope.memory.operation.base_workflow import BaseWorkflow +from memory_scope.scheme.message import Message + + +class WriteOperation(BaseOperation, BaseWorkflow): + operation_type: OPERATION_TYPE = "backend" + + def __init__(self, + chat_messages: List[Message], + max_his_msg_count: int = 0, + message_lock=None, + interval_time: int = 60, + min_count: int = 5, + **kwargs): + super().__init__(**kwargs) + + self.chat_messages: List[Message] = chat_messages + self.max_his_msg_count: int = max_his_msg_count + self.message_lock = message_lock + self.interval_time: int = interval_time + self.min_count: int = min_count + + self._operation_status_run: bool = False + self._loop_switch: bool = False + + @property + def not_memorized_size(self): + return sum([not x.memorized for x in self.chat_messages]) + + def set_memorized(self): + if self.message_lock: + with self.message_lock: + for msg in self.chat_messages: + msg.memorized = True + + def run_operation(self): + if self._operation_status_run: + return + + self._operation_status_run = True + not_memorized_size = self.not_memorized_size + if not_memorized_size < self.min_count: + return + + max_count = not_memorized_size + self.max_his_msg_count + self.context[CHAT_MESSAGES] = [x.copy() for x in self.chat_messages[-max_count:]] + self.run_workflow() + self.context.clear() + self.set_memorized() + self._operation_status_run = False + + def _loop_operation(self): + while self._loop_switch: + time.sleep(self.interval_time) + self.run_operation() + + def run_operation_backend(self): + if not self._loop_switch: + self._loop_switch = True + return G_CONTEXT.thread_pool.submit(self._loop_operation) diff --git a/memory_scope/memory/worker/base_worker.py b/memory_scope/memory/worker/base_worker.py index a47fd1ed..78bfae5e 100644 --- a/memory_scope/memory/worker/base_worker.py +++ b/memory_scope/memory/worker/base_worker.py @@ -1,73 +1,57 @@ +from abc import ABCMeta, abstractmethod from typing import Any, Dict -from utils.logger import Logger -from utils.timer import Timer +from memory_scope.utils.logger import Logger +from memory_scope.utils.timer import Timer -class BaseWorker(object): +class BaseWorker(metaclass=ABCMeta): - def __init__(self, raise_exception: bool = True, **kwargs): - super(BaseWorker, self).__init__(**kwargs) - # 异常是否继续执行 + def __init__(self, + name: str, + context: Dict[str, Any], + context_lock=None, + raise_exception: bool = True, + is_multi_thread: bool = False, + **kwargs): + + self.name: str = name + self.context: Dict[str, Any] = context + self.context_lock = context_lock self.raise_exception: bool = raise_exception - - # True 为正常运行，False会结束整个pipeline - self.continue_run: bool = True - - # 短name - self._name_simple: str = "" - - # 是否多线程环境 - self.is_multi_thread: bool = False - - # pipeline 上下文 - self.context_dict: Dict[str, Any] | None = None - self.context_lock = None - - # 日志 - self.logger: Logger = Logger.get_logger() - - # worker 参数 + self.is_multi_thread: bool = is_multi_thread self.kwargs: dict = kwargs + self.continue_run: bool = True + self.logger: Logger = Logger.get_logger() + + @abstractmethod def _run(self): raise NotImplementedError def run(self): - self.logger.info(f"----- Begin {self.name_simple} -----") - with Timer(self.name_simple, log_time=False) as t: + self.logger.info(f"----- worker_{self.name}_begin -----") + with Timer(self.name, log_time=False) as t: if self.raise_exception: self._run() else: try: self._run() except Exception as e: - self.logger.exception(f"run {self.name_simple} failed! args={e.args}") + self.logger.exception(f"run {self.name} failed! args={e.args}") - self.logger.info(f"----- End {self.name_simple} cost={t.cost_str}-----") - - def set_context_dict(self, context_dict: dict, context_lock=None): - self.context_dict = context_dict - if context_lock is not None: - self.context_lock = context_lock - self.is_multi_thread = True + self.logger.info(f"----- worker_{self.name}_end cost={t.cost_str}-----") def get_context(self, key: str, default=None): return self.context_dict.get(key, default) def set_context(self, key: str, value: Any): if self.is_multi_thread: - # add lock to multi thread with self.context_lock: self.context_dict[key] = value else: self.context_dict[key] = value def __getattr__(self, key): + # raise exception if not exists return self.kwargs[key] - - @property - def name_simple(self) -> str: - if not self._name_simple: - self._name_simple = self.__class__.__name__.replace("Worker", "") - return self._name_simple diff --git a/memory_scope/memory/worker/dummy_worker.py b/memory_scope/memory/worker/dummy_worker.py index 87ba5dbe..6b446b0c 100644 --- a/memory_scope/memory/worker/dummy_worker.py +++ b/memory_scope/memory/worker/dummy_worker.py @@ -1,6 +1,6 @@ -from memory_base_worker import MemoryBaseWorker +from memory_scope.memory.worker.base_worker import BaseWorker -class DummyWorker(MemoryBaseWorker): +class DummyWorker(BaseWorker): def _run(self): - pass \ No newline at end of file + self.logger.info("enter dummy worker!") diff --git a/memory_scope/memory/worker/memory_base_worker.py b/memory_scope/memory/worker/memory_base_worker.py deleted file mode 100644 index e8d5cc0c..00000000 --- a/memory_scope/memory/worker/memory_base_worker.py +++ /dev/null @@ -1,70 +0,0 @@ -from typing import List - -from chat.global_context import GLOBAL_CONTEXT -from constants.common_constants import MESSAGES, CHAT_NAME -from models.base_model import BaseModel -from scheme.message import Message -from storage.base_monitor import BaseMonitor -from storage.base_vector_store import BaseVectorStore -from worker.base_worker import BaseWorker - - -class MemoryBaseWorker(BaseWorker): - def __init__(self, - embedding_model: str, - generation_model: str, - rank_model: str, - **kwargs): - super(MemoryBaseWorker, self).__init__(**kwargs) - self.embedding_model_name: str = embedding_model - self.generation_model_name: str = generation_model - self.rank_model_name: str = rank_model - - self._embedding_model: BaseModel | None = None - self._generation_model: BaseModel | None = None - self._rank_model: BaseModel | None = None - - self._vector_store: BaseVectorStore | None = None - self._monitor: BaseMonitor | None = None - - @property - def messages(self) -> List[Message]: - return self.get_context(MESSAGES) - - @messages.setter - def messages(self, value): - self.set_context(MESSAGES, value) - - @property - def chat_name(self): - return self.get_context(CHAT_NAME) - - @property - def embedding_model(self): - if self._embedding_model is None: - self._embedding_model = GLOBAL_CONTEXT.model_dict.get(self.embedding_model_name) - return self._embedding_model - - @property - def generation_model(self): - if self._generation_model is None: - self._generation_model = GLOBAL_CONTEXT.model_dict.get(self.generation_model_name) - return self._generation_model - - @property - def rank_model(self): - if self._rank_model is None: - self._rank_model = GLOBAL_CONTEXT.model_dict.get(self.rank_model_name) - return self._rank_model - - @property - def vector_store(self): - if self._vector_store is None: - self._vector_store = GLOBAL_CONTEXT.vector_store - return self._vector_store - - @property - def monitor(self): - if self._monitor is None: - self._monitor = GLOBAL_CONTEXT.monitor - return self._monitor diff --git a/memory_scope/memory/workflow/backend_v1_workflow.py b/memory_scope/memory/workflow/backend_v1_workflow.py deleted file mode 100644 index 8e04bdcd..00000000 --- a/memory_scope/memory/workflow/backend_v1_workflow.py +++ /dev/null @@ -1,43 +0,0 @@ -import time - -from memory_scope.constants.common_constants import RESULT, CHAT_MESSAGES -from memory_scope.memory.workflow.base_workflow import BaseWorkflow - - -class BackendV1Workflow(BaseWorkflow): - - def __init__(self, interval_time: int, min_count: int, **kwargs): - super().__init__(**kwargs) - self.interval_time: int = interval_time - self.min_count: int = min_count - - @property - def not_memorized_size(self): - return sum([not x.memorized for x in self.chat_messages]) - - def set_memorized(self): - for msg in self.chat_messages: - msg.memorized = True - - def _loop(self): - while self.loop_switch: - time.sleep(self.interval_time) - if self.not_memorized_size < self.min_count: - continue - - self.context[CHAT_MESSAGES] = self.chat_messages - self.__call__() - self.context.clear() - self.set_memorized() - - def start_loop_run(self): - if not self.loop_switch: - self.loop_switch = True - return GLOBAL_CONTEXT.thread_pool.submit(self._thread_loop) - - def run_workflow(self): - self.context[CHAT_MESSAGES] = self.chat_messages - self.__call__() - result = self.context.get(RESULT) - self.context.clear() - return result diff --git a/memory_scope/memory/workflow/frontend_workflow.py b/memory_scope/memory/workflow/frontend_workflow.py deleted file mode 100644 index f394f6f9..00000000 --- a/memory_scope/memory/workflow/frontend_workflow.py +++ /dev/null @@ -1,12 +0,0 @@ -from memory_scope.constants.common_constants import RESULT, CHAT_MESSAGES -from memory_scope.memory.workflow.base_workflow import BaseWorkflow - - -class FrontendWorkflow(BaseWorkflow): - - def run_workflow(self): - self.context[CHAT_MESSAGES] = self.chat_messages[:1 + self.max_history_message_count] - self.__call__() - result = self.context.get(RESULT) - self.context.clear() - return result