Pydantic Validation Problem with RemoveMessage in AnyMessage Tagged Union

[maxence-oden]

Description

I am encountering a validation issue when using RemoveMessage with the AnyMessage Pydantic tagged union in the context of a StateGraph. Specifically, I want to include RemoveMessage as part of the messages field in my state object, but RemoveMessage is not tagged within AnyMessage. This leads to validation problems.

Here is a minimal reproducible example of the behavior I am trying to achieve:

from typing import Annotated
from langgraph.graph import StateGraph, START, END

from pydantic import BaseModel
from langchain_core.messages import AnyMessage, HumanMessage, RemoveMessage
from langgraph.graph import add_messages


class OverallState(BaseModel):
    messages: Annotated[AnyMessage, add_messages]


def node(state: OverallState):
    print(state.messages)
    return OverallState(messages=RemoveMessage(id="test")) # This is the problematic line
    # Alternative valid way to remove a message:
    # return {"message": RemoveMessage(id="test")}


# Build the state graph
builder = StateGraph(OverallState)
builder.add_node(node)
builder.add_edge(START, "node")
builder.add_edge("node", END)
graph = builder.compile()

# Test the graph with a valid input
graph.invoke(OverallState(messages=HumanMessage(id="test", content="hello")))

Problem

When I try to return RemoveMessage in the node function, I get the following validation error:

pydantic_core._pydantic_core.ValidationError: 1 validation error for OverallState
messages
  Input tag 'remove' found using _get_type() does not match any of the expected tags: 'ai', 'human', 'chat', 'system', 'function', 'tool', 'AIMessageChunk', 'HumanMessageChunk', 'ChatMessageChunk', 'SystemMessageChunk', 'FunctionMessageChunk', 'ToolMessageChunk' [type=union_tag_invalid, input_value=RemoveMessage(content='',..._metadata={}, id='test'), input_type=RemoveMessage]
    For further information visit https://errors.pydantic.dev/2.10/v/union_tag_invalid

I understand that RemoveMessage is not currently tagged in the AnyMessage union. However, this limitation makes it harder to work with RemoveMessage in use cases where Pydantic objects are required.

Questions

• Is it planned to include RemoveMessage in the AnyMessage union in the future?
• Is the exclusion of RemoveMessage from AnyMessage intentional because it’s primarily used with the add_messages function?
• Are there alternative approaches to incorporate RemoveMessage into the AnyMessage union or to work around this limitation?

Additional Context

Here is the version information for the relevant packages:

langgraph: 0.2.67

I’d appreciate any insights or clarification on whether this behavior is intentional or if support for RemoveMessage in AnyMessage might be added. Thank you!