From a07ec9d8c800be0283fa6332248751fa92ce1db5 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Fri, 30 Aug 2024 10:21:57 -0700
Subject: [PATCH] Handle total=false

---
 libs/langgraph/langgraph/graph/state.py |  7 ++++++-
 libs/langgraph/tests/test_state.py      | 18 +++++++++++++-----
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index 10e8ca5ff..31fdfb4e3 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -486,6 +486,10 @@ class CompiledStateGraph(CompiledGraph):
                     __root__=(self.channels[keys[0]].UpdateType, None),
                 )
             else:
+                is_total_false = (
+                    hasattr(self.builder.input, "__total__")
+                    and self.builder.input.__total__ is False
+                )
                 return create_model(  # type: ignore[call-overload]
                     self.get_name("Input"),
                     **{
@@ -493,7 +497,8 @@ class CompiledStateGraph(CompiledGraph):
                             self.channels[k].UpdateType,
                             (
                                 None
-                                if is_optional_type(self.channels[k].UpdateType)
+                                if is_total_false
+                                or is_optional_type(self.channels[k].UpdateType)
                                 else ...
                             ),
                         )
diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py
index 884895162..a3a0d2806 100644
--- a/libs/langgraph/tests/test_state.py
+++ b/libs/langgraph/tests/test_state.py
@@ -88,8 +88,9 @@ def test_state_schema_with_type_hint():
         assert c[node_name] == output_state
 
 
-def test_state_schema_optional_values():
-    class InputState(TypedDict):
+@pytest.mark.parametrize("total_", [True, False])
+def test_state_schema_optional_values(total_: bool):
+    class InputState(TypedDict, total=total_):  # type: ignore
         val1: str
         val2: Optional[str]
 
@@ -102,9 +103,16 @@ def test_state_schema_optional_values():
     graph = builder.compile()
     model = graph.input_schema
     json_schema = model.schema()
-    expected_required = {"val1"}
-    expected_optional = {"val2"}
-    assert set(json_schema["required"]) == expected_required
+
+    if total_ is False:
+        expected_required = set()
+        expected_optional = {"val2", "val1"}
+    else:
+        expected_required = {"val1"}
+
+        expected_optional = {"val2"}
+
+    assert set(json_schema.get("required", set())) == expected_required
     assert (
         set(json_schema["properties"].keys()) == expected_required | expected_optional
     )