From 22683d29ffae2981db63d7cbf811cc21b466633c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9D=92=E8=BD=A9?= Date: Tue, 13 Aug 2024 15:25:13 +0800 Subject: [PATCH] context passing --- memoryscope/core/config/arguments.py | 1 - memoryscope/core/memoryscope.py | 5 +++++ memoryscope/core/memoryscope_context.py | 14 +++++++++++++- memoryscope/core/models/base_model.py | 3 +++ memoryscope/core/utils/__init__.py | 6 ++++++ memoryscope/core/utils/logger.py | 11 ++++------- memoryscope/core/utils/singleton.py | 9 +++++++++ 7 files changed, 40 insertions(+), 9 deletions(-) create mode 100644 memoryscope/core/utils/singleton.py diff --git a/memoryscope/core/config/arguments.py b/memoryscope/core/config/arguments.py index 92ccf9e1..f409f411 100644 --- a/memoryscope/core/config/arguments.py +++ b/memoryscope/core/config/arguments.py @@ -1,7 +1,6 @@ from dataclasses import dataclass, field from typing import Literal, Dict - @dataclass class Arguments(object): language: Literal["cn", "en"] = field(default="en", metadata={"help": "support en & cn now"}) diff --git a/memoryscope/core/memoryscope.py b/memoryscope/core/memoryscope.py index 0e5219bd..d673d0ea 100644 --- a/memoryscope/core/memoryscope.py +++ b/memoryscope/core/memoryscope.py @@ -1,4 +1,5 @@ from concurrent.futures import ThreadPoolExecutor +from datetime import datetime from memoryscope.core.chat.base_memory_chat import BaseMemoryChat from memoryscope.core.config.config_manager import ConfigManager @@ -32,6 +33,10 @@ class MemoryScope(ConfigManager): self.logger.warning("If a semantic ranking model is not available, MemoryScope will use cosine similarity " "scoring as a substitute. However, the ranking effectiveness will be somewhat " "compromised.") + self.context.memory_scope_uuid = datetime.now().strftime(global_conf["logger_name_time_suffix"]) + + # set context_initialized + self.context.context_initialized = True # init memory_chat memory_chat_conf_dict = self.config["memory_chat"] diff --git a/memoryscope/core/memoryscope_context.py b/memoryscope/core/memoryscope_context.py index ac8d89fb..4de426b8 100644 --- a/memoryscope/core/memoryscope_context.py +++ b/memoryscope/core/memoryscope_context.py @@ -2,8 +2,9 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from memoryscope.enumeration.language_enum import LanguageEnum +from memoryscope.core.utils.singleton import singleton - +@singleton @dataclass class MemoryscopeContext(object): """ @@ -27,3 +28,14 @@ class MemoryscopeContext(object): worker_conf_dict: dict = field(default_factory=lambda: {}, metadata={"help": "name -> worker_conf"}) meta_data: dict = field(default_factory=lambda: {}) + + memory_scope_uuid: str = "" + + context_initialized: bool = False + +def get_ms_context(): + ms_context = MemoryscopeContext() + if ms_context.context_initialized: + return ms_context + else: + raise RuntimeError("MemoryscopeContext is not initialized yet. Please initialize it first.") diff --git a/memoryscope/core/models/base_model.py b/memoryscope/core/models/base_model.py index 69bb5c3f..4fc6befe 100644 --- a/memoryscope/core/models/base_model.py +++ b/memoryscope/core/models/base_model.py @@ -8,6 +8,8 @@ from memoryscope.core.utils.registry import Registry from memoryscope.core.utils.timer import Timer from memoryscope.enumeration.model_enum import ModelEnum from memoryscope.scheme.model_response import ModelResponse, ModelResponseGen +from memoryscope.core.memoryscope_context import MemoryscopeContext +from memoryscope.core.memoryscope_context import get_ms_context MODEL_REGISTRY = Registry("models") @@ -32,6 +34,7 @@ class BaseModel(metaclass=ABCMeta): self.retry_interval: float = retry_interval self.kwargs_filter: bool = kwargs_filter self.raise_exception: bool = raise_exception + self.context: MemoryscopeContext = get_ms_context() self.kwargs: dict = kwargs self._model: Any = None diff --git a/memoryscope/core/utils/__init__.py b/memoryscope/core/utils/__init__.py index 903b05a3..ab3e278c 100644 --- a/memoryscope/core/utils/__init__.py +++ b/memoryscope/core/utils/__init__.py @@ -15,6 +15,12 @@ from .tool_functions import ( cosine_similarity ) + +def get_context(): + from memoryscope import MemoryscopeContext + return MemoryscopeContext() + + __all__ = [ "DatetimeHandler", "Logger", diff --git a/memoryscope/core/utils/logger.py b/memoryscope/core/utils/logger.py index 92d814bd..2c9d64be 100644 --- a/memoryscope/core/utils/logger.py +++ b/memoryscope/core/utils/logger.py @@ -84,7 +84,6 @@ class Logger(logging.Logger): def format_chat_message(self, message): buf = '\n' - buf += f"++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n" buf += f"LM Input:\n" for chat_message in message.meta_data['data']['messages']: buf += chat_message.content @@ -93,13 +92,11 @@ class Logger(logging.Logger): buf += f"LM Output:\n" buf += message.message.content buf += '\n' - buf += f"++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n" buf += '\n' - return buf + return self.wrap_in_box(buf) def format_rank_message(self, model_response): buf = '\n' - buf += f"++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n" buf += f"Query Input:\n" buf += model_response.meta_data['data']['query_str'] buf += '\n' @@ -112,9 +109,8 @@ class Logger(logging.Logger): node_text = node.text buf += f"Score {score} | Rank {rank} | {node_text}\n" buf += '\n' - buf += f"++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n" buf += '\n' - return buf + return self.wrap_in_box(buf) def _add_file_handler(self): """ @@ -239,4 +235,5 @@ class Logger(logging.Logger): @staticmethod def append_timestamp(name: str) -> str: - return f"{name}_{datetime.now().strftime(r'%Y%m%d_%H%M')}" \ No newline at end of file + from memoryscope.core.memoryscope_context import get_ms_context + return f"{name}_{get_ms_context().memory_scope_uuid}" \ No newline at end of file diff --git a/memoryscope/core/utils/singleton.py b/memoryscope/core/utils/singleton.py new file mode 100644 index 00000000..b767cd1e --- /dev/null +++ b/memoryscope/core/utils/singleton.py @@ -0,0 +1,9 @@ +def singleton(cls): + _instance = {} + + def _singleton(*args, **kargs): + if cls not in _instance: + _instance[cls] = cls(*args, **kargs) + return _instance[cls] + + return _singleton \ No newline at end of file