Update Chat Model Serialization

langchain/src/callbacks/handlers/tracer.ts

@@ -6,6 +6,7 @@ import {
   RunInputs,
   RunOutputs,
 } from "../../schema/index.js";
+import { mapChatMessagesToStoredMessages } from "../../stores/message/utils.js";
 import { BaseCallbackHandler } from "../base.js";
 
 export type RunType = "llm" | "chain" | "tool";
@@ -132,10 +133,7 @@ export abstract class BaseTracer extends BaseCallbackHandler {
   ): Promise<void> {
     const execution_order = this._getExecutionOrder(parentRunId);
     const convertedMessages = messages.map((batch) =>
-      batch.map((message) => ({
-        _type: message._getType(),
-        content: message.text, // TODO: Unify serialization btwn languages
-      }))
+      mapChatMessagesToStoredMessages(batch)
     );
     const run: Run = {
       id: runId,

langchain/src/callbacks/tests/tracer.test.ts

@@ -94,7 +94,11 @@ test("Test Chat Message Run", async () => {
     execution_order: 1,
     child_execution_order: 1,
     serialized: { name: "test" },
-    inputs: { messages: [[{ _type: "human", content: "Avast" }]] },
+    inputs: {
+      messages: [
+        [{ type: "human", data: { content: "Avast", role: undefined } }],
+      ],
+    },
     run_type: "llm",
     outputs: { generations: [] },
     child_runs: [],

langchain/src/stores/message/dynamodb.ts

@@ -95,11 +95,14 @@ export class DynamoDBChatMessageHistory extends BaseListChatMessageHistory {
     const messages = items
       .map((item) => ({
         type: item.M?.type.S,
-        role: item.M?.role?.S,
-        text: item.M?.text.S,
+        data: {
+          role: item.M?.role?.S,
+          content: item.M?.text.S,
+        },
       }))
       .filter(
-        (x): x is StoredMessage => x.type !== undefined && x.text !== undefined
+        (x): x is StoredMessage =>
+          x.type !== undefined && x.data.content !== undefined
       );
     return mapStoredMessagesToChatMessages(messages);
   }
@@ -133,12 +136,12 @@ export class DynamoDBChatMessageHistory extends BaseListChatMessageHistory {
                   S: message.type,
                 },
                 text: {
-                  S: message.text,
+                  S: message.data.content,
                 },
               },
             };
-            if (message.role) {
-              dynamoSerializedMessage.M.role = { S: message.role };
+            if (message.data.role) {
+              dynamoSerializedMessage.M.role = { S: message.data.role };
             }
             return dynamoSerializedMessage;
           }),

langchain/src/stores/message/redis.ts

@@ -60,11 +60,14 @@ export class RedisChatMessageHistory extends BaseListChatMessageHistory {
     const previousMessages = orderedMessages
       .map((item) => ({
         type: item.type,
-        role: item.role,
-        text: item.text,
+        data: {
+          role: item.role,
+          content: item.text,
+        },
       }))
       .filter(
-        (x): x is StoredMessage => x.type !== undefined && x.text !== undefined
+        (x): x is StoredMessage =>
+          x.type !== undefined && x.data.content !== undefined
       );
     return mapStoredMessagesToChatMessages(previousMessages);
   }

langchain/src/stores/message/utils.ts

@@ -6,29 +6,64 @@ import {
   SystemChatMessage,
 } from "../../schema/index.js";
 
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export type AdditionalKwargs = Record<string, any>;
+export interface StoredMessageData {
+  content: string;
+  role: string | undefined;
+  additional_kwargs?: AdditionalKwargs;
+}
+
 export interface StoredMessage {
+  type: string;
+  data: StoredMessageData;
+}
+
+interface StoredMessageV1 {
   type: string;
   role: string | undefined;
   text: string;
 }
 
+export function mapV1MessageToStoredMessage(
+  message: StoredMessage | StoredMessageV1
+): StoredMessage {
+  if ("data" in message) {
+    return message as StoredMessage;
+  } else {
+    const v1Message = message as StoredMessageV1;
+    return {
+      type: v1Message.type,
+      data: {
+        content: v1Message.text,
+        role: v1Message.role,
+      },
+    };
+  }
+}
+
 export function mapStoredMessagesToChatMessages(
   messages: StoredMessage[]
 ): BaseChatMessage[] {
   return messages.map((message) => {
-    switch (message.type) {
+    const stored_msg = mapV1MessageToStoredMessage(message);
+    switch (stored_msg.type) {
       case "human":
-        return new HumanChatMessage(message.text);
+        return new HumanChatMessage(stored_msg.data.content);
       case "ai":
-        return new AIChatMessage(message.text);
+        return new AIChatMessage(stored_msg.data.content);
       case "system":
-        return new SystemChatMessage(message.text);
-      default: {
-        if (message.role === undefined) {
-          throw new Error("Role must be defined for generic messages");
+        return new SystemChatMessage(stored_msg.data.content);
+      case "chat":
+        if (stored_msg.data?.additional_kwargs?.role === undefined) {
+          throw new Error("Role must be defined for chat messages");
         }
-        return new ChatMessage(message.text, message.role);
-      }
+        return new ChatMessage(
+          stored_msg.data.content,
+          stored_msg.data.additional_kwargs.role
+        );
+      default:
+        throw new Error(`Got unexpected type: ${stored_msg.type}`);
     }
   });
 }
@@ -38,7 +73,9 @@ export function mapChatMessagesToStoredMessages(
 ): StoredMessage[] {
   return messages.map((message) => ({
     type: message._getType(),
-    role: "role" in message ? (message.role as string) : undefined,
-    text: message.text,
+    data: {
+      content: message.text,
+      role: "role" in message ? (message.role as string) : undefined,
+    },
   }));
 }

langchain/src/stores/tests/utils.test.ts

@@ -0,0 +1,32 @@
+import { test, expect } from "@jest/globals";
+import { mapV1MessageToStoredMessage } from "../message/utils.js";
+
+test("mapV1MessageToStoredMessage", () => {
+  // Test that a V1 message is upgraded.
+  // Test that a v2 message is not changed.
+  const v1Message = {
+    type: "human",
+    role: "user",
+    text: "Hello, world!",
+  };
+  const v2Message = {
+    type: "human",
+    data: {
+      content: "Hello, world!",
+      role: "user",
+    },
+  };
+  expect(mapV1MessageToStoredMessage(v1Message)).toEqual(v2Message);
+
+  const v2Message2 = {
+    type: "human",
+    data: {
+      content: "Hello, world!",
+      role: "user",
+      additional_kwargs: {
+        foo: "bar",
+      },
+    },
+  };
+  expect(mapV1MessageToStoredMessage(v2Message2)).toEqual(v2Message2);
+});