From 9bb623381b6ec608846c01fe5c7f7298ef8e521c Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Wed, 3 Jul 2024 16:13:31 -0400
Subject: [PATCH] core[minor]: update conversion utils to handle RemoveMessage
 (#23840)

---
 libs/core/langchain_core/messages/utils.py  | 12 +++++++++++-
 libs/core/tests/unit_tests/test_messages.py |  3 +++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/libs/core/langchain_core/messages/utils.py b/libs/core/langchain_core/messages/utils.py
index 109a3eb9d44..beb30ece596 100644
--- a/libs/core/langchain_core/messages/utils.py
+++ b/libs/core/langchain_core/messages/utils.py
@@ -32,6 +32,7 @@ from langchain_core.messages.base import BaseMessage, BaseMessageChunk
 from langchain_core.messages.chat import ChatMessage, ChatMessageChunk
 from langchain_core.messages.function import FunctionMessage, FunctionMessageChunk
 from langchain_core.messages.human import HumanMessage, HumanMessageChunk
+from langchain_core.messages.modifier import RemoveMessage
 from langchain_core.messages.system import SystemMessage, SystemMessageChunk
 from langchain_core.messages.tool import ToolMessage, ToolMessageChunk
 
@@ -42,7 +43,12 @@ if TYPE_CHECKING:
     from langchain_core.runnables.base import Runnable
 
 AnyMessage = Union[
-    AIMessage, HumanMessage, ChatMessage, SystemMessage, FunctionMessage, ToolMessage
+    AIMessage,
+    HumanMessage,
+    ChatMessage,
+    SystemMessage,
+    FunctionMessage,
+    ToolMessage,
 ]
 
 
@@ -113,6 +119,8 @@ def _message_from_dict(message: dict) -> BaseMessage:
         return FunctionMessage(**message["data"])
     elif _type == "tool":
         return ToolMessage(**message["data"])
+    elif _type == "remove":
+        return RemoveMessage(**message["data"])
     elif _type == "AIMessageChunk":
         return AIMessageChunk(**message["data"])
     elif _type == "HumanMessageChunk":
@@ -214,6 +222,8 @@ def _create_message_from_message_type(
         message = FunctionMessage(content=content, **kwargs)
     elif message_type == "tool":
         message = ToolMessage(content=content, **kwargs)
+    elif message_type == "remove":
+        message = RemoveMessage(**kwargs)
     else:
         raise ValueError(
             f"Unexpected message type: {message_type}. Use one of 'human',"
diff --git a/libs/core/tests/unit_tests/test_messages.py b/libs/core/tests/unit_tests/test_messages.py
index 8ab61278e70..9dfef59b5eb 100644
--- a/libs/core/tests/unit_tests/test_messages.py
+++ b/libs/core/tests/unit_tests/test_messages.py
@@ -12,6 +12,7 @@ from langchain_core.messages import (
     FunctionMessageChunk,
     HumanMessage,
     HumanMessageChunk,
+    RemoveMessage,
     SystemMessage,
     ToolCall,
     ToolCallChunk,
@@ -649,6 +650,7 @@ def test_convert_to_messages() -> None:
                 ],
             },
             {"role": "tool", "tool_call_id": "tool_id", "content": "Hi!"},
+            {"role": "remove", "id": "message_to_remove", "content": ""},
         ]
     ) == [
         SystemMessage(content="You are a helpful assistant."),
@@ -668,6 +670,7 @@ def test_convert_to_messages() -> None:
             tool_calls=[ToolCall(name="greet", args={"name": "Jane"}, id="tool_id")],
         ),
         ToolMessage(tool_call_id="tool_id", content="Hi!"),
+        RemoveMessage(id="message_to_remove"),
     ]
 
     # tuples