diff --git a/memoryscope/core/utils/logger.py b/memoryscope/core/utils/logger.py index 6a533f26..92d814bd 100644 --- a/memoryscope/core/utils/logger.py +++ b/memoryscope/core/utils/logger.py @@ -3,12 +3,20 @@ from datetime import datetime from logging.handlers import RotatingFileHandler from pathlib import Path from rich.console import Console +from rich.panel import Panel LOG_FORMAT = "%(asctime)s %(levelname)s %(threadName)s %(module)s:%(lineno)d] %(message)s" DATE_FORMAT = "%Y-%m-%d %H:%M:%S" LOGGER_DICT = {} +def rich2text(rich_table): + from rich.text import Text + console = Console(width=150) + with console.capture() as capture: + console.print(rich_table) + return '\n' + str(Text.from_ansi(capture.get())) + class Logger(logging.Logger): """ @@ -66,34 +74,13 @@ class Logger(logging.Logger): self.info(f"logger={name} is inited.") # Logs an initialization message def format_current_context(self, context): - from rich.panel import Panel - from rich.text import Text import pprint pp = pprint.PrettyPrinter() pretty_string = pp.pformat(context) - - def rich2text(rich_table): - console = Console(width=150) - with console.capture() as capture: - console.print(rich_table) - return '\n' + str(Text.from_ansi(capture.get())) - return rich2text(Panel(pretty_string, width=128)) - def format_current_memory(self, context): - from rich.panel import Panel - from rich.text import Text - import pprint - pp = pprint.PrettyPrinter() - pretty_string = pp.pformat(context) - - def rich2text(rich_table): - console = Console(width=150) - with console.capture() as capture: - console.print(rich_table) - return '\n' + str(Text.from_ansi(capture.get())) - - return rich2text(Panel(pretty_string, width=128)) + def wrap_in_box(self, context): + return rich2text(Panel(context, width=128)) def format_chat_message(self, message): buf = '\n' diff --git a/memoryscope/core/worker/memory_base_worker.py b/memoryscope/core/worker/memory_base_worker.py index 3d8fc0d7..53a72642 100644 --- a/memoryscope/core/worker/memory_base_worker.py +++ b/memoryscope/core/worker/memory_base_worker.py @@ -204,7 +204,7 @@ class MemoryBaseWorker(BaseWorker, metaclass=ABCMeta): MemoryHandler: An instance of MemoryHandler. """ if not self.has_content(MEMORY_MANAGER): - self.set_context(MEMORY_MANAGER, MemoryManager(self.memoryscope_context)) + self.set_context(MEMORY_MANAGER, MemoryManager(self.memoryscope_context, worker_name=self.name)) return self.get_context(MEMORY_MANAGER) def get_language_value(self, languages: dict | List[dict]) -> Any | List[Any]: diff --git a/memoryscope/core/worker/memory_manager.py b/memoryscope/core/worker/memory_manager.py index 7d577993..601b567b 100644 --- a/memoryscope/core/worker/memory_manager.py +++ b/memoryscope/core/worker/memory_manager.py @@ -13,7 +13,7 @@ class MemoryManager(object): The `MemoryHandler` class manages memory nodes with memory store. """ - def __init__(self, memoryscope_context: MemoryscopeContext): + def __init__(self, memoryscope_context: MemoryscopeContext, worker_name: str ="default_worker"): self.memoryscope_context: MemoryscopeContext = memoryscope_context self._memory_store: BaseMemoryStore | None = None @@ -26,6 +26,9 @@ class MemoryManager(object): self.logger = Logger.get_logger(Logger.append_timestamp("memory_manager")) + self.worker_name = worker_name + + @property def memory_store(self) -> BaseMemoryStore: """ @@ -95,10 +98,13 @@ class MemoryManager(object): self.logger.info(f"add to memory context memory id={node.memory_id} content={node.content} " f"store_status={node.store_status} action_status={node.action_status}") - self.logger.info(self.logger.format_current_memory( - [f"{node.memory_type} | {node.content}" for node in nodes] - )) - + if nodes: + self.logger.info( + self.logger.wrap_in_box( + '\n'.join([f"worker_name: {self.worker_name} | memory_type:{node.memory_type} | content:{node.content}" for node in nodes]) + ) + ) + def get_memories(self, keys: str | List[str]) -> List[MemoryNode]: """