如何扩展Langchain DynamoDBChatMessageHistory添加排序键与自定义属性

你已基于Langchain搭建好使用DynamoDB存储对话记忆的对话链，消息读写功能正常，但直接在DynamoDBChatMessageHistory构造参数中添加sortKey、sessionId和dateTimeStamp不会生效——因为该类默认不支持这些自定义配置。以下是扩展基类的具体实现方案：

原问题代码 #

const { DynamoDBChatMessageHistory } = require("@langchain/community/stores/message/dynamodb");

const MESSAGE_HISTORY = new DynamoDBChatMessageHistory(
  {
    tableName: "Conversation",
    partitionKey: "studentId", // 文档要求该键唯一，但你希望用studentId作为分区键
    studentId: studentId.toString(),
    sortKey: "sessionId", // 想要添加的排序键
    sessionId: `chat_${sessionId}`, // 未生效
    dateTimeStamp: new Date(), // 想要添加的自定义属性
    config: {
      region: "us-east-1",
      credentials: {
        accessKeyId: process.env.ACCESS_KEY_ID,
        secretAccessKey: process.env.SECRET_ACCESS_KEY,
      },
    },
  }
)

解决方案：扩展基类实现自定义逻辑 #

通过继承DynamoDBChatMessageHistory，重写构造函数和消息操作方法，实现排序键和自定义属性的写入与查询：

const { DynamoDBChatMessageHistory } = require("@langchain/community/stores/message/dynamodb");
const { DynamoDBClient, PutItemCommand, QueryCommand } = require("@aws-sdk/client-dynamodb");
const { marshall, unmarshall } = require("@aws-sdk/util-dynamodb");
const { HumanMessage, AIMessage, ChatMessage } = require("@langchain/core/messages");

class CustomDynamoDBHistory extends DynamoDBChatMessageHistory {
  constructor({
    tableName,
    partitionKey,
    partitionValue,
    sortKey,
    sortValue,
    dateTimeStamp,
    config
  }) {
    // 调用父类构造函数，传入必要参数
    super({ tableName, partitionKey, partitionValue, config });
    // 保存自定义的排序键、值和时间戳
    this.sortKey = sortKey;
    this.sortValue = sortValue;
    this.dateTimeStamp = dateTimeStamp;
    // 初始化DynamoDB客户端
    this.client = new DynamoDBClient(config);
  }

  // 重写添加消息方法，附带排序键和自定义属性
  async addMessage(message) {
    const item = {
      [this.partitionKey]: this.partitionValue,
      [this.sortKey]: this.sortValue,
      content: message.content,
      role: message._getType(),
      dateTimeStamp: this.dateTimeStamp.toISOString(),
      // 可按需添加更多自定义属性
    };

    const command = new PutItemCommand({
      TableName: this.tableName,
      Item: marshall(item),
    });

    await this.client.send(command);
  }

  // 重写获取消息方法，按排序键过滤
  async getMessages() {
    const command = new QueryCommand({
      TableName: this.tableName,
      KeyConditionExpression: "#pk = :pk AND #sk = :sk",
      ExpressionAttributeNames: {
        "#pk": this.partitionKey,
        "#sk": this.sortKey,
      },
      ExpressionAttributeValues: marshall({
        ":pk": this.partitionValue,
        ":sk": this.sortValue,
      }),
      // 可选：按时间戳倒序排列最新消息
      ScanIndexForward: false,
    });

    const response = await this.client.send(command);
    return response.Items.map((item) => {
      const unmarshalled = unmarshall(item);
      // 根据消息角色创建对应实例
      if (unmarshalled.role === "human") {
        return new HumanMessage(unmarshalled.content);
      } else if (unmarshalled.role === "ai") {
        return new AIMessage(unmarshalled.content);
      }
      return new ChatMessage(unmarshalled.content, unmarshalled.role);
    });
  }
}

// 使用自定义类初始化对话记忆
const MESSAGE_HISTORY = new CustomDynamoDBHistory({
  tableName: "Conversation",
  partitionKey: "studentId",
  partitionValue: studentId.toString(),
  sortKey: "sessionId",
  sortValue: `chat_${sessionId}`,
  dateTimeStamp: new Date(),
  config: {
    region: "us-east-1",
    credentials: {
      accessKeyId: process.env.ACCESS_KEY_ID,
      secretAccessKey: process.env.SECRET_ACCESS_KEY,
    },
  },
});

关键说明 #

1. 类继承设计：保留父类核心逻辑，仅修改需要自定义的消息读写流程，避免重复造轮子。
2. 构造函数扩展：新增接收排序键、排序值和时间戳参数，将其保存为实例属性供后续方法调用。
3. 写入逻辑重写：手动构建DynamoDB写入条目，强制包含分区键、排序键、消息元数据和自定义时间戳。
4. 查询逻辑重写：使用QueryCommand替代默认查询，通过分区键+排序键的复合条件精准定位会话消息。
5. 表结构匹配：确保你的DynamoDB表Conversation已配置复合主键：分区键studentId，排序键sessionId，同时预留dateTimeStamp属性字段。

内容的提问来源于stack exchange，提问作者Dinuda Yaggahavita