Passing additional_kwargs to messages being stored by memory in ConversationalRetrievalQAChain

[rysloan4]

I have a use case where I'd like to store the timestamp of the messages being sent in a chat app. I'd like to pass the information through the additional_kwargs argument in BaseMessage and can do so manually like:

const userMessage = new HumanMessage({
      content: 'Some Content'
      additional_kwargs: {
        timestamp: new Date().getTime()
      }
    })

However, within the context of a ConversationalRetrievalQAChain I can't figure out a way to specify additional_kwargs. Alternately, if there was a way for the chain to simple read from the BufferMemory, I could just manage inserting messages outside the chain. Any thoughts would be greatly appreciated.

-Best-

[dosubot[bot]]

🤖

Hello,

Thank you for reaching out with your question. Currently, the LangChainJS framework does not support the additionalKwargs field in the BufferMemory configuration. The BufferMemory class constructor only accepts the following fields: chatHistory, returnMessages, inputKey, outputKey, humanPrefix, aiPrefix, and memoryKey. There is no mention or use of additionalKwargs in the current implementation.

However, you can extend the BufferMemory class to support additional kwargs. Here's a simple example of how you might do this:

class ExtendedBufferMemory extends BufferMemory {
  constructor(config) {
    super(config);
  }

  insertMessage(message) {
    message.additional_kwargs = {
      timestamp: new Date().getTime()
    };
    super.insertMessage(message);
  }
}

Then, you can use ExtendedBufferMemory instead of BufferMemory when initializing your ConversationalRetrievalQAChain:

const chain = ConversationalRetrievalQAChain.fromLLM(
  slowerModel,
  vectorStore.asRetriever(),
  {
    returnSourceDocuments: true,
    memory: new ExtendedBufferMemory({
      memoryKey: "chat_history",
      inputKey: "question",
      outputKey: "text",
      returnMessages: true,
    }),
    questionGeneratorChainOptions: {
      llm: fasterModel,
    },
  }
);

This will add a timestamp to every message that is inserted into the memory. Please note that this is a simple example and you might need to adjust it to fit your specific needs.

I hope this helps! If you have any other questions, feel free to ask.

Best regards.

Sources

• langchain/src/chains/tests/conversational_retrieval_chain.int.test.ts
• examples/src/chains/conversational_qa_built_in_memory.ts
• langchain/src/memory/buffer_memory.ts

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[rysloan4]

simplest solution was to do the following:

import { type BaseMessage } from 'langchain/schema'
import { RedisChatMessageHistory } from 'langchain/stores/message/ioredis'

class ExtendedRedisChatMemory extends RedisChatMessageHistory {
  async addMessage (message: BaseMessage): Promise<void> {
    message.additional_kwargs = {
      timestamp: new Date().getTime()
    }
    await super.addMessage(message)
  }
}

export default ExtendedRedisChatMemory

Which is a bit of a hack. Ideally I could pass in the timestamp of the message being sent from the client for the HumanMessage as opposed to the system time.

[rollsrobby]

Hi, did you ever manage to solve this using the timestamp of the message? I'm facing a similar problem, but I can't solve it using your proposal since the value I need to save comes from the output of an agent...

[stevenmilstein]

I have a similar question as @rollsrobby .

Ideally, I'd like to pass a parameter (not derived like timestamp) that will update a new column in the Xata Messages table, but I'll settle for persisting it in additionalKwargs.

In my case, I was thinking about extending BufferMemory because the new parameter would be a composable key with sessionId

[stevenmilstein]

With the help of ChatGPT, here's the solution that works for me.

After adding column userId to my xata messages table I extended XataChatMessageHistory to treat sessionId and userId as a composite key as follows:

import { XataChatMessageHistory } from "@langchain/community/stores/message/xata";
import {
  mapChatMessagesToStoredMessages,
  mapStoredMessagesToChatMessages,
} from "@langchain/core/messages";

/**
 * @name XataChatMessageHistoryExtended
 * @description Extended XataChatMessageHistory to add userId parameter for messages table
 */
class XataChatMessageHistoryExtended extends XataChatMessageHistory {
  // Add your custom properties and methods here

  constructor(fields) {
    super(fields);
    // Additional initialization for your custom class
    this.userId = fields.userId;
  }

  // Override the getMessage method to include userId
  async getMessages() {
    await this.ensureTable();
    const records = await this.client.db[this.table]
      .filter({ sessionId: this.sessionId, userId: this.userId }) // Include userId in the filter
      .sort("xata.createdAt", "asc")
      .getAll();

    const rawStoredMessages = records.map((message) => {
      const data = {
        content: message.content,
        additional_kwargs: JSON.parse(message.additionalKwargs),
      };
      if (message.role) {
        data.role = message.role;
      }
      if (message.name) {
        data.name = message.name;
      }

      return {
        type: message.type,
        data,
      };
    });
    return mapStoredMessagesToChatMessages(rawStoredMessages);
  }

  // Override the addMessage method to include userId
  async addMessage(message) {
    try {
      await this.ensureTable();
      const messageToAdd = mapChatMessagesToStoredMessages([message]);
      await this.client.db[this.table].create({
        sessionId: this.sessionId,
        userId: this.userId, // Include the userId
        type: messageToAdd[0].type,
        content: messageToAdd[0].data.content,
        role: messageToAdd[0].data.role,
        name: messageToAdd[0].data.name,
        additionalKwargs: JSON.stringify(
          messageToAdd[0].data.additional_kwargs
        ),
      });
    } catch (err) {
      console.log(err);
      throw err;
    }
  }

  // Override the getMessage method to include userId
  async clear() {
    await this.ensureTable();
    const records = await this.client.db[this.table]
      .select(["id"])
      .filter({ sessionId: this.sessionId, userId: this.userId }) // Include userId in the filter
      .getAll();
    const ids = records.map((m) => m.id);
    await this.client.db[this.table].delete(ids);
  }
}

Then I used XataChatMessageHistoryExtended instead of XataChatMessageHistory