From 7ecad39ecbcf6fa5e815e55c89877369b1a733ca Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 15 Jan 2025 21:15:05 -0500 Subject: [PATCH 1/4] x --- libs/langgraph/tests/test_pregel.py | 101 ++++++++++++++++++++++++---- 1 file changed, 87 insertions(+), 14 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 67536deed..fdbd032c4 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,3 +1,4 @@ +import asyncio import enum import json import logging @@ -35,6 +36,20 @@ from langchain_core.runnables import ( from langsmith import traceable from pytest_mock import MockerFixture from syrupy import SnapshotAssertion +from tests.agents import AgentAction, AgentFinish +from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence +from tests.conftest import ( + ALL_CHECKPOINTERS_SYNC, + ALL_STORES_SYNC, + REGULAR_CHECKPOINTERS_SYNC, + SHOULD_CHECK_SNAPSHOTS, +) +from tests.memory_assert import MemorySaverAssertCheckpointMetadata +from tests.messages import ( + _AnyIdAIMessage, + _AnyIdHumanMessage, + _AnyIdToolMessage, +) from typing_extensions import TypedDict from langgraph.channels.base import BaseChannel @@ -67,20 +82,6 @@ from langgraph.types import ( StreamWriter, interrupt, ) -from tests.agents import AgentAction, AgentFinish -from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence -from tests.conftest import ( - ALL_CHECKPOINTERS_SYNC, - ALL_STORES_SYNC, - REGULAR_CHECKPOINTERS_SYNC, - SHOULD_CHECK_SNAPSHOTS, -) -from tests.memory_assert import MemorySaverAssertCheckpointMetadata -from tests.messages import ( - _AnyIdAIMessage, - _AnyIdHumanMessage, - _AnyIdToolMessage, -) logger = logging.getLogger(__name__) @@ -5491,3 +5492,75 @@ def test_double_interrupt_subgraph( "invoke_sub_agent": {"input": True}, }, ] + + +def test_sync_streaming_with_functional_api() -> None: + """Test streaming with functional API. + + This test verifies that we're able to stream results as they're being generated + rather than have all the results arrive at once after the graph has completed. + + The time of arrival between the two updates corresponding to the two `slow` tasks + should be greater than the time delay between the two tasks. + """ + + time_delay = 0.01 + + @task() + def slow() -> dict: + time.sleep(time_delay) # Simulate a delay of 10 ms + return {"tic": time.time()} + + @entrypoint() + def graph(inputs: dict) -> list: + first = slow().result() + second = slow().result() + return [first, second] + + arrival_times = [] + + for chunk in graph.stream({}): + if "slow" not in chunk: # We'll just look at the updates from `slow` + continue + arrival_times.append(time.time()) + + assert len(arrival_times) == 2 + delta = arrival_times[1] - arrival_times[0] + # Delta cannot be less than 10 ms if it is streaming as results are generated. + assert delta > time_delay + + +async def test_async_streaming_with_functional_api() -> None: + """Test streaming with functional API. + + This test verifies that we're able to stream results as they're being generated + rather than have all the results arrive at once after the graph has completed. + + The time of arrival between the two updates corresponding to the two `slow` tasks + should be greater than the time delay between the two tasks. + """ + + time_delay = 0.01 + + @task() + async def slow() -> dict: + await asyncio.sleep(time_delay) # Simulate a delay of 10 ms + return {"tic": time.time()} + + @entrypoint() + async def graph(inputs: dict) -> list: + first = await slow() + second = await slow() + return [first, second] + + arrival_times = [] + + async for chunk in graph.astream({}): + if "slow" not in chunk: # We'll just look at the updates from `slow` + continue + arrival_times.append(time.time()) + + assert len(arrival_times) == 2 + delta = arrival_times[1] - arrival_times[0] + # Delta cannot be less than 10 ms if it is streaming as results are generated. + assert delta > time_delay From f9c25bba07cb7c84738c67258930739f1a71c78e Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 15 Jan 2025 21:15:38 -0500 Subject: [PATCH 2/4] x --- libs/langgraph/tests/test_pregel.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index fdbd032c4..4eae525d1 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -36,20 +36,6 @@ from langchain_core.runnables import ( from langsmith import traceable from pytest_mock import MockerFixture from syrupy import SnapshotAssertion -from tests.agents import AgentAction, AgentFinish -from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence -from tests.conftest import ( - ALL_CHECKPOINTERS_SYNC, - ALL_STORES_SYNC, - REGULAR_CHECKPOINTERS_SYNC, - SHOULD_CHECK_SNAPSHOTS, -) -from tests.memory_assert import MemorySaverAssertCheckpointMetadata -from tests.messages import ( - _AnyIdAIMessage, - _AnyIdHumanMessage, - _AnyIdToolMessage, -) from typing_extensions import TypedDict from langgraph.channels.base import BaseChannel @@ -82,6 +68,20 @@ from langgraph.types import ( StreamWriter, interrupt, ) +from tests.agents import AgentAction, AgentFinish +from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence +from tests.conftest import ( + ALL_CHECKPOINTERS_SYNC, + ALL_STORES_SYNC, + REGULAR_CHECKPOINTERS_SYNC, + SHOULD_CHECK_SNAPSHOTS, +) +from tests.memory_assert import MemorySaverAssertCheckpointMetadata +from tests.messages import ( + _AnyIdAIMessage, + _AnyIdHumanMessage, + _AnyIdToolMessage, +) logger = logging.getLogger(__name__) From 145220f2a81aafe4e0a0eca105379ba087484724 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 18:29:34 -0800 Subject: [PATCH 3/4] Fix --- libs/langgraph/langgraph/func/__init__.py | 1 + libs/langgraph/langgraph/pregel/__init__.py | 20 ++++++- libs/langgraph/tests/test_pregel.py | 65 +++++---------------- libs/langgraph/tests/test_pregel_async.py | 56 ++++++++++++++---- 4 files changed, 79 insertions(+), 63 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index d625be560..4749c075f 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -148,6 +148,7 @@ def entrypoint( output_channels=END, stream_channels=END, stream_mode=stream_mode, + stream_eager=True, checkpointer=checkpointer, store=store, ) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 102e68be8..6671bf70a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -203,6 +203,10 @@ class Pregel(PregelProtocol): stream_mode: StreamMode = "values" """Mode to stream output, defaults to 'values'.""" + stream_eager: bool = False + """Whether to force emitting stream events eagerly, automatically turned on + for stream_mode "messages" and "custom".""" + output_channels: Union[str, Sequence[str]] stream_channels: Optional[Union[str, Sequence[str]]] = None @@ -242,6 +246,7 @@ class Pregel(PregelProtocol): channels: Optional[dict[str, Union[BaseChannel, ManagedValueSpec]]], auto_validate: bool = True, stream_mode: StreamMode = "values", + stream_eager: bool = False, output_channels: Union[str, Sequence[str]], stream_channels: Optional[Union[str, Sequence[str]]] = None, interrupt_after_nodes: Union[All, Sequence[str]] = (), @@ -259,6 +264,7 @@ class Pregel(PregelProtocol): self.nodes = nodes self.channels = channels or {} self.stream_mode = stream_mode + self.stream_eager = stream_eager self.output_channels = output_channels self.stream_channels = stream_channels self.interrupt_after_nodes = interrupt_after_nodes @@ -1655,7 +1661,12 @@ class Pregel(PregelProtocol): if subgraphs: loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream # enable concurrent streaming - if subgraphs or "messages" in stream_modes or "custom" in stream_modes: + if ( + self.stream_eager + or subgraphs + or "messages" in stream_modes + or "custom" in stream_modes + ): # we are careful to have a single waiter live at any one time # because on exit we increment semaphore count by exactly 1 waiter: Optional[concurrent.futures.Future] = None @@ -1886,7 +1897,12 @@ class Pregel(PregelProtocol): stream_put, stream_modes ) # enable concurrent streaming - if subgraphs or "messages" in stream_modes or "custom" in stream_modes: + if ( + self.stream_eager + or subgraphs + or "messages" in stream_modes + or "custom" in stream_modes + ): def get_waiter() -> asyncio.Task[None]: return aioloop.create_task(stream.wait()) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 4eae525d1..f35442bdf 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,4 +1,3 @@ -import asyncio import enum import json import logging @@ -2168,10 +2167,10 @@ def test_in_one_fan_out_state_graph_waiting_edge( @workflow.add_node def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -2308,10 +2307,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( docs: Annotated[list[str], sorted_add] def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -2741,11 +2740,11 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( docs: Annotated[list[str], sorted_add] def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyzer_one(data: State) -> State: time.sleep(0.1) - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -2831,10 +2830,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: docs: Annotated[list[str], sorted_add] def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -2904,10 +2903,10 @@ def test_callable_in_conditional_edges_with_no_path_map() -> None: query: str def rewrite(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyze(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} class ChooseAnalyzer: def __call__(self, data: State) -> str: @@ -2930,10 +2929,10 @@ def test_function_in_conditional_edges_with_no_path_map() -> None: query: str def rewrite(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyze(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def choose_analyzer(data: State) -> str: return "analyzer" @@ -2966,13 +2965,13 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None: docs: Annotated[list[str], sorted_add] def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def retriever_picker(data: State) -> list[str]: return ["analyzer_one", "retriever_two"] def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -5528,39 +5527,3 @@ def test_sync_streaming_with_functional_api() -> None: delta = arrival_times[1] - arrival_times[0] # Delta cannot be less than 10 ms if it is streaming as results are generated. assert delta > time_delay - - -async def test_async_streaming_with_functional_api() -> None: - """Test streaming with functional API. - - This test verifies that we're able to stream results as they're being generated - rather than have all the results arrive at once after the graph has completed. - - The time of arrival between the two updates corresponding to the two `slow` tasks - should be greater than the time delay between the two tasks. - """ - - time_delay = 0.01 - - @task() - async def slow() -> dict: - await asyncio.sleep(time_delay) # Simulate a delay of 10 ms - return {"tic": time.time()} - - @entrypoint() - async def graph(inputs: dict) -> list: - first = await slow() - second = await slow() - return [first, second] - - arrival_times = [] - - async for chunk in graph.astream({}): - if "slow" not in chunk: # We'll just look at the updates from `slow` - continue - arrival_times.append(time.time()) - - assert len(arrival_times) == 2 - delta = arrival_times[1] - arrival_times[0] - # Delta cannot be less than 10 ms if it is streaming as results are generated. - assert delta > time_delay diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 2f7d7a1c4..bdbec78e2 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -4294,10 +4294,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) - docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -4384,10 +4384,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -4801,11 +4801,11 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def analyzer_one(data: State) -> State: await asyncio.sleep(0.1) - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -4895,10 +4895,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -4979,13 +4979,13 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> N docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def retriever_picker(data: State) -> list[str]: return ["analyzer_one", "retriever_two"] async def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -6875,3 +6875,39 @@ async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: "invoke_sub_agent": {"input": True}, }, ] + + +async def test_async_streaming_with_functional_api() -> None: + """Test streaming with functional API. + + This test verifies that we're able to stream results as they're being generated + rather than have all the results arrive at once after the graph has completed. + + The time of arrival between the two updates corresponding to the two `slow` tasks + should be greater than the time delay between the two tasks. + """ + + time_delay = 0.01 + + @task() + async def slow() -> dict: + await asyncio.sleep(time_delay) # Simulate a delay of 10 ms + return {"tic": asyncio.get_running_loop().time()} + + @entrypoint() + async def graph(inputs: dict) -> list: + first = await slow() + second = await slow() + return [first, second] + + arrival_times = [] + + async for chunk in graph.astream({}): + if "slow" not in chunk: # We'll just look at the updates from `slow` + continue + arrival_times.append(asyncio.get_running_loop().time()) + + assert len(arrival_times) == 2 + delta = arrival_times[1] - arrival_times[0] + # Delta cannot be less than 10 ms if it is streaming as results are generated. + assert delta > time_delay From 0e4dbb4c62dc9542ff6a5728fea5d3a3390c9531 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 18:36:49 -0800 Subject: [PATCH 4/4] Guard --- libs/langgraph/tests/test_pregel_async.py | 81 ++++++----------------- 1 file changed, 21 insertions(+), 60 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index bdbec78e2..9e984d214 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -89,6 +89,11 @@ logger = logging.getLogger(__name__) pytestmark = pytest.mark.anyio +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + async def test_checkpoint_errors() -> None: class FaultyGetCheckpointer(MemorySaver): @@ -501,10 +506,7 @@ async def test_node_cancellation_on_other_node_exception_two() -> None: await graph.ainvoke(1) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_dynamic_interrupt(checkpointer_name: str) -> None: class State(TypedDict): @@ -678,10 +680,7 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: ) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: class SubgraphState(TypedDict): @@ -872,10 +871,7 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: ) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_copy_checkpoint(checkpointer_name: str) -> None: class State(TypedDict): @@ -1079,10 +1075,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: ) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_node_not_cancelled_on_other_node_interrupted( checkpointer_name: str, @@ -2442,10 +2435,7 @@ async def test_send_sequences(checkpointer_name: str) -> None: ] -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2493,10 +2483,7 @@ async def test_imp_task(checkpointer_name: str) -> None: assert mapper_calls == 2 -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task_cancel(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2547,10 +2534,7 @@ async def test_imp_task_cancel(checkpointer_name: str) -> None: assert mapper_cancels == 2 -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_sync_from_async(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2583,10 +2567,7 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None: ] -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_stream_order(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -6117,10 +6098,7 @@ async def test_parent_command(checkpointer_name: str) -> None: ) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interrupt_subgraph(checkpointer_name: str): class State(TypedDict): @@ -6153,10 +6131,7 @@ async def test_interrupt_subgraph(checkpointer_name: str): assert await graph.ainvoke(Command(resume="bar"), thread1) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interrupt_multiple(checkpointer_name: str): class State(TypedDict): @@ -6220,10 +6195,7 @@ async def test_interrupt_multiple(checkpointer_name: str): ] -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interrupt_loop(checkpointer_name: str): class State(TypedDict): @@ -6508,10 +6480,7 @@ async def test_parallel_node_execution(): assert duration < 3.0 -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_multiple_interrupt_state_persistence(checkpointer_name: str) -> None: """Test that state is preserved correctly across multiple interrupts.""" @@ -6692,10 +6661,7 @@ async def test_multiple_updates() -> None: ] -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_falsy_return_from_task(checkpointer_name: str) -> None: """Test with a falsy return from a task.""" @@ -6717,10 +6683,7 @@ async def test_falsy_return_from_task(checkpointer_name: str) -> None: await graph.ainvoke(Command(resume="123"), configurable) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None: """Test multiple interrupts with an imperative API.""" @@ -6760,10 +6723,7 @@ async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None: assert counter == 3 -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: class AgentState(TypedDict): @@ -6877,6 +6837,7 @@ async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: ] +@NEEDS_CONTEXTVARS async def test_async_streaming_with_functional_api() -> None: """Test streaming with functional API.