Update Chat Model Serializtion

langchain-ai · May 11, 2023 · 5c7efd3 · 5c7efd3
1 parent 32c1c33
commit 5c7efd3
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 25 deletions.
diff --git a/langchain/src/callbacks/handlers/tracer.ts b/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,

diff --git a/langchain/src/stores/message/dynamodb.ts b/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;
           }),

diff --git a/langchain/src/stores/message/redis.ts b/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);
   }

diff --git a/langchain/src/stores/message/utils.ts b/langchain/src/stores/message/utils.ts
@@ -6,10 +6,17 @@ 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;
-  role: string | undefined;
-  text: string;
+  data: StoredMessageData;
 }
 
 export function mapStoredMessagesToChatMessages(
@@ -18,17 +25,21 @@ export function mapStoredMessagesToChatMessages(
   return messages.map((message) => {
     switch (message.type) {
       case "human":
-        return new HumanChatMessage(message.text);
+        return new HumanChatMessage(message.data.content);
       case "ai":
-        return new AIChatMessage(message.text);
+        return new AIChatMessage(message.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(message.data.content);
+      case "chat":
+        if (message.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(
+          message.data.content,
+          message.data.additional_kwargs.role
+        );
+      default:
+        throw new Error(`Got unexpected type: ${message.type}`);
     }
   });
 }
@@ -38,7 +49,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,
+    },
   }));
 }