Resolve type annotation forward refs when parsing StateGraph schema type

This commit is contained in:
Nuno Campos
2024-04-16 15:05:04 -07:00
parent f91fb184a1
commit bdf0def798
+7 -9
View File
@@ -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]: