From 4239e6fdff41358758b0698fa1f716e86dc7f1ab Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 12 Jun 2024 10:18:22 -0400 Subject: [PATCH] (RFC) add support for deleting messages --- langgraph/graph/message.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/langgraph/graph/message.py b/langgraph/graph/message.py index 6e2a1729e..6908028df 100644 --- a/langgraph/graph/message.py +++ b/langgraph/graph/message.py @@ -1,5 +1,5 @@ import uuid -from typing import Annotated, TypedDict, Union +from typing import Annotated, Literal, TypedDict, Union from langchain_core.messages import ( AnyMessage, @@ -7,12 +7,21 @@ 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): + id: str + type: Literal["delete"] = "delete" + + def __init__(self, **kwargs): + return super().__init__("delete-message", **kwargs) + + def add_messages(left: Messages, right: Messages) -> Messages: """Merges two lists of messages, updating existing messages by ID. @@ -79,8 +88,16 @@ def add_messages(left: Messages, right: Messages) -> Messages: merged = left.copy() for m in right: if (existing_idx := left_idx_by_id.get(m.id)) is not None: - merged[existing_idx] = m + if isinstance(m, DeleteMessage): + del merged[existing_idx] + else: + 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}')" + ) + merged.append(m) return merged