[dev] add operation

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

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):

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

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)

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)

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

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
+        self.logger.info("enter dummy worker!")

memory_scope/memory/worker/memory_base_worker.py

@@ -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

memory_scope/memory/workflow/backend_v1_workflow.py

@@ -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

memory_scope/memory/workflow/frontend_workflow.py

@@ -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