From ed33749cf68c3b138924c4ed0df1e6c61abe9796 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Mon, 5 Jan 2026 22:39:10 +0800 Subject: [PATCH] feat(memory): add memory node schema and utilities --- reme_ai/core/enumeration/__init__.py | 2 + reme_ai/core/enumeration/memory_type.py | 25 +++ reme_ai/core/schema/__init__.py | 2 + reme_ai/core/schema/memory_node.py | 213 +++++++++++++++++++++++ reme_ai/core/utils/__init__.py | 3 +- reme_ai/core/utils/llm_utils.py | 23 ++- reme_ai/core/utils/{timer.py => time.py} | 10 ++ 7 files changed, 274 insertions(+), 4 deletions(-) create mode 100644 reme_ai/core/enumeration/memory_type.py create mode 100644 reme_ai/core/schema/memory_node.py rename reme_ai/core/utils/{timer.py => time.py} (89%) diff --git a/reme_ai/core/enumeration/__init__.py b/reme_ai/core/enumeration/__init__.py index 3323bdc9..7202a949 100644 --- a/reme_ai/core/enumeration/__init__.py +++ b/reme_ai/core/enumeration/__init__.py @@ -3,6 +3,7 @@ from .chunk_enum import ChunkEnum from .http_enum import HttpEnum from .json_schema_enum import JsonSchemaEnum +from .memory_type import MemoryType from .registry_enum import RegistryEnum from .role import Role @@ -10,6 +11,7 @@ __all__ = [ "ChunkEnum", "HttpEnum", "JsonSchemaEnum", + "MemoryType", "RegistryEnum", "Role", ] diff --git a/reme_ai/core/enumeration/memory_type.py b/reme_ai/core/enumeration/memory_type.py new file mode 100644 index 00000000..22d35481 --- /dev/null +++ b/reme_ai/core/enumeration/memory_type.py @@ -0,0 +1,25 @@ +"""Memory type enumeration for the three-layer memory architecture.""" + +from enum import Enum + + +class MemoryType(str, Enum): + """ + Three-layer memory architecture for agent memory management. + + Layer 1 - High-level Abstraction Memory: + - IDENTITY: Self-cognition (identity, personality, current state) + - PERSONAL: Person-specific memory (preferences and context about specific individuals) + - PROCEDURAL: Procedural memory (how-to knowledge, e.g., 4 steps to write financial reports) + - TOOL: Tool memory (tool usage patterns, success rates, token consumption, latency) + + Layer 2 - Summary Memory (Compressed): Summarized digest of raw message history + Layer 3 - History Memory (Raw): Raw message history + """ + + IDENTITY = "identity" + PERSONAL = "personal" + PROCEDURAL = "procedural" + TOOL = "tool" + SUMMARY = "summary" + HISTORY = "history" diff --git a/reme_ai/core/schema/__init__.py b/reme_ai/core/schema/__init__.py index f55f054e..b7b73719 100644 --- a/reme_ai/core/schema/__init__.py +++ b/reme_ai/core/schema/__init__.py @@ -1,5 +1,6 @@ """schema""" +from .memory_node import MemoryNode from .message import ContentBlock, Message, Trajectory from .request import Request from .response import Response @@ -19,6 +20,7 @@ from .tool_call import ToolAttr, ToolCall from .vector_node import VectorNode __all__ = [ + "MemoryNode", "ContentBlock", "EmbeddingModelConfig", "FlowConfig", diff --git a/reme_ai/core/schema/memory_node.py b/reme_ai/core/schema/memory_node.py new file mode 100644 index 00000000..b8b6c92a --- /dev/null +++ b/reme_ai/core/schema/memory_node.py @@ -0,0 +1,213 @@ +"""Memory schema module for the ReMe AI system. + +This module defines the MemoryNode class for storing and retrieving +memories in the ReMe system. +""" + +import hashlib +import json +from typing import Any + +from pydantic import BaseModel, Field, model_validator + +from ..enumeration import MemoryType +from ..schema import VectorNode +from ..utils import get_now_time + +# Length of the memory ID (first N characters of SHA-256 hash) +MEMORY_ID_LENGTH: int = 16 + + +class MemoryNode(BaseModel): + """Memory node for storing memories in the ReMe system. + + Attributes: + memory_id: Unique identifier, auto-generated from content hash. + memory_type: Type of memory (e.g., SUMMARY, PERSONAL). + memory_target: Target or topic this memory relates to. + when_to_use: Condition description for vector retrieval. + content: Actual memory content. + ref_memory_id: Reference to related raw history memory. + time_created: Creation timestamp. + time_modified: Last modification timestamp. + author: Author or source of this memory. + score: Relevance or importance score. + metadata: Additional metadata for extensibility. + """ + + memory_id: str = Field(default="", description="Unique memory identifier") + memory_type: MemoryType = Field(default=..., description="Type of memory") + memory_target: str = Field(default="", description="Target or topic of the memory") + when_to_use: str = Field(default="", description="Condition description for vector retrieval") + content: str = Field(default="", description="Actual memory content") + ref_memory_id: str = Field(default="", description="Reference to related raw history memory ID") + + time_created: str = Field(default_factory=get_now_time, description="Creation timestamp") + time_modified: str = Field(default_factory=get_now_time, description="Last modification timestamp") + author: str = Field(default="", description="Author or source of the memory") + score: float = Field(default=0, description="Relevance or importance score") + + metadata: dict[str, Any] = Field(default_factory=dict, description="Additional metadata") + + def _update_modified_time(self) -> "MemoryNode": + """Update time_modified to current timestamp. + + Returns: + Self: Returns self for method chaining. + """ + self.time_modified = get_now_time() + return self + + def _update_memory_id(self) -> "MemoryNode": + """Generate memory_id from SHA-256 hash of content. + + Takes the first MEMORY_ID_LENGTH characters of the hash. + + Returns: + Self: Returns self for method chaining. + """ + if not self.content: + return self + + hash_obj = hashlib.sha256(self.content.encode("utf-8")) + hex_dig = hash_obj.hexdigest() + self.memory_id = hex_dig[:MEMORY_ID_LENGTH] + return self + + @model_validator(mode="after") + def _update_after_init(self) -> "MemoryNode": + """Post-initialization validator. + + Auto-generates memory_id from content if not provided. + + Returns: + Self: Returns self for method chaining. + """ + if not self.memory_id: + self._update_memory_id() + return self + + def __setattr__(self, name: str, value): + """Auto-update timestamps and memory_id when content or when_to_use changes. + + Args: + name: Attribute name being set. + value: New value for the attribute. + """ + should_update: bool = name in ("when_to_use", "content") and getattr(self, name, None) != value + super().__setattr__(name, value) + if should_update: + self._update_modified_time() + if name == "content": + self._update_memory_id() + + def to_vector_node(self) -> VectorNode: + """Convert to VectorNode for vector storage. + + When when_to_use is set, use it as vector content and store content in metadata. + When when_to_use is empty, use content as vector content directly. + + Returns: + VectorNode: Vector node representation of this memory. + """ + # Build base metadata (shared fields) + metadata: dict[str, Any] = { + "memory_type": self.memory_type.value, + "memory_target": self.memory_target, + "ref_memory_id": self.ref_memory_id, + "time_created": self.time_created, + "time_modified": self.time_modified, + "author": self.author, + "score": self.score, + **self.metadata, + } + + if self.when_to_use: + # Use when_to_use for vector embedding, store content in metadata + vector_content = self.when_to_use + metadata["content"] = self.content + else: + # Use content directly for vector embedding + vector_content = self.content + + return VectorNode( + vector_id=self.memory_id, + content=vector_content, + metadata=metadata, + ) + + def format_memory(self) -> str: + """Format memory as human-readable string. + + Returns: + str: Formatted string with when_to_use, content, and ref_memory_id. + """ + parts: list[str] = [ + f"memory_id={self.memory_id}" f"modified_time={self.time_modified}", + ] + + if self.when_to_use: + parts.append(self.when_to_use) + + if self.content: + parts.append(self.content) + + if self.metadata: + parts.append(f"metadata={json.dumps(self.metadata, ensure_ascii=False)}") + + if self.ref_memory_id: + parts.append(f"history_memory.ref_memory_id={self.ref_memory_id}") + + return " ".join(parts) + + @classmethod + def from_vector_node(cls, node: VectorNode) -> "MemoryNode": + """Reconstruct MemoryNode from VectorNode. + + Reverses the to_vector_node conversion: + - If metadata contains 'content': node.content -> when_to_use, metadata['content'] -> content + - Otherwise: node.content -> content, when_to_use remains empty + + Args: + node: VectorNode containing memory data. + + Returns: + Self: Reconstructed MemoryNode instance. + + Raises: + ValueError: If memory_type in metadata is invalid. + """ + metadata = node.metadata.copy() + memory_type_str = metadata.pop("memory_type", None) + + try: + memory_type: MemoryType = MemoryType(memory_type_str) + except ValueError as e: + raise ValueError( + f"Invalid memory_type '{memory_type_str}' in VectorNode metadata. " + f"Valid types are: {[t.value for t in MemoryType]}", + ) from e + + # Restore when_to_use and content based on metadata structure + if "content" in metadata: + # Original had when_to_use set + when_to_use = node.content + content = metadata.pop("content", "") + else: + # Original had empty when_to_use + when_to_use = "" + content = node.content + + return cls( + memory_id=node.vector_id, + memory_type=memory_type, + memory_target=metadata.pop("memory_target", ""), + when_to_use=when_to_use, + content=content, + ref_memory_id=metadata.pop("ref_memory_id", ""), + time_created=metadata.pop("time_created", ""), + time_modified=metadata.pop("time_modified", ""), + author=metadata.pop("author", ""), + score=metadata.pop("score", 0), + metadata=metadata, + ) diff --git a/reme_ai/core/utils/__init__.py b/reme_ai/core/utils/__init__.py index 3d3bdf9a..a3f018de 100644 --- a/reme_ai/core/utils/__init__.py +++ b/reme_ai/core/utils/__init__.py @@ -13,7 +13,7 @@ from .mcp_client import MCPClient from .pydantic_config_parser import PydanticConfigParser from .pydantic_utils import create_pydantic_model from .singleton import singleton -from .timer import timer +from .time import timer, get_now_time __all__ = [ "CacheHandler", @@ -34,4 +34,5 @@ __all__ = [ "create_pydantic_model", "singleton", "timer", + "get_now_time", ] diff --git a/reme_ai/core/utils/llm_utils.py b/reme_ai/core/utils/llm_utils.py index ae0e763f..2362a278 100644 --- a/reme_ai/core/utils/llm_utils.py +++ b/reme_ai/core/utils/llm_utils.py @@ -4,19 +4,27 @@ import json import re from ..enumeration import Role -from ..schema import Message +from ..schema import Message, MemoryNode def format_messages(messages: list[Message | dict], enable_system: bool = False) -> str: """Formats a list of messages into a single string, optionally filtering system roles.""" formatted_lines = [] - for message in messages: + for i, message in enumerate(messages): if isinstance(message, dict): message = Message(**message) if not enable_system and message.role is Role.SYSTEM: continue - formatted_lines.append(message.format_message()) + formatted_lines.append( + message.format_message( + index=i, + add_time=True, + use_name=True, + add_reasoning=True, + add_tools=True, + ), + ) return "\n".join(formatted_lines) @@ -38,3 +46,12 @@ def extract_content(text: str, language_tag: str = "json", greedy: bool = False) result = None return result + + +def deduplicate_memories(memories: list[MemoryNode]) -> list[MemoryNode]: + """Deduplicates a list of memories by memory ID.""" + seen_memories: dict[str, MemoryNode] = {} + for memory in memories: + if memory.memory_id not in seen_memories: + seen_memories[memory.memory_id] = memory + return list(seen_memories.values()) diff --git a/reme_ai/core/utils/timer.py b/reme_ai/core/utils/time.py similarity index 89% rename from reme_ai/core/utils/timer.py rename to reme_ai/core/utils/time.py index f03225fa..541feab7 100644 --- a/reme_ai/core/utils/timer.py +++ b/reme_ai/core/utils/time.py @@ -2,6 +2,7 @@ Utility module for timing function execution with log metadata preservation. """ +import datetime import functools import inspect import time @@ -13,6 +14,15 @@ from loguru import logger F = TypeVar("F", bound=Callable[..., Any]) +def get_now_time() -> str: + """Get current timestamp in YYYY-MM-DD HH:MM:SS format. + + Returns: + str: Current timestamp string in format 'YYYY-MM-DD HH:MM:SS'. + """ + return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + + def timer(func: F) -> F: """ Decorator that logs execution time and patches log records with original function metadata.