From 0a178dcb705f1e8d5cc3cb98b7d8f4acd70d31de Mon Sep 17 00:00:00 2001
From: Brace Sproul <braceasproul@gmail.com>
Date: Wed, 19 Jun 2024 16:54:11 -0700
Subject: [PATCH] core[minor]: Add ID field to messages (#5817)

* core[minor]: Add ID field to messages

* added id to openai

* revert oai change
---
 langchain-core/src/messages/ai.ts       |  1 +
 langchain-core/src/messages/base.ts     | 12 ++++++++++++
 langchain-core/src/messages/chat.ts     |  1 +
 langchain-core/src/messages/function.ts |  1 +
 langchain-core/src/messages/human.ts    |  1 +
 langchain-core/src/messages/system.ts   |  1 +
 langchain-core/src/messages/tool.ts     |  1 +
 7 files changed, 18 insertions(+)

diff --git a/langchain-core/src/messages/ai.ts b/langchain-core/src/messages/ai.ts
index 7a9504038af9..aedd46f6dc75 100644
--- a/langchain-core/src/messages/ai.ts
+++ b/langchain-core/src/messages/ai.ts
@@ -244,6 +244,7 @@ export class AIMessageChunk extends BaseMessageChunk {
         chunk.response_metadata
       ),
       tool_call_chunks: [],
+      id: this.id ?? chunk.id,
     };
     if (
       this.tool_call_chunks !== undefined ||
diff --git a/langchain-core/src/messages/base.ts b/langchain-core/src/messages/base.ts
index 1135c75b84f0..c99fc9ebb016 100644
--- a/langchain-core/src/messages/base.ts
+++ b/langchain-core/src/messages/base.ts
@@ -106,6 +106,11 @@ export type BaseMessageFields = {
   /** Response metadata. For example: response headers, logprobs, token counts. */
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   response_metadata?: Record<string, any>;
+  /**
+   * An optional unique identifier for the message. This should ideally be
+   * provided by the provider/model which created the message.
+   */
+  id?: string;
 };
 
 export function mergeContent(
@@ -170,6 +175,12 @@ export abstract class BaseMessage
   /** Response metadata. For example: response headers, logprobs, token counts. */
   response_metadata: NonNullable<BaseMessageFields["response_metadata"]>;
 
+  /**
+   * An optional unique identifier for the message. This should ideally be
+   * provided by the provider/model which created the message.
+   */
+  id?: string;
+
   /** The type of the message. */
   abstract _getType(): MessageType;
 
@@ -200,6 +211,7 @@ export abstract class BaseMessage
     this.content = fields.content;
     this.additional_kwargs = fields.additional_kwargs;
     this.response_metadata = fields.response_metadata;
+    this.id = fields.id;
   }
 
   toDict(): StoredMessage {
diff --git a/langchain-core/src/messages/chat.ts b/langchain-core/src/messages/chat.ts
index 837660f398aa..fb71956e6f8d 100644
--- a/langchain-core/src/messages/chat.ts
+++ b/langchain-core/src/messages/chat.ts
@@ -90,6 +90,7 @@ export class ChatMessageChunk extends BaseMessageChunk {
         chunk.response_metadata
       ),
       role: this.role,
+      id: this.id ?? chunk.id,
     });
   }
 }
diff --git a/langchain-core/src/messages/function.ts b/langchain-core/src/messages/function.ts
index 4f755f25840a..12cd9c5d3693 100644
--- a/langchain-core/src/messages/function.ts
+++ b/langchain-core/src/messages/function.ts
@@ -69,6 +69,7 @@ export class FunctionMessageChunk extends BaseMessageChunk {
         chunk.response_metadata
       ),
       name: this.name ?? "",
+      id: this.id ?? chunk.id,
     });
   }
 }
diff --git a/langchain-core/src/messages/human.ts b/langchain-core/src/messages/human.ts
index c0c3f54fc2fa..a64ad32b7e2a 100644
--- a/langchain-core/src/messages/human.ts
+++ b/langchain-core/src/messages/human.ts
@@ -43,6 +43,7 @@ export class HumanMessageChunk extends BaseMessageChunk {
         this.response_metadata,
         chunk.response_metadata
       ),
+      id: this.id ?? chunk.id,
     });
   }
 }
diff --git a/langchain-core/src/messages/system.ts b/langchain-core/src/messages/system.ts
index 1c7adcdb95e8..5d88fdea5c8d 100644
--- a/langchain-core/src/messages/system.ts
+++ b/langchain-core/src/messages/system.ts
@@ -43,6 +43,7 @@ export class SystemMessageChunk extends BaseMessageChunk {
         this.response_metadata,
         chunk.response_metadata
       ),
+      id: this.id ?? chunk.id,
     });
   }
 }
diff --git a/langchain-core/src/messages/tool.ts b/langchain-core/src/messages/tool.ts
index 48324d878f04..3375cfd8572d 100644
--- a/langchain-core/src/messages/tool.ts
+++ b/langchain-core/src/messages/tool.ts
@@ -88,6 +88,7 @@ export class ToolMessageChunk extends BaseMessageChunk {
         chunk.response_metadata
       ),
       tool_call_id: this.tool_call_id,
+      id: this.id ?? chunk.id,
     });
   }
 }