Support optional types from typeddict

In the input schema

libs/langgraph/langgraph/graph/state.py

@@ -20,9 +20,7 @@ from typing import (
 from langchain_core.pydantic_v1 import BaseModel
 from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.base import RunnableLike
-from langchain_core.runnables.utils import (
-    create_model,
-)
+from langchain_core.runnables.utils import create_model
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
@@ -33,14 +31,7 @@ from langgraph.channels.named_barrier_value import NamedBarrierValue
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.constants import NS_END, NS_SEP, TAG_HIDDEN
 from langgraph.errors import InvalidUpdateError
-from langgraph.graph.graph import (
-    END,
-    START,
-    Branch,
-    CompiledGraph,
-    Graph,
-    Send,
-)
+from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send
 from langgraph.managed.base import (
     ChannelKeyPlaceholder,
     ChannelTypePlaceholder,
@@ -53,7 +44,7 @@ from langgraph.pregel.read import ChannelRead, PregelNode
 from langgraph.pregel.types import All, RetryPolicy
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
-from langgraph.utils import RunnableCallable, coerce_to_runnable
+from langgraph.utils import RunnableCallable, coerce_to_runnable, is_optional_type
 
 logger = logging.getLogger(__name__)
 
@@ -498,7 +489,14 @@ class CompiledStateGraph(CompiledGraph):
                 return create_model(  # type: ignore[call-overload]
                     self.get_name("Input"),
                     **{
-                        k: (self.channels[k].UpdateType, None)
+                        k: (
+                            self.channels[k].UpdateType,
+                            (
+                                None
+                                if is_optional_type(self.channels[k].UpdateType)
+                                else ...
+                            ),
+                        )
                         for k in self.builder.schemas[self.builder.input]
                         if isinstance(self.channels[k], BaseChannel)
                     },

libs/langgraph/langgraph/utils.py

@@ -4,7 +4,15 @@ import inspect
 import sys
 from contextvars import copy_context
 from functools import partial, wraps
-from typing import Any, AsyncIterator, Awaitable, Callable, Optional
+from typing import (
+    Any,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Optional,
+    Union,
+    get_origin,
+)
 
 from langchain_core.runnables.base import (
     Runnable,
@@ -34,8 +42,6 @@ except ImportError:
 class StrEnum(str, enum.Enum):
     """A string enum."""
 
-    pass
-
 
 class RunnableCallable(Runnable):
     """A much simpler version of RunnableLambda that requires sync and async functions."""
@@ -183,3 +189,20 @@ def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnab
             f"Expected a Runnable, callable or dict."
             f"Instead got an unsupported type: {type(thing)}"
         )
+
+
+def is_optional_type(type_: Any) -> bool:
+    """Check if a type is Optional."""
+
+    if hasattr(type_, "__origin__") and hasattr(type_, "__args__"):
+        origin = get_origin(type_)
+        if origin is Optional:
+            return True
+        if origin is Union:
+            return any(
+                arg is type(None) or is_optional_type(arg) for arg in type_.__args__
+            )
+        return origin is None
+    if hasattr(type_, "__bound__") and type_.__bound__ is not None:
+        return is_optional_type(type_.__bound__)
+    return type_ is None

libs/langgraph/tests/test_state.py

@@ -1,5 +1,5 @@
 from typing import Annotated as Annotated2
-from typing import Any
+from typing import Any, Optional
 
 import pytest
 from langchain_core.runnables import RunnableConfig
@@ -86,3 +86,25 @@ def test_state_schema_with_type_hint():
     for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
         node_name = actions[i].__name__
         assert c[node_name] == output_state
+
+
+def test_state_schema_optional_values():
+    class InputState(TypedDict):
+        val1: str
+        val2: Optional[str]
+
+    class State(InputState):
+        val4: dict
+
+    builder = StateGraph(State, input=InputState)
+    builder.add_node("n", lambda x: x)
+    builder.add_edge("__start__", "n")
+    graph = builder.compile()
+    model = graph.input_schema
+    json_schema = model.schema()
+    expected_required = {"val1"}
+    expected_optional = {"val2"}
+    assert set(json_schema["required"]) == expected_required
+    assert (
+        set(json_schema["properties"].keys()) == expected_required | expected_optional
+    )

libs/langgraph/tests/test_utils.py

@@ -1,7 +1,18 @@
 import functools
 import sys
 import uuid
-from typing import TypedDict
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    ForwardRef,
+    List,
+    Literal,
+    Optional,
+    TypedDict,
+    TypeVar,
+    Union,
+)
 from unittest.mock import patch
 
 import langsmith
@@ -9,7 +20,7 @@ import pytest
 
 from langgraph.graph import END, StateGraph
 from langgraph.graph.graph import CompiledGraph
-from langgraph.utils import is_async_callable, is_async_generator
+from langgraph.utils import is_async_callable, is_async_generator, is_optional_type
 
 
 def test_is_async() -> None:
@@ -119,3 +130,51 @@ async def test_runnable_callable_tracing_nested_async(rt_graph: CompiledGraph) -
             with langsmith.tracing_context(enabled=True):
                 res = await rt_graph.ainvoke({"foo": 1})
     assert isinstance(res["node_run_id"], uuid.UUID)
+
+
+def test_is_optional_type():
+    assert is_optional_type(None)
+    assert not is_optional_type(type(None))
+    assert is_optional_type(Optional[list])
+    assert not is_optional_type(int)
+    assert is_optional_type(Optional[Literal[1, 2, 3]])
+    assert not is_optional_type(Literal[1, 2, 3])
+    assert is_optional_type(Optional[List[int]])
+    assert is_optional_type(Optional[Dict[str, int]])
+    assert not is_optional_type(List[Optional[int]])
+    assert is_optional_type(Union[Optional[str], Optional[int]])
+    assert is_optional_type(
+        Union[
+            Union[Optional[str], Optional[int]], Union[Optional[float], Optional[dict]]
+        ]
+    )
+    assert not is_optional_type(Union[Union[str, int], Union[float, dict]])
+
+    assert is_optional_type(Union[int, None])
+    assert is_optional_type(Union[str, None, int])
+    assert is_optional_type(Union[None, str, int])
+    assert not is_optional_type(Union[int, str])
+
+    assert not is_optional_type(Any)  # Do we actually want this?
+    assert is_optional_type(Optional[Any])
+
+    class MyClass:
+        pass
+
+    assert is_optional_type(Optional[MyClass])
+    assert not is_optional_type(MyClass)
+    assert is_optional_type(Optional[ForwardRef("MyClass")])
+    assert not is_optional_type(ForwardRef("MyClass"))
+
+    assert is_optional_type(Optional[Union[List[int], Dict[str, Optional[int]]]])
+    assert not is_optional_type(Union[List[int], Dict[str, Optional[int]]])
+
+    assert is_optional_type(Optional[Callable[[int], str]])
+    assert not is_optional_type(Callable[[int], Optional[str]])
+
+    T = TypeVar("T")
+    assert is_optional_type(Optional[T])
+    assert not is_optional_type(T)
+
+    U = TypeVar("U", bound=Optional[T])
+    assert is_optional_type(U)