From 2fccd0958534f6266c53a331b6717c73b62e83ab Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 21 Sep 2024 17:52:17 -0700 Subject: [PATCH] Fix --- libs/langgraph/langgraph/pregel/algo.py | 3 +++ libs/langgraph/langgraph/pregel/loop.py | 11 ++++++----- libs/langgraph/tests/test_tracing_interops.py | 2 +- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index f9b0096c8..3c98a248b 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -30,6 +30,7 @@ from langgraph.checkpoint.base import ( from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINTER, + CONFIG_KEY_GRAPH_COUNT, CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_TASK_ID, @@ -429,6 +430,7 @@ def prepare_single_task( manager.get_child(f"graph:step:{step}") if manager else None ), configurable={ + CONFIG_KEY_GRAPH_COUNT: 0, CONFIG_KEY_TASK_ID: task_id, # deque.extend is thread-safe CONFIG_KEY_SEND: partial( @@ -539,6 +541,7 @@ def prepare_single_task( else None ), configurable={ + CONFIG_KEY_GRAPH_COUNT: 0, CONFIG_KEY_TASK_ID: task_id, # deque.extend is thread-safe CONFIG_KEY_SEND: partial( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f5dba4f9b..a40b5f8f6 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -221,11 +221,12 @@ class PregelLoop: self.config = patch_configurable( self.config, {"checkpoint_ns": "", "checkpoint_id": None} ) - if config["configurable"].get(CONFIG_KEY_GRAPH_COUNT, 0) > 0: - raise ValueError("Detected multiple subgraphs called in a single node.") - else: - # mutate config so that sibling subgraphs can be detected - self.config["configurable"][CONFIG_KEY_GRAPH_COUNT] = 1 + if self.is_nested: + if config["configurable"].get(CONFIG_KEY_GRAPH_COUNT, 0) > 0: + raise ValueError("Detected multiple subgraphs called in a single node.") + else: + # mutate config so that sibling subgraphs can be detected + self.config["configurable"][CONFIG_KEY_GRAPH_COUNT] = 1 if ( CONFIG_KEY_CHECKPOINT_MAP in self.config["configurable"] and self.config["configurable"].get("checkpoint_ns") diff --git a/libs/langgraph/tests/test_tracing_interops.py b/libs/langgraph/tests/test_tracing_interops.py index a3dae6420..5b458394b 100644 --- a/libs/langgraph/tests/test_tracing_interops.py +++ b/libs/langgraph/tests/test_tracing_interops.py @@ -55,7 +55,7 @@ def wait_for( raise ValueError(f"Callable did not return within {total_time}") -@pytest.skip("This test times out in CI") +@pytest.mark.skip("This test times out in CI") async def test_nested_tracing(): lt_py_311 = sys.version_info < (3, 11) mock_client = _get_mock_client()