From e05cef773203bed133ffdc41cdc85c3c11f3deb8 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Tue, 24 Jun 2025 08:17:22 -0400 Subject: [PATCH] support invokable signature --- libs/langgraph/langgraph/graph/state.py | 13 +++++++++++-- libs/langgraph/langgraph/typing.py | 12 +++--------- libs/langgraph/tests/test_type_checking.py | 15 +++++++++++++++ 3 files changed, 29 insertions(+), 11 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 151b1b278..b44d47586 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -148,6 +148,15 @@ class _NodeWithConfigWriterStore(Protocol[StateT_contra]): ) -> Any: ... +class _Invokable(Protocol[StateT_contra]): + def invoke( + self, + input: StateT_contra, + config: RunnableConfig | None = None, + **kwargs: Any, + ) -> Any: ... + + # TODO: we probably don't want to explicitly support the config / store signatures once # we move to adding a context arg. Maybe what we do is we add support for kwargs with param spec # this is purely for typing purposes though, so can easily change in the coming weeks. @@ -160,7 +169,7 @@ StateNode: TypeAlias = Union[ _NodeWithConfigWriter[StateT_contra], _NodeWithConfigStore[StateT_contra], _NodeWithConfigWriterStore[StateT_contra], - Runnable[StateT_contra, Any], + _Invokable[StateT_contra], ] @@ -536,7 +545,7 @@ class StateGraph(Generic[StateT, InputT, OutputT]): if input_schema is not None: self._add_schema(input_schema) self.nodes[node] = StateNodeSpec( - coerce_to_runnable(action, name=node, trace=False), + coerce_to_runnable(action, name=node, trace=False), # type: ignore[arg-type] metadata, input=input_schema or self.state_schema, retry_policy=retry_policy, diff --git a/libs/langgraph/langgraph/typing.py b/libs/langgraph/langgraph/typing.py index 01ea27adf..7c8346d37 100644 --- a/libs/langgraph/langgraph/typing.py +++ b/libs/langgraph/langgraph/typing.py @@ -1,7 +1,5 @@ from __future__ import annotations -from typing import Union - from typing_extensions import TypeVar from langgraph._typing import StateLike @@ -19,12 +17,8 @@ InputT = TypeVar("InputT", bound=StateLike, default=StateT) Defaults to `StateT`. """ -ResolvedInputT = TypeVar("ResolvedInputT", bound=StateLike) -"""Type variable used to represent the resolved input to a state graph. +OutputT = TypeVar("OutputT", bound=StateLike, default=StateT) +"""Type variable used to represent the output of a state graph. -No default. +Defaults to `StateT`. """ - - -OutputT = TypeVar("OutputT", bound=Union[StateLike, None], default=StateT) -"""Type variable used to represent the output of a state graph.""" diff --git a/libs/langgraph/tests/test_type_checking.py b/libs/langgraph/tests/test_type_checking.py index 0b7ee0679..09d29bd6d 100644 --- a/libs/langgraph/tests/test_type_checking.py +++ b/libs/langgraph/tests/test_type_checking.py @@ -103,3 +103,18 @@ def test_input_state_specified() -> None: new_graph.invoke({"something": 1}) new_graph.invoke({"something": 2, "info": ["hello", "world"]}) # type: ignore[arg-type] + + +def test_invokeable_node_signature() -> None: + class State(TypedDict): + info: Annotated[list[str], add] + + graph_builder = StateGraph(State) + + class RunnableIsh: + def invoke( + self, input: State, config: RunnableConfig | None = None, **kwargs: Any + ) -> dict[str, str]: + return {} + + graph_builder.add_node("runnable", RunnableIsh())