From 7a1aad5fb35e2b43ff5b154274ffc991bb075865 Mon Sep 17 00:00:00 2001 From: Francis Gajitos Date: Sat, 30 Nov 2024 15:18:42 -0500 Subject: [PATCH] feat: Fields of BaseModels can be updated by aliases --- libs/langgraph/langgraph/graph/state.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 0684fb29c..2a8e4ae56 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -641,6 +641,13 @@ def _get_state_key(input: Union[None, dict, Any], *, key: str) -> Any: if input is None: return SKIP_WRITE elif isinstance(input, dict): + if hasattr(self.builder.output, 'model_fields'): + model_fields = self.builder.output.model_fields + # Extract the alias for the key if it exists + if key in model_fields: + alias = model_fields[key].alias + if alias in input: + return input.get(alias, SKIP_WRITE) if all(k not in output_keys for k in input): raise InvalidUpdateError( f"Expected node {node_key} to update at least one of {output_keys}, got {input}" @@ -823,6 +830,8 @@ def _get_state_reader( def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: + if hasattr(schema, '__fields__') and any(field.alias and field.alias != name for name, field in schema.__fields__.items()): + input = {schema.model_fields[key].alias if key in schema.model_fields else key: value for key, value in input.items()} return schema(**input)