From 3b4fcf6ee6b84265a962a4df317d64f0b1c2a719 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Sun, 9 Jun 2024 16:33:23 -0700
Subject: [PATCH] In Binop channels unwrap types before trying to instantiate -
 this makes List/Sequence annotations behave same as list

---
 langgraph/channels/binop.py | 35 ++++++++++++++++++--
 tests/test_pregel.py        | 65 +++++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 2 deletions(-)

diff --git a/langgraph/channels/binop.py b/langgraph/channels/binop.py
index 6cef26cee..a45b71f9f 100644
--- a/langgraph/channels/binop.py
+++ b/langgraph/channels/binop.py
@@ -1,5 +1,15 @@
+import collections.abc
 from contextlib import contextmanager
-from typing import Callable, Generator, Generic, Optional, Sequence, Type
+from typing import (
+    Callable,
+    Generator,
+    Generic,
+    NotRequired,
+    Optional,
+    Required,
+    Sequence,
+    Type,
+)
 
 from typing_extensions import Self
 
@@ -7,6 +17,17 @@ from langgraph.channels.base import BaseChannel, Value
 from langgraph.errors import EmptyChannelError
 
 
+# Adapted from typing_extensions
+def _strip_extras(t):
+    """Strips Annotated, Required and NotRequired from a given type."""
+    if hasattr(t, "__origin__"):
+        return _strip_extras(t.__origin__)
+    if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired):
+        return _strip_extras(t.__args__[0])
+
+    return t
+
+
 class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
     """Stores the result of applying a binary operator to the current value and each new value.
 
@@ -18,8 +39,18 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
     """
 
     def __init__(self, typ: Type[Value], operator: Callable[[Value, Value], Value]):
-        self.typ = typ
         self.operator = operator
+        # keep the type exposed by ValueType/UpdateType as-is
+        self.typ = typ
+        # special forms from typing or collections.abc are not instantiable
+        # so we need to replace them with their concrete counterparts
+        typ = _strip_extras(typ)
+        if typ in (collections.abc.Sequence, collections.abc.MutableSequence):
+            typ = list
+        if typ in (collections.abc.Set, collections.abc.MutableSet):
+            typ = set
+        if typ in (collections.abc.Mapping, collections.abc.MutableMapping):
+            typ = dict
         try:
             self.value = typ()
         except Exception:
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index d5ebffcbd..712d260f5 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -10,6 +10,7 @@ from typing import (
     Any,
     Dict,
     Generator,
+    List,
     Literal,
     Optional,
     Sequence,
@@ -170,6 +171,70 @@ def test_graph_validation() -> None:
         graph.compile()
 
 
+def test_reducer_before_first_node() -> None:
+    from langchain_core.messages import HumanMessage
+
+    class State(TypedDict):
+        hello: str
+        messages: Annotated[list[str], add_messages]
+
+    def node_a(state: State) -> State:
+        assert state == {
+            "hello": "there",
+            "messages": [HumanMessage(content="hello", id=AnyStr())],
+        }
+
+    builder = StateGraph(State)
+    builder.add_node("a", node_a)
+    builder.set_entry_point("a")
+    builder.set_finish_point("a")
+    graph = builder.compile()
+    assert graph.invoke({"hello": "there", "messages": "hello"}) == {
+        "hello": "there",
+        "messages": [HumanMessage(content="hello", id=AnyStr())],
+    }
+
+    class State(TypedDict):
+        hello: str
+        messages: Annotated[List[str], add_messages]
+
+    def node_a(state: State) -> State:
+        assert state == {
+            "hello": "there",
+            "messages": [HumanMessage(content="hello", id=AnyStr())],
+        }
+
+    builder = StateGraph(State)
+    builder.add_node("a", node_a)
+    builder.set_entry_point("a")
+    builder.set_finish_point("a")
+    graph = builder.compile()
+    assert graph.invoke({"hello": "there", "messages": "hello"}) == {
+        "hello": "there",
+        "messages": [HumanMessage(content="hello", id=AnyStr())],
+    }
+
+    class State(TypedDict):
+        hello: str
+        messages: Annotated[Sequence[str], add_messages]
+
+    def node_a(state: State) -> State:
+        assert state == {
+            "hello": "there",
+            "messages": [HumanMessage(content="hello", id=AnyStr())],
+        }
+
+    builder = StateGraph(State)
+    builder.add_node("a", node_a)
+    builder.set_entry_point("a")
+    builder.set_finish_point("a")
+    graph = builder.compile()
+    assert graph.invoke({"hello": "there", "messages": "hello"}) == {
+        "hello": "there",
+        "messages": [HumanMessage(content="hello", id=AnyStr())],
+    }
+
+
 def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")