From 4a60eaf32f327f8d81120729ff9b151f9aa87319 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 Jan 2025 16:26:23 -0800 Subject: [PATCH] Guard and cache calls to find_subgraph_pregel - The result of these doesnt change once a node is created, and it's fairly expensive to run, so great thing to cache - There's a variety of errors that can come from inspecting the source code of a function (part of what this does) so adding a catch-all try-except block as this should be best-effort, not crash your graph --- libs/langgraph/langgraph/pregel/__init__.py | 10 +++---- libs/langgraph/langgraph/pregel/algo.py | 2 ++ libs/langgraph/langgraph/pregel/debug.py | 3 +-- libs/langgraph/langgraph/pregel/read.py | 29 ++++++++++++++++----- libs/langgraph/langgraph/pregel/utils.py | 2 +- libs/langgraph/langgraph/types.py | 2 ++ libs/langgraph/tests/test_pregel.py | 2 ++ 7 files changed, 36 insertions(+), 14 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index a30d2fe58..327373bcb 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -97,7 +97,7 @@ from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner -from langgraph.pregel.utils import find_subgraph_pregel, get_new_channel_versions +from langgraph.pregel.utils import get_new_channel_versions from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore @@ -429,7 +429,7 @@ class Pregel(PregelProtocol): def get_subgraphs( self, *, namespace: Optional[str] = None, recurse: bool = False - ) -> Iterator[tuple[str, Pregel]]: + ) -> Iterator[tuple[str, PregelProtocol]]: for name, node in self.nodes.items(): # filter by prefix if namespace is not None: @@ -437,7 +437,7 @@ class Pregel(PregelProtocol): continue # find the subgraph, if any - graph = cast(Optional[Pregel], find_subgraph_pregel(node.bound)) + graph = node.subgraphs[0] if node.subgraphs else None # if found, yield recursively if graph: @@ -446,7 +446,7 @@ class Pregel(PregelProtocol): return # we found it, stop searching if namespace is None: yield name, graph - if recurse: + if recurse and isinstance(graph, Pregel): if namespace is not None: namespace = namespace[len(name) + 1 :] yield from ( @@ -458,7 +458,7 @@ class Pregel(PregelProtocol): async def aget_subgraphs( self, *, namespace: Optional[str] = None, recurse: bool = False - ) -> AsyncIterator[tuple[str, Pregel]]: + ) -> AsyncIterator[tuple[str, PregelProtocol]]: for name, node in self.get_subgraphs(namespace=namespace, recurse=recurse): yield name, node diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 5663cfa5f..03b2af6f4 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -630,6 +630,7 @@ def prepare_single_task( task_id, task_path[:3], writers=proc.flat_writers, + subgraphs=proc.subgraphs, ) else: return PregelTask(task_id, packet.node, task_path[:3]) @@ -754,6 +755,7 @@ def prepare_single_task( task_id, task_path[:3], writers=proc.flat_writers, + subgraphs=proc.subgraphs, ) else: return PregelTask(task_id, name, task_path[:3]) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index cc398a599..8429fd538 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -30,7 +30,6 @@ from langgraph.constants import ( TAG_HIDDEN, ) from langgraph.pregel.io import read_channels -from langgraph.pregel.utils import find_subgraph_pregel from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot from langgraph.utils.config import patch_checkpoint_map @@ -157,7 +156,7 @@ def map_debug_checkpoint( task_states: dict[str, Union[RunnableConfig, StateSnapshot]] = {} for task in tasks: - if not find_subgraph_pregel(task.proc): + if not task.subgraphs: continue # assemble checkpoint_ns for this task diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 6733258dc..02e466475 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -22,7 +22,9 @@ from langchain_core.runnables.base import Input, Other, coerce_to_runnable from langchain_core.runnables.utils import ConfigurableFieldSpec from langgraph.constants import CONF, CONFIG_KEY_READ +from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.retry import RetryPolicy +from langgraph.pregel.utils import find_subgraph_pregel from langgraph.pregel.write import ChannelWrite from langgraph.utils.config import merge_configs from langgraph.utils.runnable import RunnableCallable, RunnableSeq @@ -145,6 +147,9 @@ class PregelNode(Runnable): metadata: Optional[Mapping[str, Any]] """Metadata to attach to the node for tracing.""" + subgraphs: Sequence[PregelProtocol] + """Subgraphs used by the node.""" + def __init__( self, *, @@ -165,9 +170,21 @@ class PregelNode(Runnable): self.retry_policy = retry_policy self.tags = tags self.metadata = metadata + if self.bound is not DEFAULT_BOUND: + try: + subgraph = find_subgraph_pregel(self.bound) + except Exception: + subgraph = None + if subgraph: + self.subgraphs = [subgraph] + else: + self.subgraphs = [] + else: + self.subgraphs = [] def copy(self, update: dict[str, Any]) -> PregelNode: attrs = {**self.__dict__, **update} + attrs.pop("subgraphs") return PregelNode(**attrs) @cached_property @@ -205,12 +222,12 @@ class PregelNode(Runnable): return self.bound def join(self, channels: Sequence[str]) -> PregelNode: - assert isinstance(channels, list) or isinstance( - channels, tuple - ), "channels must be a list or tuple" - assert isinstance( - self.channels, dict - ), "all channels must be named when using .join()" + assert isinstance(channels, list) or isinstance(channels, tuple), ( + "channels must be a list or tuple" + ) + assert isinstance(self.channels, dict), ( + "all channels must be named when using .join()" + ) return self.copy( update=dict( channels={ diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index 66464ef9a..f484664a0 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -26,7 +26,7 @@ def get_new_channel_versions( return new_versions -def find_subgraph_pregel(candidate: Runnable) -> Optional[Runnable]: +def find_subgraph_pregel(candidate: Runnable) -> Optional[PregelProtocol]: from langgraph.pregel import Pregel candidates: list[Runnable] = [candidate] diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 68bb9830e..b52949cd0 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -25,6 +25,7 @@ from typing_extensions import Self from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata if TYPE_CHECKING: + from langgraph.pregel.protocol import PregelProtocol from langgraph.store.base import BaseStore @@ -155,6 +156,7 @@ class PregelExecutableTask(NamedTuple): path: tuple[Union[str, int, tuple], ...] scheduled: bool = False writers: Sequence[Runnable] = () + subgraphs: Sequence["PregelProtocol"] = () class StateSnapshot(NamedTuple): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index b08bc88dc..a54e77889 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6256,6 +6256,8 @@ def test_merging_updates_command_parent(): ), ((), {"node_3": {"bar": ["node_3"]}}), ] + + def test_entrypoint_output_schema_with_return_and_save() -> None: """Test output schema inference with entrypoint.final."""