MastraAI•6d ago

How to dictate assistant message ID when using Agent + Memory?

We have a use case where we need to persist some data after generation in our DB, and want to associate it with the assistant message ID. It would be best if we could dictate that ID to begin with. But it would be okay if we could also just retrieve it from agent.stream. We're on 0.21.0. Using generateMessageId does not seem to work for persistence. I could not find a Memory class property that does this either. We're using the Postgres storage layer for Memory.

4 Replies

ArielOP•6d ago

We landed on a pretty ugly workaround of patching saveMessages, generating the memory object inside a closure, where assistantMessageId is a param, and patching it before saving.

  function patchedSaveMessages(
    args: SaveMessagesV1Args,
  ): Promise<MastraMessageV1[]>;
  function patchedSaveMessages(
    args: SaveMessagesV2Args,
  ): Promise<MastraMessageV2[]>;
  async function patchedSaveMessages(
    args: SaveMessagesArgs,
  ): Promise<MastraMessageV1[] | MastraMessageV2[]> {
    console.log("### CALLED PATCHED SAVE MESSAGES ###", args);
    const nextMessages = args.messages.map((message) => {
      const typedMessage = message as MastraMessageV2;

      if (typedMessage.role === "assistant" && assistantMessageId) {
        typedMessage.id = assistantMessageId;
      }

  function patchedSaveMessages(
    args: SaveMessagesV1Args,
  ): Promise<MastraMessageV1[]>;
  function patchedSaveMessages(
    args: SaveMessagesV2Args,
  ): Promise<MastraMessageV2[]>;
  async function patchedSaveMessages(
    args: SaveMessagesArgs,
  ): Promise<MastraMessageV1[] | MastraMessageV2[]> {
    console.log("### CALLED PATCHED SAVE MESSAGES ###", args);
    const nextMessages = args.messages.map((message) => {
      const typedMessage = message as MastraMessageV2;

      if (typedMessage.role === "assistant" && assistantMessageId) {
        typedMessage.id = assistantMessageId;
      }

if there's a cleaner way, I'd love to know.

Mastra Triager•5d ago

📝 Created GitHub issue: https://github.com/mastra-ai/mastra/issues/8952

GitHub

[DISCORD:1428430169646960721] How to dictate assistant message ID w...

This issue was created from Discord post: https://discord.com/channels/1309558646228779139/1428430169646960721 We have a use case where we need to persist some data after generation in our DB, and ...

_roamin_•5d ago

Hi @Ariel ! You can get the generated assistant message id from the response.uiMessages property :

const result = await agent.stream("...", {
  memory: {
    resource: "...",
    thread: "...",
  }
});

const response = await result.response

console.dir(response.uiMessages, { depth: Infinity });

const result = await agent.stream("...", {
  memory: {
    resource: "...",
    thread: "...",
  }
});

const response = await result.response

console.dir(response.uiMessages, { depth: Infinity });

ArielOP•4d ago

is this going to be the same ID as the ID stored in the Mastra memory DB?

Gaming

Programming

How to dictate assistant message ID when using Agent + Memory?

Did you find this page helpful?