记忆栈：构建高效的对话记忆系统

创作时间:

作者:

@小白创作中心

记忆栈：构建高效的对话记忆系统

引用

CSDN

https://m.blog.csdn.net/weixin_36829761/article/details/141032266

在现代人工智能应用中，记忆的构建与管理是一个至关重要的环节。随着用户交互的复杂性增加，如何有效地保存和检索对话历史成为了一个技术挑战。本文将探讨一个基于栈结构的记忆系统实现，名为 StackMemory，它使用了 Python 的数据类和 TinyDB 库来实现一个内存或持久化存储的对话记忆系统。

栈节点：结构与功能

StackMemoryNode 类是我们记忆系统的基本单元。每个节点都代表一次对话的一个片段，包含以下属性：

role：表示节点的角色，可以是 user、system、root 或 assistant。
type：节点的类型，分为 text 和 list，指示内容的格式。
content：节点的具体内容，可以是文本信息或 JSON 格式的列表。
node_id：节点的唯一标识符。
parent：指向父节点的 ID，构成树状结构。
childrens：一个列表，存储当前节点的子节点 ID。

from dataclasses import dataclass
from typing import List, Union
import json

@dataclass
class StackMemoryNode:
    role: str
    type: str = None
    content: str = None
    node_id: int = None
    parent: int = None
    childrens: List[int] = None

    def __post_init__(self):
        assert self.role in ['user', 'system', 'root', 'assistant'], self.role
        self.childrens = self.childrens if self.childrens else []
        if self.type is None:
            self.type = 'text'

通过定义节点的结构，我们能够灵活地管理对话的记录。每个节点不仅可以保存对话内容，还能记录其在对话树中的位置，从而方便后续的检索和操作。

记忆栈的构建

StackMemory 类是整个记忆栈的管理器，负责创建、更新和删除节点。它的初始化方法可以选择使用内存存储或持久化存储：

from tinydb import TinyDB, MemoryStorage

class StackMemory:
    def __init__(self, serialize_path='./memory.json'):
        if serialize_path is not None:
            self.db = TinyDB(serialize_path)
        else:
            self.db = TinyDB(storage=MemoryStorage)

在初始化过程中，系统会检查是否已有存储的节点，并加载这些节点。如果没有节点，则会创建一个根节点。

节点的添加与删除

在 StackMemory 中，添加节点的方法分为几种场景，包括在根节点下添加新节点、在现有节点之后添加节点，以及在父节点内部添加新节点。

添加节点示例

假设我们需要在当前对话中添加一条用户消息，我们可以使用如下方法：

def add_message(self, role, message: Union[str, list]):
    assert role in ['user', 'system', 'assistant'], role
    type = 'text'
    if isinstance(message, list):
        type = 'list'
        message = json.dumps(message)
    new_node = StackMemoryNode(role=role, content=message, type=type)
    if self.next_position == 'after':
        self.add_node_after(self.current_node, new_node)
    else:
        self.add_node_in(self.current_node, new_node)
    self.next_position = 'after'
    self.set_current_node(new_node)
    return new_node.node_id

通过调用 add_message 方法，我们可以将用户的消息以节点的形式添加到对话树中，确保每次交互的内容被准确记录。

节点的检索与更新

在对话过程中，用户可能需要查看当前节点的所有相关信息。这可以通过 get_related_messages_for_node 方法实现。该方法会返回与当前节点相关的所有消息，包括其父节点的消息和兄弟节点的消息。

def get_related_messages_for_node(self, node: StackMemoryNode):
    nodes_with_position = self.get_related_nodes_for_node(node)
    messages = [_parse_node(node) for position, node in nodes_with_position]
    return messages

该方法利用了节点间的关系，确保在显示消息时能够提供完整的上下文。

深度与层级管理

为了管理对话的深度与层级，StackMemory 提供了 get_node_level 和 get_level 方法。它们能够返回一个节点在树中的层级，帮助开发者了解当前对话的结构。

def get_node_level(self, node: StackMemoryNode):
    if node.is_root():
        return 0
    else:
        return self.get_node_level(self.get_node_parent(node)) + 1

这种层级管理使得在复杂对话中，开发者能够快速定位到特定的对话节点。