From bdf0def7981e0f0b09b439e084245d94a351b5d3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 16 Apr 2024 15:05:04 -0700 Subject: [PATCH] Resolve type annotation forward refs when parsing StateGraph schema type --- langgraph/graph/state.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index bc42d0370..86e764504 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -1,7 +1,7 @@ import logging from functools import partial from inspect import signature -from typing import Any, Optional, Sequence, Type, Union +from typing import Any, Optional, Sequence, Type, Union, get_type_hints from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.base import RunnableLike @@ -286,15 +286,13 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: def _get_channels(schema: Type[dict]) -> dict[str, BaseChannel]: if not hasattr(schema, "__annotations__"): - return { - "__root__": _get_channel(schema), - } + return {"__root__": _get_channel(schema)} - channels: dict[str, BaseChannel] = {} - for name, typ in schema.__annotations__.items(): - channels[name] = _get_channel(typ) - - return channels + return { + name: _get_channel(typ) + for name, typ in get_type_hints(schema, include_extras=True).items() + if name != "__slots__" + } def _get_channel(annotation: Any) -> Optional[BaseChannel]: