From af6552a17e5aa86876d6a880fa64bf25dce3c7d6 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Thu, 3 Apr 2025 13:43:55 +0200 Subject: [PATCH] Move to langgraph/graph --- .../langgraph/graph}/ui.py | 34 ++++++++++++++----- 1 file changed, 25 insertions(+), 9 deletions(-) rename libs/{prebuilt/langgraph/prebuilt => langgraph/langgraph/graph}/ui.py (85%) diff --git a/libs/prebuilt/langgraph/prebuilt/ui.py b/libs/langgraph/langgraph/graph/ui.py similarity index 85% rename from libs/prebuilt/langgraph/prebuilt/ui.py rename to libs/langgraph/langgraph/graph/ui.py index 711b85e45..504711be9 100644 --- a/libs/prebuilt/langgraph/prebuilt/ui.py +++ b/libs/langgraph/langgraph/graph/ui.py @@ -84,7 +84,6 @@ def push_ui_message( ) """ - writer = get_stream_writer() config = get_config() @@ -147,7 +146,7 @@ def remove_ui_message(id: str, *, state_key: str = "ui") -> RemoveUIMessage: return evt -def reduce_ui_messages( +def add_ui_messages( left: Union[list[AnyUIMessage], AnyUIMessage], right: Union[list[AnyUIMessage], AnyUIMessage], ) -> list[AnyUIMessage]: @@ -180,11 +179,28 @@ def reduce_ui_messages( if not isinstance(right, list): right = [right] - new_state = left.copy() - for m in right: - if m.get("type") == "remove-ui": - new_state = [m for m in new_state if m.get("id") != m.get("id")] - else: - new_state.append(m) + # merge messages + merged = left.copy() + merged_by_id = {m.get("id"): i for i, m in enumerate(merged)} + ids_to_remove = set() - return new_state + for msg in right: + msg_id = msg.get("id") + + if (existing_idx := merged_by_id.get(msg_id)) is not None: + if msg.get("type") == "remove-ui": + ids_to_remove.add(msg_id) + else: + ids_to_remove.discard(msg_id) + merged[existing_idx] = msg + else: + if msg.get("type") == "remove-ui": + raise ValueError( + f"Attempting to delete an UI message with an ID that doesn't exist ('{msg_id}')" + ) + + merged_by_id[msg_id] = len(merged) + merged.append(msg) + + merged = [m for m in merged if m.get("id") not in ids_to_remove] + return merged