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 e9ac31502..877ddcf4e 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -204,6 +204,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 @@ -243,6 +247,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]] = (), @@ -260,6 +265,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 @@ -1631,7 +1637,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 @@ -1861,7 +1872,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 76a067533..64b5076de 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2166,10 +2166,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"]} @@ -2306,10 +2306,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"]} @@ -2739,11 +2739,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"]} @@ -2829,10 +2829,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"]} @@ -2902,10 +2902,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: @@ -2928,10 +2928,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" @@ -2964,13 +2964,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"]} @@ -5490,3 +5490,39 @@ 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 diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a6fab066f..558a93709 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: @@ -4293,10 +4274,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"]} @@ -4383,10 +4364,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"]} @@ -4800,11 +4781,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"]} @@ -4894,10 +4875,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"]} @@ -4978,13 +4959,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"]} @@ -6116,10 +6097,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): @@ -6152,10 +6130,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): @@ -6219,10 +6194,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): @@ -6507,10 +6479,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.""" @@ -6691,10 +6660,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.""" @@ -6716,10 +6682,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.""" @@ -6759,10 +6722,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): @@ -6874,3 +6834,40 @@ async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: "invoke_sub_agent": {"input": True}, }, ] + + +@NEEDS_CONTEXTVARS +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