From 65b64c13312fdc69dd2b8df4b50673171892a528 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Fri, 30 Aug 2024 09:35:06 -0700 Subject: [PATCH] Support optional types from typeddict In the input schema --- libs/langgraph/langgraph/graph/state.py | 24 +++++----- libs/langgraph/langgraph/utils.py | 29 ++++++++++-- libs/langgraph/tests/test_state.py | 24 +++++++++- libs/langgraph/tests/test_utils.py | 63 ++++++++++++++++++++++++- 4 files changed, 121 insertions(+), 19 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 1416d14d0..10e8ca5ff 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -20,9 +20,7 @@ from typing import ( from langchain_core.pydantic_v1 import BaseModel from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.base import RunnableLike -from langchain_core.runnables.utils import ( - create_model, -) +from langchain_core.runnables.utils import create_model from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate @@ -33,14 +31,7 @@ from langgraph.channels.named_barrier_value import NamedBarrierValue from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import NS_END, NS_SEP, TAG_HIDDEN from langgraph.errors import InvalidUpdateError -from langgraph.graph.graph import ( - END, - START, - Branch, - CompiledGraph, - Graph, - Send, -) +from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send from langgraph.managed.base import ( ChannelKeyPlaceholder, ChannelTypePlaceholder, @@ -53,7 +44,7 @@ from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All, RetryPolicy from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.utils import RunnableCallable, coerce_to_runnable +from langgraph.utils import RunnableCallable, coerce_to_runnable, is_optional_type logger = logging.getLogger(__name__) @@ -498,7 +489,14 @@ class CompiledStateGraph(CompiledGraph): return create_model( # type: ignore[call-overload] self.get_name("Input"), **{ - k: (self.channels[k].UpdateType, None) + k: ( + self.channels[k].UpdateType, + ( + None + if is_optional_type(self.channels[k].UpdateType) + else ... + ), + ) for k in self.builder.schemas[self.builder.input] if isinstance(self.channels[k], BaseChannel) }, diff --git a/libs/langgraph/langgraph/utils.py b/libs/langgraph/langgraph/utils.py index 64cb21128..9aa1c257c 100644 --- a/libs/langgraph/langgraph/utils.py +++ b/libs/langgraph/langgraph/utils.py @@ -4,7 +4,15 @@ import inspect import sys from contextvars import copy_context from functools import partial, wraps -from typing import Any, AsyncIterator, Awaitable, Callable, Optional +from typing import ( + Any, + AsyncIterator, + Awaitable, + Callable, + Optional, + Union, + get_origin, +) from langchain_core.runnables.base import ( Runnable, @@ -34,8 +42,6 @@ except ImportError: class StrEnum(str, enum.Enum): """A string enum.""" - pass - class RunnableCallable(Runnable): """A much simpler version of RunnableLambda that requires sync and async functions.""" @@ -183,3 +189,20 @@ def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnab f"Expected a Runnable, callable or dict." f"Instead got an unsupported type: {type(thing)}" ) + + +def is_optional_type(type_: Any) -> bool: + """Check if a type is Optional.""" + + if hasattr(type_, "__origin__") and hasattr(type_, "__args__"): + origin = get_origin(type_) + if origin is Optional: + return True + if origin is Union: + return any( + arg is type(None) or is_optional_type(arg) for arg in type_.__args__ + ) + return origin is None + if hasattr(type_, "__bound__") and type_.__bound__ is not None: + return is_optional_type(type_.__bound__) + return type_ is None diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 4088e38c9..884895162 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -1,5 +1,5 @@ from typing import Annotated as Annotated2 -from typing import Any +from typing import Any, Optional import pytest from langchain_core.runnables import RunnableConfig @@ -86,3 +86,25 @@ def test_state_schema_with_type_hint(): for i, c in enumerate(graph.stream(input_state, stream_mode="updates")): node_name = actions[i].__name__ assert c[node_name] == output_state + + +def test_state_schema_optional_values(): + class InputState(TypedDict): + val1: str + val2: Optional[str] + + class State(InputState): + val4: dict + + builder = StateGraph(State, input=InputState) + builder.add_node("n", lambda x: x) + builder.add_edge("__start__", "n") + graph = builder.compile() + model = graph.input_schema + json_schema = model.schema() + expected_required = {"val1"} + expected_optional = {"val2"} + assert set(json_schema["required"]) == expected_required + assert ( + set(json_schema["properties"].keys()) == expected_required | expected_optional + ) diff --git a/libs/langgraph/tests/test_utils.py b/libs/langgraph/tests/test_utils.py index 1e3bf8ea1..d5bfb53c7 100644 --- a/libs/langgraph/tests/test_utils.py +++ b/libs/langgraph/tests/test_utils.py @@ -1,7 +1,18 @@ import functools import sys import uuid -from typing import TypedDict +from typing import ( + Any, + Callable, + Dict, + ForwardRef, + List, + Literal, + Optional, + TypedDict, + TypeVar, + Union, +) from unittest.mock import patch import langsmith @@ -9,7 +20,7 @@ import pytest from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph -from langgraph.utils import is_async_callable, is_async_generator +from langgraph.utils import is_async_callable, is_async_generator, is_optional_type def test_is_async() -> None: @@ -119,3 +130,51 @@ async def test_runnable_callable_tracing_nested_async(rt_graph: CompiledGraph) - with langsmith.tracing_context(enabled=True): res = await rt_graph.ainvoke({"foo": 1}) assert isinstance(res["node_run_id"], uuid.UUID) + + +def test_is_optional_type(): + assert is_optional_type(None) + assert not is_optional_type(type(None)) + assert is_optional_type(Optional[list]) + assert not is_optional_type(int) + assert is_optional_type(Optional[Literal[1, 2, 3]]) + assert not is_optional_type(Literal[1, 2, 3]) + assert is_optional_type(Optional[List[int]]) + assert is_optional_type(Optional[Dict[str, int]]) + assert not is_optional_type(List[Optional[int]]) + assert is_optional_type(Union[Optional[str], Optional[int]]) + assert is_optional_type( + Union[ + Union[Optional[str], Optional[int]], Union[Optional[float], Optional[dict]] + ] + ) + assert not is_optional_type(Union[Union[str, int], Union[float, dict]]) + + assert is_optional_type(Union[int, None]) + assert is_optional_type(Union[str, None, int]) + assert is_optional_type(Union[None, str, int]) + assert not is_optional_type(Union[int, str]) + + assert not is_optional_type(Any) # Do we actually want this? + assert is_optional_type(Optional[Any]) + + class MyClass: + pass + + assert is_optional_type(Optional[MyClass]) + assert not is_optional_type(MyClass) + assert is_optional_type(Optional[ForwardRef("MyClass")]) + assert not is_optional_type(ForwardRef("MyClass")) + + assert is_optional_type(Optional[Union[List[int], Dict[str, Optional[int]]]]) + assert not is_optional_type(Union[List[int], Dict[str, Optional[int]]]) + + assert is_optional_type(Optional[Callable[[int], str]]) + assert not is_optional_type(Callable[[int], Optional[str]]) + + T = TypeVar("T") + assert is_optional_type(Optional[T]) + assert not is_optional_type(T) + + U = TypeVar("U", bound=Optional[T]) + assert is_optional_type(U)