From 28e0e7f3aebf3bccdd546b35b557f5919302a0e2 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Thu, 22 May 2025 12:59:39 +0200 Subject: [PATCH 1/2] feat(langgraph): push_messages should directly write to the state --- libs/langgraph/langgraph/graph/message.py | 11 ++++++++++- libs/langgraph/langgraph/graph/ui.py | 5 +++-- libs/langgraph/tests/test_messages_state.py | 12 +++++------- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index 16077eeb7..2e2682fe8 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -23,6 +23,7 @@ from langchain_core.messages import ( ) from typing_extensions import TypedDict +from langgraph.constants import CONF, CONFIG_KEY_SEND from langgraph.graph.state import StateGraph Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation] @@ -298,8 +299,13 @@ def _format_messages(messages: Sequence[BaseMessage]) -> list[BaseMessage]: def push_message( message: Union[MessageLikeRepresentation, BaseMessageChunk], + *, + state_key: Optional[str] = "messages", ) -> AnyMessage: - """Write a message manually to the `messages` / `messages-tuple` stream mode.""" + """Write a message manually to the `messages` / `messages-tuple` stream mode. + + Will automatically write to the channel specified in the `state_key` unless `state_key` is `None`. + """ from langchain_core.callbacks.base import ( BaseCallbackHandler, @@ -334,4 +340,7 @@ def push_message( ) stream_handler._emit(message_meta, message, dedupe=False) + if state_key: + config[CONF][CONFIG_KEY_SEND]([(state_key, message)]) + return message diff --git a/libs/langgraph/langgraph/graph/ui.py b/libs/langgraph/langgraph/graph/ui.py index 08dc6fac7..a2e66f91a 100644 --- a/libs/langgraph/langgraph/graph/ui.py +++ b/libs/langgraph/langgraph/graph/ui.py @@ -54,7 +54,7 @@ def push_ui_message( id: Optional[str] = None, metadata: Optional[dict[str, Any]] = None, message: Optional[AnyMessage] = None, - state_key: str = "ui", + state_key: Optional[str] = "ui", merge: bool = False, ) -> UIMessage: """Push a new UI message to update the UI state. @@ -111,7 +111,8 @@ def push_ui_message( } writer(evt) - config[CONF][CONFIG_KEY_SEND]([(state_key, evt)]) + if state_key: + config[CONF][CONFIG_KEY_SEND]([(state_key, evt)]) return evt diff --git a/libs/langgraph/tests/test_messages_state.py b/libs/langgraph/tests/test_messages_state.py index 4315cd25e..a611dc6b6 100644 --- a/libs/langgraph/tests/test_messages_state.py +++ b/libs/langgraph/tests/test_messages_state.py @@ -342,13 +342,11 @@ def test_push_messages_in_graph(): with pytest.raises(ValueError, match="Message ID is required"): push_message(AIMessage(content="No ID")) - return { - "messages": [ - push_message(AIMessage(content="First", id="1")), - push_message(HumanMessage(content="Second", id="2")), - push_message(AIMessage(content="Third", id="3")), - ] - } + push_message(AIMessage(content="First", id="1")) + push_message(HumanMessage(content="Second", id="2")) + push_message(AIMessage(content="Third", id="3")) + + return {"messages": []} builder = StateGraph(MessagesState) builder.add_node(chat) From 10c6ed73205f23cf134833b6f4e56cd4e38de8f0 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Thu, 22 May 2025 19:35:44 +0200 Subject: [PATCH 2/2] Code review --- libs/langgraph/tests/test_messages_state.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/libs/langgraph/tests/test_messages_state.py b/libs/langgraph/tests/test_messages_state.py index a611dc6b6..a481123a4 100644 --- a/libs/langgraph/tests/test_messages_state.py +++ b/libs/langgraph/tests/test_messages_state.py @@ -346,8 +346,6 @@ def test_push_messages_in_graph(): push_message(HumanMessage(content="Second", id="2")) push_message(AIMessage(content="Third", id="3")) - return {"messages": []} - builder = StateGraph(MessagesState) builder.add_node(chat) builder.add_edge(START, "chat")