diff --git a/memory_scope/chat_v2/cli_memory_chat.py b/memory_scope/chat_v2/cli_memory_chat.py index 1cf6db9e..0b39c914 100644 --- a/memory_scope/chat_v2/cli_memory_chat.py +++ b/memory_scope/chat_v2/cli_memory_chat.py @@ -61,7 +61,7 @@ class CliMemoryChat(BaseMemoryChat): def run(self): - op_description_dict: Dict[str, str] = self.memory_service.get_op_description_dict() + op_description_dict: Dict[str, str] = self.memory_service.op_description_dict self.USER_COMMANDS.update({f"/{k}": v for k, v in op_description_dict.items()}) console = Console() @@ -92,14 +92,14 @@ def run(self): questionary.print(f" {desc}") elif query in op_description_dict: if not args: - result = self.memory_service.do_operation(op_name=query) + result = self.memory_service.operate(op_name=query) questionary.print(result) elif args[0].isdigit(): refresh_time = int(args[0]) while True: time.sleep(refresh_time) - result = self.memory_service.do_operation(op_name=query) + result = self.memory_service.operate(op_name=query) questionary.print(result) else: console.print("unknown command received. Please try again!") diff --git a/memory_scope/memory/service/base_memory_service.py b/memory_scope/memory/service/base_memory_service.py index 05ee4b54..84288b01 100644 --- a/memory_scope/memory/service/base_memory_service.py +++ b/memory_scope/memory/service/base_memory_service.py @@ -1,6 +1,8 @@ +import threading from abc import ABCMeta, abstractmethod from typing import List, Dict +from memory_scope.memory.operation.base_operation import BaseOperation from memory_scope.scheme.message import Message from memory_scope.utils.logger import Logger @@ -8,22 +10,32 @@ from memory_scope.utils.logger import Logger class BaseMemoryService(metaclass=ABCMeta): def __init__(self, read_memory_key: str = "read_memory", **kwargs): self.read_memory_key: str = read_memory_key + + self._operation_dict: Dict[str, BaseOperation] = {} + self._op_description_dict: Dict[str, str] = {} + + self.chat_messages: List[Message] = [] + self.message_lock = threading.Lock + self.logger = Logger.get_logger() self.kwargs = kwargs - def submit_messages(self, messages: List[Message] | Message): + def add_messages(self, messages: List[Message] | Message): pass def prepare_service(self): pass @abstractmethod - def do_operation(self, op_name: str): - pass + def operate(self, op_name: str): + raise NotImplementedError - @abstractmethod - def get_op_description_dict(self) -> Dict[str, str]: - pass + @property + def op_description_dict(self) -> Dict[str, str]: + 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 def read_memory(self): - return self.do_operation(self.read_memory_key) + assert self.read_memory_key in self._operation_dict, f"op={self.read_memory_key} is not inited!" + return self.operate(self.read_memory_key) diff --git a/memory_scope/memory/service/chat_memory_service.py b/memory_scope/memory/service/chat_memory_service.py index 5bb9895c..820d1420 100644 --- a/memory_scope/memory/service/chat_memory_service.py +++ b/memory_scope/memory/service/chat_memory_service.py @@ -1,7 +1,5 @@ -import threading from typing import List, Dict -from memory_scope.memory.operation.base_operation import BaseOperation from memory_scope.memory.service.base_memory_service import BaseMemoryService from memory_scope.scheme.message import Message from memory_scope.utils.tool_functions import init_instance_by_config @@ -15,29 +13,24 @@ class ChatMemoryService(BaseMemoryService): contextual_msg_count: int = 6, **kwargs): super().__init__(**kwargs) - - self.op_dict: Dict[str, BaseOperation] = self._init_operation(memory_operations) self.history_msg_count: int = history_msg_count self.contextual_msg_count: int = contextual_msg_count assert self.history_msg_count >= self.contextual_msg_count - self.chat_messages: List[Message] = [] - self.message_lock = threading.Lock + self._init_operation(memory_operations) def _init_operation(self, memory_operations: Dict[str, dict]): - op_dict: Dict[str, BaseOperation] = {} for name, operation_config in memory_operations.items(): - if name in self.op_dict: + if name in self._operation_dict: self.logger.warning(f"memory operation={name} is repeated!") continue - self.op_dict[name] = init_instance_by_config(config=operation_config, - name=name, - chat_messages=self.chat_messages, - message_lock=self.message_lock, - contextual_msg_count=self.contextual_msg_count) - return op_dict + self._operation_dict[name] = init_instance_by_config(config=operation_config, + name=name, + chat_messages=self.chat_messages, + message_lock=self.message_lock, + contextual_msg_count=self.contextual_msg_count) - def submit_messages(self, messages: List[Message] | Message): + def add_messages(self, messages: List[Message] | Message): if isinstance(messages, Message): messages = [messages] @@ -49,16 +42,13 @@ class ChatMemoryService(BaseMemoryService): self.chat_messages.pop(0) def prepare_service(self): - for _, operation in self.op_dict.items(): + for _, operation in self._operation_dict.items(): operation.init_workflow() if operation.operation_type == "backend": operation.run_operation_backend() - def do_operation(self, op_name: str): - if op_name not in self.op_dict: + def operate(self, op_name: str): + if op_name not in self._operation_dict: self.logger.warning(f"op_name={op_name} is not inited!") return - return self.op_dict[op_name].run_operation() - - def get_op_description_dict(self) -> Dict[str, str]: - return {k: v.description for k, v in self.op_dict.items()} + return self._operation_dict[op_name].run_operation()