From 2c908f1557c3ef3e408a37bcb31085bf8643082c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 11 Mar 2025 17:32:54 -0700 Subject: [PATCH] Avoid validating pydantic state models when we can - 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 | 39 +++++++++++++++------ libs/langgraph/langgraph/pregel/__init__.py | 8 +++++ libs/langgraph/langgraph/pregel/loop.py | 29 +++++++++++++-- 3 files changed, 63 insertions(+), 13 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 48267c5df..285542bdc 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -626,6 +626,9 @@ class StateGraph(Graph): compiled = CompiledStateGraph( builder=self, config_type=self.config_schema, + input_model=self.input + if isclass(self.input) and issubclass(self.input, (BaseModel, BaseModelV1)) + else None, nodes={}, channels={ **self.channels, @@ -812,11 +815,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 +930,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[[Type[Any], dict[str, Any]], dict[str, 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) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 6948847d2..362c1459c 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index e654f17df..1ba444dad 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/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[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[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,11 @@ 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 + 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 +632,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 +674,17 @@ 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 # update config if not self.is_nested: self.config = patch_configurable( @@ -840,10 +859,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[BaseModel] = None, debug: bool = False, ) -> None: super().__init__( input, + input_model=input_model, stream=stream, config=config, checkpointer=checkpointer, @@ -979,10 +1000,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[BaseModel] = None, debug: bool = False, ) -> None: super().__init__( input, + input_model=input_model, stream=stream, config=config, checkpointer=checkpointer,