diff --git a/memory_scope/chat/base_memory_chat.py b/memory_scope/chat/base_memory_chat.py index f15386a3..e15d4d81 100644 --- a/memory_scope/chat/base_memory_chat.py +++ b/memory_scope/chat/base_memory_chat.py @@ -1,20 +1,41 @@ from abc import ABCMeta, abstractmethod - from memory_scope.memory.service.base_memory_service import BaseMemoryService class BaseMemoryChat(metaclass=ABCMeta): + """ + An abstract base class representing a chat system integrated with memory services. + It outlines the method to initiate a chat session leveraging memory data, which concrete subclasses must implement. + """ @abstractmethod def chat_with_memory(self, query: str): """ - :param query: - :return: + Initiates a chat interaction using the memory service, with the provided query as input. + + Args: + query (str): The user's query or message to start the chat. + + Returns: + This method should return the chat response generated after processing the query + with the associated memory context. The actual return type and content are defined by the implementing subclass. """ @property def memory_service(self) -> BaseMemoryService: + """ + Abstract property to access the memory service. + + Raises: + NotImplementedError: This method should be implemented in a subclass. + """ raise NotImplementedError def run(self): + """ + Abstract method to run the chat system. + + This method should contain the logic to initiate and manage the chat process, + utilizing the memory service as needed. It must be implemented by subclasses. + """ pass diff --git a/memory_scope/chat/cli_memory_chat.py b/memory_scope/chat/cli_memory_chat.py index 646fcc64..b3f3dbf2 100644 --- a/memory_scope/chat/cli_memory_chat.py +++ b/memory_scope/chat/cli_memory_chat.py @@ -18,11 +18,15 @@ from memory_scope.utils.tool_functions import char_logo class CliMemoryChat(BaseMemoryChat): + """ + Command-line interface for chatting with an AI that integrates memory functionality. + Allows users to interact, manage chat history, adjust streaming settings, and view commands' help. + """ USER_COMMANDS = { - "exit": "exit the CLI", - "clear": "clear commands", - "help": "get cli commands help", - "stream": "get stream response" + "exit": "Exit the CLI.", + "clear": "Clear the command history.", + "help": "Display available CLI commands and their descriptions.", + "stream": "Toggle between getting streamed responses from the model." } def __init__(self, @@ -32,7 +36,21 @@ class CliMemoryChat(BaseMemoryChat): human_name: str = DEFAULT_HUMAN_NAME[G_CONTEXT.language], assistant_name: str = "AI", **kwargs): + """ + Initializes the CLI chat instance with specified services, models, and personalized settings. + Args: + memory_service (str | BaseMemoryService): The memory service to be used for storing conversation history. + generation_model (str | BaseModel): The model responsible for generating AI responses. + stream (bool, optional): Flag indicating whether responses should be streamed. Defaults to True. + human_name (str, optional): The name assigned to the human user. Defaults to a language-specific default. + assistant_name (str, optional): The name of the AI assistant. Defaults to "AI". + **kwargs: Additional keyword arguments for flexibility or future extensions. + + Side Effects: + - Updates global context with human and AI names. + - Initializes logging for the instance. + """ self._memory_service: BaseMemoryService | str = memory_service self._generation_model: BaseModel | str = generation_model self.stream: bool = stream @@ -51,25 +69,61 @@ class CliMemoryChat(BaseMemoryChat): @property def prompt_handler(self) -> PromptHandler: + """ + Lazy initialization property for the prompt handler. + + This property ensures that the `_prompt_handler` attribute is only instantiated when it is first accessed. + It uses the current file's path and additional keyword arguments for configuration. + + Returns: + PromptHandler: An instance of the PromptHandler configured for this CLI session. + """ if self._prompt_handler is None: self._prompt_handler = PromptHandler(__file__, **self.kwargs) return self._prompt_handler def print_logo(self): + """ + Prints the logo of the CLI application to the console. + + The logo is composed of multiple lines, which are iterated through + and printed one by one to provide a visual identity for the chat interface. + """ for line in self._logo: print(line) @property def memory_service(self) -> BaseMemoryService: + """ + Property to access the memory service. If the service is initially set as a string, + it will be looked up in the global context's memory service dictionary, initialized, + and then returned as an instance of `BaseMemoryService`. Ensures the memory service + is properly started before use. + + Returns: + BaseMemoryService: The active memory service instance. + + Raises: + ValueError: If the memory service string reference is not found in the global context's dictionary. + """ if isinstance(self._memory_service, str): if self._memory_service not in G_CONTEXT.memory_service_dict: raise ValueError("Missing declaration of memory_service in yaml configuration: " + self._memory_service) self._memory_service = G_CONTEXT.memory_service_dict[self._memory_service] - self._memory_service.start_service() + self._memory_service.start_service() # ⭐ Initialize and start the memory service return self._memory_service @property def generation_model(self) -> BaseModel: + """ + Property to get the generation model. If the model is set as a string, it will be resolved from the global context's model dictionary. + + Raises: + ValueError: If the model string is not found in the global context's model dictionary. + + Returns: + BaseModel: The actual generation model instance. + """ if isinstance(self._generation_model, str): if self._generation_model not in G_CONTEXT.model_dict: raise ValueError(f"Missing declaration of generation model in yaml config: {self._generation_model}") @@ -77,12 +131,30 @@ class CliMemoryChat(BaseMemoryChat): return self._generation_model def chat_with_memory(self, query: str, remember_response: bool = False) -> ModelResponse | ModelResponseGen: + """ + Engages in a conversation with the AI model, utilizing conversation memory. + The function sends the user's query, incorporates conversation history and memory, + and optionally remembers the AI's response based on the user's preference. + + Args: + query (str): The user's input or query for the AI. + remember_response (bool, optional): Flag indicating whether to save the AI's response to memory. + Defaults to False. + + Returns: + - ModelResponse: In non-streaming mode, returns the complete AI response. + - ModelResponseGen: In streaming mode, returns a generator yielding AI response parts. + + Side Effects: + - Updates the conversation memory with the user's query and (optionally) the AI's response. + - Retrieves and includes historical messages and memory content in the conversation context. + """ new_message: Message = Message(role=MessageRoleEnum.USER.value, role_name=self.human_name, content=query) self.memory_service.add_messages(new_message) messages: List[Message] = [] - # add memory to system prompt + # Incorporate memory into the system prompt if available system_prompt = self.prompt_handler.system_prompt memories: str = self.memory_service.read_memory() if memories: @@ -90,45 +162,67 @@ class CliMemoryChat(BaseMemoryChat): system_prompt = "\n".join([x.strip() for x in [system_prompt, memory_prompt, memories]]) messages.append(Message(role=MessageRoleEnum.SYSTEM, content=system_prompt)) - # add history messages + # Include past conversation history in the message list history_messages = self.memory_service.read_message() if history_messages: messages.extend(history_messages) - # add new_message + # Append the current user's message to the conversation context messages.append(new_message) self.logger.info(f"messages={messages}") - # call LLM. in stream mode, return generator. in non-stream mode, return response. + # Invoke the Language Model with the constructed message context, respecting streaming setting generated = self.generation_model.call(messages=messages, stream=self.stream) - # in non-stream mode, remember the response if user demand to do so. + # In non-streaming interactions, explicitly save the AI's reply to memory if instructed if remember_response: - assert not self.stream + assert not self.stream # Ensure we're not in streaming mode when remembering responses generated.message.role_name = self.assistant_name self.memory_service.add_messages(generated.message) - # return response or generator + # Return the AI's response directly or as a generator based on the streaming mode return generated @staticmethod def parse_query_command(query: str): - query_split = query.lstrip("/").lower().split(" ") - command = query_split[0] - args = query_split[1:] - kwargs = {} + """ + Parses the user's input query command, separating it into the command and its associated keyword arguments. + + Args: + query (str): The raw input string from the user which includes the command and its arguments. + + Returns: + tuple: A tuple containing the command (str) as the first element and a dictionary (kwargs) of keyword arguments as the second element. + """ + query_split = query.lstrip("/").lower().split(" ") # Split and preprocess the input command + command = query_split[0] # Extract the command + args = query_split[1:] # Extract the arguments following the command + kwargs = {} # Initialize dictionary to hold keyword arguments + for arg in args: - if not args: + if not args: # Skip if no arguments exist (unnecessary check due to prior assignment, but retained as per original) continue - arg_split = arg.split("=") - if len(arg_split) >= 2: - k = arg_split[0] - v = arg_split[1] - if k and v: + arg_split = arg.split("=") # Split argument into key-value pair + if len(arg_split) >= 2: # Ensure there's both a key and value + k = arg_split[0] # Extract key + v = arg_split[1] # Extract value + if k and v: # Only add to kwargs if both key and value are non-empty kwargs[k] = v - return command, kwargs + + return command, kwargs # Return the parsed command and keyword arguments def process_commands(self, query: str) -> bool: + """ + Parses and executes commands from user input in the CLI chat interface. + Supports operations like exiting, clearing screen, showing help, toggling stream mode, + executing predefined memory operations, and handling unknown commands. + + Args: + query (str): The user's input command string. + + Returns: + bool: Indicates whether to continue running the CLI after processing the command. + """ continue_run = True command, kwargs = self.parse_query_command(query) @@ -170,6 +264,13 @@ class CliMemoryChat(BaseMemoryChat): return continue_run def run(self): + """ + Runs the CLI chat loop, which handles user input, processes commands, + communicates with the AI model, manages conversation memory, and controls + the chat session including streaming responses, command execution, and error handling. + + The loop continues until the user explicitly chooses to exit. + """ self.print_logo() self.USER_COMMANDS.update(self.memory_service.op_description_dict) @@ -181,16 +282,18 @@ class CliMemoryChat(BaseMemoryChat): query: str = query.strip() - # handle cli / commands with memory ops + # Handle special commands prefixed with '/' if query.startswith("/"): if self.process_commands(query=query): continue else: break + # Print prompt for AI's response questionary.print("> ", end="", style="fg:yellow") questionary.print(f"{self.assistant_name}: ", end="", style="bold") + # Fetch and display AI's response, with support for streaming if self.stream: model_response = None for model_response in self.chat_with_memory(query=query): @@ -201,18 +304,20 @@ class CliMemoryChat(BaseMemoryChat): model_response = self.chat_with_memory(query=query) questionary.print(model_response.message.content) - # add response to memory + # Append AI's response to the conversation memory model_response.message.role_name = self.assistant_name self.memory_service.add_messages(model_response.message) except KeyboardInterrupt: + # Handle user interruption and confirm exit questionary.print("User interrupt occurred.") - is_exit = questionary.confirm("continue exit").unsafe_ask() + is_exit = questionary.confirm("Continue exit?").unsafe_ask() if is_exit: self.memory_service.stop_service() break except Exception as e: + # Log and handle any unanticipated exceptions import traceback traceback.print_exc() self.logger.exception(f"An exception occurred when running cli memory chat. args={e.args}.") diff --git a/memory_scope/constants/common_constants.py b/memory_scope/constants/common_constants.py index c27aca87..b4dd15ef 100644 --- a/memory_scope/constants/common_constants.py +++ b/memory_scope/constants/common_constants.py @@ -1,33 +1,89 @@ +# common_constants.py +# This module defines constants used as keys throughout the application to maintain a consistent reference +# for data structures related to workflow management, chat interactions, context storage, memory operations, +# node processing, and temporal inference functionalities. + WORKFLOW_NAME = "workflow_name" +""" +The constant represents the key for the name of the workflow in the application context. +""" RESULT = "result" +""" +Indicates the key for storing the result of an operation or processing within the application. +""" CHAT_MESSAGES = "chat_messages" +""" +Used as the key for accessing or storing a collection of chat messages in the application's data model. +""" CONTEXT_MEMORY_DICT = "context_memory_dict" +""" +Denotes the key for a dictionary that holds contextual memory information, which might be utilized for maintaining conversation context. +""" CHAT_KWARGS = "chat_kwargs" +""" +Key for passing keyword arguments specifically related to chat functionalities. +""" QUERY_WITH_TS = "query_with_ts" +""" +Refers to a query operation that includes a timestamp, used when retrieving data with temporal consideration. +""" RETRIEVE_MEMORY_NODES = "retrieve_memory_nodes" +""" +Specifies the action of retrieving nodes from the memory, often used in the context of information retrieval processes. +""" RANKED_MEMORY_NODES = "ranked_memory_nodes" +""" +Indicates a collection of memory nodes that have been ranked, typically by relevance or other criteria. +""" NOT_REFLECTED_NODES = "not_reflected_nodes" +""" +Represents nodes that have not been incorporated or reflected back into the system, such as unprocessed updates. +""" NOT_UPDATED_NODES = "not_updated_nodes" +""" +Identifies nodes that were not updated during a processing cycle, useful for tracking changes. +""" EXTRACT_TIME_DICT = "extract_time_dict" +""" +A key pointing to a dictionary used for extracting or storing time-related information extracted from data nodes. +""" NEW_OBS_NODES = "new_obs_nodes" +""" +Refers to nodes representing new observations added to the system. +""" NEW_OBS_WITH_TIME_NODES = "new_obs_with_time_nodes" +""" +Similar to 'new_obs_nodes', but specifically denotes these observations are associated with timestamps. +""" INSIGHT_NODES = "insight_nodes" +""" +Key for nodes that encapsulate insights derived from data analysis or processing stages. +""" TODAY_NODES = "today_nodes" +""" +Indicates nodes relevant or generated on the current day, assisting in daily summaries or time-sensitive operations. +""" MERGE_OBS_NODES = "merge_obs_nodes" +""" +Specifies nodes that are candidates for merging, often to consolidate similar or redundant observations. +""" TIME_INFER = "time_infer" +""" +Involves the process of inferring time information from data, crucial for temporal understanding within the application. +""" diff --git a/memory_scope/constants/language_constants.py b/memory_scope/constants/language_constants.py index 0999e2fb..9dfeaa10 100644 --- a/memory_scope/constants/language_constants.py +++ b/memory_scope/constants/language_constants.py @@ -1,5 +1,7 @@ from memory_scope.enumeration.language_enum import LanguageEnum +# This dictionary maps languages to lists of words related to datetime expressions. +# It aids in recognizing and processing datetime mentions in text, enhancing the system's ability to understand temporal context across different languages. DATATIME_WORD_LIST = { LanguageEnum.CN: [ "天", @@ -86,6 +88,7 @@ DATATIME_WORD_LIST = { ] } +# A mapping of weekdays for each supported language, facilitating calendar-related operations and understanding within the application. WEEKDAYS = { LanguageEnum.CN: [ "周一", @@ -107,42 +110,49 @@ WEEKDAYS = { ] } +# Constants for the word 'none' in different languages NONE_WORD = { LanguageEnum.CN: "无", LanguageEnum.EN: "none" } +# Constants for the word 'repeated' in different languages REPEATED_WORD = { LanguageEnum.CN: "重复", LanguageEnum.EN: "repeated" } +# Constants for the word 'contradictory' in different languages CONTRADICTORY_WORD = { LanguageEnum.CN: "矛盾", LanguageEnum.EN: "contradictory" } +# Constants for the phrase 'included' in different languages INCLUDED_WORD = { LanguageEnum.CN: "被包含", LanguageEnum.EN: "included" } +# Constants for the symbol ':' in different languages' representations COLON_WORD = { LanguageEnum.CN: "：", LanguageEnum.EN: ":" } - +# Constants for the symbol ',' in different languages' representations COMMA_WORD = { LanguageEnum.CN: "，", LanguageEnum.EN: "," } +# Default human name placeholders for different languages DEFAULT_HUMAN_NAME = { LanguageEnum.CN: "用户", LanguageEnum.EN: "user" } +# Mapping of datetime terms from natural language to standardized keys for each supported language DATATIME_KEY_MAP = { LanguageEnum.CN: { "年": "year", @@ -160,6 +170,7 @@ DATATIME_KEY_MAP = { } } +# Phrase for indicating inferred time in different languages TIME_INFER_WORD = { LanguageEnum.CN: "推断时间", LanguageEnum.EN: "Inference time" diff --git a/memory_scope/enumeration/language_enum.py b/memory_scope/enumeration/language_enum.py index 9c0423f1..7a34ed1b 100644 --- a/memory_scope/enumeration/language_enum.py +++ b/memory_scope/enumeration/language_enum.py @@ -2,6 +2,12 @@ from enum import Enum class LanguageEnum(str, Enum): + """ + An enumeration representing supported languages. + + Members: + - CN: Represents the Chinese language. + - EN: Represents the English language. + """ CN = "cn" - EN = "en" diff --git a/memory_scope/enumeration/memory_status_enum.py b/memory_scope/enumeration/memory_status_enum.py index 60703d95..de4c929d 100644 --- a/memory_scope/enumeration/memory_status_enum.py +++ b/memory_scope/enumeration/memory_status_enum.py @@ -2,12 +2,18 @@ from enum import Enum class MemoryNodeStatus(str, Enum): + """ + Enumeration representing various statuses of a memory node. + + Each status reflects a different state of the node in terms of its lifecycle or content: + - NEW: Indicates a newly created node. + - MODIFIED: Signifies that the node has been altered. + - CONTENT_MODIFIED: Specifies changes in the actual content of the node. + - ACTIVE: Denotes that the node is currently in use or accessible. + - EXPIRED: Implies that the node is no longer valid or needed. + """ NEW = "new" - MODIFIED = "modified" - CONTENT_MODIFIED = "content_modified" - ACTIVE = "active" - EXPIRED = "expired" diff --git a/memory_scope/enumeration/memory_type_enum.py b/memory_scope/enumeration/memory_type_enum.py index 42a4e413..1f1343a8 100644 --- a/memory_scope/enumeration/memory_type_enum.py +++ b/memory_scope/enumeration/memory_type_enum.py @@ -2,10 +2,16 @@ from enum import Enum class MemoryTypeEnum(str, Enum): + """ + Defines an enumeration for different types of memory categories. + + Each member represents a distinct type of memory content: + - CONVERSATION: Represents conversation-based memories. + - OBSERVATION: Denotes observational memories. + - INSIGHT: Indicates insightful memories derived from analysis. + - OBS_CUSTOMIZED: Customized observational memories. + """ CONVERSATION = "conversation" - OBSERVATION = "observation" - INSIGHT = "insight" - OBS_CUSTOMIZED = "obs_customized" diff --git a/memory_scope/enumeration/message_role_enum.py b/memory_scope/enumeration/message_role_enum.py index 669ebe01..a3c91c17 100644 --- a/memory_scope/enumeration/message_role_enum.py +++ b/memory_scope/enumeration/message_role_enum.py @@ -2,8 +2,15 @@ from enum import Enum class MessageRoleEnum(str, Enum): - USER = "user" + """ + Enumeration for different message roles within a conversation context. + + This enumeration includes predefined roles such as User, Assistant, and System, + which can be used to categorize messages in chat interfaces, AI interactions, or + any system that involves distinct participant roles. + """ + USER = "user" # Represents a message sent by the user. - ASSISTANT = "assistant" + ASSISTANT = "assistant" # Represents a response or action performed by an assistant. - SYSTEM = "system" + SYSTEM = "system" # Represents system-level messages or actions. diff --git a/memory_scope/enumeration/model_enum.py b/memory_scope/enumeration/model_enum.py index d1d9aecc..ff9aef75 100644 --- a/memory_scope/enumeration/model_enum.py +++ b/memory_scope/enumeration/model_enum.py @@ -2,8 +2,14 @@ from enum import Enum class ModelEnum(str, Enum): + """ + An enumeration representing different types of models used within the system. + + Members: + GENERATION_MODEL: Represents a model responsible for generating content. + EMBEDDING_MODEL: Represents a model tasked with creating embeddings, typically used for transforming data into a numerical form suitable for machine learning tasks. + RANK_MODEL: Denotes a model that specializes in ranking, often used to order items based on relevance or importance. + """ GENERATION_MODEL = "generation_model" - EMBEDDING_MODEL = "embedding_model" - RANK_MODEL = "rank_model" diff --git a/memory_scope/memory/operation/base_backend_operation.py b/memory_scope/memory/operation/base_backend_operation.py index 0cfc43a2..f14b977e 100644 --- a/memory_scope/memory/operation/base_backend_operation.py +++ b/memory_scope/memory/operation/base_backend_operation.py @@ -7,9 +7,20 @@ from memory_scope.utils.logger import Logger class BaseBackendOperation(BaseOperation): + """ + BaseBackendOperation serves as an abstract base class for defining backend operations within a specified time interval. + It manages operation status, loop control, and integrates with a logging facility and a global context for thread management. + """ operation_type: OPERATION_TYPE = "backend" def __init__(self, interval_time: int, **kwargs): + """ + Initializes the BaseBackendOperation instance with an interval time for recurring operations. + + Args: + interval_time (int): The time interval in seconds at which the operation should run. + **kwargs: Additional keyword arguments passed to the parent class's initializer. + """ super(BaseBackendOperation, self).__init__(**kwargs) self.interval_time: int = interval_time @@ -22,9 +33,31 @@ class BaseBackendOperation(BaseOperation): @abstractmethod def _run_operation(self, **kwargs): + """ + Abstract method to define the logic of the operation executed by the backend. + + This method needs to be implemented by any subclass of BaseBackendOperation. + It serves as the core execution unit for backend-specific tasks. + + Args: + **kwargs: Arbitrary keyword arguments that might be necessary for the operation. + + Raises: + NotImplementedError: If the method is not overridden in a subclass. + """ raise NotImplementedError def run_operation(self, **kwargs): + """ + Executes the operation defined by `_run_operation` method with given keyword arguments, + while managing the operation status and exception handling. + + Args: + **kwargs: Arbitrary keyword arguments to be passed to `_run_operation`. + + Returns: + The result of the `_run_operation` method if no exception occurs, otherwise None. + """ if self._operation_status_run: return @@ -39,6 +72,10 @@ class BaseBackendOperation(BaseOperation): return result def _loop_operation(self): + """ + Loops until _loop_switch is False, sleeping for 1 second in each interval. + At each interval, it checks if _loop_switch is still True, and if so, executes the operation. + """ while self._loop_switch: for _ in range(self.interval_time): if self._loop_switch: @@ -49,9 +86,16 @@ class BaseBackendOperation(BaseOperation): self.run_operation() def run_operation_backend(self): + """ + Initiates the background operation loop if it's not already running. + Sets the _loop_switch to True and submits the _loop_operation to a thread from the global thread pool. + """ if not self._loop_switch: self._loop_switch = True self._run_thread = G_CONTEXT.thread_pool.submit(self._loop_operation) def stop_operation_backend(self): + """ + Stops the background operation loop by setting the _loop_switch to False. + """ self._loop_switch = False diff --git a/memory_scope/memory/operation/base_operation.py b/memory_scope/memory/operation/base_operation.py index 64cc23c1..5c7df589 100644 --- a/memory_scope/memory/operation/base_operation.py +++ b/memory_scope/memory/operation/base_operation.py @@ -5,22 +5,64 @@ OPERATION_TYPE = Literal["frontend", "backend"] class BaseOperation(metaclass=ABCMeta): + """ + An abstract base class representing an operation that can be categorized as either frontend or backend. + + Attributes: + operation_type (OPERATION_TYPE): Specifies the type of operation, defaulting to "frontend". + name (str): The name of the operation. + description (str): A description of the operation. + kwargs (dict): Additional keyword arguments for operation configuration. + """ + operation_type: OPERATION_TYPE = "frontend" def __init__(self, name: str, description: str = "", **kwargs): + """ + Initializes a new instance of the BaseOperation. + + Args: + name (str): The name identifying the operation. + description (str): An optional description detailing the operation's purpose or behavior. + **kwargs: Arbitrary keyword arguments for custom settings or parameters. + """ self.name: str = name self.description: str = description self.kwargs: dict = kwargs def init_workflow(self, **kwargs): + """ + Initialize the workflow with additional keyword arguments if needed. + + Args: + **kwargs: Additional parameters for initializing the workflow. + """ pass @abstractmethod def run_operation(self, **kwargs): + """ + Abstract method to define the operation to be run. + Subclasses must implement this method. + + Args: + **kwargs: Keyword arguments for running the operation. + + Raises: + NotImplementedError: If the subclass does not implement this method. + """ raise NotImplementedError def run_operation_backend(self): + """ + Placeholder method for running an operation specific to the backend. + Intended to be overridden by subclasses if backend operations are required. + """ pass def stop_operation_backend(self): + """ + Placeholder method to stop any ongoing backend operations. + Should be implemented in subclasses where backend operations are managed. + """ pass diff --git a/memory_scope/memory/operation/base_workflow.py b/memory_scope/memory/operation/base_workflow.py index 6353ab6d..0541c7a7 100644 --- a/memory_scope/memory/operation/base_workflow.py +++ b/memory_scope/memory/operation/base_workflow.py @@ -37,21 +37,34 @@ class BaseWorkflow(object): self._print_workflow() def _parse_workflow(self): - # re-match e.g., [a|b],c,[d,e,f|g,h],j + """ + Parses the workflow string to configure worker threads and organizes them into execution order. + + The workflow string format supports complex configurations with optional multi-threading indications. + E.g., `[task1,task2|task3],task4` denotes task1 and task2 can run in parallel to task3, followed by task4. + + Returns: + List[List[List[str]]]: A nested list representing the execution plan, including parallel groups and tasks. + """ + # Regular expression to match components of the workflow, handling both plain items and grouped items. pattern = r'(\[[^\]]*\]|[^,]+)' + # Find all matches in the workflow string based on the pattern. workflow_split = re.findall(pattern, self.workflow) + for workflow_part in workflow_split: # e.g., [d,e,f|g,h] workflow_part = workflow_part.strip() if '[' in workflow_part or ']' in workflow_part: workflow_part = workflow_part.replace('[', '').replace(']', '') - # e.g., ["d,e,f", "g,h"] + # Split the part by '|' to identify potential parallel task groups. line_split = [x.strip() for x in workflow_part.split("|") if x] + + # Skip if no valid tasks are identified after splitting. if len(line_split) <= 0: continue - # is under multi thread cond + # Determine if the current part involves multi-threading based on the number of groups. is_multi_thread: bool = len(line_split) > 1 # e.g., ["d","e","f"] @@ -62,27 +75,56 @@ class BaseWorkflow(object): # add workers for sub_item in sub_split: self.worker_dict[sub_item] = is_multi_thread + + # Append the parsed and structured tasks to the workflow execution plan. self.workflow_worker_list.append(line_split_split) + + # Return the fully constructed workflow execution plan. return self.workflow_worker_list def _print_workflow(self): + """ + Prints the workflow stages in a structured format. Each stage of the workflow + is detailed with its constituent parts, either single elements or grouped + elements separated by ' | '. + + The method iterates over the workflow parts, handling both singular steps + and parallel steps (where elements are zipped together). + """ self.logger.info(f"----- workflow.{self.name}.print.begin -----") i: int = 0 for workflow_part in self.workflow_worker_list: if len(workflow_part) == 1: + # Handles workflow parts with single elements for w in workflow_part[0]: self.logger.info(f"stage{i}: {w}") i += 1 else: + # Handles workflow parts with multiple parallel elements (zipped) for w_zip in zip_longest(*workflow_part, fillvalue="-"): self.logger.info(f"stage{i}: {' | '.join(w_zip)}") i += 1 + # Skips placeholder '-' used for uneven lists in zip_longest for w in w_zip: if w == "-": continue self.logger.info(f"----- workflow.{self.name}.print.end -----") def init_workers(self, is_backend: bool = False, **kwargs): + """ + Initializes worker instances based on the configuration for each worker defined in `G_CONTEXT.worker_config`. + Each worker can be set to run in a multi-threaded mode depending on the `is_backend` flag or the worker's individual configuration. + + Args: + is_backend (bool, optional): A flag indicating whether the workers should be initialized in a backend context. Defaults to False. + **kwargs: Additional keyword arguments to be passed during worker initialization. + + Raises: + RuntimeError: If a worker mentioned in `self.worker_dict` does not exist in `G_CONTEXT.worker_config`. + + Note: + This method modifies `self.worker_dict` in-place, replacing the keys with actual worker instances. + """ for name in list(self.worker_dict.keys()): if name not in G_CONTEXT.worker_config: raise RuntimeError(f"worker={name} is not exists in worker_config!") @@ -106,18 +148,31 @@ class BaseWorkflow(object): return True def run_workflow(self): + """ + Executes the workflow by orchestrating the steps defined in `self.workflow_worker_list`. + This method supports both sequential and parallel execution of sub-workflows based on the structure of `self.workflow_worker_list`. + + If a workflow part consists of a single item, it is executed sequentially. For parts with multiple items, + they are submitted for parallel execution using a thread pool. The workflow will stop if any sub-workflow returns False. + """ self.logger.info(f"----- workflow.{self.name}.begin -----") with Timer(self.name, log_time=False) as t: self.context[WORKFLOW_NAME] = self.name + + # Iterate over each part of the workflow for workflow_part in self.workflow_worker_list: + # Sequential execution for single-item parts if len(workflow_part) == 1: if not self._run_sub_workflow(workflow_part[0]): break + # Parallel execution for multi-item parts else: t_list = [] + # Submit tasks to the thread pool for sub_workflow in workflow_part: t_list.append(G_CONTEXT.thread_pool.submit(self._run_sub_workflow, sub_workflow)) - + + # Check results; if any task returns False, stop the workflow flag = True for future in as_completed(t_list): if not future.result(): @@ -125,4 +180,5 @@ class BaseWorkflow(object): break if not flag: break + self.logger.info(f"----- workflow.{self.name}.end cost={t.cost_str}-----") diff --git a/memory_scope/memory/operation/read_memory.py b/memory_scope/memory/operation/read_memory.py index e272ea4a..d6288fed 100644 --- a/memory_scope/memory/operation/read_memory.py +++ b/memory_scope/memory/operation/read_memory.py @@ -22,13 +22,30 @@ class ReadMemory(BaseWorkflow, BaseOperation): self.his_msg_count: int = his_msg_count def init_workflow(self, **kwargs): + """ + Initializes the workflow by setting up workers with provided keyword arguments. + + Args: + **kwargs: Arbitrary keyword arguments to be passed during worker initialization. + """ self.init_workers(**kwargs) def run_operation(self, **kwargs): - self.context.clear() - max_count = 1 + self.his_msg_count + """ + Executes the main operation of reading recent chat messages, initializing workflow, + and returning the result of the workflow execution. + + Args: + **kwargs: Additional keyword arguments used in the operation context. + + Returns: + Any: The result obtained from executing the workflow. + """ + self.context.clear() # Clear the previous operation context + max_count = 1 + self.his_msg_count # Determine the number of historical messages to include + # Include the most recent messages in the operation context self.context[CHAT_MESSAGES] = [x.copy(deep=True) for x in self.chat_messages[-max_count:]] - self.context[CHAT_KWARGS] = kwargs - self.run_workflow() - result = self.context.get(RESULT) + self.context[CHAT_KWARGS] = kwargs # Add additional arguments to the context + self.run_workflow() # Execute the workflow with the prepared context + result = self.context.get(RESULT) # Retrieve the result from the context after workflow execution return result diff --git a/memory_scope/memory/operation/summary_memory.py b/memory_scope/memory/operation/summary_memory.py index 8afef7f4..abb471c2 100644 --- a/memory_scope/memory/operation/summary_memory.py +++ b/memory_scope/memory/operation/summary_memory.py @@ -5,16 +5,44 @@ from memory_scope.memory.operation.base_workflow import BaseWorkflow class SummaryMemory(BaseWorkflow, BaseBackendOperation): + """ + A class that combines workflow management and backend operations to process and summarize data. + + This class inherits functionality from both `BaseWorkflow` for managing workflow steps and + `BaseBackendOperation` for executing backend-specific tasks. + """ operation_type: OPERATION_TYPE = "backend" def __init__(self, **kwargs): - super().__init__(**kwargs) - BaseBackendOperation.__init__(self, **kwargs) + """ + Initializes the SummaryMemory instance, setting up both workflow and backend operation capabilities. + + Args: + **kwargs: Additional keyword arguments used in initializing the parent classes. + """ + super().__init__(**kwargs) # Initialize the BaseWorkflow part of the instance + BaseBackendOperation.__init__(self, **kwargs) # Initialize the BaseBackendOperation part def init_workflow(self, **kwargs): + """ + Initializes the workflow with backend settings using provided keyword arguments. + + Args: + **kwargs: Additional keyword arguments to initialize the workflow. + """ self.init_workers(is_backend=True, **kwargs) def _run_operation(self, **kwargs): + """ + Executes an operation within the workflow by clearing the context, + setting chat arguments, running the workflow, and returning the result. + + Args: + **kwargs: Keyword arguments necessary for the operation, including chat parameters. + + Returns: + Any: The result obtained after executing the workflow. + """ self.context.clear() self.context[CHAT_KWARGS] = kwargs self.run_workflow() diff --git a/memory_scope/memory/operation/write_memory.py b/memory_scope/memory/operation/write_memory.py index e42482ac..b66066db 100644 --- a/memory_scope/memory/operation/write_memory.py +++ b/memory_scope/memory/operation/write_memory.py @@ -27,6 +27,12 @@ class WriteMemory(BaseWorkflow, BaseBackendOperation): @property def not_memorized_size(self): + """ + Calculates the count of chat messages that have not been memorized. + + Returns: + int: The total count of chat messages which are not marked as memorized. + """ return sum([not x.memorized for x in self.chat_messages]) def set_memorized(self): @@ -36,9 +42,29 @@ class WriteMemory(BaseWorkflow, BaseBackendOperation): msg.memorized = True def init_workflow(self, **kwargs): + """ + Initializes the workflow by setting up workers, considering backend-specific parameters. + + Args: + **kwargs: Additional keyword arguments passed for initializing workers. + """ self.init_workers(is_backend=True, **kwargs) def _run_operation(self, **kwargs): + """ + Executes an operation after preparing the chat context, checking message memory status, + and updating workflow status accordingly. + + If the number of not-memorized messages is less than the contextual message count, + the operation is skipped. Otherwise, it sets up the chat context, runs the workflow, + captures the result, and updates the memory status. + + Args: + **kwargs: Keyword arguments for chat operation configuration. + + Returns: + Any: The result obtained from running the workflow. + """ self.context.clear() self.context[CHAT_KWARGS] = kwargs not_memorized_size = self.not_memorized_size diff --git a/memory_scope/memory/service/base_memory_service.py b/memory_scope/memory/service/base_memory_service.py index 5076920f..c37c46c2 100644 --- a/memory_scope/memory/service/base_memory_service.py +++ b/memory_scope/memory/service/base_memory_service.py @@ -8,11 +8,27 @@ from memory_scope.utils.logger import Logger class BaseMemoryService(metaclass=ABCMeta): + """ + An abstract base class for managing memory operations within a multi-threaded context. + It sets up the infrastructure for operation handling, message storage, and synchronization, + along with logging capabilities and customizable configurations. + """ + def __init__(self, memory_operations: Dict[str, dict], read_memory_key: str = "read_memory", read_message_key: str = "read_message", **kwargs): + """ + Initializes the BaseMemoryService with operation definitions, keys for memory access, + and additional keyword arguments for flexibility. + + Args: + memory_operations (Dict[str, dict]): A dictionary defining available memory operations. + read_memory_key (str): The key indicating a read memory operation. Defaults to "read_memory". + read_message_key (str): The key for reading messages. Defaults to "read_message". + **kwargs: Additional parameters to customize service behavior. + """ self.memory_operations: Dict[str, dict] = memory_operations self.read_memory_key: str = read_memory_key self.read_message_key: str = read_message_key @@ -27,6 +43,20 @@ class BaseMemoryService(metaclass=ABCMeta): @abstractmethod def _init_operation(self, memory_operations: Dict[str, dict]): + """ + Initializes the memory operations with a given dictionary of operations. + + This method is to be implemented by subclasses to set up or configure + the memory operations based on the provided dictionary. + + Args: + memory_operations (Dict[str, dict]): A dictionary containing configuration + details for each memory operation. + + Raises: + NotImplementedError: This exception is raised to indicate that the method + needs to be overridden in the subclass. + """ raise NotImplementedError @abstractmethod @@ -34,25 +64,68 @@ class BaseMemoryService(metaclass=ABCMeta): raise NotImplementedError def start_service(self, **kwargs): + """ + This method is intended to initiate the service with provided keyword arguments, + preparing the necessary resources for executing operations. + + Args: + **kwargs: Additional keyword arguments used to configure the service upon startup. + """ pass @abstractmethod def do_operation(self, op_name: str, **kwargs): + """ + Abstract method defining the interface for executing a specific operation by its name. + This method must be implemented by subclasses to provide the actual operation logic. + + Args: + op_name (str): The name identifying the operation to be performed. + **kwargs: Additional keyword arguments required for the operation execution. + + Raises: + NotImplementedError: This exception is raised when the method is not overridden in a subclass. + """ raise NotImplementedError @property def op_description_dict(self) -> Dict[str, str]: + """ + Property to retrieve a dictionary mapping operation keys to their descriptions. + Lazily initializes the dictionary on first access. + + Returns: + Dict[str, str]: A dictionary where keys are operation identifiers and values are their descriptions. + """ 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): + """ + Executes the operation associated with reading memory. + Asserts that the operation for reading memory has been initialized. + + Returns: + Any: The result of the read memory operation. + """ assert self.read_memory_key in self._operation_dict, f"op={self.read_memory_key} is not inited!" return self.do_operation(self.read_memory_key) def read_message(self): + """ + Executes the operation associated with reading messages. + Asserts that the operation for reading messages has been initialized. + + Returns: + Any: The result of the read message operation. + """ assert self.read_message_key in self._operation_dict, f"op={self.read_message_key} is not inited!" return self.do_operation(self.read_message_key) def stop_service(self): + """ + Placeholder method to stop the service. + Intended to be overridden by subclasses to define specific shutdown logic. + """ pass diff --git a/memory_scope/memory/service/chat_memory_service.py b/memory_scope/memory/service/chat_memory_service.py index c8aa6188..1142541f 100644 --- a/memory_scope/memory/service/chat_memory_service.py +++ b/memory_scope/memory/service/chat_memory_service.py @@ -15,12 +15,24 @@ class ChatMemoryService(BaseMemoryService): self._init_operation(self.memory_operations) def _init_operation(self, memory_operations: Dict[str, dict]): + """ + Initializes memory operations based on the provided configuration dictionary. + Ensures that each operation is only initialized once by checking for duplicates. + + Args: + memory_operations (Dict[str, dict]): A dictionary where keys are operation names + and values are configuration dictionaries for each operation. + + Note: + Logs a warning if an attempt is made to initialize an operation with a name that already exists. + Logs an info message upon successful initialization of each operation. + """ for name, operation_config in memory_operations.items(): if name in self._operation_dict: self.logger.warning(f"memory operation={name} is repeated!") continue - self._operation_dict[name] = init_instance_by_config( + self._operation_dict[name] = init_instance_by_config( # ⭐ Initialize operation instance by its config config=operation_config, name=name, chat_messages=self.chat_messages, @@ -29,29 +41,65 @@ class ChatMemoryService(BaseMemoryService): self.logger.info(f"service={self.__class__.__name__} init operation={name}") def add_messages(self, messages: List[Message] | Message): + """ + Adds a single message or a list of messages to the chat history, ensuring the message list + remains sorted by creation time and does not exceed the maximum history message count. + + Args: + messages (List[Message] | Message): A single message instance or a list of message instances + to be added to the chat history. + """ + # If a single message is provided, convert it into a list for uniform processing if isinstance(messages, Message): messages = [messages] + # Sort the messages by their creation time to maintain chronological order messages = sorted(messages, key=lambda x: x.time_created) + + # Append the sorted messages to the chat history self.chat_messages.extend(messages) + + # If the chat history exceeds the allowed message count, remove the oldest messages if len(self.chat_messages) > self.history_msg_count: gap_size = len(self.chat_messages) - self.history_msg_count for _ in range(gap_size): self.chat_messages.pop(0) def start_service(self, **kwargs): + """ + Initializes and starts backend operations defined in `_operation_dict`. + + Args: + **kwargs: Additional keyword arguments passed to each operation's initialization. + """ for _, operation in self._operation_dict.items(): - operation.init_workflow(**kwargs) + operation.init_workflow(**kwargs) # Initialize workflow for each operation if operation.operation_type == "backend": - operation.run_operation_backend() + operation.run_operation_backend() # Run backend operations def do_operation(self, op_name: str, **kwargs): + """ + Executes a specific operation by its name with provided keyword arguments. + + Args: + op_name (str): The name of the operation to execute. + **kwargs: Keyword arguments for the operation's execution. + + Returns: + The result of the operation execution, if any. Otherwise, None. + + Raises: + Warning: If the operation name is not initialized in `_operation_dict`. + """ if op_name not in self._operation_dict: - self.logger.warning(f"op_name={op_name} is not inited!") + self.logger.warning(f"op_name={op_name} is not inited!") # Warn if operation not initialized return - return self._operation_dict[op_name].run_operation(**kwargs) + return self._operation_dict[op_name].run_operation(**kwargs) # Execute the operation def stop_service(self): + """ + Stops all backend operations that are currently running. + """ for _, operation in self._operation_dict.items(): if operation.operation_type == "backend": - operation.stop_operation_backend() + operation.stop_operation_backend() # Stop backend operations diff --git a/memory_scope/memory/worker/dummy_worker.py b/memory_scope/memory/worker/dummy_worker.py index 2b488b7e..cabda23d 100644 --- a/memory_scope/memory/worker/dummy_worker.py +++ b/memory_scope/memory/worker/dummy_worker.py @@ -6,9 +6,17 @@ from memory_scope.memory.worker.base_worker import BaseWorker class DummyWorker(BaseWorker): def _run(self): + """ + Executes the dummy worker's run logic by logging workflow entry, capturing the current timestamp, + file path, and setting the result context with details about the workflow execution. + + This method utilizes the BaseWorker's capabilities to interact with the workflow context. + """ workflow_name = self.get_context(WORKFLOW_NAME) chat_kwargs = self.get_context(CHAT_KWARGS) - self.logger.info(f"enter workflow={workflow_name}.dummy_worker!") + self.logger.info(f"Entering workflow={workflow_name}.dummy_worker!") + # Records the current timestamp as an integer ts = int(datetime.datetime.now().timestamp()) + # Retrieves the current file's path file_path = __file__ self.set_context(RESULT, f"test {workflow_name} kwargs={chat_kwargs} file_path={file_path} \nts={ts}") diff --git a/memory_scope/memory/worker/memory_base_worker.py b/memory_scope/memory/worker/memory_base_worker.py index 32448044..db885e95 100644 --- a/memory_scope/memory/worker/memory_base_worker.py +++ b/memory_scope/memory/worker/memory_base_worker.py @@ -20,6 +20,18 @@ class MemoryBaseWorker(BaseWorker, metaclass=ABCMeta): generation_model: str = "", rank_model: str = "", **kwargs): + """ + Initializes the MemoryBaseWorker with specified models and configurations. + + Args: + embedding_model (str): Identifier or instance of the embedding model used for transforming text. + generation_model (str): Identifier or instance of the text generation model. + rank_model (str): Identifier or instance of the ranking model for sorting or prioritizing data. + **kwargs: Additional keyword arguments passed to the parent class initializer. + + The constructor also initializes key attributes related to memory store, monitoring, + user and target identification, and a prompt handler, setting them up for later use. + """ super(MemoryBaseWorker, self).__init__(**kwargs) self._embedding_model: BaseModel | str = embedding_model @@ -35,47 +47,122 @@ class MemoryBaseWorker(BaseWorker, metaclass=ABCMeta): @property def chat_messages(self) -> List[Message]: + """ + Getter property to retrieve the list of chat messages from the context. + + Returns: + List[Message]: A list of Message objects representing the chat messages. + """ return self.get_context(CHAT_MESSAGES) @chat_messages.setter + def chat_messages(self, messages: List[Message]) -> None: + """ + Setter property to update the list of chat messages in the context. + + Args: + messages (List[Message]): A list of Message objects to set as the new chat messages. + """ def chat_messages(self, value): + """ + Sets the context for chat messages with the provided value. + + Args: + value: The value to be set for the chat messages context. The type of `value` is inferred from the usage context. + """ self.set_context(CHAT_MESSAGES, value) @property def chat_kwargs(self) -> Dict[str, str]: + """ + Retrieves the chat keyword arguments from the context. + + This property getter fetches the chat-related parameters stored in the context, + which are used to configure how chat interactions are handled. + + Returns: + Dict[str, str]: A dictionary containing the chat keyword arguments. + """ return self.get_context(CHAT_KWARGS) @property def embedding_model(self) -> BaseModel: + """ + Property to get the embedding model. If the model is currently stored as a string, + it will be replaced with the actual model instance from the global context's model dictionary. + + Returns: + BaseModel: The embedding model used for converting text into vector representations. + """ if isinstance(self._embedding_model, str): - self._embedding_model = G_CONTEXT.model_dict[self._embedding_model] + self._embedding_model = G_CONTEXT.model_dict[self._embedding_model] # ⭐ Retrieve the actual model instance when the attribute is a string reference return self._embedding_model @property def generation_model(self) -> BaseModel: + """ + Property to access the generation model. If the model is stored as a string, + it retrieves the actual model instance from the global context's model dictionary. + + Returns: + BaseModel: The model used for text generation. + """ if isinstance(self._generation_model, str): - self._generation_model = G_CONTEXT.model_dict[self._generation_model] + self._generation_model = G_CONTEXT.model_dict[self._generation_model] # ⭐ Retrieve the model instance if currently a string reference return self._generation_model @property def rank_model(self) -> BaseModel: + """ + Property to access the rank model. If the stored rank model is a string, it fetches the actual model instance + from the global context's model dictionary before returning it. + + Returns: + BaseModel: The rank model instance used for ranking tasks within the conversation management system. + """ if isinstance(self._rank_model, str): - self._rank_model = G_CONTEXT.model_dict[self._rank_model] + self._rank_model = G_CONTEXT.model_dict[self._rank_model] # Fetch model instance if string reference return self._rank_model @property def memory_store(self) -> BaseMemoryStore: + """ + Property to access the memory store. If not initialized, it fetches the memory store from the global context. + + Returns: + BaseMemoryStore: The memory store instance associated with this worker. + """ if self._memory_store is None: self._memory_store = G_CONTEXT.memory_store return self._memory_store @property def contex_memory_dict(self) -> Dict[str, MemoryNode]: + """ + Retrieves the context memory dictionary. If it does not exist, initializes it first. + + Returns: + Dict[str, MemoryNode]: The dictionary storing context memory nodes. + """ if not self.has_content(CONTEXT_MEMORY_DICT): - self.set_context(CONTEXT_MEMORY_DICT, {}) + self.set_context(CONTEXT_MEMORY_DICT, {}) # Initialize context memory dict if not present return self.get_context(CONTEXT_MEMORY_DICT) def get_memories(self, keys: str | List[str]) -> List[MemoryNode]: + """ + Retrieves memory nodes associated with the given keys. + + This method accepts a single key or a list of keys. For each key, it fetches the + associated memory IDs from the context. If memory IDs are found, they are used to + collect the corresponding MemoryNode objects from the contex_memory_dict. The final + result is a list of unique MemoryNode values, avoiding duplicates. + + Args: + keys (str | List[str]): The key or list of keys to retrieve memories for. + + Returns: + List[MemoryNode]: A list of MemoryNode objects associated with the input keys. + """ memories: Dict[str, MemoryNode] = {} if isinstance(keys, str): keys = [keys] @@ -87,6 +174,19 @@ class MemoryBaseWorker(BaseWorker, metaclass=ABCMeta): return list(memories.values()) def set_memories(self, key: str, nodes: MemoryNode | List[MemoryNode], log_repeat: bool = True): + """ + Stores or updates multiple memory nodes in the context, optionally logging if a memory ID is repeated. + + Args: + key (str): The key under which to categorize the memory nodes in the context. + nodes (MemoryNode | List[MemoryNode]): A single MemoryNode instance or a list of MemoryNode instances to be set. + log_repeat (bool, optional): If True, logs a warning when a memory ID is encountered more than once. Defaults to True. + + Notes: + - If 'nodes' is None, it is treated as an empty list. + - If 'nodes' is a single MemoryNode instance, it is converted into a list containing that single node. + - Existing memory nodes with duplicate IDs are skipped, with an optional warning logged. + """ if nodes is None: nodes = [] elif isinstance(nodes, MemoryNode): @@ -103,51 +203,104 @@ class MemoryBaseWorker(BaseWorker, metaclass=ABCMeta): self.set_context(key, [n.memory_id for n in nodes]) def save_memories(self, keys: str | List[str] = None): + """ + Saves memories from the context to the memory store. If no keys are provided, + all memories are saved and the context is cleared. If keys are provided, only + the associated memories are saved and removed from the context. + + Args: + keys (str | List[str], optional): The keys identifying which memories to save. + If None, all memories are saved. Defaults to None. + """ if keys is None: - self.memory_store.update_memories(list(self.contex_memory_dict.values())) + self.memory_store.update_memories(list(self.contex_memory_dict.values())) # Save all memories and clear context self.contex_memory_dict.clear() return if isinstance(keys, str): keys = [keys] - ids: Set[str] = Set[str]() + ids: Set[str] = set() # corrected type annotation for Python typing for key in keys: t_ids: List[str] = self.get_context(key) if t_ids: ids.update(t_ids) - nodes = [self.contex_memory_dict.pop(_) for _ in ids] - self.memory_store.update_memories(nodes) + nodes = [self.contex_memory_dict.pop(_) for _ in ids] # Remove and collect nodes by IDs + self.memory_store.update_memories(nodes) # Save collected nodes to memory store @property def monitor(self) -> BaseMonitor: + """ + Property to access the monitoring component. If not initialized, it fetches + the global monitor. + + Returns: + BaseMonitor: The monitoring component instance. + """ if self._monitor is None: self._monitor = G_CONTEXT.monitor return self._monitor @property def user_name(self) -> str: + """ + Property to get the user name from the meta_data of the global context. + If not set initially, it retrieves the 'assistant_name' as the user name. + + Returns: + str: The name of the user. + """ if self._user_name is None: self._user_name = G_CONTEXT.meta_data["assistant_name"] return self._user_name @property def target_name(self) -> str: + """ + Retrieves the target name, initializing it from meta_data if not set. + + Returns: + str: The human-readable name of the target. + """ if self._target_name is None: self._target_name = G_CONTEXT.meta_data["human_name"] return self._target_name @property def prompt_handler(self) -> PromptHandler: + """ + Lazily initializes and returns the PromptHandler instance. + + Returns: + PromptHandler: An instance of PromptHandler initialized with specific file path and keyword arguments. + """ if self._prompt_handler is None: self._prompt_handler = PromptHandler(self.FILE_PATH, **self.kwargs) return self._prompt_handler def __getattr__(self, key: str): + """ + Custom attribute access to directly retrieve values from kwargs. + + Args: + key (str): The attribute key to look up in kwargs. + + Returns: + Any: The value associated with the key in kwargs. + """ return self.kwargs[key] @staticmethod def get_language_value(languages: dict | list[dict]) -> Any | list[Any]: + """ + Retrieves the value(s) corresponding to the current language context. + + Args: + languages (dict | list[dict]): A dictionary or list of dictionaries containing language-keyed values. + + Returns: + Any | list[Any]: The value or list of values matching the current language setting. + """ if isinstance(languages, list): return [x[G_CONTEXT.language] for x in languages] return languages[G_CONTEXT.language] diff --git a/memory_scope/memory/worker/read/extract_time_worker.py b/memory_scope/memory/worker/read/extract_time_worker.py index 4a0d68d9..60ee0f41 100644 --- a/memory_scope/memory/worker/read/extract_time_worker.py +++ b/memory_scope/memory/worker/read/extract_time_worker.py @@ -9,19 +9,33 @@ from memory_scope.utils.tool_functions import prompt_to_msg class ExtractTimeWorker(MemoryBaseWorker): + """ + A specialized worker class designed to identify and extract time-related information + from text generated by an LLM, translating date-time keywords based on the set language, + and storing this extracted data within a shared context. + """ + EXTRACT_TIME_PATTERN = r'-\s*(\S+)：(\d+)' FILE_PATH: str = __file__ def _run(self): + """ + Executes the primary logic of identifying and extracting time data from an LLM's response. + + This method first checks if the input query contains any datetime keywords. If not, it logs and returns. + It then constructs a prompt with contextual information including formatted timestamps and calls the LLM. + The response is parsed for time-related data using regex, translated via a language-specific key map, + and the resulting time data is stored in the shared context. + """ query, query_timestamp = self.get_context(QUERY_WITH_TS) - # find datetime keyword + # Identify if the query contains datetime keywords contain_datetime = DatetimeHandler.has_time_word(query) if not contain_datetime: self.logger.info(f"contain_datetime={contain_datetime}") return - # prepare prompt + # Prepare the prompt with necessary contextual details query_time_str = DatetimeHandler(dt=query_timestamp).string_format(self.prompt_handler.time_string_format) system_prompt = self.prompt_handler.extract_time_system few_shot = self.prompt_handler.extract_time_few_shot.format(user_name=self.target_name) @@ -29,15 +43,15 @@ class ExtractTimeWorker(MemoryBaseWorker): extract_time_message = prompt_to_msg(system_prompt=system_prompt, few_shot=few_shot, user_query=user_query) self.logger.info(f"extract_time_message={extract_time_message}") - # call llm + # Invoke the LLM to generate a response response = self.generation_model.call(messages=extract_time_message, top_k=self.generation_model_top_k) - # if empty, return + # Handle empty or unsuccessful responses if not response.status or not response.message.content: return response_text = response.message.content - # re-match time info to dict + # Extract time information from the LLM's response using regex extract_time_dict: Dict[str, str] = {} matches = re.findall(self.EXTRACT_TIME_PATTERN, response_text) key_map: dict = self.get_language_value(DATATIME_KEY_MAP) diff --git a/memory_scope/memory/worker/read/fuse_rerank_worker.py b/memory_scope/memory/worker/read/fuse_rerank_worker.py index 1282d6c7..3910e0fa 100644 --- a/memory_scope/memory/worker/read/fuse_rerank_worker.py +++ b/memory_scope/memory/worker/read/fuse_rerank_worker.py @@ -38,27 +38,40 @@ class FuseRerankWorker(MemoryBaseWorker): return match_event_flag, match_msg_flag def _run(self): - # parse input + """ + Executes the reranking process on memory nodes considering their scores, types, and temporal relevance. + + This method performs the following steps: + 1. Retrieves extraction time data and a list of ranked memory nodes from the worker's context. + 2. Reranks nodes based on a combination of their original rank score, type, and temporal alignment with extracted events/messages. + 3. Selects the top-K reranked nodes according to the predefined threshold. + 4. Optionally infuses inferred time information into the content of selected nodes. + 5. Logs reranking details and formats the final list of memories for output. + """ + # Parse input parameters from the worker's context extract_time_dict: Dict[str, str] = self.get_context(EXTRACT_TIME_DICT) memory_node_list: List[MemoryNode] = self.get_memories(RANKED_MEMORY_NODES) + + # Check if memory nodes are available; warn and return if not if not memory_node_list: - self.logger.warning(f"ranked memory nodes is empty!") + self.logger.warning("Ranked memory nodes list is empty.") return - # get reranked nodes + # Perform reranking based on score, type, and time relevance reranked_memory_nodes = [] for node in memory_node_list: + # Skip nodes below the fuse score threshold if node.score_rank < self.fuse_score_threshold: continue - - # memory type ratio + + # Calculate type-based adjustment factor type_ratio: float = self.fuse_ratio_dict.get(node.memory_type, 0.1) - - # memory fuse time ratio + + # Determine time relevance adjustment factor match_event_flag, match_msg_flag = self.match_node_time(extract_time_dict=extract_time_dict, node=node) fuse_time_ratio: float = self.fuse_time_ratio if match_event_flag or match_msg_flag else 1.0 - - # fuse rerank score + + # Apply reranking score adjustments node.score_rerank = node.score_rank * type_ratio * fuse_time_ratio reranked_memory_nodes.append(node) @@ -68,11 +81,13 @@ class FuseRerankWorker(MemoryBaseWorker): key=lambda x: x.score_rerank, reverse=True)[: self.fuse_rerank_top_k] for node in reranked_memory_nodes: + # Log reranking details including flags for event and message matches f_event = int(node.meta_data["match_event_flag"]) f_msg = int(node.meta_data["match_msg_flag"]) - self.logger.info(f"rerank_stage: content={node.content} score={node.score_rerank} " - f"f_event={f_event} f_msg={f_msg}") - + self.logger.info(f"Rerank Stage: Content={node.content}, Score={node.score_rerank}, " + f"Event Flag={f_event}, Message Flag={f_msg}") + + # Infuse time inference if relevant flags are set content = node.content if f_event or f_msg: time_infer = DatetimeHandler.format_time_by_extract_time(extract_time_dict=extract_time_dict, @@ -81,4 +96,5 @@ class FuseRerankWorker(MemoryBaseWorker): content = f"{time_infer}{self.get_language_value(COLON_WORD)}{content}" memories.append(content) + # Set the final list of formatted memories back into the worker's context self.set_context(RESULT, "\n".join(memories)) diff --git a/memory_scope/memory/worker/read/retrieve_memory_worker.py b/memory_scope/memory/worker/read/retrieve_memory_worker.py index 26e1bb29..bef21cef 100644 --- a/memory_scope/memory/worker/read/retrieve_memory_worker.py +++ b/memory_scope/memory/worker/read/retrieve_memory_worker.py @@ -9,9 +9,25 @@ from memory_scope.utils.timer import timer class RetrieveMemoryWorker(MemoryBaseWorker): + """ + Retrieves memories based on specified criteria such as status, type, and timestamp. + Processes these memories concurrently, sorts them by similarity, and logs the activity, + facilitating efficient memory retrieval operations within a given scope. + """ @timer def retrieve_from_observation(self, query: str) -> List[MemoryNode]: + """ + Retrieves memory nodes from observation based on a query, considering active memories + with specific types. If the retrieval limit is not set, an empty list is returned. + + Args: + query (str): The query string used to filter and rank the memory nodes. + + Returns: + List[MemoryNode]: A list of MemoryNode objects that match the query criteria, + sorted by their relevance. Returns an empty list if no retrieval limit is configured. + """ if not self.retrieve_obs_top_k: return [] @@ -21,12 +37,24 @@ class RetrieveMemoryWorker(MemoryBaseWorker): "status": MemoryNodeStatus.ACTIVE.value, "memory_type": [MemoryTypeEnum.OBSERVATION.value, MemoryTypeEnum.OBS_CUSTOMIZED.value], } + # ⭐ Retrieve memories matching the query, filtered by the specified conditions, + # limited to a certain number, and sorted by relevance. return self.memory_store.retrieve_memories(query=query, top_k=self.retrieve_obs_top_k, filter_dict=filter_dict) @timer def retrieve_from_insight_and_profile(self, query: str) -> List[MemoryNode]: + """ + Retrieves memories marked as insights from the database based on a query, filtered by user, target, and set to active status. + + Args: + query (str): The search query to match against the insights. + + Returns: + List[MemoryNode]: A list of MemoryNode objects that match the query criteria, limited by 'retrieve_ins_pf_top_k'. + Returns an empty list if 'retrieve_ins_pf_top_k' is not set. + """ if not self.retrieve_ins_pf_top_k: return [] @@ -36,6 +64,7 @@ class RetrieveMemoryWorker(MemoryBaseWorker): "status": MemoryNodeStatus.ACTIVE.value, "memory_type": MemoryTypeEnum.INSIGHT.value, } + # ⭐ Retrieve insights matching the query, filtered, and limited by top_k return self.memory_store.retrieve_memories(query=query, top_k=self.retrieve_ins_pf_top_k, filter_dict=filter_dict) @@ -56,6 +85,20 @@ class RetrieveMemoryWorker(MemoryBaseWorker): filter_dict=filter_dict) def _run(self): + """ + Executes the main retrieval流程 for memories. It fetches the query from the context, initiates concurrent tasks + to retrieve memories from observations, insights, and expired sources, collects the results, sorts them by + similarity score, logs the details, and finally sets the retrieved memory nodes. + + The method follows these steps: + 1. Retrieves the query from the worker's context. + 2. Submits tasks to asynchronously retrieve memories from various sources. + 3. Gathers the results from all submitted tasks. + 4. Logs the total number of collected memory nodes. + 5. Sorts the memory nodes based on their similarity scores in descending order. + 6. Logs detailed information about each memory node. + 7. Stores the processed memory nodes for further use. + """ query, _ = self.get_context(QUERY_WITH_TS) self.submit_thread_task(self.retrieve_from_observation, query=query) self.submit_thread_task(self.retrieve_from_insight_and_profile, query=query) diff --git a/memory_scope/memory/worker/read/semantic_rank_worker.py b/memory_scope/memory/worker/read/semantic_rank_worker.py index e6d283b1..1bfb8bdd 100644 --- a/memory_scope/memory/worker/read/semantic_rank_worker.py +++ b/memory_scope/memory/worker/read/semantic_rank_worker.py @@ -6,13 +6,32 @@ from memory_scope.scheme.memory_node import MemoryNode class SemanticRankWorker(MemoryBaseWorker): + """ + The `SemanticRankWorker` class processes queries by retrieving memory nodes, + removing duplicates, ranking them based on semantic relevance using a model, + assigning scores, sorting the nodes, and storing the ranked nodes back, + while logging relevant information. + """ def _run(self): + """ + Executes the primary workflow of the SemanticRankWorker which includes: + - Retrieving query and timestamp from context. + - Fetching memory nodes. + - Removing duplicate nodes. + - Ranking nodes semantically. + - Assigning scores to nodes. + - Sorting nodes by score. + - Saving the ranked nodes back with logging. + + If no memory nodes are retrieved or if the ranking model fails, + appropriate warnings are logged. + """ # query query, _ = self.get_context(QUERY_WITH_TS) memory_node_list: List[MemoryNode] = self.get_memories(RETRIEVE_MEMORY_NODES) if not memory_node_list: - self.logger.warning(f"retrieve memory nodes is empty!") + self.logger.warning("Retrieve memory nodes is empty!") return # drop repeated @@ -26,11 +45,16 @@ class SemanticRankWorker(MemoryBaseWorker): # set score for idx, score in response.rank_scores.items(): if idx >= len(memory_node_list): - self.logger.warning(f"idx={idx} exceeds the maximum length of rank_scores!") + self.logger.warning(f"Idx={idx} exceeds the maximum length of rank_scores!") continue memory_node_list[idx].score_rank = score + # sort by score memory_node_list = sorted(memory_node_list, key=lambda n: n.score_rank, reverse=True) + + # log ranked nodes for node in memory_node_list: - self.logger.info(f"rank_stage: content={node.content} score={node.score_rank}") + self.logger.info(f"Rank stage: Content={node.content}, Score={node.score_rank}") + + # save ranked nodes back to memory self.set_memories(RANKED_MEMORY_NODES, memory_node_list, log_repeat=False) diff --git a/memory_scope/memory/worker/read/set_query_worker.py b/memory_scope/memory/worker/read/set_query_worker.py index 4dec4bbb..e58cf5fa 100644 --- a/memory_scope/memory/worker/read/set_query_worker.py +++ b/memory_scope/memory/worker/read/set_query_worker.py @@ -5,17 +5,32 @@ from memory_scope.memory.worker.memory_base_worker import MemoryBaseWorker class SetQueryWorker(MemoryBaseWorker): + """ + The `SetQueryWorker` class is responsible for setting a query and its associated timestamp + into the context, utilizing either provided chat parameters or details from the most recent + chat message. + """ def _run(self): - query = "_" - query_timestamp = int(datetime.datetime.now().timestamp()) + """ + Executes the worker's primary function, which involves determining the query and its + timestamp, then storing these values within the context. + If 'query' is found within `self.chat_kwargs`, it is considered as the query input. + Otherwise, the content of the last message in `self.chat_messages` is used as the query, + along with its creation timestamp. + """ + query = "_" # Default query value + query_timestamp = int(datetime.datetime.now().timestamp()) # Current timestamp as default + + # Check if a specific 'query' has been provided via chat kwargs if "query" in self.chat_kwargs: - # cli test query query = self.chat_kwargs["query"] + # If no explicit query is given, use the content of the latest chat message elif self.chat_messages: query = self.chat_messages[-1].content query_timestamp = self.chat_messages[-1].time_created + # Store the determined query and its timestamp in the context self.set_context(QUERY_WITH_TS, (query, query_timestamp)) diff --git a/memory_scope/memory/worker/summary/get_reflection_subject_worker.py b/memory_scope/memory/worker/summary/get_reflection_subject_worker.py index 1d367a02..4854ea31 100644 --- a/memory_scope/memory/worker/summary/get_reflection_subject_worker.py +++ b/memory_scope/memory/worker/summary/get_reflection_subject_worker.py @@ -12,13 +12,27 @@ from memory_scope.utils.tool_functions import prompt_to_msg class GetReflectionSubjectWorker(MemoryBaseWorker): + """ + A specialized worker class responsible for retrieving unreflected memory nodes, + generating reflection prompts with current insights, invoking an LLM for fresh insights, + parsing the LLM responses, forming new insight nodes, and updating memory statuses accordingly. + """ FILE_PATH: str = __file__ def new_insight_node(self, insight_key: str) -> MemoryNode: - dt_handler = DatetimeHandler() - meta_data = {k: str(v) for k, v in dt_handler.dt_info_dict.items()} + """ + Creates a new MemoryNode for an insight with the given key, enriched with current datetime metadata. - return MemoryNode(user_name=self.user_name, + Args: + insight_key (str): The unique identifier for the insight. + + Returns: + MemoryNode: A new MemoryNode instance representing the insight, marked as new and of type INSIGHT. + """ + dt_handler = DatetimeHandler() + meta_data = {k: str(v) for k, v in dt_handler.dt_info_dict.items()} # ⭐ Prepare metadata with current datetime info + + return MemoryNode(user_name=self.user_name, # ⭐ Populate MemoryNode attributes target_name=self.target_name, meta_data=meta_data, key=insight_key, @@ -26,21 +40,33 @@ class GetReflectionSubjectWorker(MemoryBaseWorker): status=MemoryNodeStatus.NEW.value) def _run(self): + """ + Executes the main logic of reflecting on unaudited memory nodes to derive new insights. + + Steps include: + - Retrieving unaudited memory nodes. + - Checking the count against a threshold to decide whether to proceed. + - Compiling a list of existing insight keys. + - Generating a reflection prompt with system message, few-shot examples, and user queries. + - Calling the language model for new insights. + - Parsing the model's responses for new insight keys. + - Creating new insight nodes and updating the memory status accordingly. + """ not_reflected_nodes: List[MemoryNode] = self.get_memories(NOT_REFLECTED_NODES) insight_nodes: List[MemoryNode] = self.get_memories(INSIGHT_NODES) - # count + # Count unaudited nodes not_reflected_count = len(not_reflected_nodes) if not_reflected_count <= self.reflect_obs_cnt_threshold: - self.logger.info(f"not_reflected_count={not_reflected_count} is not enough, stop.") + self.logger.info(f"not_reflected_count={not_reflected_count} is not enough, stopping process.") self.continue_run = False return - # get profile_keys + # Compile existing insight keys exist_keys: List[str] = [n.key for n in insight_nodes] self.logger.info(f"exist_keys={exist_keys}") - # gen reflect prompt + # Generate reflection prompt components user_query_list = [n.content for n in not_reflected_nodes] system_prompt = self.prompt_handler.get_reflection_subject_system.format( user_name=self.target_name, @@ -51,21 +77,23 @@ class GetReflectionSubjectWorker(MemoryBaseWorker): exist_keys=self.get_language_value(COMMA_WORD).join(exist_keys), user_query="\n".join(user_query_list)) + # Construct and log reflection message reflect_message = prompt_to_msg(system_prompt=system_prompt, few_shot=few_shot, user_query=user_query) self.logger.info(f"reflect_message={reflect_message}") - # # call LLM + # Invoke Language Model for new insights response = self.generation_model.call(messages=reflect_message, top_k=self.generation_model_top_k) - # return if empty + # Handle empty response if not response.status or not response.message.content: return - # parse text & save + # Parse LLM response for new insight keys and update memory new_insight_keys = ResponseTextParser(response.message.content).parse_v2(self.__class__.__name__) if new_insight_keys: for insight_key in new_insight_keys: insight_nodes.append(self.new_insight_node(insight_key)) + # Mark unaudited nodes as reflected for node in not_reflected_nodes: node.obs_reflected = True diff --git a/memory_scope/memory/worker/summary/long_contra_repeat_worker.py b/memory_scope/memory/worker/summary/long_contra_repeat_worker.py index ece8d65c..26a370d4 100644 --- a/memory_scope/memory/worker/summary/long_contra_repeat_worker.py +++ b/memory_scope/memory/worker/summary/long_contra_repeat_worker.py @@ -11,21 +11,51 @@ from memory_scope.utils.tool_functions import prompt_to_msg class LongContraRepeatWorker(MemoryBaseWorker): + """ + Manages and updates memory entries within a conversation scope by identifying + and handling contradictions or redundancies. It extends the base MemoryBaseWorker + to provide specialized functionality for long conversations with potential + contradictory or repetitive statements. + """ FILE_PATH: str = __file__ def retrieve_similar_content(self, node: MemoryNode) -> (MemoryNode, List[MemoryNode]): + """ + Retrieves memory nodes with content similar to the given node, filtering by user, target, status, and memory type. + Only returns nodes whose similarity score meets or exceeds the predefined threshold. + + Args: + node (MemoryNode): The reference node used to find similar content in memory. + + Returns: + Tuple[MemoryNode, List[MemoryNode]]: A tuple containing the original node and a list of similar nodes + that passed the similarity threshold. + """ filter_dict = { "user_name": self.user_name, "target_name": self.target_name, "status": MemoryNodeStatus.ACTIVE.value, "memory_type": [MemoryTypeEnum.OBSERVATION.value, MemoryTypeEnum.OBS_CUSTOMIZED.value] } + # Retrieve memories similar to the node's content, limited by top_k and filtered by filter_dict retrieve_nodes = self.memory_store.retrieve_memories(query=node.content, top_k=self.long_contra_repeat_top_k, filter_dict=filter_dict) + # Filter retrieved nodes based on the similarity threshold return node, [n for n in retrieve_nodes if n.score_similar >= self.long_contra_repeat_threshold] def _run(self): + """ + Executes the primary routine of the LongContraRepeatWorker. This involves: + 1. Retrieving not updated memory nodes. + 2. Gathering similar content for these nodes. + 3. Organizing observed nodes and generating a prompt for the language model. + 4. Calling the language model to process the prompt and receive a response. + 5. Parsing the model's response to update memory node statuses. + 6. Saving the modified memory nodes. + + The process helps in maintaining conversation coherence by resolving contradictions and redundancies. + """ not_updated_nodes: List[MemoryNode] = self.get_memories(NOT_UPDATED_NODES) for node in not_updated_nodes: self.submit_thread_task(fn=self.retrieve_similar_content, node=node) @@ -59,20 +89,20 @@ class LongContraRepeatWorker(MemoryBaseWorker): user_query=user_query) self.logger.info(f"long_contra_repeat_message={long_contra_repeat_message}") - # call llm + # Invokes the language model for processing the constructed prompt response = self.generation_model.call(messages=long_contra_repeat_message, top_k=self.generation_model_top_k) - # return if empty + # Handles the case where the model's response is empty if not response or not response.message.content: return - # parse text + # Parses the model's response text to identify updates for memory nodes idx_obs_info_list = ResponseTextParser(response.message.content).parse_v1(self.__class__.__name__) if len(idx_obs_info_list) <= 0: self.logger.warning("idx_obs_info_list is empty!") return - # add merged obs + # Processes parsed information to update memory nodes' statuses merge_obs_nodes: List[MemoryNode] = [] for idx_obs_info in idx_obs_info_list: if not idx_obs_info: diff --git a/memory_scope/memory/worker/summary/update_insight_worker.py b/memory_scope/memory/worker/summary/update_insight_worker.py index a3a367d5..ad7ef1e8 100644 --- a/memory_scope/memory/worker/summary/update_insight_worker.py +++ b/memory_scope/memory/worker/summary/update_insight_worker.py @@ -11,34 +11,60 @@ from memory_scope.utils.tool_functions import prompt_to_msg class UpdateInsightWorker(MemoryBaseWorker): + """ + This class is responsible for updating insights in a memory system. It filters insight nodes + based on their association with observed nodes, utilizes a ranking model to prioritize them, + generates refreshed insights via an LLM, and manages node statuses and content updates, + incorporating features for concurrent execution and logging. + """ FILE_PATH: str = __file__ def filter_obs_nodes(self, insight_node: MemoryNode, obs_nodes: List[MemoryNode]) -> (MemoryNode, List[MemoryNode], float): + """ + Filters observed nodes based on their relevance to a given insight node using a ranking model. + + Args: + insight_node (MemoryNode): The insight node used as the basis for filtering. + obs_nodes (List[MemoryNode]): A list of observed nodes to be filtered. + + Returns: + tuple: A tuple containing: + - The original insight node. + - A list of filtered observed nodes that are relevant to the insight node. + - The maximum relevance score among the filtered nodes. + """ max_score: float = 0 filtered_nodes: List[MemoryNode] = [] + # Check if insight node key or value is empty and log a warning if not insight_node.key or not insight_node.value: self.logger.warning(f"insight_key={insight_node.key} insight_value={insight_node.value} is empty!") return insight_node, filtered_nodes, max_score + # Call the ranking model to get scores for each observed node's content against the insight key response = self.rank_model.call(query=insight_node.key, documents=[x.content for x in obs_nodes]) if not response.status: return insight_node, filtered_nodes, max_score - # find nodes related to query + # Iterate over the ranked scores to filter nodes for index, score in response.rank_scores.items(): node = obs_nodes[index] + # Determine if the node should be kept based on the threshold keep_flag = score >= self.update_insight_threshold if keep_flag: filtered_nodes.append(node) max_score = max(max_score, score) + # Log information about each node's processing self.logger.info(f"insight_key={insight_node.key} insight_value={insight_node.value} " f"score={score} keep_flag={keep_flag}") + # Warn if no nodes were filtered if not filtered_nodes: self.logger.warning(f"update_insight={insight_node.key} filtered_nodes is empty!") + + # Return the original insight node, the list of filtered nodes, and the highest score return insight_node, filtered_nodes, max_score def update_insight_node(self, insight_node: MemoryNode, insight_value: str): @@ -56,10 +82,20 @@ class UpdateInsightWorker(MemoryBaseWorker): return insight_node def update_insight(self, insight_node: MemoryNode, filtered_nodes: List[MemoryNode]) -> MemoryNode: - self.logger.info(f"update_insight insight_key={insight_node.key} insight_value={insight_node.value} " - f"doc.size={len(filtered_nodes)}") + """ + Updates the insight value of a given MemoryNode based on the context from a list of filtered MemoryNodes. - # gen prompt + Args: + insight_node (MemoryNode): The MemoryNode whose insight value needs to be updated. + filtered_nodes (List[MemoryNode]): A list of MemoryNodes used as context for updating the insight. + + Returns: + MemoryNode: The updated MemoryNode with potentially revised insight value. + """ + self.logger.info(f"Updating insight for key={insight_node.key}, value={insight_node.value}, " + f"with {len(filtered_nodes)} documents considered.") + + # Generate the prompt for updating insight user_query_list = [n.content for n in filtered_nodes] system_prompt = self.prompt_handler.update_insight_system.foramt(user_name=self.target_name) few_shot = self.prompt_handler.update_insight_few_shot.foramt(user_name=self.target_name) @@ -68,13 +104,14 @@ class UpdateInsightWorker(MemoryBaseWorker): insight_key=insight_node.key, insight_key_value=insight_node.key + self.get_language_value(COLON_WORD) + insight_node.value) + # Construct the message for LLM interaction update_insight_message = prompt_to_msg(system_prompt=system_prompt, few_shot=few_shot, user_query=user_query) - self.logger.info(f"update_insight_message={update_insight_message}") + self.logger.info(f"Generated insight update message: {update_insight_message}") - # call LLM + # Call the Language Model for insight update response = self.generation_model.call(messages=update_insight_message, top_k=self.generation_model_top_k) - # return if empty + # Handle empty or invalid responses if not response.status or not response.message.content: return insight_node @@ -101,14 +138,29 @@ class UpdateInsightWorker(MemoryBaseWorker): return insight_node def _run(self): + """ + Executes the main routine of the UpdateInsightWorker. This involves filtering and updating insight nodes + based on their association with observed nodes. It processes nodes in batches, selects the top nodes + according to a scoring mechanism, and then initiates tasks to update these insights using an LLM. Finally, + it updates the status of processed nodes and gathers the results from all threads. + + Steps include: + 1. Retrieve lists of insight, not updated, and not reflected nodes from memory. + 2. Filter and process active insight nodes with respective not updated nodes. + 3. Sort processed results by score and select the top N. + 4. Submit tasks to update insights for the selected nodes. + 5. Gather the results of all update tasks. + 6. Mark processed nodes as updated in memory. + """ insight_nodes: List[MemoryNode] = self.get_memories(INSIGHT_NODES) not_updated_nodes: List[MemoryNode] = self.get_memories(NOT_UPDATED_NODES) not_reflected_nodes: List[MemoryNode] = self.get_memories(NOT_REFLECTED_NODES) if not insight_nodes: - self.logger.warning("insight_nodes is empty, stop.") + self.logger.warning("insight_nodes is empty, stopping processing.") return + # Process active insight nodes with corresponding not updated nodes for node in insight_nodes: if node.status == MemoryNodeStatus.ACTIVE.value: self.submit_thread_task(fn=self.filter_obs_nodes, @@ -128,11 +180,11 @@ class UpdateInsightWorker(MemoryBaseWorker): result_list.append(result) result_sorted = sorted(result_list, key=lambda x: x[2], reverse=True)[: self.update_insight_max_thread] - # submit llm update task + # Submit tasks to update insights for the top nodes for insight_node, filtered_nodes, _ in result_sorted: self.submit_thread_task(fn=self.update_insight, insight_node=insight_node, filtered_nodes=filtered_nodes) - # get result + # Gather the final results from all update tasks self.gather_thread_result() for node in not_updated_nodes: diff --git a/memory_scope/memory/worker/write/contra_repeat_worker.py b/memory_scope/memory/worker/write/contra_repeat_worker.py index 1f4e6e03..94d81601 100644 --- a/memory_scope/memory/worker/write/contra_repeat_worker.py +++ b/memory_scope/memory/worker/write/contra_repeat_worker.py @@ -1,5 +1,4 @@ from typing import List - from memory_scope.constants.common_constants import NEW_OBS_NODES, NEW_OBS_WITH_TIME_NODES, MERGE_OBS_NODES, TODAY_NODES from memory_scope.constants.language_constants import NONE_WORD, CONTRADICTORY_WORD, INCLUDED_WORD from memory_scope.enumeration.memory_status_enum import MemoryNodeStatus @@ -10,9 +9,34 @@ from memory_scope.utils.tool_functions import prompt_to_msg class ContraRepeatWorker(MemoryBaseWorker): + """ + The `ContraRepeatWorker` class specializes in processing memory nodes to identify and handle + contradictory and repetitive information. It extends the base functionality of `MemoryBaseWorker`. + + Responsibilities: + - Collects observation nodes from various memory categories. + - Constructs a prompt with these observations for language model analysis. + - Parses the model's response to detect contradictions or redundancies. + - Adjusts the status of memory nodes based on the analysis. + - Persists the updated node statuses back into memory. + """ FILE_PATH: str = __file__ def _run(self): + """ + Executes the primary routine of the ContraRepeatWorker which involves fetching memory nodes, + constructing a prompt, querying a language model, parsing the response to identify nodes for merging, + updating node statuses, and saving the updated nodes back to memory. + + Steps: + 1. Retrieves new observation nodes and nodes observed on the current day. + 2. Optionally combines today's nodes with the new ones, sorts, and limits the list by a predefined count. + 3. Constructs a prompt using the combined nodes, system prompt, and a few-shot example. + 4. Queries a language model with the constructed prompt. + 5. Parses the model's response to identify nodes to merge or exclude based on contradiction or redundancy. + 6. Updates the status of nodes accordingly. + 7. Persists the changes back to memory storage. + """ all_obs_nodes: List[MemoryNode] = self.get_memories([NEW_OBS_NODES, NEW_OBS_WITH_TIME_NODES]) if not all_obs_nodes: self.logger.info("all_obs_nodes is empty!") @@ -58,7 +82,7 @@ class ContraRepeatWorker(MemoryBaseWorker): if not obs_content_list: continue - # [6, skipping classes] + # Expecting a pair [index, flag] if len(obs_content_list) != 2: self.logger.warning(f"obs_content_list={obs_content_list} is invalid!") continue diff --git a/memory_scope/memory/worker/write/get_observation_with_time_worker.py b/memory_scope/memory/worker/write/get_observation_with_time_worker.py index 98dfd0c4..5e1a3459 100644 --- a/memory_scope/memory/worker/write/get_observation_with_time_worker.py +++ b/memory_scope/memory/worker/write/get_observation_with_time_worker.py @@ -1,5 +1,4 @@ from typing import List - from memory_scope.constants.common_constants import NEW_OBS_WITH_TIME_NODES from memory_scope.constants.language_constants import COLON_WORD from memory_scope.memory.worker.write.get_observation_worker import GetObservationWorker @@ -9,30 +8,66 @@ from memory_scope.utils.tool_functions import prompt_to_msg class GetObservationWithTimeWorker(GetObservationWorker): + """ + A specialized worker class that extends GetObservationWorker functionality to handle + retrieval of observations which include associated timestamp information from chat messages. + """ FILE_PATH: str = __file__ OBS_STORE_KEY: str = NEW_OBS_WITH_TIME_NODES def filter_messages(self) -> List[Message]: + """ + Filters the chat messages to only include those which contain time-related keywords. + + Returns: + List[Message]: A list of filtered messages that mention time. + """ filter_messages = [] for msg in self.chat_messages: + # Checks if the message content has any time reference words if DatetimeHandler.has_time_word(query=msg.content): filter_messages.append(msg) return filter_messages def build_message(self, filter_messages: List[Message]) -> List[Message]: + """ + Constructs a message for obtaining observations with timestamps based on filtered chat messages. + + This method processes each filtered message to append a timestamp formatted per the specified format. + It then organizes these timestamped queries into a structured prompt that includes a system prompt, + few-shot examples, and the concatenated user queries, tailored to a target individual with a given language setting. + + Args: + filter_messages (List[Message]): A list of Message objects that have been filtered for processing. + + Returns: + List[Message]: A list containing the newly constructed Message object for further interaction. + """ user_query_list = [] for i, msg in enumerate(filter_messages): + # Create a DatetimeHandler instance for each message's timestamp and format it dt_handler = DatetimeHandler(dt=msg.time_created) dt = dt_handler.string_format(self.prompt_handler.time_string_format) + # Append formatted timestamp-query pairs to the user_query_list user_query_list.append(f"{i} {dt} {self.target_name}{self.get_language_value(COLON_WORD)}{msg.content}") + # Construct the system prompt with the count of observations system_prompt = self.prompt_handler.get_observation_with_time_system.format(num_obs=len(user_query_list), user_name=self.target_name) + + # Retrieve the few-shot examples for the prompt few_shot = self.prompt_handler.get_observation_with_time_few_shot.format(user_name=self.target_name) + + # Format the user query section with the concatenated list of timestamped queries user_query = self.prompt_handler.get_observation_with_time_user_query.format( user_query="\n".join(user_query_list), user_name=self.target_name) + # Assemble the final message for observation retrieval obtain_obs_message = prompt_to_msg(system_prompt=system_prompt, few_shot=few_shot, user_query=user_query) + + # Log the constructed message for debugging purposes self.logger.info(f"obtain_obs_message={obtain_obs_message}") + + # Return the newly created message return obtain_obs_message diff --git a/memory_scope/memory/worker/write/get_observation_worker.py b/memory_scope/memory/worker/write/get_observation_worker.py index 72ed081c..4360e9d4 100644 --- a/memory_scope/memory/worker/write/get_observation_worker.py +++ b/memory_scope/memory/worker/write/get_observation_worker.py @@ -51,21 +51,55 @@ class GetObservationWorker(MemoryBaseWorker): return filter_messages def build_message(self, filter_messages: List[Message]) -> List[Message]: + """ + Constructs a formatted message for observation based on input messages, incorporating system prompts, + few-shot examples, and user queries. + + Args: + filter_messages (List[Message]): A list of messages filtered for observation processing. + + Returns: + List[Message]: A list containing the constructed message ready for observation. + """ user_query_list = [] for i, msg in enumerate(filter_messages): + # Construct each user query item with index, target name, and message content user_query_list.append(f"{i} {self.target_name}{self.get_language_value(COLON_WORD)}{msg.content}") + # Format the system prompt with the number of observations and target name system_prompt = self.prompt_handler.get_observation_system.format(num_obs=len(user_query_list), user_name=self.target_name) + + # Incorporate few-shot examples into the prompt with the target name few_shot = self.prompt_handler.get_observation_few_shot.format(user_name=self.target_name) + + # Assemble the user query part of the prompt with the list of formatted user queries user_query = self.prompt_handler.get_observation_user_query.format(user_query="\n".join(user_query_list), user_name=self.target_name) + # Combine system prompt, few-shot, and user query into a single message for obtaining observations obtain_obs_message = prompt_to_msg(system_prompt=system_prompt, few_shot=few_shot, user_query=user_query) + + # Log the constructed observation message self.logger.info(f"obtain_obs_message={obtain_obs_message}") + + # Return the processed message(s) for further steps in the observation workflow return obtain_obs_message def _run(self): + """ + Processes chat messages to extract observations, inferring timestamps and content relevance, + and stores the extracted information as MemoryNode objects within the conversation memory. + + Steps: + 1. Filters messages based on predefined criteria. + 2. Constructs a message for the language model to generate observations. + 3. Calls the language model to predict observation details. + 4. Parses the model's response to extract observation lists. + 5. Validates and structures each observed event into MemoryNode objects. + 6. Stores these MemoryNodes in the conversation memory under a specific key. + """ + # Filters messages and constructs an input message for the language model filter_messages = self.filter_messages() if not filter_messages: self.logger.warning("get obs filter_messages is empty!") @@ -73,36 +107,37 @@ class GetObservationWorker(MemoryBaseWorker): obtain_obs_message = self.build_message(filter_messages) - # call LLM + # Generates observations using the language model response = self.generation_model.call(messages=obtain_obs_message, top_k=self.generation_model_top_k) - - # return if empty if not response.status or not response.message.content: return + response_text = response.message.content - # parse text + # Parses the generated text to extract observation indices, times, contents, and keywords idx_obs_list = ResponseTextParser(response_text).parse_v1(self.__class__.__name__) if len(idx_obs_list) <= 0: self.logger.warning("idx_obs_list is empty!") return - # gene new obs nodes + # Processes each extracted observation to create MemoryNode objects new_obs_nodes: List[MemoryNode] = [] for obs_content_list in idx_obs_list: if not obs_content_list: continue - # [1, In June 2022, the user will travel to Hangzhou for tourism, tourism] + # Expected format: [index, time_inference, observation_content, keywords] if len(obs_content_list) != 4: self.logger.warning(f"obs_content_list={obs_content_list} is invalid!") continue idx, time_infer, obs_content, keywords = obs_content_list + # Skips processing if content indicates no meaningful observation if obs_content in self.get_language_value([NONE_WORD, REPEATED_WORD]): continue + # Validates index format if not idx.isdigit(): self.logger.warning(f"idx={idx} is invalid!") continue @@ -110,15 +145,17 @@ class GetObservationWorker(MemoryBaseWorker): if time_infer == self.get_language_value(NONE_WORD): time_infer = "" - # index number needs to be corrected to -1 + # Adjusts index to zero-based and checks validity against filtered messages idx = int(idx) - 1 if idx >= len(filter_messages): self.logger.warning(f"idx={idx} is invalid! filter_messages.size={len(filter_messages)}") continue + # Creates a MemoryNode for the validated observation and adds it to the list new_obs_nodes.append(self.add_observation(message=filter_messages[idx], time_infer=time_infer, obs_content=obs_content, keywords=keywords)) + # Stores the extracted and structured observations in the conversation memory self.set_memories(self.OBS_STORE_KEY, new_obs_nodes) diff --git a/memory_scope/memory/worker/write/info_filter_worker.py b/memory_scope/memory/worker/write/info_filter_worker.py index 44fd52eb..97c04c50 100644 --- a/memory_scope/memory/worker/write/info_filter_worker.py +++ b/memory_scope/memory/worker/write/info_filter_worker.py @@ -10,11 +10,26 @@ from memory_scope.utils.tool_functions import prompt_to_msg class InfoFilterWorker(MemoryBaseWorker): """ - This worker will filter and modify `self.chat_messages`, preserving only the messages that contain important information. + This worker filters and modifies the chat message history (`self.chat_messages`) by retaining only the messages + that include significant information. It then constructs a prompt from these filtered messages, utilizes an AI + model to process this prompt, parses the AI's generated response to allocate scores, and ultimately retains + messages in `self.chat_messages` based on these assigned scores. """ FILE_PATH: str = __file__ def _run(self): + """ + Filters user messages in the chat, generates a prompt incorporating these messages, + utilizes an LLM to process the prompt, parses the LLM's response to score each message, + and updates `self.chat_messages` to only include messages with designated scores. + + This method executes the following steps: + 1. Filters out non-user messages and truncates long messages. + 2. Constructs a prompt with user messages for LLM input. + 3. Calls the LLM model with the constructed prompt. + 4. Parses the LLM's response to extract message scores. + 5. Retains messages in `self.chat_messages` based on their scores. + """ # filter user msg info_messages: List[Message] = [] for msg in self.chat_messages: diff --git a/memory_scope/memory/worker/write/load_memory_worker.py b/memory_scope/memory/worker/write/load_memory_worker.py index cf209118..5704677b 100644 --- a/memory_scope/memory/worker/write/load_memory_worker.py +++ b/memory_scope/memory/worker/write/load_memory_worker.py @@ -87,9 +87,17 @@ class LoadMemoryWorker(MemoryBaseWorker): self.set_memories(TODAY_NODES, nodes) def _run(self): - mock_query = "-" + """ + Initiates asynchronous tasks to retrieve various types of memory data including + not reflected, not updated, insights, and data from today. After submitting all tasks, + it waits for their completion by calling `gather_thread_result`. + + This method serves as the controller for data retrieval operations, enhancing efficiency + by handling tasks concurrently. + """ + mock_query = "-" # Placeholder query self.submit_thread_task(self.retrieve_not_reflected_memory, query=mock_query) self.submit_thread_task(self.retrieve_not_updated_memory, query=mock_query) self.submit_thread_task(self.retrieve_insight_memory, query=mock_query) self.submit_thread_task(self.retrieve_today_memory) - self.gather_thread_result() + self.gather_thread_result() # Waits for all submitted tasks to complete diff --git a/memory_scope/models/dummy_generation_model.py b/memory_scope/models/dummy_generation_model.py index d4a0acd6..c56c5549 100644 --- a/memory_scope/models/dummy_generation_model.py +++ b/memory_scope/models/dummy_generation_model.py @@ -11,14 +11,34 @@ from memory_scope.scheme.model_response import ModelResponse, ModelResponseGen class DummyGenerationModel(BaseModel): + """ + The `DummyGenerationModel` class serves as a placeholder model for generating responses. + It processes input prompts or sequences of messages, adapting them into a structure compatible + with chat interfaces. It also facilitates the generation of mock (dummy) responses for testing, + supporting both immediate and streamed output. + """ m_type: ModelEnum = ModelEnum.GENERATION_MODEL class DummyModel: + """ + An inner class representing the dummy model placeholder. + """ pass MODEL_REGISTRY.register("dummy_generation", DummyModel) def before_call(self, **kwargs): + """ + Prepares the input data before making a call to the model's generate function. + Accepts either a 'prompt' or a list of 'messages'. If both are provided or missing, + a RuntimeError is raised. Transforms the input into a standardized format for processing. + + Args: + **kwargs: Arbitrary keyword arguments including 'prompt' or 'messages'. + + Raises: + RuntimeError: If neither 'prompt' nor 'messages' is provided, or both are provided. + """ prompt: str = kwargs.pop("prompt", "") messages: List[Message] | List[dict] = kwargs.pop("messages", []) @@ -30,12 +50,27 @@ class DummyGenerationModel(BaseModel): else: self.data = {"messages": [ChatMessage(role=msg.role, content=msg.content) for msg in messages]} else: - raise RuntimeError("prompt and messages is both empty!") + raise RuntimeError("Both 'prompt' and 'messages' are empty!") def after_call(self, model_response: ModelResponse, stream: bool = False, **kwargs) -> ModelResponse | ModelResponseGen: + """ + Processes the model's response post-call, optionally streaming the output or returning it as a whole. + + This method modifies the input `model_response` by resetting its message content and, based on the `stream` + parameter, either yields the response in a generated stream or returns the complete response directly. + + Args: + model_response (ModelResponse): The initial response object to be processed. + stream (bool, optional): Flag indicating whether to stream the response. Defaults to False. + **kwargs: Additional keyword arguments (not used in this implementation). + + Returns: + ModelResponse | ModelResponseGen: If `stream` is True, a generator yielding updated `ModelResponse` objects; + otherwise, a modified `ModelResponse` object with the complete content. + """ model_response.message = Message(role=MessageRoleEnum.ASSISTANT, content="") call_result = ["-" for _ in range(10)] @@ -44,20 +79,43 @@ class DummyGenerationModel(BaseModel): for delta in call_result: model_response.message.content += delta model_response.delta = delta - time.sleep(0.1) + time.sleep(0.1) # ⭐ Introduce a delay to simulate streaming yield model_response return gen() else: - model_response.message.content = "".join(call_result) + model_response.message.content = "".join(call_result) # ⭐ Concatenate results for non-streaming return model_response def _call(self, stream: bool = False, **kwargs) -> ModelResponse | ModelResponseGen: + """ + Generates a dummy response based on the input data, supporting both immediate + and streamed response types. + + Args: + stream (bool, optional): If True, indicates the response should be generated + in a streaming manner. Defaults to False. + **kwargs: Additional keyword arguments not used in this dummy implementation. + + Returns: + Union[ModelResponse, ModelResponseGen]: A dummy response object or a generator + object capable of streaming responses. + """ assert "prompt" in self.data or "messages" in self.data results = ModelResponse(m_type=self.m_type) return results async def _async_call(self, **kwargs) -> ModelResponse: + """ + Asynchronous version of `_call`, providing the same functionality but designed + to be used in asynchronous contexts. + + Args: + **kwargs: Additional keyword arguments not used in this dummy implementation. + + Returns: + ModelResponse: A dummy response object suitable for asynchronous use. + """ assert "prompt" in self.data or "messages" in self.data results = ModelResponse(m_type=self.m_type) return results diff --git a/memory_scope/models/llama_index_embedding_model.py b/memory_scope/models/llama_index_embedding_model.py index bbecbce9..eaa51b49 100644 --- a/memory_scope/models/llama_index_embedding_model.py +++ b/memory_scope/models/llama_index_embedding_model.py @@ -8,8 +8,23 @@ from memory_scope.scheme.model_response import ModelResponse class LlamaIndexEmbeddingModel(BaseModel): + """ + Manages text embeddings utilizing the DashScopeEmbedding within the LlamaIndex framework, + facilitating embedding operations for both sync and async modes, inheriting from BaseModel. + """ m_type: ModelEnum = ModelEnum.EMBEDDING_MODEL + @classmethod + def register_model(cls, model_name: str, model_class: type): + """ + Registers a new embedding model class with the model registry. + + Args: + model_name (str): The name to register the model under. + model_class (type): The class of the model to register. + """ + MODEL_REGISTRY.register(model_name, model_class) + MODEL_REGISTRY.register("dashscope_embedding", DashScopeEmbedding) def before_call(self, **kwargs): @@ -34,14 +49,32 @@ class LlamaIndexEmbeddingModel(BaseModel): def _call(self, **kwargs) -> ModelResponse: """ - :param kwargs: - :return: + Executes a synchronous call to generate embeddings for the input data. + + This method utilizes the `get_text_embedding_batch` method of the encapsulated model, + passing the processed data from `self.data`. The result is then packaged into a + `ModelResponse` object with the model type specified by `self.m_type`. + + Args: + **kwargs: Additional keyword arguments that might be used in the embedding process. + + Returns: + ModelResponse: An object containing the embedding results and the model type. """ return ModelResponse(m_type=self.m_type, raw=self.model.get_text_embedding_batch(**self.data)) async def _async_call(self, **kwargs) -> ModelResponse: """ - :param kwargs: - :return: + Executes an asynchronous call to generate embeddings for the input data. + + Similar to `_call`, but uses the asynchronous `aget_text_embedding_batch` method + of the model. It handles the input data asynchronously and packages the result + within a `ModelResponse` instance. + + Args: + **kwargs: Additional keyword arguments for the embedding process, if any. + + Returns: + ModelResponse: An object encapsulating the embedding output and the model's type. """ return ModelResponse(m_type=self.m_type, raw=await self.model.aget_text_embedding_batch(**self.data)) diff --git a/memory_scope/models/llama_index_generation_model.py b/memory_scope/models/llama_index_generation_model.py index 462f88c7..cebd9f85 100644 --- a/memory_scope/models/llama_index_generation_model.py +++ b/memory_scope/models/llama_index_generation_model.py @@ -11,11 +11,32 @@ from memory_scope.scheme.model_response import ModelResponse, ModelResponseGen class LlamaIndexGenerationModel(BaseModel): + """ + This class represents a generation model within the LlamaIndex framework, + capable of processing input prompts or message histories, selecting an appropriate + language model service from a registry, and generating text responses, with support + for both streaming and non-streaming modes. It encapsulates logic for formatting + these interactions within the context of a memory scope management system. + """ + m_type: ModelEnum = ModelEnum.GENERATION_MODEL MODEL_REGISTRY.register("dashscope_generation", DashScope) def before_call(self, **kwargs): + """ + Prepares the input data before making a call to the language model. + It accepts either a 'prompt' directly or a list of 'messages'. + If 'prompt' is provided, it sets the data accordingly. + If 'messages' are provided, it constructs a list of ChatMessage objects from the list. + Raises an error if neither 'prompt' nor 'messages' are supplied. + + Args: + **kwargs: Arbitrary keyword arguments including 'prompt' and 'messages'. + + Raises: + RuntimeError: When both 'prompt' and 'messages' inputs are not provided. + """ prompt: str = kwargs.pop("prompt", "") messages: List[Message] | List[dict] = kwargs.pop("messages", []) @@ -27,7 +48,7 @@ class LlamaIndexGenerationModel(BaseModel): else: self.data = {"messages": [ChatMessage(role=msg.role, content=msg.content) for msg in messages]} else: - raise RuntimeError("prompt and messages is both empty!") + raise RuntimeError("prompt and messages are both empty!") def after_call(self, model_response: ModelResponse, @@ -71,6 +92,20 @@ class LlamaIndexGenerationModel(BaseModel): return results async def _async_call(self, **kwargs) -> ModelResponse: + """ + Asynchronously calls the language model with the provided prompt or message history, + and packages the raw response into a ModelResponse object. + + This method checks if the input data contains a 'prompt' or 'messages' key to decide + which method to call on the model instance. It uses 'acomplete' for simple prompts and + 'achat' for chat-based message histories. + + Args: + **kwargs: Additional keyword arguments that might be used in the model call. + + Returns: + ModelResponse: An object containing the raw response from the language model. + """ assert "prompt" in self.data or "messages" in self.data results = ModelResponse(m_type=self.m_type) @@ -78,5 +113,6 @@ class LlamaIndexGenerationModel(BaseModel): response = await self.model.acomplete(**self.data) else: response = await self.model.achat(**self.data) + results.raw = response return results diff --git a/memory_scope/models/llama_index_rank_model.py b/memory_scope/models/llama_index_rank_model.py index e5cefa12..dc88621f 100644 --- a/memory_scope/models/llama_index_rank_model.py +++ b/memory_scope/models/llama_index_rank_model.py @@ -10,11 +10,24 @@ from memory_scope.scheme.model_response import ModelResponse class LlamaIndexRankModel(BaseModel): + """ + The LlamaIndexRankModel class is designed to rerank documents according to their relevance + to a provided query, utilizing the DashScope Rerank model. It transforms document lists + and queries into a compatible format for ranking, manages the ranking process, and allocates + rank scores to individual documents. + """ m_type: ModelEnum = ModelEnum.RANK_MODEL MODEL_REGISTRY.register("dashscope_rank", DashScopeRerank) def before_call(self, **kwargs) -> None: + """ + Prepares necessary data before the ranking call by extracting the query and documents, + ensuring they are valid, and initializing nodes with dummy scores. + + Args: + **kwargs: Keyword arguments containing 'query' and 'documents'. + """ query: str = kwargs.pop("query", "") documents: List[str] = kwargs.pop("documents", []) if isinstance(documents, str): @@ -22,13 +35,24 @@ class LlamaIndexRankModel(BaseModel): assert query and documents, f"query or documents is empty! query={query}, documents={len(documents)}" - # using -1.0 as dummy scores + # Using -1.0 as dummy scores nodes = [NodeWithScore(node=Node(text=doc), score=-1.0) for doc in documents] self._get_documents_mapping(documents) self.data = {"nodes": nodes, "query_str": query} def after_call(self, model_response: ModelResponse, **kwargs) -> ModelResponse: + """ + Processes the model response post-ranking, assigning calculated rank scores to each document + based on their index in the original document list. + + Args: + model_response (ModelResponse): The initial response from the ranking model. + **kwargs: Additional keyword arguments (unused). + + Returns: + ModelResponse: Updated response with rank scores assigned to documents. + """ if not model_response.rank_scores: model_response.rank_scores = {} @@ -39,12 +63,36 @@ class LlamaIndexRankModel(BaseModel): return model_response def _call(self, **kwargs) -> ModelResponse: + """ + Executes the ranking process by passing prepared data to the model's postprocessing method. + + Args: + **kwargs: Keyword arguments (unused). + + Returns: + ModelResponse: A response object encapsulating the ranked nodes. + """ return ModelResponse(m_type=self.m_type, raw=self.model.postprocess_nodes(**self.data)) async def _async_call(self, **kwargs) -> ModelResponse: + """ + Asynchronous wrapper for the `_call` method, maintaining the same functionality. + + Args: + **kwargs: Keyword arguments (unused). + + Returns: + ModelResponse: A response object encapsulating the ranked nodes. + """ return self._call(**kwargs) def _get_documents_mapping(self, documents): + """ + Generates a mapping of each document to its index within the provided document list. + + Args: + documents (List[str]): The list of documents. + """ self.documents_map = {} for idx, doc in enumerate(documents): self.documents_map[doc] = idx diff --git a/memory_scope/scheme/memory_node.py b/memory_scope/scheme/memory_node.py index b4072e66..96ccd9d2 100644 --- a/memory_scope/scheme/memory_node.py +++ b/memory_scope/scheme/memory_node.py @@ -6,6 +6,10 @@ from pydantic import Field, BaseModel class MemoryNode(BaseModel): + """ + Represents a memory node with comprehensive attributes to store memory information including unique ID, user details, + content, metadata, scoring metrics, and status indicators. Automatically handles timestamp conversion to date format during initialization. + """ memory_id: str = Field(default_factory=lambda: uuid4().hex, description="unique id for memory") user_name: str = Field("", description="the user who owns the memory") diff --git a/memory_scope/scheme/message.py b/memory_scope/scheme/message.py index 43308850..e6196d4e 100644 --- a/memory_scope/scheme/message.py +++ b/memory_scope/scheme/message.py @@ -5,15 +5,26 @@ from pydantic import Field, BaseModel class Message(BaseModel): + """ + Represents a structured message object with details about the sender, content, and metadata. + + Attributes: + role (str): The role of the message sender (e.g., 'user', 'assistant', 'system'). + role_name (str): Optional name associated with the role of the message sender. + content (str): The actual content or text of the message. + time_created (int): Timestamp indicating when the message was created. + memorized (bool): Flag to indicate if the message has been saved or remembered. + meta_data (Dict[str, str]): Additional data or context attached to the message. + """ role: str = Field(..., description="The role of the message sender (user, assistant, system)") - role_name: str = Field("", description="role name") + role_name: str = Field("", description="Name describing the role of the message sender") - content: str = Field(..., description="The body of the message") + content: str = Field(..., description="The primary content of the message") time_created: int = Field(int(datetime.datetime.now().timestamp()), - description="Timestamp when the message was created") + description="Timestamp marking the message creation time") - memorized: bool = Field(False, description="indicate whether message is memorized") + memorized: bool = Field(False, description="Indicates if the message is flagged for memory retention") - meta_data: Dict[str, str] = Field({}, description="meta data for msg") + meta_data: Dict[str, str] = Field({}, description="Supplementary data attached to the message") diff --git a/memory_scope/storage/base_memory_store.py b/memory_scope/storage/base_memory_store.py index 9b041ba9..cf2ae698 100644 --- a/memory_scope/storage/base_memory_store.py +++ b/memory_scope/storage/base_memory_store.py @@ -5,6 +5,10 @@ from memory_scope.scheme.memory_node import MemoryNode class BaseMemoryStore(metaclass=ABCMeta): + """ + An abstract base class defining the interface for a memory store which handles memory nodes. + It outlines essential operations like retrieval, updating, flushing, and closing of memory scopes. + """ @abstractmethod def retrieve_memories(self, query: str, top_k: int, filter_dict: Dict[str, List[str]]) -> List[MemoryNode]: @@ -13,17 +17,29 @@ class BaseMemoryStore(metaclass=ABCMeta): @abstractmethod def update_memories(self, nodes: MemoryNode | List[MemoryNode]): """ - status: - 1. new: emb & insert - 2. modified: update - 3. content_modified: emb & update - 4. active: do nothing - 5. expired: update + Updates the memories based on their status: + - New: Embeds and inserts the memory node. + - Modified: Directly updates the memory node. + - Content Modified: Embeds and then updates the memory node. + - Active: No action required. + - Expired: Updates the memory node. + + Args: + nodes (MemoryNode | List[MemoryNode]): A single memory node or a list of memory nodes to be updated. """ + pass def flush(self): + """ + Flushes any pending memory updates or operations to ensure data consistency. + This method should be overridden by subclasses to provide the specific flushing mechanism. + """ pass @abstractmethod def close(self): + """ + Closes the memory store, releasing any resources associated with it. + Subclasses must implement this method to define how the memory store is properly closed. + """ pass diff --git a/memory_scope/storage/base_monitor.py b/memory_scope/storage/base_monitor.py index 05465cd3..7a4521c0 100644 --- a/memory_scope/storage/base_monitor.py +++ b/memory_scope/storage/base_monitor.py @@ -2,6 +2,10 @@ from abc import ABCMeta, abstractmethod class BaseMonitor(metaclass=ABCMeta): + """ + An abstract base class defining the interface for monitor classes. + Subclasses should implement the methods defined here to provide concrete monitoring behavior. + """ def __init__(self, **kwargs): pass @@ -9,17 +13,35 @@ class BaseMonitor(metaclass=ABCMeta): @abstractmethod def add(self): """ - :return: + Abstract method to add data or events to the monitor. + This method should be implemented by subclasses to define how data is added into the monitoring system. + + :return: None """ @abstractmethod def add_token(self): """ - :return: + Abstract method to add a token or a specific type of identifier to the monitor. + Subclasses should implement this to specify how tokens are managed within the monitoring context. + + :return: None """ def flush(self): + """ + Method to flush any buffered data in the monitor. + Intended to ensure that all pending recorded data is processed or written out. + + :return: None + """ pass def close(self): + """ + Method to close the monitor, performing necessary cleanup operations. + This could include releasing resources, closing files, or any other termination tasks. + + :return: None + """ pass diff --git a/memory_scope/storage/dummy_memory_store.py b/memory_scope/storage/dummy_memory_store.py index 40802d77..645b9d8f 100644 --- a/memory_scope/storage/dummy_memory_store.py +++ b/memory_scope/storage/dummy_memory_store.py @@ -6,19 +6,69 @@ from memory_scope.storage.base_memory_store import BaseMemoryStore class DummyMemoryStore(BaseMemoryStore): + """ + Placeholder implementation of a memory storage system interface. Defines methods for querying, updating, + and closing memory nodes with asynchronous capabilities, leveraging an embedding model for potential + semantic retrieval. Actual storage operations are not implemented. + """ def __init__(self, embedding_model: BaseModel, **kwargs): + """ + Initializes the DummyMemoryStore with an embedding model and additional keyword arguments. + + Args: + embedding_model (BaseModel): The model used to embed data for potential similarity-based retrieval. + **kwargs: Additional keyword arguments for configuration or future expansion. + """ self.embedding_model: BaseModel = embedding_model self.kwargs = kwargs def retrieve_memories(self, query: str, top_k: int, filter_dict: Dict[str, List[str]]) -> List[MemoryNode]: + """ + Retrieves a list of MemoryNode objects that are most relevant to the query, + considering a filter dictionary for additional constraints. The number of nodes returned + is limited by top_k. + + Args: + query (str): The query string used to find relevant memories. + top_k (int): The maximum number of MemoryNode objects to return. + filter_dict (Dict[str, List[str]]): A dictionary with keys representing filter fields + and values as lists of strings for filtering criteria. + + Returns: + List[MemoryNode]: A list of MemoryNode objects sorted by relevance to the query, + limited to top_k items. + """ pass async def a_retrieve_memories(self, query: str, top_k: int, filter_dict: Dict[str, List[str]]) -> List[MemoryNode]: + """ + Asynchronously retrieves a list of MemoryNode objects that best match the query, + respecting a filter dictionary, with the result size capped at top_k. + + Args: + query (str): The text to search for in memory nodes. + top_k (int): Maximum number of nodes to return. + filter_dict (Dict[str, List[str]]): Filters to apply on memory nodes. + + Returns: + List[MemoryNode]: A list of up to top_k MemoryNode objects matching the criteria. + """ pass def update_memories(self, nodes: MemoryNode | List[MemoryNode]): + """ + Updates the stored memories with a single MemoryNode or a list of MemoryNode objects. + + Args: + nodes (MemoryNode | List[MemoryNode]): A single MemoryNode or a collection of MemoryNode objects + to be updated in the memory store. + """ pass def close(self): + """ + Closes the memory store, releasing any resources it holds. This method should be called + when the memory store is no longer needed. + """ pass diff --git a/memory_scope/storage/dummy_monitor.py b/memory_scope/storage/dummy_monitor.py index 6c754ae5..735f643e 100644 --- a/memory_scope/storage/dummy_monitor.py +++ b/memory_scope/storage/dummy_monitor.py @@ -2,11 +2,29 @@ from memory_scope.storage.base_monitor import BaseMonitor class DummyMonitor(BaseMonitor): + """ + DummyMonitor serves as a placeholder or mock class extending BaseMonitor, + providing empty method bodies for 'add', 'add_token', and 'close' operations. + This can be used for testing or in situations where a full monitor implementation is not required. + """ + def add(self): + """ + Placeholder for adding data to the monitor. + This method currently does nothing. + """ pass def add_token(self): + """ + Placeholder for adding a token to the monitored data. + This method currently does nothing. + """ pass def close(self): + """ + Placeholder for closing the monitor and performing any necessary cleanup. + This method currently does nothing. + """ pass diff --git a/memory_scope/storage/llama_index_es_memory_store.py b/memory_scope/storage/llama_index_es_memory_store.py index 01f33121..8af77e11 100644 --- a/memory_scope/storage/llama_index_es_memory_store.py +++ b/memory_scope/storage/llama_index_es_memory_store.py @@ -12,11 +12,25 @@ from memory_scope.utils.logger import Logger class _AsyncDenseVectorStrategy(AsyncDenseVectorStrategy): + """ + Custom asynchronous dense vector strategy extending LlamaIndex's ElasticsearchStore's strategy. + This strategy enables hybrid search combining KNN queries with text queries and supports customizable ranking functions. + """ + def _hybrid(self, query: str, knn: Dict[str, Any], filter: List[Dict[str, Any]], top_k: int) -> Dict[str, Any]: - # Add a query to the knn query. - # RRF is used to even the score from the knn query and text query - # RRF has two optional parameters: {'rank_constant':int, 'window_size':int} - # https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html + """ + Constructs a hybrid query body combining KNN search with a text query, and applies filters. + + Args: + query (str): The text query to be combined with the KNN results. + knn (Dict[str, Any]): The KNN query part specifying the vector search parameters. + filter (List[Dict[str, Any]]): A list of filters to apply to the search. + top_k (int): The number of top results to retrieve. + + Returns: + Dict[str, Any]: The constructed query body for Elasticsearch to perform the hybrid search. + """ + # Combines KNN query with a text query and applies optional RRF ranking for result balancing query_body = { "knn": knn, "query": { @@ -35,6 +49,7 @@ class _AsyncDenseVectorStrategy(AsyncDenseVectorStrategy): }, } + # Configures Rank-Risk Function (RRF) if enabled or specified, to balance scores between KNN and text matches if isinstance(self.rrf, Dict): query_body["rank"] = {"rrf": self.rrf} elif isinstance(self.rrf, bool) and self.rrf is True: @@ -98,15 +113,21 @@ class _ElasticsearchStore(ElasticsearchStore): def _to_elasticsearch_filter(standard_filters: Dict[str, List[str]]) -> Dict[str, Any]: """ - Convert standard filters to Elasticsearch filter. + Converts the provided standard Llama-index filters into an Elasticsearch compatible filter format. + + This function processes each key-value pair in the input dictionary. If the value is a list, + it constructs a 'should' clause with multiple 'term' sub-clauses for each item in the list, + requiring at least one to match. If the value is not a list, it forms a 'must' clause with a single 'term' + sub-clause. The resulting structure is nested within a 'bool' clause which is the standard way to combine + boolean logic in Elasticsearch queries. Args: - standard_filters: Standard Llama-index filters. + standard_filters (Dict[str, List[str]]): A dictionary where keys represent filter fields and values are + either single values or lists of values to filter by. Returns: - Elasticsearch filter. + Dict[str, Any]: An Elasticsearch query filter dictionary ready to be used in a query. """ - result = { "bool": {} } @@ -140,6 +161,10 @@ def _to_elasticsearch_filter(standard_filters: Dict[str, List[str]]) -> Dict[str import ray ray.init(ignore_reinit_error=True) + +# The following decorator '@ray.remote' is used to define a function or class that should be executed remotely +# by Ray. This facilitates parallel and distributed computation. However, due to the instruction constraints, +# no modification or additional explanation is provided for this part. @ray.remote class _LlamaIndexEsMemoryStore(BaseMemoryStore): def __init__(self, @@ -175,9 +200,24 @@ class _LlamaIndexEsMemoryStore(BaseMemoryStore): return [self._text_node_2_memory_node(n) for n in text_nodes] def insert(self, node: MemoryNode): + """ + Inserts a MemoryNode into the Elasticsearch store by converting it to aTextNode. + + Args: + node (MemoryNode): The MemoryNode to be inserted into the store. + """ self.index.insert_nodes([self._memory_node_2_text_node(node)]) def delete(self, node: MemoryNode): + """ + Deletes a MemoryNode from the Elasticsearch store based on its memory_id. + + Args: + node (MemoryNode): The MemoryNode to be deleted, identified by its memory_id. + + Returns: + bool: The result of the deletion operation, typically True if successful. + """ memory_id = node.memory_id return self.es_store.delete(memory_id) @@ -190,6 +230,12 @@ class _LlamaIndexEsMemoryStore(BaseMemoryStore): self.update(node) def close(self): + """ + Closes the Elasticsearch store, releasing any resources associated with it. + + This method ensures that the connection to the Elasticsearch instance is properly closed, + which is a good practice to prevent resource leaks when you're done interacting with the store. + """ self.es_store.close() def update_memories(self, nodes: MemoryNode | List[MemoryNode]): @@ -236,12 +282,30 @@ class _LlamaIndexEsMemoryStore(BaseMemoryStore): @staticmethod def _memory_node_2_text_node(memory_node: MemoryNode) -> TextNode: + """ + Converts a MemoryNode object into a TextNode object. + + Args: + memory_node (MemoryNode): The MemoryNode to be converted. + + Returns: + TextNode: The converted TextNode object with the content and metadata from the MemoryNode. + """ return TextNode(id_=memory_node.memory_id, text=memory_node.content, metadata=memory_node.model_dump(exclude={"content"})) @staticmethod def _text_node_2_memory_node(text_node: NodeWithScore) -> MemoryNode: + """ + Converts a NodeWithScore object into a MemoryNode object. + + Args: + text_node (NodeWithScore): The NodeWithScore to be converted. + + Returns: + MemoryNode: The converted MemoryNode object with the text and metadata from the NodeWithScore. + """ return MemoryNode(content=text_node.text, **text_node.metadata) @@ -271,21 +335,63 @@ class LlamaIndexEsMemoryStore(): def update(self, node: MemoryNode): return ray.get(self.proxy_obj.update.remote(node)) - def update_batch(self, nodes: List[MemoryNode]): + def update_batch(self, nodes: List[MemoryNode]) -> Any: + """ + Updates a batch of memory nodes asynchronously using Ray. + + Args: + nodes (List[MemoryNode]): A list of MemoryNode objects to be updated. + + Returns: + Any: The result from the remote task once completed. + """ return ray.get(self.proxy_obj.update_batch.remote(nodes)) - def close(self): + def close(self) -> Any: + """ + Closes the Elasticsearch memory store asynchronously using Ray. + + Returns: + Any: The result from the remote task once completed. + """ return ray.get(self.proxy_obj.close.remote()) - def update_memories(self, nodes: MemoryNode | List[MemoryNode]): + def update_memories(self, nodes: MemoryNode | List[MemoryNode]) -> Any: + """ + Updates one or more memory nodes asynchronously using Ray. + + Args: + nodes (MemoryNode | List[MemoryNode]): A single MemoryNode or a list of MemoryNode objects to be updated. + + Returns: + Any: The result from the remote task once completed. + """ return ray.get(self.proxy_obj.update_memories.remote(nodes)) @staticmethod def _memory_node_2_text_node(memory_node: MemoryNode) -> TextNode: + """ + Converts a MemoryNode object into a TextNode object. + + Args: + memory_node (MemoryNode): The MemoryNode to convert. + + Returns: + TextNode: The converted TextNode object with content and metadata. + """ return TextNode(id_=memory_node.memory_id, text=memory_node.content, metadata=memory_node.model_dump(exclude={"content"})) @staticmethod def _text_node_2_memory_node(text_node: NodeWithScore) -> MemoryNode: + """ + Converts a TextNode (with score) into a MemoryNode object. + + Args: + text_node (NodeWithScore): The TextNode to convert, which includes a 'score' attribute. + + Returns: + MemoryNode: The converted MemoryNode object with content and metadata. + """ return MemoryNode(content=text_node.text, **text_node.metadata) diff --git a/memory_scope/storage/llama_index_es_memory_store_sync.py b/memory_scope/storage/llama_index_es_memory_store_sync.py index 38919028..727463c6 100644 --- a/memory_scope/storage/llama_index_es_memory_store_sync.py +++ b/memory_scope/storage/llama_index_es_memory_store_sync.py @@ -84,15 +84,21 @@ class _AsyncDenseVectorStrategy(AsyncDenseVectorStrategy): def _to_elasticsearch_filter(standard_filters: Dict[str, List[str]]) -> Dict[str, Any]: """ - Convert standard filters to Elasticsearch filter. + Converts standard Llama-index filters into a format compatible with Elasticsearch. + + This function transforms dictionary-based filters, where each key represents a field and + the value is a list of strings, into an Elasticsearch query structure. It supports both + list values (interpreted as 'should' clauses for OR logic) and single values (interpreted + as 'must' clauses for AND logic). Args: - standard_filters: Standard Llama-index filters. + standard_filters (Dict[str, List[str]]): A dictionary containing filter criteria, + where keys are field names and values are lists of strings or single string values + representing filter values. Returns: - Elasticsearch filter. + Dict[str, Any]: A dictionary structured as an Elasticsearch filter query. """ - result = { "bool": {} } @@ -108,8 +114,8 @@ def _to_elasticsearch_filter(standard_filters: Dict[str, List[str]]) -> Dict[str } } ) - result['bool'].update({"should": operands}) - result['bool'].update({"minimum_should_match": 1}) + result['bool'].update({"should": operands}) # ⭐ Add 'should' clause for OR logic + result['bool'].update({"minimum_should_match": 1}) # Ensure at least one 'should' match else: operand = [{ "term": { @@ -119,9 +125,9 @@ def _to_elasticsearch_filter(standard_filters: Dict[str, List[str]]) -> Dict[str } }] if "must" in result['bool']: - result['bool']['must'].extend(operand) + result['bool']['must'].extend(operand) # Extend existing 'must' clause for AND logic else: - result['bool'].update({"must": operand}) + result['bool'].update({"must": operand}) # Initialize 'must' clause if not present return result @@ -192,9 +198,26 @@ class LlamaIndexEsMemoryStoreSync(BaseMemoryStore): self.update(node) def close(self): + """ + Closes the Elasticsearch store, releasing any resources associated with it. + """ self.es_store.close() def update_memories(self, nodes: MemoryNode | List[MemoryNode]): + """ + Processes a list of MemoryNodes to update the Elasticsearch store based on their statuses. + NEW nodes are inserted, CONTENT_MODIFIED and MODIFIED nodes are updated (with embeddings for the former), + and EXPIRED nodes are deleted from the store. + + Args: + nodes (MemoryNode | List[MemoryNode]): A single MemoryNode or a list of MemoryNodes to be processed. + + Note: + - NEW nodes transition to ACTIVE after insertion. + - CONTENT_MODIFIED and MODIFIED nodes are re-embedded and transitioned to ACTIVE. + - EXPIRED nodes are removed from the store. + - Batch processing for insertion and deletion is planned but not yet implemented (TODOs). + """ if not nodes: self.logger.warning("empty nodes!") return @@ -203,7 +226,7 @@ class LlamaIndexEsMemoryStoreSync(BaseMemoryStore): if isinstance(nodes, MemoryNode): nodes = [nodes] - # emb & insert new memories + # Embed and insert new memories # TODO batch insert new_memories = [n for n in nodes if n.status == MemoryNodeStatus.NEW.value] if new_memories: @@ -211,7 +234,7 @@ class LlamaIndexEsMemoryStoreSync(BaseMemoryStore): n.status = MemoryNodeStatus.ACTIVE.value self.insert(n) - # emb & update new memories + # Embed and update content modified memories (overwriting existing) # TODO insert overwrite c_modified_memories = [n for n in nodes if n.status == MemoryNodeStatus.CONTENT_MODIFIED.value] if c_modified_memories: @@ -220,7 +243,7 @@ class LlamaIndexEsMemoryStoreSync(BaseMemoryStore): self.delete(n) self.insert(n) - # update new memories + # Update modified memories without re-embedding # TODO no emb modified_memories = [n for n in nodes if n.status == MemoryNodeStatus.MODIFIED.value] if modified_memories: @@ -229,19 +252,36 @@ class LlamaIndexEsMemoryStoreSync(BaseMemoryStore): self.delete(n) self.insert(n) - # set memories expired + # Set and remove expired memories expired_memories = [n for n in nodes if n.status == MemoryNodeStatus.EXPIRED.value] if expired_memories: for n in expired_memories: self.delete(n) - self.insert(n) @staticmethod def _memory_node_2_text_node(memory_node: MemoryNode) -> TextNode: + """ + Converts a MemoryNode object into a TextNode object. + + Args: + memory_node (MemoryNode): The MemoryNode to be converted. + + Returns: + TextNode: The converted TextNode with content and metadata from the MemoryNode. + """ return TextNode(id_=memory_node.memory_id, text=memory_node.content, metadata=memory_node.model_dump(exclude={"content"})) @staticmethod def _text_node_2_memory_node(text_node: NodeWithScore) -> MemoryNode: + """ + Converts a NodeWithScore object into a MemoryNode object. + + Args: + text_node (NodeWithScore): The NodeWithScore to be converted, typically retrieved from search results. + + Returns: + MemoryNode: The converted MemoryNode with text and metadata from the NodeWithScore. + """ return MemoryNode(content=text_node.text, **text_node.metadata) diff --git a/memory_scope/storage/llama_index_sync_elasticsearch.py b/memory_scope/storage/llama_index_sync_elasticsearch.py index 89ec5191..cd0f7bce 100644 --- a/memory_scope/storage/llama_index_sync_elasticsearch.py +++ b/memory_scope/storage/llama_index_sync_elasticsearch.py @@ -82,15 +82,20 @@ def get_elasticsearch_client( def _to_elasticsearch_filter(standard_filters: MetadataFilters) -> Dict[str, Any]: """ - Convert standard filters to Elasticsearch filter. + Transforms Llama-index standard filters into an Elasticsearch-compatible filter structure. + + This function supports both single-term filters and multiple operands combined + with a boolean 'should' clause for more complex queries. Args: - standard_filters: Standard Llama-index filters. + standard_filters (MetadataFilters): An instance of MetadataFilters containing + the filtering criteria to be applied. Returns: - Elasticsearch filter. + Dict[str, Any]: A dictionary representing the Elasticsearch filter query. """ if len(standard_filters.legacy_filters()) == 1: + # For a single filter term, construct a simple term filter. filter = standard_filters.legacy_filters()[0] return { "term": { @@ -100,6 +105,8 @@ def _to_elasticsearch_filter(standard_filters: MetadataFilters) -> Dict[str, Any } } else: + # When multiple filters are present, create a boolean 'should' clause + # with each individual filter as an operand. operands = [] for filter in standard_filters.legacy_filters(): operands.append( @@ -115,10 +122,21 @@ def _to_elasticsearch_filter(standard_filters: MetadataFilters) -> Dict[str, Any def _to_llama_similarities(scores: List[float]) -> List[float]: + """ + Converts a list of similarity scores into a normalized form for LlamaIndex compatibility. + The normalization involves an exponential transformation based on the maximum score in the list. + + Args: + scores (List[float]): A list of raw similarity scores. + + Returns: + List[float]: A list of normalized similarity scores suitable for LlamaIndex. + """ if scores is None or len(scores) == 0: return [] scores_to_norm: np.ndarray = np.array(scores) + # Normalize scores by subtracting the max score and applying the exponential function return np.exp(scores_to_norm - np.max(scores_to_norm)).tolist() @@ -303,7 +321,12 @@ class SyncElasticsearchStore(BasePydanticVectorStore): @property def client(self) -> Any: - """Get async elasticsearch client.""" + """ + Get the asynchronous Elasticsearch client. + + Returns: + Any: The asynchronous Elasticsearch client instance configured for this store. + """ return self._store.client def close(self) -> None: @@ -317,23 +340,25 @@ class SyncElasticsearchStore(BasePydanticVectorStore): **add_kwargs: Any, ) -> List[str]: """ - Add nodes to Elasticsearch index. + Adds a list of nodes, each containing embeddings, to an Elasticsearch index. + Optionally creates the index if it does not already exist. Args: - nodes: List of nodes with embeddings. - create_index_if_not_exists: Optional. Whether to create - the Elasticsearch index if it - doesn't already exist. - Defaults to True. + nodes (List[BaseNode]): A list of node objects, each encapsulating an embedding. + create_index_if_not_exists (bool, optional): + A flag indicating whether to create the Elasticsearch index if it's not present. + Defaults to True. Returns: - List of node IDs that were added to the index. + List[str]: A list of node IDs that have been successfully added to the index. Raises: - ImportError: If elasticsearch['async'] python package is not installed. - BulkIndexError: If AsyncElasticsearch async_bulk indexing fails. + ImportError: If the 'elasticsearch[async]' Python package is not installed. + BulkIndexError: If there is a failure during the asynchronous bulk indexing with AsyncElasticsearch. + + Note: + This method delegates the actual operation to the `sync_add` method. """ - return self.sync_add(nodes, create_index_if_not_exists=create_index_if_not_exists) def sync_add( @@ -344,38 +369,46 @@ class SyncElasticsearchStore(BasePydanticVectorStore): **add_kwargs: Any, ) -> List[str]: """ - Asynchronous method to add nodes to Elasticsearch index. + Asynchronously adds a list of nodes, each containing an embedding, to the Elasticsearch index. + + This method processes each node to extract its ID, embedding, text content, and metadata, + preparing them for batch insertion into the index. It ensures the index is created if not present + and respects the dimensionality of the embeddings for consistency. Args: - nodes: List of nodes with embeddings. - create_index_if_not_exists: Optional. Whether to create - the AsyncElasticsearch index if it - doesn't already exist. - Defaults to True. + nodes (List[BaseNode]): A list of node objects, each encapsulating an embedding. + create_index_if_not_exists (bool, optional): A flag indicating whether to create the Elasticsearch + index if it does not already exist. Defaults to True. + **add_kwargs (Any): Additional keyword arguments passed to the underlying add_texts method + for customization during the indexing process. Returns: - List of node IDs that were added to the index. + List[str]: A list of node IDs that were successfully added to the index. Raises: - ImportError: If elasticsearch python package is not installed. - BulkIndexError: If AsyncElasticsearch async_bulk indexing fails. + ImportError: If the Elasticsearch Python client is not installed. + BulkIndexError: If there's a failure during the asynchronous bulk indexing operation. """ if len(nodes) == 0: return [] - embeddings: List[List[float]] = [] - texts: List[str] = [] - metadatas: List[dict] = [] - ids: List[str] = [] + # Extract necessary components from each node + embeddings: List[List[float]] = [] # Embedding vectors + texts: List[str] = [] # Textual contents of nodes + metadatas: List[dict] = [] # Metadata associated with nodes + ids: List[str] = [] # Unique identifiers for nodes + for node in nodes: - ids.append(node.node_id) - embeddings.append(node.get_embedding()) - texts.append(node.get_content(metadata_mode=MetadataMode.NONE)) - metadatas.append(node_to_metadata_dict(node, remove_text=True)) + ids.append(node.node_id) # Node identifier + embeddings.append(node.get_embedding()) # Node's embedding vector + texts.append(node.get_content(metadata_mode=MetadataMode.NONE)) # Node's raw text content + metadatas.append(node_to_metadata_dict(node, remove_text=True)) # Convert node to metadata dictionary + # Initialize the number of dimensions in the store if not set if not self._store.num_dimensions: - self._store.num_dimensions = len(embeddings[0]) + self._store.num_dimensions = len(embeddings[0]) # Set based on the first node's embedding size + # Add the prepared data to the Elasticsearch index asynchronously return self._store.add_texts( texts=texts, metadatas=metadatas, @@ -387,33 +420,45 @@ class SyncElasticsearchStore(BasePydanticVectorStore): def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: """ - Delete node from Elasticsearch index. - + Deletes a node from the Elasticsearch index using the provided reference document ID. + + Optionally, extra keyword arguments can be supplied to customize the deletion behavior, + which are passed directly to Elasticsearch's `delete_by_query` operation. + Args: - ref_doc_id: ID of the node to delete. - delete_kwargs: Optional. Additional arguments to - pass to Elasticsearch delete_by_query. - + ref_doc_id (str): The unique identifier of the node/document to be deleted. + delete_kwargs (Any): Additional keyword arguments for Elasticsearch's + `delete_by_query`. These might include query filters, + timeouts, or other operational configurations. + Raises: - Exception: If Elasticsearch delete_by_query fails. + Exception: If the deletion operation via Elasticsearch's `delete_by_query` fails. + + Note: + This method internally calls a synchronous delete method (`sync_delete`) + to execute the deletion operation against Elasticsearch. """ return self.sync_delete(ref_doc_id, **delete_kwargs) def sync_delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: """ - Async delete node from Elasticsearch index. + Synchronously deletes a node from the Elasticsearch index based on the reference document ID. Args: - ref_doc_id: ID of the node to delete. - delete_kwargs: Optional. Additional arguments to - pass to AsyncElasticsearch delete_by_query. + ref_doc_id (str): The unique identifier of the node/document to be deleted. + delete_kwargs (Any): Optional keyword arguments to be passed + to the delete_by_query operation of AsyncElasticsearch, + allowing for additional customization of the deletion process. Raises: - Exception: If AsyncElasticsearch delete_by_query fails. + Exception: If the deletion operation via AsyncElasticsearch's delete_by_query fails. + + Note: + The function directly uses '_id' field to match the document for deletion instead of 'metadata.ref_doc_id', + ensuring targeted removal based on the document's unique identifier within Elasticsearch. """ - # return self._store.delete( - # query={"term": {"metadata.ref_doc_id": ref_doc_id}}, **delete_kwargs - # ) + # The original commented line suggests an alternative query using 'metadata.ref_doc_id', + # but the active code line performs the deletion based on '_id', which typically aligns with 'ref_doc_id'. return self._store.delete(query={"term": {"_id": ref_doc_id}}, **delete_kwargs) def query( @@ -426,23 +471,25 @@ class SyncElasticsearchStore(BasePydanticVectorStore): **kwargs: Any, ) -> VectorStoreQueryResult: """ - Query index for top k most similar nodes. + Executes a query against the Elasticsearch index to retrieve the top k most similar nodes + based on the input query embedding. Supports customization of the query process and + application of Elasticsearch filters. Args: - query (List[float]): query embedding - custom_query: Optional. custom query function that takes in the es query - body and returns a modified query body. - This can be used to add additional query - parameters to the Elasticsearch query. - es_filter: Optional. Elasticsearch filter to apply to the - query. If filter is provided in the query, - this filter will be ignored. + query (VectorStoreQuery): The query containing the embedding and other parameters. + custom_query (Callable[[Dict, Union[VectorStoreQuery, None]], Dict], optional): + An optional custom function to modify the Elasticsearch query body, allowing for + additional query parameters or logic. Defaults to None. + es_filter (Optional[List[Dict]], optional): An optional Elasticsearch filter list to + apply to the query. If a filter is directly included in the `query`, this argument + will not be used. Defaults to None. + **kwargs (Any): Additional keyword arguments that might be used in the query process. Returns: - VectorStoreQueryResult: Result of the query. + VectorStoreQueryResult: The result of the query operation, including the most similar nodes. Raises: - Exception: If Elasticsearch query fails. + Exception: If an error occurs during the Elasticsearch query execution. """ return self.sync_query(query, custom_query, es_filter, **kwargs) @@ -457,24 +504,27 @@ class SyncElasticsearchStore(BasePydanticVectorStore): **kwargs: Any, ) -> VectorStoreQueryResult: """ - Asynchronous query index for top k most similar nodes. + Asynchronously queries the Elasticsearch index for the top k most similar nodes + based on the provided query embedding. Supports custom query modifications + and application of Elasticsearch filters. Args: - query_embedding (VectorStoreQuery): query embedding - custom_query: Optional. custom query function that takes in the es query - body and returns a modified query body. - This can be used to add additional query - parameters to the AsyncElasticsearch query. - es_filter: Optional. AsyncElasticsearch filter to apply to the - query. If filter is provided in the query, - this filter will be ignored. + query (VectorStoreQuery): The query containing the embedding and other details. + custom_query (Callable[[Dict, Union[VectorStoreQuery, None]], Dict], optional): + A custom function to modify the Elasticsearch query body. Defaults to None. + es_filter (List[Dict], optional): Additional filters to apply during the query. + If filters are present in the query, these filters will not be used. Defaults to None. Returns: - VectorStoreQueryResult: Result of the query. + VectorStoreQueryResult: The result of the query, including nodes, their IDs, + and similarity scores. Raises: - Exception: If AsyncElasticsearch query fails. + Exception: If the Elasticsearch query encounters an error. + Note: + The mode of the query must align with the retrieval strategy set for this store. + In case of legacy metadata, a warning is logged and nodes are constructed accordingly. """ _mode_must_match_retrieval_strategy(query.mode, self.retrieval_strategy) @@ -501,6 +551,7 @@ class SyncElasticsearchStore(BasePydanticVectorStore): node_id = hit["_id"] try: + # Attempt to parse metadata using the standard method node = metadata_dict_to_node(metadata) node.text = text except Exception: diff --git a/memory_scope/utils/datetime_handler.py b/memory_scope/utils/datetime_handler.py index f274991d..a9c9bf7e 100644 --- a/memory_scope/utils/datetime_handler.py +++ b/memory_scope/utils/datetime_handler.py @@ -8,9 +8,28 @@ from memory_scope.utils.logger import Logger class DatetimeHandler(object): + """ + Handles operations related to datetime such as parsing, extraction, and formatting, + with support for both Chinese and English contexts including weekday names and + specialized text parsing for date components. + """ + logger = Logger.get_logger() def __init__(self, dt: datetime.datetime | str | int | float = None): + """ + Initialize the DatetimeHandler instance with a datetime object, string, integer, or float representation of a timestamp. + If no argument is provided, the current time is used. + + Args: + dt (datetime.datetime | str | int | float, optional): + The datetime to be handled. Can be a datetime object, a timestamp string, or a numeric timestamp. + Defaults to None, which sets the instance to the current datetime. + + Attributes: + _dt (datetime.datetime): The internal datetime representation of the input. + _dt_info_dict (dict | None): A dictionary containing parsed datetime information, initialized as None. + """ if isinstance(dt, str | int | float): if isinstance(dt, str): dt = float(dt) @@ -23,6 +42,15 @@ class DatetimeHandler(object): self._dt_info_dict: dict | None = None def _parse_dt_info(self): + """ + Parses the datetime object (_dt) into a dictionary containing detailed date and time components, + including language-specific weekday representation. + + Returns: + dict: A dictionary with keys representing date and time parts such as 'year', 'month', + 'day', 'hour', 'minute', 'second', 'week', and 'weekday' with respective values. + The 'weekday' value is translated based on the current language context. + """ return { "year": self._dt.year, "month": self._dt.month, @@ -36,12 +64,36 @@ class DatetimeHandler(object): @property def dt_info_dict(self): + """ + Property method to get the dictionary containing parsed datetime information. + If not already parsed, it triggers the parsing process using `_parse_dt_info`. + + Returns: + dict: A dictionary with parsed datetime information. + """ if self._dt_info_dict is None: self._dt_info_dict = self._parse_dt_info() return self._dt_info_dict @classmethod def extract_date_parts_cn(cls, input_string: str) -> dict: + """ + Extracts date components from a Chinese text string into a dictionary. + + This method identifies year, month, day, weekday, and hour components within the input + string based on predefined patterns. It supports relative terms like '每' (every) and + translates weekday names into numeric representations. + + Args: + input_string (str): The Chinese text containing date and time information. + + Returns: + dict: A dictionary with keys 'year', 'month', 'day', 'weekday', and 'hour', + each holding the corresponding extracted value. If a component is not found, + it will not be included in the dictionary. For relative terms like '每' (every), + the value is set to -1. + + """ # Extending our pattern to handle every/每 as a possible value. patterns = { "year": r"(\d+|每)年", @@ -67,6 +119,19 @@ class DatetimeHandler(object): @classmethod def extract_date_parts_en(cls, input_string: str) -> dict: + """ + Extracts various components of a date (year, month, day, etc.) from an input string based on English formats. + + This method employs regex patterns to identify and parse different date and time elements within the provided text. + It supports extraction of year, month name, day, 12-hour and 24-hour time formats, and weekdays. + + Args: + input_string (str): The string from which to extract date and time components. + + Returns: + dict: A dictionary containing the extracted date parts with default values of -1 where components are not found. + Keys include 'year', 'month', 'day', 'hour', 'minute', 'second', and 'weekday'. + """ date_info = { "year": -1, "month": -1, @@ -98,26 +163,27 @@ class DatetimeHandler(object): "Monday": 1, "Tuesday": 2, "Wednesday": 3, "Thursday": 4, "Friday": 5, "Saturday": 6, "Sunday": 7 } + # Attempt to match full date (day month year) day_month_year_match = re.search(patterns["day_month_year"], input_string) if day_month_year_match: date_info["year"] = int(day_month_year_match.group("year")) date_info["month"] = month_mapping[day_month_year_match.group("month")] date_info["day"] = int(day_month_year_match.group("day")) - # Extract month and day without year + # If year wasn't found, try matching day and month without year elif date_info["year"] == -1: day_month_match = re.search(patterns["day_month"], input_string) if day_month_match: date_info["month"] = month_mapping[day_month_match.group("month")] date_info["day"] = int(day_month_match.group("day")) - # Extract year + # Extract year if not already found if date_info["year"] == -1: year_match = re.search(patterns["year"], input_string) if year_match: date_info["year"] = int(year_match.group(0)) - # Extract month + # Extract month if not already found if date_info["month"] == -1: month_match = re.search(patterns["month"], input_string) if month_match: @@ -133,7 +199,7 @@ class DatetimeHandler(object): hour = 0 date_info["hour"] = hour - # Extract weekday + # Identify weekday for week_day, value in weekday_mapping.items(): if week_day in input_string: date_info["weekday"] = value @@ -143,6 +209,19 @@ class DatetimeHandler(object): @classmethod def extract_date_parts(cls, input_string: str) -> dict: + """ + Extracts various date components from the input string based on the current language context. + + This method dynamically selects a language-specific function to parse the input string and extract + date parts such as year, month, day, etc. If the function for the current language context does not exist, + a warning is logged and an empty dictionary is returned. + + Args: + input_string (str): The string containing date information to be parsed. + + Returns: + dict: A dictionary containing extracted date components, or an empty dictionary if parsing fails. + """ func_name = f"extract_date_parts_{G_CONTEXT.language.value}" if not hasattr(cls, func_name): cls.logger.warning(f"language={G_CONTEXT.language.value} needs to complete extract_date_parts func!") @@ -151,6 +230,20 @@ class DatetimeHandler(object): @classmethod def format_time_by_extract_time_cn(cls, extract_time_dict: Dict[str, str], meta_data: Dict[str, str]) -> str: + """ + Formats a time string based on extracted time elements with Chinese context. + + This method constructs a time string using Chinese characters for year, month, and day, + and includes special handling for weekdays. If a value is '-1', it is replaced with '每' + to denote a recurring event in the formatted string. + + Args: + extract_time_dict (Dict[str, str]): A dictionary mapping time elements to their respective placeholders. + meta_data (Dict[str, str]): A dictionary containing the actual values for the time elements. + + Returns: + str: The formatted time string in Chinese context. + """ cn_key_dict = {"year": "年", "month": "月", "day": "日", "weekday": ""} format_time_str = "" for key, value_cn in cn_key_dict.items(): @@ -166,6 +259,20 @@ class DatetimeHandler(object): @classmethod def format_time_by_extract_time(cls, extract_time_dict: Dict[str, str], meta_data: Dict[str, str]) -> str: + """ + Formats the time based on extracted time components and additional metadata, considering the current language context. + + This method dynamically selects a language-specific function to format the time. If the function for the current + language context does not exist, a warning is logged and an empty string is returned. + + Args: + extract_time_dict (Dict[str, str]): A dictionary containing extracted time components like year, month, etc. + meta_data (Dict[str, str]): Additional contextual metadata that might be used in formatting. + + Returns: + str: The formatted time string according to the current language settings, or an empty string if the formatting + function is not implemented for the current language. + """ func_name = f"format_time_by_extract_time_{G_CONTEXT.language.value}" if not hasattr(cls, func_name): cls.logger.warning(f"language={G_CONTEXT.language.value} needs to complete format_time_by_extract_time func!") @@ -174,6 +281,15 @@ class DatetimeHandler(object): @classmethod def has_time_word(cls, query: str) -> bool: + """ + Check if the input query contains any datetime-related words based on the current language context. + + Args: + query (str): The input string to check for datetime words. + + Returns: + bool: True if the query contains at least one datetime word, False otherwise. + """ contain_datetime = False # TODO use re for datetime_word in DATATIME_WORD_LIST[G_CONTEXT.language]: @@ -182,12 +298,36 @@ class DatetimeHandler(object): break return contain_datetime - def datetime_format(self, dt_format: str = "%Y%m%d"): + def datetime_format(self, dt_format: str = "%Y%m%d") -> str: + """ + Format the stored datetime object into a string based on the provided format. + + Args: + dt_format (str, optional): The datetime format string. Defaults to "%Y%m%d". + + Returns: + str: The formatted datetime string. + """ return self._dt.strftime(dt_format) - def string_format(self, string_format: str): + def string_format(self, string_format: str) -> str: + """ + Format the datetime information stored in the instance using a custom string format. + + Args: + string_format (str): A format string where placeholders are keys from `dt_info_dict`. + + Returns: + str: The formatted string with datetime information inserted. + """ return string_format.format(**self.dt_info_dict) @property def timestamp(self) -> int: + """ + Get the timestamp representation of the stored datetime. + + Returns: + int: The timestamp value of the datetime. + """ return int(self._dt.timestamp()) diff --git a/memory_scope/utils/logger.py b/memory_scope/utils/logger.py index 77f32042..bad19709 100644 --- a/memory_scope/utils/logger.py +++ b/memory_scope/utils/logger.py @@ -21,6 +21,22 @@ class Logger(logging.Logger): dir_path: str = "log", max_bytes: int = 1024 * 1024 * 1024, backup_count: int = 10): + """ + Initializes the Logger instance, setting up handlers for console and/or file logging based on provided parameters. + + Args: + name (str): Identifier for the logger. + level (int, optional): Logging level. Defaults to logging.INFO. + format_style (str, optional): Log message format. Defaults to LOG_FORMAT constant. + date_format_style (str, optional): Date format for logs. Defaults to DATE_FORMAT constant. + to_stream (bool, optional): Enables console logging. Defaults to True. + to_file (bool, optional): Enables file logging. Defaults to True. + file_mode (str, optional): File open mode. Defaults to 'w'. + file_type (str, optional): Log file extension type. Defaults to 'log'. + dir_path (str, optional): Directory for log files. Defaults to 'log'. + max_bytes (int, optional): Maximum log file size before rotation. Defaults to 1GB. + backup_count (int, optional): Number of rotated log files to retain. Defaults to 10. + """ super(Logger, self).__init__(name, level) self.formatter = logging.Formatter(format_style, date_format_style) @@ -37,53 +53,120 @@ class Logger(logging.Logger): self.trace_id: str = "" if self.to_stream: - self._add_stream_handler() + self._add_stream_handler() # Adds a handler to output logs to the console if self.to_file: - self._add_file_handler() + self._add_file_handler() # Adds a handler to output logs to a file - self.info(f"logger={name} is inited.") + self.info(f"logger={name} is inited.") # Logs an initialization message def _add_file_handler(self): - file_path = Path().joinpath(self.dir_path, f"{self.name}.{self.file_type}") - file_path.parent.mkdir(exist_ok=True) - file_name = file_path.as_posix() + """ + Adds a file handler to the logger which logs messages to a rotating file. + The file is stored in a specified directory with a name derived from the logger's name and type. + The file handler is set up to rotate when it reaches a certain size and keeps a defined number of backups. + + This method ensures the directory exists before creating the file handler and sets the formatter + for consistent log message formatting. + """ + file_path = Path().joinpath(self.dir_path, f"{self.name}.{self.file_type}") + file_path.parent.mkdir(exist_ok=True) # Ensure the directory exists + file_name = file_path.as_posix() # Get the absolute path as a string + + # Instantiate a rotating file handler with specified parameters file_handler = RotatingFileHandler( filename=file_name, - maxBytes=self.max_bytes, - backupCount=self.backup_count, - encoding="utf-8") - file_handler.setFormatter(self.formatter) - self.addHandler(file_handler) + maxBytes=self.max_bytes, # Maximum size of the log file before rotation + backupCount=self.backup_count, # Number of backup files to keep + encoding="utf-8") # Set the encoding to UTF-8 + file_handler.setFormatter(self.formatter) # Apply the logger's formatter to the handler + self.addHandler(file_handler) # Add the file handler to this logger instance def _add_stream_handler(self): + """ + Adds a stream handler to the logger for console output. The handler is configured + with the logger's formatter and set to use UTF-8 encoding. + """ stream_handler = logging.StreamHandler() - stream_handler.setFormatter(self.formatter) - stream_handler.encoding = 'utf-8' - self.addHandler(stream_handler) + stream_handler.setFormatter(self.formatter) # Configure the handler with the logger's formatter + stream_handler.encoding = 'utf-8' # Set the handler's encoding to UTF-8 + self.addHandler(stream_handler) # Add the handler to the logger def close(self): + """ + Closes all handlers associated with this logger instance. + + This method iterates over the handlers attached to the logger and + calls their `close` method to ensure that any system resources used + by the handlers are freed properly. + """ for handler in self.handlers: - handler.close() + handler.close() # ⭐ Close each handler to release resources def clear(self): + """ + Clears all handlers from the logger. + """ self.handlers.clear() def set_trace_id(self, trace_id: str): + """ + Sets the trace ID for the logger. If the provided trace ID is longer than 8 characters, + it will be truncated to the first 8 characters. + + Args: + trace_id (str): The trace identifier to be associated with the logs. + """ self.trace_id: str = trace_id if len(self.trace_id) >= 8: self.trace_id = self.trace_id[:8] def makeRecord(self, name, level, fn, lno, msg, args, exc_info, func=None, extra=None, sinfo=None): + """ + Creates a log record with additional trace_id included in the extra information. + + This method extends the default behavior of creating a log record by adding + a trace_id from the logger instance to the record's extra data, allowing + for traceability within logged data. + + Args: + name (str): The name of the logger. + level (int): The logging level of the record. + fn (str): The name of the function containing the logging call. + lno (int): The line number at which the logging call was made. + msg (str): The logged message, before formatting. + args (tuple): The arguments to the log message. + exc_info (tuple): Exception information or None. + func (function): The function where the logging call was made. + extra (dict): Additional information for the log record, can be None. + sinfo (str): Stack trace information or None. + + Returns: + logging.LogRecord: The created log record with potentially enriched 'extra' field. + """ if extra is None: extra = {} if self.trace_id: - extra["trace_id"] = self.trace_id + extra["trace_id"] = self.trace_id # ⭐ Include trace_id from the logger in the log record extra data return super().makeRecord(name, level, fn, lno, msg, args, exc_info, func, extra, sinfo) @classmethod def get_logger(cls, name: str = None, **kwargs): + """ + Retrieves or creates a logger instance with the specified name and configurations. + + If no name is provided, it defaults to the first registered logger's name or 'default' if none exist. + This method ensures that only one logger instance exists per name by reusing existing instances + stored in `LOGGER_DICT`. + + Args: + name (str, optional): The name of the logger. Defaults to None, which triggers auto-naming logic. + **kwargs: Additional keyword arguments to configure the logger. + + Returns: + Logger: The requested or newly created logger instance. + """ if name is None: if LOGGER_DICT: name = list(LOGGER_DICT.keys())[0] diff --git a/memory_scope/utils/prompt_handler.py b/memory_scope/utils/prompt_handler.py index b4f1dc90..d0aec970 100644 --- a/memory_scope/utils/prompt_handler.py +++ b/memory_scope/utils/prompt_handler.py @@ -8,8 +8,21 @@ from memory_scope.utils.global_context import G_CONTEXT class PromptHandler(object): + """ + The `PromptHandler` class manages prompt messages by loading them from YAML or JSON files and dictionaries, + supporting language selection based on a global context, and providing dictionary-like access to the prompt messages. + """ def __init__(self, class_path: str, prompt_file: str = "", prompt_dict: dict = None, **kwargs): + """ + Initializes the PromptHandler with paths to prompt sources and additional keyword arguments. + + Args: + class_path (str): The path to the class where prompts are utilized. + prompt_file (str, optional): The path to an external file containing prompts. Defaults to "". + prompt_dict (dict, optional): A dictionary directly containing prompt definitions. Defaults to None. + **kwargs: Additional keyword arguments that might be used in prompt handling. + """ self._class_path: str = class_path self._prompt_dict: Dict[str, str] = {} self.kwargs = kwargs @@ -26,7 +39,19 @@ class PromptHandler(object): @staticmethod def file_path_completion(file_path: str) -> str: + """ + Attempts to complete the given file path by appending either a `.yaml` or `.json` extension + based on the existence of the respective file. If neither exists, an exception is raised. + Args: + file_path (str): The base path of the file to be completed. + + Returns: + str: The completed file path with the appropriate extension. + + Raises: + RuntimeError: If neither the `.yaml` nor `.json` file exists at the given path. + """ if file_path.endswith(".yaml") or file_path.endswith(".json"): return file_path @@ -39,21 +64,43 @@ class PromptHandler(object): raise RuntimeError(f"{file_path}/yaml/json is not exists!") def add_prompt_file(self, file_path: str): + """ + Adds prompt messages from a YAML or JSON file to the internal dictionary. + + This method supports loading prompts from files ending with '.yaml' or '.json'. + It uses the respective libraries to parse the content and merge it into the current prompt dictionary. + + Args: + file_path (str): The path to the YAML or JSON file containing the prompts. + """ file_path = self.file_path_completion(file_path) prompt_dict = {} if file_path.endswith(".yaml"): + # Load prompts from a YAML file with open(file_path) as f: prompt_dict = yaml.load(f, yaml.FullLoader) elif file_path.endswith(".json"): - with open(f"{file_path}.json") as f: + # Load prompts from a JSON file (corrected file handling) + with open(file_path) as f: prompt_dict = json.load(f) + # Merge the loaded prompts into the existing dictionary self.add_prompt_dict(prompt_dict) def add_prompt_dict(self, prompt_dict: dict): + """ + Adds prompt messages from a dictionary, ensuring each message has a valid entry for the current language. + + Args: + prompt_dict (dict): A dictionary where keys represent prompt identifiers and values are nested dictionaries + containing language-specific prompt messages. + + Raises: + RuntimeError: If a prompt message for the current language is not found. + """ for key, language_dict in prompt_dict.items(): prompts = language_dict.get(G_CONTEXT.language) if not prompts: @@ -62,13 +109,44 @@ class PromptHandler(object): @property def prompt_dict(self): + """ + Retrieves the internal dictionary containing all prompt messages. + + Returns: + dict: The dictionary of prompt messages with keys as identifiers and values as prompt strings. + """ return self._prompt_dict def __getitem__(self, key: str): + """ + Enables accessing prompt messages using dictionary-like indexing. + + Args: + key (str): The identifier for the prompt message. + + Returns: + str: The prompt message corresponding to the given key. + """ return self._prompt_dict[key] def __setitem__(self, key: str, value: str): + """ + Allows setting prompt messages using dictionary-like item assignment. + + Args: + key (str): The identifier for the prompt message. + value (str): The new prompt message content. + """ self._prompt_dict[key] = value def __getattr__(self, key: str): + """ + Overrides attribute access to provide prompt messages dynamically. + + Args: + key (str): The identifier for the prompt message attempted to access as an attribute. + + Returns: + str: The prompt message corresponding to the given attribute-like key. + """ return self._prompt_dict[key] diff --git a/memory_scope/utils/registry.py b/memory_scope/utils/registry.py index 9d7996ea..420ac5c3 100644 --- a/memory_scope/utils/registry.py +++ b/memory_scope/utils/registry.py @@ -6,7 +6,22 @@ from typing import Dict, Any, List class Registry(object): + """ + A registry to manage and instantiate various modules by their names, ensuring the uniqueness of registered entries. + It supports both individual and bulk registration of modules, as well as retrieval of modules by name. + + Attributes: + name (str): The name of the registry. + module_dict (Dict[str, Any]): A dictionary holding registered modules where keys are module names and values are the modules themselves. + """ + def __init__(self, name: str): + """ + Initializes the Registry with a given name. + + Args: + name (str): The name to identify this registry. + """ self.name: str = name self.module_dict: Dict[str, Any] = {} @@ -20,14 +35,35 @@ class Registry(object): self.module_dict[module_name] = module def batch_register(self, modules: List[Any] | Dict[str, Any]): + """ + Registers multiple modules in the registry in a single call. Accepts either a list of modules or a dictionary mapping names to modules. + + Args: + modules (List[Any] | Dict[str, Any]): A list of modules or a dictionary mapping module names to the modules. + + Raises: + NotImplementedError: If the input is neither a list nor a dictionary. + """ if isinstance(modules, list): module_name_dict = {m.__name__: m for m in modules} elif isinstance(modules, dict): module_name_dict = modules else: - raise NotImplementedError + raise NotImplementedError("Input must be a list or a dictionary.") self.module_dict.update(module_name_dict) def __getitem__(self, module_name: str): + """ + Retrieves a registered module by its name using index notation. + + Args: + module_name (str): The name of the module to retrieve. + + Returns: + The registered module corresponding to the given name. + + Raises: + AssertionError: If the specified module is not found in the registry. + """ assert module_name in self.module_dict, f"{module_name} not found in {self.name}" return self.module_dict[module_name] diff --git a/memory_scope/utils/response_text_parser.py b/memory_scope/utils/response_text_parser.py index d665b458..fbcb0889 100644 --- a/memory_scope/utils/response_text_parser.py +++ b/memory_scope/utils/response_text_parser.py @@ -4,11 +4,22 @@ from memory_scope.utils.logger import Logger class ResponseTextParser(object): - pattern_v1 = re.compile(r"<(.*?)>") + """ + The `ResponseTextParser` class is designed to process and parse response texts. It provides methods to extract specific + patterns from the text and filter out unnecessary information, while also logging the processing steps and outcomes. + """ + + pattern_v1 = re.compile(r"<(.*?)>") # Regular expression pattern to match content within angle brackets def __init__(self, response_text: str): - self.response_text: str = response_text.strip() - self.logger: Logger = Logger.get_logger() + """ + Initializes the `ResponseTextParser` instance with the provided response text and sets up a logger. + + Args: + response_text (str): The raw response text that needs to be parsed and processed. + """ + self.response_text: str = response_text.strip() # Strips leading and trailing whitespace from the response text + self.logger: Logger = Logger.get_logger() # Initializes a logger instance for logging parsing activities def parse_v1(self, prefix: str = ""): result = [] diff --git a/memory_scope/utils/timer.py b/memory_scope/utils/timer.py index 9e035302..62071be2 100644 --- a/memory_scope/utils/timer.py +++ b/memory_scope/utils/timer.py @@ -4,8 +4,21 @@ from memory_scope.utils.logger import Logger class Timer(object): + """ + A class used to measure the execution time of code blocks. It supports logging the elapsed time and can be customized + to display time in seconds or milliseconds. + """ def __init__(self, name: str, log_time: bool = True, use_ms: bool = True, **kwargs): + """ + Initializes the Timer object with a name, logging preference, time unit preference, and additional keyword arguments. + + Args: + name (str): The name associated with this timer instance, often used in logs. + log_time (bool, optional): Determines if the elapsed time should be logged. Defaults to True. + use_ms (bool, optional): Specifies whether to use milliseconds as the time unit in logs. Defaults to True. + **kwargs: Additional keyword arguments that might be utilized by the logger or other components. + """ self.name: str = name self.log_time: bool = log_time self.use_ms: bool = use_ms @@ -20,16 +33,26 @@ class Timer(object): @classmethod def kwargs_to_str(cls, float_precision: int = 4, **kwargs): + """ + Converts keyword arguments into a formatted string, with floats controlled by a precision setting. + + Args: + float_precision (int, optional): The number of decimal places for floating point numbers. Defaults to 4. + **kwargs: Arbitrary keyword arguments to be converted into strings. + + Returns: + str: A single string composed of the keyword arguments and their values, separated by spaces. + """ line_list = [] for k, v in kwargs.items(): if isinstance(v, float): float_style = f".{float_precision}f" - line = f"{k}={v:{float_style}}" + line = f"{k}={v:{float_style}}" # Format float value with specified precision else: - line = f"{k}={v}" + line = f"{k}={v}" # Keep other types as is line_list.append(line) - return " ".join(line_list) + return " ".join(line_list) # Join all parts into a single string with spaces def __enter__(self): self.t_start = time.time() @@ -37,6 +60,15 @@ class Timer(object): return self def __exit__(self, exc_type, exc_val, exc_tb): + """ + Records the end time of the timed code block and calculates the elapsed time. + Logs the time cost if logging is enabled, with optional message customization. + + Args: + exc_type: The exception type (unused). + exc_val: The exception value (unused). + exc_tb: The traceback (unused). + """ self.t_end = time.time() self.cost = self.t_end - self.t_start if self.use_ms: @@ -57,6 +89,13 @@ class Timer(object): @property def cost_str(self): + """ + Returns a string representation of the time cost, formatted as seconds or milliseconds + based on the `use_ms` attribute. + + Returns: + A string indicating the time cost in the chosen unit (seconds or milliseconds). + """ if self.use_ms: return f"{self.cost:.1f}ms" else: @@ -64,7 +103,26 @@ class Timer(object): def timer(func): + """ + A decorator function that measures the execution time of the wrapped function. + + Args: + func (Callable): The function to be wrapped and timed. + + Returns: + Callable: The wrapper function that includes timing functionality. + """ def wrapper(*args, **kwargs): + """ + The wrapper function that manages the timing of the original function. + + Args: + *args: Variable length argument list for the decorated function. + **kwargs: Arbitrary keyword arguments for the decorated function. + + Returns: + Any: The result of the decorated function. + """ with Timer(name=func.__name__, **kwargs): return func(*args, **kwargs) diff --git a/memory_scope/utils/tool_functions.py b/memory_scope/utils/tool_functions.py index 98f27f0e..a3571fb5 100644 --- a/memory_scope/utils/tool_functions.py +++ b/memory_scope/utils/tool_functions.py @@ -25,6 +25,15 @@ def underscore_to_camelcase(name: str, is_first_title: bool = True): def camelcase_to_underscore(name: str): + """ + Converts a CamelCase string to underscore_notation. + + Args: + name (str): The CamelCase formatted string to be converted. + + Returns: + str: The converted string in underscore_notation. + """ return re.sub(r'(? List[Message]: + """ + Converts input strings into a structured list of message objects suitable for AI interactions. + + Args: + system_prompt (str): The system-level instruction or context. + few_shot (str): An example or demonstration input, often used for illustrating expected behavior. + user_query (str): The actual user query or prompt to be processed. + + Returns: + List[Message]: A list of Message objects, each representing a part of the conversation setup. + """ return [ - Message(role=MessageRoleEnum.SYSTEM.value, content=system_prompt.strip()), + Message(role=MessageRoleEnum.SYSTEM.value, content=system_prompt.strip()), # System message Message(role=MessageRoleEnum.USER.value, - content="\n".join([x.strip() for x in [few_shot, system_prompt, user_query]])) + content="\n".join([x.strip() for x in [few_shot, system_prompt, user_query]])) # User message combining few shot, system prompt, and user query ] @@ -106,12 +126,31 @@ def char_logo(words: str, seed: int = time.time_ns(), color=None): def md5_hash(input_string: str): + """ + Computes the MD5 hash of a given input string. + + Args: + input_string (str): The string for which the MD5 hash needs to be computed. + + Returns: + str: The hexadecimal representation of the MD5 hash. + """ m = hashlib.md5() m.update(input_string.encode('utf-8')) return m.hexdigest() def contains_keyword(text, keywords): + """ + Checks if the given text contains any of the specified keywords, ignoring case. + + Args: + text (str): The text to search within. + keywords (List[str]): A list of keywords to look for in the text. + + Returns: + bool: True if any keyword is found in the text, False otherwise. + """ escaped_keywords = map(re.escape, keywords) pattern = re.compile('|'.join(escaped_keywords), re.IGNORECASE) return pattern.search(text) is not None