diff --git a/langgraph/graph/message.py b/langgraph/graph/message.py index 6908028df..26a9cfb29 100644 --- a/langgraph/graph/message.py +++ b/langgraph/graph/message.py @@ -1,3 +1,4 @@ +from itertools import filterfalse import uuid from typing import Annotated, Literal, TypedDict, Union @@ -7,19 +8,17 @@ from langchain_core.messages import ( convert_to_messages, message_chunk_to_message, ) -from langchain_core.messages.base import BaseMessage from langgraph.graph.state import StateGraph -Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation] - -class DeleteMessage(BaseMessage): +class MessageModifier(TypedDict): id: str - type: Literal["delete"] = "delete" + action: Literal["delete"] = "delete" - def __init__(self, **kwargs): - return super().__init__("delete-message", **kwargs) + +Message = Union[MessageLikeRepresentation, MessageModifier] +Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation] def add_messages(left: Messages, right: Messages) -> Messages: @@ -73,6 +72,10 @@ def add_messages(left: Messages, right: Messages) -> Messages: left = [left] if not isinstance(right, list): right = [right] + + is_modifier = lambda m: isinstance(m, dict) and m.get("action") + message_modifiers = list(filter(is_modifier, right)) + right = list(filterfalse(is_modifier, right)) # coerce to message left = [message_chunk_to_message(m) for m in convert_to_messages(left)] right = [message_chunk_to_message(m) for m in convert_to_messages(right)] @@ -84,21 +87,26 @@ def add_messages(left: Messages, right: Messages) -> Messages: if m.id is None: m.id = str(uuid.uuid4()) # merge - left_idx_by_id = {m.id: i for i, m in enumerate(left)} + existing_idx_by_id = {m.id: i for i, m in enumerate(left)} merged = left.copy() for m in right: - if (existing_idx := left_idx_by_id.get(m.id)) is not None: - if isinstance(m, DeleteMessage): - del merged[existing_idx] - else: - merged[existing_idx] = m + if (existing_idx := existing_idx_by_id.get(m.id)) is not None: + merged[existing_idx] = m else: - if isinstance(m, DeleteMessage): - raise ValueError( - f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')" - ) - + existing_idx_by_id[m.id] = len(merged) merged.append(m) + + for modifier in message_modifiers: + if (existing_idx := existing_idx_by_id.get(modifier["id"])) is None: + raise ValueError( + f"Attempting to modify a message with an ID that doesn't exist ('{modifier['id']}')" + ) + + if modifier["action"] == "delete": + del merged[existing_idx] + else: + raise ValueError(f"Unsupported modifier action '{modifier['action']}'") + return merged