Avoid validating pydantic state models when we can (#3782)

- When a pydantic input schema isued but dict input is passed in
validate it once after running hidden START node. If the input is an
instance of the input model we skip validation altogether
- When entering each node we need to create a standalone instance of the
state class, but we can now skip validation, as it's now run once
elsewhere

libs/langgraph/langgraph/graph/state.py

@@ -626,6 +626,11 @@ class StateGraph(Graph):
         compiled = CompiledStateGraph(
             builder=self,
             config_type=self.config_schema,
+            input_model=self.input
+            if len(self.channels) > 1
+            and isclass(self.input)
+            and issubclass(self.input, (BaseModel, BaseModelV1))
+            else None,
             nodes={},
             channels={
                 **self.channels,
@@ -812,11 +817,7 @@ class CompiledStateGraph(CompiledGraph):
                 # read state keys and managed values
                 channels=(list(input_values) if is_single_input else input_values),
                 # coerce state dict to schema class (eg. pydantic model)
-                mapper=(
-                    None
-                    if is_single_input or issubclass(input_schema, dict)
-                    else partial(_coerce_state, input_schema)
-                ),
+                mapper=_pick_mapper(list(input_values), input_schema),
                 writers=[
                     # publish to this channel and state keys
                     ChannelWrite(
@@ -931,14 +932,34 @@ def _get_state_reader(
         select=select[0] if select == ["__root__"] else select,
         fresh=True,
         # coerce state dict to schema class (eg. pydantic model)
-        mapper=(
-            None
-            if state_keys == ["__root__"] or issubclass(schema, dict)
-            else partial(_coerce_state, schema)
-        ),
+        mapper=_pick_mapper(state_keys, schema),
     )
 
 
+def _pick_mapper(
+    state_keys: Sequence[str], schema: Type[Any]
+) -> Optional[Callable[[Any], Any]]:
+    if state_keys == ["__root__"]:
+        return None
+    if issubclass(schema, dict):
+        return None
+    if issubclass(schema, BaseModel):
+        return partial(_coerce_state_pydantic, schema)
+    if issubclass(schema, BaseModelV1):
+        return partial(_coerce_state_pydantic_v1, schema)
+    return partial(_coerce_state, schema)
+
+
+def _coerce_state_pydantic(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
+    return schema.model_construct(**input)
+
+
+def _coerce_state_pydantic_v1(
+    schema: Type[Any], input: dict[str, Any]
+) -> dict[str, Any]:
+    return schema.construct(**input)
+
+
 def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
     return schema(**input)
 

libs/langgraph/langgraph/pregel/__init__.py

@@ -496,6 +496,8 @@ class Pregel(PregelProtocol):
 
     config_type: Optional[Type[Any]] = None
 
+    input_model: Optional[Type[BaseModel]] = None
+
     config: Optional[RunnableConfig] = None
 
     name: str = "LangGraph"
@@ -519,6 +521,7 @@ class Pregel(PregelProtocol):
         store: Optional[BaseStore] = None,
         retry_policy: Optional[RetryPolicy] = None,
         config_type: Optional[Type[Any]] = None,
+        input_model: Optional[Type[BaseModel]] = None,
         config: Optional[RunnableConfig] = None,
         name: str = "LangGraph",
     ) -> None:
@@ -537,6 +540,7 @@ class Pregel(PregelProtocol):
         self.store = store
         self.retry_policy = retry_policy
         self.config_type = config_type
+        self.input_model = input_model
         self.config = config
         self.name = name
         if auto_validate:
@@ -650,6 +654,8 @@ class Pregel(PregelProtocol):
     def get_input_schema(
         self, config: Optional[RunnableConfig] = None
     ) -> Type[BaseModel]:
+        if self.input_model is not None:
+            return self.input_model
         config = merge_configs(self.config, config)
         if isinstance(self.input_channels, str):
             return super().get_input_schema(config)
@@ -1967,6 +1973,7 @@ class Pregel(PregelProtocol):
                 )
             with SyncPregelLoop(
                 input,
+                input_model=self.input_model,
                 stream=StreamProtocol(stream.put, stream_modes),
                 config=config,
                 store=store,
@@ -2257,6 +2264,7 @@ class Pregel(PregelProtocol):
                 )
             async with AsyncPregelLoop(
                 input,
+                input_model=self.input_model,
                 stream=StreamProtocol(stream.put_nowait, stream_modes),
                 config=config,
                 store=store,

libs/langgraph/langgraph/pregel/loop.py

@@ -23,6 +23,7 @@ from typing import (
 
 from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
 from langchain_core.runnables import RunnableConfig
+from pydantic import BaseModel
 from typing_extensions import ParamSpec, Self
 
 from langgraph.channels.base import BaseChannel
@@ -125,6 +126,7 @@ P = ParamSpec("P")
 
 INPUT_DONE = object()
 INPUT_RESUMING = object()
+INPUT_SHOULD_VALIDATE = object()
 SPECIAL_CHANNELS = (ERROR, INTERRUPT, SCHEDULED)
 
 
@@ -139,6 +141,7 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
 
 class PregelLoop(LoopProtocol):
     input: Optional[Any]
+    input_model: Optional[Type[BaseModel]]
     checkpointer: Optional[BaseCheckpointSaver]
     nodes: Mapping[str, PregelNode]
     specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
@@ -202,6 +205,7 @@ class PregelLoop(LoopProtocol):
         interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
         interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
+        input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
     ) -> None:
         super().__init__(
@@ -212,6 +216,7 @@ class PregelLoop(LoopProtocol):
             store=store,
         )
         self.input = input
+        self.input_model = input_model
         self.checkpointer = checkpointer
         self.nodes = nodes
         self.specs = specs
@@ -395,7 +400,7 @@ class PregelLoop(LoopProtocol):
         if self.status != "pending":
             raise RuntimeError("Cannot tick when status is no longer 'pending'")
 
-        if self.input not in (INPUT_DONE, INPUT_RESUMING):
+        if self.input not in (INPUT_DONE, INPUT_RESUMING, INPUT_SHOULD_VALIDATE):
             self._first(input_keys=input_keys)
         elif self.to_interrupt:
             # if we need to interrupt, do so
@@ -425,6 +430,13 @@ class PregelLoop(LoopProtocol):
             # apply writes to managed values
             for key, values in mv_writes.items():
                 self._update_mv(key, values)
+            # validate input if requested
+            if self.input is INPUT_SHOULD_VALIDATE:
+                self.input = INPUT_DONE
+                # validate
+                cast(Type[BaseModel], self.input_model)(
+                    **read_channels(self.channels, self.stream_keys)
+                )
             # produce values output
             self._emit(
                 "values", map_output_values, self.output_keys, writes, self.channels
@@ -622,6 +634,8 @@ class PregelLoop(LoopProtocol):
             self._emit(
                 "values", map_output_values, self.output_keys, True, self.channels
             )
+            # set flag
+            self.input = INPUT_RESUMING
         # map inputs to channel updates
         elif input_writes := deque(map_input(input_keys, self.input)):
             # TODO shouldn't these writes be passed to put_writes too?
@@ -662,10 +676,19 @@ class PregelLoop(LoopProtocol):
             assert not mv_writes, "Can't write to SharedValues in graph input"
             # save input checkpoint
             self._put_checkpoint({"source": "input", "writes": dict(input_writes)})
+            # set flag
+            if (
+                self.input_model is not None
+                and not isinstance(self.input, self.input_model)
+                and not isinstance(self.stream_keys, str)
+            ):
+                self.input = INPUT_SHOULD_VALIDATE
+            else:
+                self.input = INPUT_DONE
         elif CONFIG_KEY_RESUMING not in configurable:
             raise EmptyInputError(f"Received no input for {input_keys}")
-        # done with input
-        self.input = INPUT_RESUMING if is_resuming else INPUT_DONE
+        else:
+            self.input = INPUT_DONE
         # update config
         if not self.is_nested:
             self.config = patch_configurable(
@@ -840,10 +863,12 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
     ) -> None:
         super().__init__(
             input,
+            input_model=input_model,
             stream=stream,
             config=config,
             checkpointer=checkpointer,
@@ -979,10 +1004,12 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
     ) -> None:
         super().__init__(
             input,
+            input_model=input_model,
             stream=stream,
             config=config,
             checkpointer=checkpointer,