From 2c908f1557c3ef3e408a37bcb31085bf8643082c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 11 Mar 2025 17:32:54 -0700 Subject: [PATCH 1/4] 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, From 14c22418539d3ffaf65b0516b4a747d7150d171f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 11 Mar 2025 17:44:24 -0700 Subject: [PATCH 2/4] Lint --- libs/langgraph/langgraph/graph/state.py | 2 +- libs/langgraph/langgraph/pregel/loop.py | 12 +++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 285542bdc..45564334f 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -936,7 +936,7 @@ def _get_state_reader( def _pick_mapper( state_keys: Sequence[str], schema: Type[Any] -) -> Optional[Callable[[Type[Any], dict[str, Any]], dict[str, Any]]]: +) -> Optional[Callable[[Any], Any]]: if state_keys == ["__root__"]: return None if issubclass(schema, dict): diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 1ba444dad..169943718 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -141,7 +141,7 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol: class PregelLoop(LoopProtocol): input: Optional[Any] - input_model: Optional[BaseModel] + input_model: Optional[Type[BaseModel]] checkpointer: Optional[BaseCheckpointSaver] nodes: Mapping[str, PregelNode] specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]] @@ -205,7 +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, + input_model: Optional[Type[BaseModel]] = None, debug: bool = False, ) -> None: super().__init__( @@ -434,7 +434,9 @@ class PregelLoop(LoopProtocol): if self.input is INPUT_SHOULD_VALIDATE: self.input = INPUT_DONE # validate - self.input_model(**read_channels(self.channels, self.stream_keys)) + 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 @@ -859,7 +861,7 @@ 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, + input_model: Optional[Type[BaseModel]] = None, debug: bool = False, ) -> None: super().__init__( @@ -1000,7 +1002,7 @@ 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, + input_model: Optional[Type[BaseModel]] = None, debug: bool = False, ) -> None: super().__init__( From e566ed4b3fd04962fddd3ae95959a671645d652d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 11 Mar 2025 17:51:13 -0700 Subject: [PATCH 3/4] Fix py 3.9 - isclass and issubclass disagree on whether something like list[str] is a class --- libs/langgraph/langgraph/graph/state.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 45564334f..1b5ef0191 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -627,7 +627,9 @@ class StateGraph(Graph): builder=self, config_type=self.config_schema, input_model=self.input - if isclass(self.input) and issubclass(self.input, (BaseModel, BaseModelV1)) + if len(self.channels) > 1 + and isclass(self.input) + and issubclass(self.input, (BaseModel, BaseModelV1)) else None, nodes={}, channels={ From 1d3926af27464d7a9678dd9ca175c7766e8dbd8e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 11 Mar 2025 18:13:36 -0700 Subject: [PATCH 4/4] Fix kafka --- libs/langgraph/langgraph/pregel/loop.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 169943718..66f0a60f5 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -687,6 +687,8 @@ class PregelLoop(LoopProtocol): self.input = INPUT_DONE elif CONFIG_KEY_RESUMING not in configurable: raise EmptyInputError(f"Received no input for {input_keys}") + else: + self.input = INPUT_DONE # update config if not self.is_nested: self.config = patch_configurable(