From 098d5ec403aedff38428a608c99e66200eb9d341 Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Mon, 3 Feb 2025 11:25:20 -0500
Subject: [PATCH] langgraph: update RunnableLike to support injected kwargs
 (#3288)

Fixes #3257
---
 libs/langgraph/langgraph/graph/graph.py    |  3 +--
 libs/langgraph/langgraph/graph/state.py    |  3 +--
 libs/langgraph/langgraph/utils/runnable.py | 17 ++++++++++++++---
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py
index dce0be20c..07ddd8381 100644
--- a/libs/langgraph/langgraph/graph/graph.py
+++ b/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__)
 
diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index 51b12c08e..ae7c188ec 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/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__)
 
diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py
index f3f4c4b79..742cc0bff 100644
--- a/libs/langgraph/langgraph/utils/runnable.py
+++ b/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."""