langgraph: update RunnableLike to support injected kwargs (#3288)

Fixes #3257

libs/langgraph/langgraph/graph/graph.py

@@ -19,7 +19,6 @@ from typing import (
 )
 
 from langchain_core.runnables import Runnable
-from langchain_core.runnables.base import RunnableLike
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.runnables.graph import Graph as DrawableGraph
 from langchain_core.runnables.graph import Node as DrawableNode
@@ -40,7 +39,7 @@ from langgraph.pregel import Channel, Pregel
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.types import All, Checkpointer
-from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
+from langgraph.utils.runnable import RunnableCallable, RunnableLike, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
 

libs/langgraph/langgraph/graph/state.py

@@ -22,7 +22,6 @@ from typing import (
 )
 
 from langchain_core.runnables import Runnable, RunnableConfig
-from langchain_core.runnables.base import RunnableLike
 from pydantic import BaseModel
 from pydantic.v1 import BaseModel as BaseModelV1
 from typing_extensions import Self
@@ -60,7 +59,7 @@ from langgraph.store.base import BaseStore
 from langgraph.types import All, Checkpointer, Command, RetryPolicy
 from langgraph.utils.fields import get_field_default
 from langgraph.utils.pydantic import create_model
-from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
+from langgraph.utils.runnable import RunnableCallable, RunnableLike, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
 

libs/langgraph/langgraph/utils/runnable.py

@@ -23,17 +23,19 @@ from langchain_core.runnables.base import (
     Runnable,
     RunnableConfig,
     RunnableLambda,
-    RunnableLike,
     RunnableParallel,
     RunnableSequence,
 )
+from langchain_core.runnables.base import (
+    RunnableLike as LCRunnableLike,
+)
 from langchain_core.runnables.config import (
     run_in_executor,
     var_child_runnable_config,
 )
-from langchain_core.runnables.utils import Input
+from langchain_core.runnables.utils import Input, Output
 from langchain_core.tracers._streaming import _StreamingCallbackHandler
-from typing_extensions import TypeGuard
+from typing_extensions import Concatenate, ParamSpec, TypeGuard
 
 from langgraph.constants import (
     CONF,
@@ -130,6 +132,15 @@ Each tuple contains:
 VALID_KINDS = (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
 
 
+P = ParamSpec("P")  # to handle injected kwargs like `writer` / `store`
+
+RunnableLike = Union[
+    LCRunnableLike,
+    Callable[Concatenate[Input, P], Output],
+    Callable[Concatenate[Input, P], Awaitable[Output]],
+]
+
+
 class RunnableCallable(Runnable):
     """A much simpler version of RunnableLambda that requires sync and async functions."""