From a184600915b2759cb603e88bd3c9712b585b0dc5 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 1 Apr 2024 18:44:06 -0700
Subject: [PATCH] MessageGraph now accepts same shorthand message formats as
 langchain-core

---
 langgraph/graph/message.py | 12 ++++++----
 tests/test_pregel.py       | 47 +++++++++++++++++++++++++++++---------
 2 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/langgraph/graph/message.py b/langgraph/graph/message.py
index b718913b1..42b7b2a49 100644
--- a/langgraph/graph/message.py
+++ b/langgraph/graph/message.py
@@ -1,7 +1,11 @@
 import uuid
 from typing import Annotated, Union
 
-from langchain_core.messages import AnyMessage, message_chunk_to_message
+from langchain_core.messages import (
+    AnyMessage,
+    convert_to_messages,
+    message_chunk_to_message,
+)
 
 from langgraph.graph.state import StateGraph
 
@@ -14,6 +18,9 @@ def add_messages(left: Messages, right: Messages) -> Messages:
         left = [left]
     if not isinstance(right, list):
         right = [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)]
     # assign missing ids
     for m in left:
         if m.id is None:
@@ -21,9 +28,6 @@ def add_messages(left: Messages, right: Messages) -> Messages:
     for m in right:
         if m.id is None:
             m.id = str(uuid.uuid4())
-    # coerce to message
-    left = [message_chunk_to_message(m) for m in left]
-    right = [message_chunk_to_message(m) for m in right]
     # merge
     left_idx_by_id = {m.id: i for i, m in enumerate(left)}
     merged = left.copy()
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index acf92d053..e3cc6c203 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -2357,10 +2357,7 @@ def test_message_graph(
     config = {"configurable": {"thread_id": "1"}}
 
     assert [
-        c
-        for c in app_w_interrupt.stream(
-            HumanMessage(content="what is weather in sf"), config
-        )
+        c for c in app_w_interrupt.stream(("human", "what is weather in sf"), config)
     ] == [
         {
             "agent": AIMessage(
@@ -2466,7 +2463,7 @@ def test_message_graph(
 
     app_w_interrupt.update_state(
         config,
-        AIMessage(content="answer", id="ai2"),
+        AIMessage(content="answer", id="ai2"),  # replace existing message
     )
 
     # replaces message even if object identity is different, as long as id is the same
@@ -2504,12 +2501,7 @@ def test_message_graph(
     config = {"configurable": {"thread_id": "2"}}
     model.i = 0  # reset the llm
 
-    assert [
-        c
-        for c in app_w_interrupt.stream(
-            HumanMessage(content="what is weather in sf"), config
-        )
-    ] == [
+    assert [c for c in app_w_interrupt.stream("what is weather in sf", config)] == [
         {
             "agent": AIMessage(
                 content="",
@@ -2651,6 +2643,39 @@ def test_message_graph(
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
     )
 
+    # add an extra message as if it came from "action" node
+    app_w_interrupt.update_state(config, ("ai", "an extra message"), as_node="action")
+
+    # extra message is coerced BaseMessge and appended
+    # now the next node is "agent" per the graph edges
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values=[
+            HumanMessage(
+                content="what is weather in sf",
+                id=AnyStr(),
+            ),
+            AIMessage(
+                content="",
+                additional_kwargs={
+                    "function_call": {
+                        "name": "search_api",
+                        "arguments": '"a different query"',
+                    }
+                },
+                id="ai1",
+            ),
+            FunctionMessage(
+                content="result for a different query",
+                name="search_api",
+                id=AnyStr(),
+            ),
+            AIMessage(content="answer", id="ai2"),
+            AIMessage(content="an extra message", id=AnyStr()),
+        ],
+        next=("agent",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
+    )
+
 
 def test_in_one_fan_out_out_one_graph_state() -> None:
     def sorted_add(x: list[str], y: list[str]) -> list[str]: