diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index 7aa0a6f8f..f484664a0 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -1,6 +1,7 @@ from typing import Optional from langchain_core.runnables import RunnableLambda, RunnableSequence +from langchain_core.runnables.utils import get_function_nonlocals from langgraph.checkpoint.base import ChannelVersions from langgraph.pregel.protocol import PregelProtocol @@ -41,5 +42,16 @@ def find_subgraph_pregel(candidate: Runnable) -> Optional[PregelProtocol]: candidates.extend(c.steps) elif isinstance(c, RunnableLambda): candidates.extend(c.deps) + elif isinstance(c, RunnableCallable): + if c.func is not None: + candidates.extend( + nl.__self__ if hasattr(nl, "__self__") else nl + for nl in get_function_nonlocals(c.func) + ) + elif c.afunc is not None: + candidates.extend( + nl.__self__ if hasattr(nl, "__self__") else nl + for nl in get_function_nonlocals(c.afunc) + ) return None