diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index be43a021a..ca7f787c3 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import concurrent.futures +import time from collections import defaultdict, deque from functools import partial from typing import ( @@ -898,21 +899,56 @@ class Pregel( map_debug_tasks(step, next_tasks), ) - futures = [ - executor.submit(run_with_retry, task, self.retry_policy) - for task in next_tasks - ] - # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks - done, inflight = concurrent.futures.wait( - futures, - return_when=concurrent.futures.FIRST_EXCEPTION, - timeout=self.step_timeout, + # yield updates/debug output as each task finishes + futures = { + executor.submit(run_with_retry, task, self.retry_policy): task + for task in next_tasks + } + end_time = ( + self.step_timeout + time.monotonic() + if self.step_timeout + else None ) + while futures: + done, inflight = concurrent.futures.wait( + futures, + return_when=concurrent.futures.FIRST_COMPLETED, + timeout=max(0, end_time - time.monotonic()) + if end_time + else None, + ) + for fut in done: + task = futures.pop(fut) + if fut.exception() is not None: + # we got an exception, break out of while loop + # exception will be handle in panic_or_proceed + futures.clear() + else: + # yield updates output for the finished task + if "updates" in stream_modes: + yield from _with_mode( + "updates", + isinstance(stream_mode, list), + map_output_updates(output_keys, [task]), + ) + if "debug" in stream_modes: + yield from _with_mode( + "debug", + isinstance(stream_mode, list), + map_debug_task_results( + step, [task], self.stream_channels_list + ), + ) + else: + # remove references to loop vars + del fut, task # panic on failure or timeout _panic_or_proceed(done, inflight, step) + # don't keep futures around in memory longer than needed + del done, inflight, futures # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() @@ -927,21 +963,7 @@ class Pregel( # apply writes to channels _apply_writes(checkpoint, channels, pending_writes) - # yield current value or updates - if "updates" in stream_modes: - yield from _with_mode( - "updates", - isinstance(stream_mode, list), - map_output_updates(output_keys, next_tasks), - ) - if "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_task_results( - step, next_tasks, self.stream_channels_list - ), - ) + # yield values output if "values" in stream_modes: yield from _with_mode( "values", @@ -1030,6 +1052,7 @@ class Pregel( None, ) try: + loop = asyncio.get_event_loop() bg: list[asyncio.Task] = [] if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") @@ -1201,23 +1224,58 @@ class Pregel( ): yield chunk - futures = [ - asyncio.create_task( - arun_with_retry(task, self.retry_policy, do_stream) - ) - for task in next_tasks - ] - # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks - done, inflight = await asyncio.wait( - futures, - return_when=asyncio.FIRST_EXCEPTION, - timeout=self.step_timeout, + # yield updates/debug output as each task finishes + futures = { + asyncio.create_task( + arun_with_retry(task, self.retry_policy, do_stream) + ): task + for task in next_tasks + } + end_time = ( + self.step_timeout + loop.time() if self.step_timeout else None ) + while futures: + done, inflight = await asyncio.wait( + futures, + return_when=asyncio.FIRST_COMPLETED, + timeout=max(0, end_time - loop.time()) + if end_time + else None, + ) + for fut in done: + task = futures.pop(fut) + if fut.exception() is not None: + # we got an exception, break out of while loop + # exception will be handle in panic_or_proceed + futures.clear() + else: + # yield updates output for the finished task + if "updates" in stream_modes: + for chunk in _with_mode( + "updates", + isinstance(stream_mode, list), + map_output_updates(output_keys, [task]), + ): + yield chunk + if "debug" in stream_modes: + for chunk in _with_mode( + "debug", + isinstance(stream_mode, list), + map_debug_task_results( + step, [task], self.stream_channels_list + ), + ): + yield chunk + else: + # remove references to loop vars + del fut, task # panic on failure or timeout _panic_or_proceed(done, inflight, step) + # don't keep futures around in memory longer than needed + del done, inflight, futures # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() @@ -1232,23 +1290,7 @@ class Pregel( # apply writes to channels _apply_writes(checkpoint, channels, pending_writes) - # yield current value or updates - if "updates" in stream_modes: - for chunk in _with_mode( - "updates", - isinstance(stream_mode, list), - map_output_updates(output_keys, next_tasks), - ): - yield chunk - if "debug" in stream_modes: - for chunk in _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_task_results( - step, next_tasks, self.stream_channels_list - ), - ): - yield chunk + # yield current values if "values" in stream_modes: for chunk in _with_mode( "values", @@ -1602,6 +1644,7 @@ def _prepare_next_tasks( "langgraph_step": step, "langgraph_node": packet.node, "langgraph_triggers": [TASKS], + "langgraph_task_idx": len(tasks), } }, ), @@ -1672,6 +1715,7 @@ def _prepare_next_tasks( "langgraph_step": step, "langgraph_node": name, "langgraph_triggers": triggers, + "langgraph_task_idx": len(tasks), } }, ), diff --git a/langgraph/pregel/debug.py b/langgraph/pregel/debug.py index c042567a5..7fb294903 100644 --- a/langgraph/pregel/debug.py +++ b/langgraph/pregel/debug.py @@ -66,7 +66,7 @@ def map_debug_tasks( step: int, tasks: list[PregelExecutableTask] ) -> Iterator[DebugOutputTask]: ts = datetime.now(timezone.utc).isoformat() - for idx, (name, input, _, _, config, triggers) in enumerate(tasks): + for name, input, _, _, config, triggers in tasks: if config is not None and TAG_HIDDEN in config.get("tags", []): continue @@ -75,7 +75,9 @@ def map_debug_tasks( "timestamp": ts, "step": step, "payload": { - "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, idx)))), + "id": str( + uuid5(TASK_NAMESPACE, json.dumps((name, step, config["metadata"]))) + ), "name": name, "input": input, "triggers": triggers, @@ -89,7 +91,7 @@ def map_debug_task_results( stream_channels_list: Sequence[str], ) -> Iterator[DebugOutputTaskResult]: ts = datetime.now(timezone.utc).isoformat() - for idx, (name, _, _, writes, config, _) in enumerate(tasks): + for name, _, _, writes, config, _ in tasks: if config is not None and TAG_HIDDEN in config.get("tags", []): continue @@ -98,7 +100,9 @@ def map_debug_task_results( "timestamp": ts, "step": step, "payload": { - "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, idx)))), + "id": str( + uuid5(TASK_NAMESPACE, json.dumps((name, step, config["metadata"]))) + ), "name": name, "result": [w for w in writes if w[0] in stream_channels_list], }, diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 0b6979606..fc1fc38ea 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -4,7 +4,7 @@ from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union from langchain_core.runnables.utils import AddableDict from langgraph.channels.base import BaseChannel, EmptyChannelError -from langgraph.constants import TAG_HIDDEN, TASKS +from langgraph.constants import TAG_HIDDEN from langgraph.pregel.log import logger from langgraph.pregel.types import PregelExecutableTask @@ -104,42 +104,33 @@ def map_output_updates( ] if isinstance(output_channels, str): if updated := [ - (triggers == [TASKS], node, value) - for node, _, _, writes, _, triggers in output_tasks + (node, value) + for node, _, _, writes, _, _ in output_tasks for chan, value in writes if chan == output_channels ]: grouped = defaultdict(list) - for from_packet, node, value in updated: - if from_packet: - grouped[node].append(value) - for from_packet, node, value in updated: - if not from_packet: - if grouped[node]: - grouped[node].append(value) - else: - grouped[node] = value + for node, value in updated: + grouped[node].append(value) + for node, value in grouped.items(): + if len(value) == 1: + grouped[node] = value[0] yield AddableUpdatesDict(grouped) else: if updated := [ ( - triggers == [TASKS], node, {chan: value for chan, value in writes if chan in output_channels}, ) - for node, _, _, writes, _, triggers in output_tasks + for node, _, _, writes, _, _ in output_tasks if any(chan in output_channels for chan, _ in writes) ]: grouped = defaultdict(list) - for from_packet, node, value in updated: - if from_packet: - grouped[node].append(value) - for from_packet, node, value in updated: - if not from_packet: - if grouped[node]: - grouped[node].append(value) - else: - grouped[node] = value + for node, value in updated: + grouped[node].append(value) + for node, value in grouped.items(): + if len(value) == 1: + grouped[node] = value[0] yield AddableUpdatesDict(grouped) diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index a3826314a..e47d94b57 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -16,7 +16,7 @@ class PregelExecutableTask(NamedTuple): input: Any proc: Runnable writes: deque[tuple[str, Any]] - config: Optional[RunnableConfig] + config: RunnableConfig triggers: list[str] diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 500637b1d..4997b62fe 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -504,6 +504,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: two = ( Channel.subscribe_to("inbox") | RunnableLambda(add_one).batch + | RunnablePassthrough(lambda _: time.sleep(0.1)) | Channel.write_to("output").batch ) @@ -534,7 +535,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: ] assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="updates")] == [ - {"one": {"inbox": 3}, "two": {"output": 13}}, + {"one": {"inbox": 3}}, + {"two": {"output": 13}}, {"two": {"output": 4}}, ] assert [*app.stream({"input": 2, "inbox": 12})] == [ @@ -547,7 +549,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "9379da35-ae1c-5a7b-8556-7ce22a1f8fde", + "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923", "name": "one", "input": 2, "triggers": ["input"], @@ -558,7 +560,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "49ac8f60-4ff2-5cdd-a319-66bbd9837e5a", + "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e", "name": "two", "input": [12], "triggers": ["inbox"], @@ -569,7 +571,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "9379da35-ae1c-5a7b-8556-7ce22a1f8fde", + "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923", "name": "one", "result": [("inbox", 3)], }, @@ -579,7 +581,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "49ac8f60-4ff2-5cdd-a319-66bbd9837e5a", + "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e", "name": "two", "result": [("output", 13)], }, @@ -589,7 +591,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "b97f26c1-a34b-51e0-884e-44a41a3a3b47", + "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62", "name": "two", "input": [3], "triggers": ["inbox"], @@ -600,7 +602,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "b97f26c1-a34b-51e0-884e-44a41a3a3b47", + "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62", "name": "two", "result": [("output", 4)], }, @@ -3475,12 +3477,12 @@ def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ), @@ -3506,6 +3508,7 @@ def test_state_graph_packets() -> None: return END def tools_node(tool_call: ToolCall, config: RunnableConfig) -> AgentState: + time.sleep(tool_call["args"].get("idx", 0) / 10) output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) return { "messages": ToolMessage( @@ -3563,12 +3566,12 @@ def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ), @@ -3610,16 +3613,14 @@ def test_state_graph_packets() -> None: }, }, { - "tools": [ - { - "messages": ToolMessage( - content="result for query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ) - } - ] + "tools": { + "messages": ToolMessage( + content="result for query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } }, { "agent": { @@ -3630,36 +3631,36 @@ def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ) } }, { - "tools": [ - { - "messages": ToolMessage( - content="result for another", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call234", - ) - }, - { - "messages": ToolMessage( - content="result for a third one", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call567", - ), - }, - ] + "tools": { + "messages": ToolMessage( + content="result for another", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call234", + ) + }, + }, + { + "tools": { + "messages": ToolMessage( + content="result for a third one", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call567", + ), + }, }, {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, ] @@ -3792,16 +3793,14 @@ def test_state_graph_packets() -> None: assert [c for c in app_w_interrupt.stream(None, config)] == [ { - "tools": [ - { - "messages": ToolMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ) - } - ] + "tools": { + "messages": ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } }, { "agent": { @@ -3812,12 +3811,12 @@ def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ) @@ -3856,12 +3855,12 @@ def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ), @@ -3882,12 +3881,12 @@ def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ) @@ -5520,6 +5519,11 @@ def test_in_one_fan_out_out_one_graph_state() -> None: return {"query": f'query: {data["query"]}'} def retriever_one(data: State) -> State: + # timer ensures stream output order is stable + # also, it confirms that the update order is not dependent on finishing order + # instead being defined by the order of the nodes/edges in the graph definition + # ie. stable between invocations + time.sleep(0.1) return {"docs": ["doc1", "doc2"]} def retriever_two(data: State) -> State: @@ -5552,10 +5556,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None: assert [*app.stream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "retriever_two": {"docs": ["doc3", "doc4"]}, - "retriever_one": {"docs": ["doc1", "doc2"]}, - }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -5587,7 +5589,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "03dadab4-fb41-5308-a8a4-6eeb9ef7b9aa", + "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf", "name": "rewrite_query", "input": { "query": "what is weather in sf", @@ -5606,7 +5608,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "03dadab4-fb41-5308-a8a4-6eeb9ef7b9aa", + "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf", "name": "rewrite_query", "result": [("query", "query: what is weather in sf")], }, @@ -5620,7 +5622,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "96f499e2-e203-5a13-9259-08cb62f4a2e5", + "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b", "name": "retriever_one", "input": { "query": "query: what is weather in sf", @@ -5638,7 +5640,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "6b344a90-a061-5f17-8714-51f0cf67cf01", + "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2", "name": "retriever_two", "input": { "query": "query: what is weather in sf", @@ -5651,10 +5653,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ), ( "updates", - { - "retriever_one": {"docs": ["doc1", "doc2"]}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, ), ( "debug", @@ -5663,12 +5662,16 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "96f499e2-e203-5a13-9259-08cb62f4a2e5", - "name": "retriever_one", - "result": [("docs", ["doc1", "doc2"])], + "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2", + "name": "retriever_two", + "result": [("docs", ["doc3", "doc4"])], }, }, ), + ( + "updates", + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ), ( "debug", { @@ -5676,9 +5679,9 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "6b344a90-a061-5f17-8714-51f0cf67cf01", - "name": "retriever_two", - "result": [("docs", ["doc3", "doc4"])], + "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b", + "name": "retriever_one", + "result": [("docs", ["doc1", "doc2"])], }, }, ), @@ -5696,7 +5699,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "0dda6269-4ce3-5b98-9cea-d40737a68500", + "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0", "name": "qa", "input": { "query": "query: what is weather in sf", @@ -5715,7 +5718,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "0dda6269-4ce3-5b98-9cea-d40737a68500", + "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0", "name": "qa", "result": [("answer", "doc1,doc2,doc3,doc4")], }, @@ -5996,7 +5999,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "d6e87693-41fb-58f5-8e0d-ee9ab46890b5", + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", "name": "prepare", "input": {"my_key": "value", "market": "DE"}, "triggers": ["start:prepare"], @@ -6007,7 +6010,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "d6e87693-41fb-58f5-8e0d-ee9ab46890b5", + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", "name": "prepare", "result": [("my_key", " prepared")], }, @@ -6044,7 +6047,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "b1826010-0028-5aa7-abd2-ed24984614ea", + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", "name": "tool_two_slow", "input": {"my_key": "value prepared", "market": "DE"}, "triggers": ["branch:prepare:condition:tool_two_slow"], @@ -6055,7 +6058,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "b1826010-0028-5aa7-abd2-ed24984614ea", + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", "name": "tool_two_slow", "result": [("my_key", " slow")], }, @@ -6092,7 +6095,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "a22dbd2d-f136-57f0-a86a-bc2c234ffcb1", + "id": "ceada3c5-5f25-59e4-9ea5-544599ce1d2f", "name": "finish", "input": {"my_key": "value prepared slow", "market": "DE"}, "triggers": ["branch:prepare:condition:then"], @@ -6103,7 +6106,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "a22dbd2d-f136-57f0-a86a-bc2c234ffcb1", + "id": "ceada3c5-5f25-59e4-9ea5-544599ce1d2f", "name": "finish", "result": [("my_key", " finished")], }, @@ -6378,6 +6381,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> return {"docs": ["doc1", "doc2"]} def retriever_two(data: State) -> State: + time.sleep(0.1) # to ensure stream order return {"docs": ["doc3", "doc4"]} def qa(data: State) -> State: @@ -6407,10 +6411,8 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> assert [*app.stream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -6425,10 +6427,8 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) ] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] @@ -6464,6 +6464,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( return {"docs": ["doc1", "doc2"]} def retriever_two(data: State) -> State: + time.sleep(0.1) return {"docs": ["doc3", "doc4"]} def qa(data: State) -> State: @@ -6499,10 +6500,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( assert [*app.stream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -6517,10 +6516,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) ] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] @@ -6563,6 +6560,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( return {"docs": ["doc1", "doc2"]} def retriever_two(data: State) -> State: + time.sleep(0.1) return {"docs": ["doc3", "doc4"]} def qa(data: State) -> State: @@ -6604,10 +6602,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( assert [*app.stream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -6622,10 +6618,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) ] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] @@ -6653,12 +6647,14 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: return {"query": f'query: {data["query"]}'} def analyzer_one(data: State) -> State: + time.sleep(0.1) return {"query": f'analyzed: {data["query"]}'} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} def retriever_two(data: State) -> State: + time.sleep(0.2) return {"docs": ["doc3", "doc4"]} def qa(data: State) -> State: @@ -6693,11 +6689,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: assert [*app.stream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - "qa": {"answer": ""}, - }, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -6712,11 +6706,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) ] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - "qa": {"answer": ""}, - }, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] @@ -6750,6 +6742,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: return {"docs": ["doc1", "doc2"]} def retriever_two(data: State) -> State: + time.sleep(0.1) return {"docs": ["doc3", "doc4"]} def qa(data: State) -> State: @@ -6791,21 +6784,17 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: assert [*app.stream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, { "analyzer_one": { "query": "analyzed: query: analyzed: query: what is weather in sf" - }, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, - { - "retriever_one": {"docs": ["doc1", "doc2"]}, + } }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, ] @@ -6838,6 +6827,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None: return {"docs": ["doc1", "doc2"]} def retriever_two(data: State) -> State: + time.sleep(0.1) return {"docs": ["doc3", "doc4"]} def qa(data: State) -> State: @@ -6878,21 +6868,17 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None: assert [*app.stream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, { "analyzer_one": { "query": "analyzed: query: analyzed: query: what is weather in sf" - }, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, - { - "retriever_one": {"docs": ["doc1", "doc2"]}, + } }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 47192c0c7..944a6bef8 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -99,7 +99,38 @@ async def test_node_cancellation_on_other_node_exception() -> None: graph = builder.compile() with pytest.raises(ValueError, match="I am bad"): - await graph.ainvoke(1) + # This will raise ValueError, not TimeoutError + await asyncio.wait_for(graph.ainvoke(1), 0.5) + + assert inner_task_cancelled + + +async def test_step_timeout_on_stream_hang() -> None: + inner_task_cancelled = False + + async def awhile(input: Any) -> None: + try: + await asyncio.sleep(1.5) + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + + async def alittlewhile(input: Any) -> None: + await asyncio.sleep(0.6) + return "1" + + builder = Graph() + builder.add_node(awhile) + builder.add_node(alittlewhile) + builder.set_conditional_entry_point(lambda _: ["awhile", "alittlewhile"], then=END) + graph = builder.compile() + graph.step_timeout = 1 + + with pytest.raises(asyncio.CancelledError): + async for chunk in graph.astream(1, stream_mode="updates"): + assert chunk == {"alittlewhile": {"alittlewhile": "1"}} + await asyncio.sleep(0.6) assert inner_task_cancelled @@ -492,7 +523,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: assert [ c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="updates") ] == [ - {"one": {"inbox": 3}, "two": {"output": 13}}, + {"one": {"inbox": 3}}, + {"two": {"output": 13}}, {"two": {"output": 4}}, ] assert [c async for c in app.astream({"input": 2, "inbox": 12})] == [ @@ -507,7 +539,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "9379da35-ae1c-5a7b-8556-7ce22a1f8fde", + "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923", "name": "one", "input": 2, "triggers": ["input"], @@ -518,7 +550,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "49ac8f60-4ff2-5cdd-a319-66bbd9837e5a", + "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e", "name": "two", "input": [12], "triggers": ["inbox"], @@ -529,7 +561,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "9379da35-ae1c-5a7b-8556-7ce22a1f8fde", + "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923", "name": "one", "result": [("inbox", 3)], }, @@ -539,7 +571,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "49ac8f60-4ff2-5cdd-a319-66bbd9837e5a", + "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e", "name": "two", "result": [("output", 13)], }, @@ -549,7 +581,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "b97f26c1-a34b-51e0-884e-44a41a3a3b47", + "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62", "name": "two", "input": [3], "triggers": ["inbox"], @@ -560,7 +592,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "b97f26c1-a34b-51e0-884e-44a41a3a3b47", + "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62", "name": "two", "result": [("output", 4)], }, @@ -3099,12 +3131,12 @@ async def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ), @@ -3120,8 +3152,11 @@ async def test_state_graph_packets() -> None: else: return END - def tools_node(tool_call: ToolCall, config: RunnableConfig) -> AgentState: - output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) + async def tools_node(tool_call: ToolCall, config: RunnableConfig) -> AgentState: + await asyncio.sleep(tool_call["args"].get("idx", 0) / 10) + output = await tools_by_name[tool_call["name"]].ainvoke( + tool_call["args"], config + ) return { "messages": ToolMessage( content=output, name=tool_call["name"], tool_call_id=tool_call["id"] @@ -3180,12 +3215,12 @@ async def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ), @@ -3227,16 +3262,14 @@ async def test_state_graph_packets() -> None: }, }, { - "tools": [ - { - "messages": ToolMessage( - content="result for query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ) - } - ] + "tools": { + "messages": ToolMessage( + content="result for query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } }, { "agent": { @@ -3247,36 +3280,36 @@ async def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ) } }, { - "tools": [ - { - "messages": ToolMessage( - content="result for another", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call234", - ) - }, - { - "messages": ToolMessage( - content="result for a third one", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call567", - ), - }, - ] + "tools": { + "messages": ToolMessage( + content="result for another", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call234", + ) + }, + }, + { + "tools": { + "messages": ToolMessage( + content="result for a third one", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call567", + ), + }, }, {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, ] @@ -3410,16 +3443,14 @@ async def test_state_graph_packets() -> None: assert [c async for c in app_w_interrupt.astream(None, config)] == [ { - "tools": [ - { - "messages": ToolMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ) - } - ] + "tools": { + "messages": ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } }, { "agent": { @@ -3430,12 +3461,12 @@ async def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ) @@ -3474,12 +3505,12 @@ async def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ), @@ -3502,12 +3533,12 @@ async def test_state_graph_packets() -> None: { "id": "tool_call234", "name": "search_api", - "args": {"query": "another"}, + "args": {"query": "another", "idx": 0}, }, { "id": "tool_call567", "name": "search_api", - "args": {"query": "a third one"}, + "args": {"query": "a third one", "idx": 1}, }, ], ) @@ -3955,12 +3986,13 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: class State(TypedDict, total=False): query: str answer: str - docs: Annotated[list[str], sorted_add] + docs: Annotated[list[str], operator.add] async def rewrite_query(data: State) -> State: return {"query": f'query: {data["query"]}'} async def retriever_one(data: State) -> State: + await asyncio.sleep(0.1) return {"docs": ["doc1", "doc2"]} async def retriever_two(data: State) -> State: @@ -3993,10 +4025,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "retriever_two": {"docs": ["doc3", "doc4"]}, - "retriever_one": {"docs": ["doc1", "doc2"]}, - }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -4034,7 +4064,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "03dadab4-fb41-5308-a8a4-6eeb9ef7b9aa", + "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf", "name": "rewrite_query", "input": { "query": "what is weather in sf", @@ -4053,7 +4083,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "03dadab4-fb41-5308-a8a4-6eeb9ef7b9aa", + "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf", "name": "rewrite_query", "result": [("query", "query: what is weather in sf")], }, @@ -4067,7 +4097,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "96f499e2-e203-5a13-9259-08cb62f4a2e5", + "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b", "name": "retriever_one", "input": { "query": "query: what is weather in sf", @@ -4085,7 +4115,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "6b344a90-a061-5f17-8714-51f0cf67cf01", + "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2", "name": "retriever_two", "input": { "query": "query: what is weather in sf", @@ -4098,10 +4128,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ), ( "updates", - { - "retriever_one": {"docs": ["doc1", "doc2"]}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, ), ( "debug", @@ -4110,12 +4137,16 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "96f499e2-e203-5a13-9259-08cb62f4a2e5", - "name": "retriever_one", - "result": [("docs", ["doc1", "doc2"])], + "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2", + "name": "retriever_two", + "result": [("docs", ["doc3", "doc4"])], }, }, ), + ( + "updates", + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ), ( "debug", { @@ -4123,9 +4154,9 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "6b344a90-a061-5f17-8714-51f0cf67cf01", - "name": "retriever_two", - "result": [("docs", ["doc3", "doc4"])], + "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b", + "name": "retriever_one", + "result": [("docs", ["doc1", "doc2"])], }, }, ), @@ -4143,7 +4174,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "0dda6269-4ce3-5b98-9cea-d40737a68500", + "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0", "name": "qa", "input": { "query": "query: what is weather in sf", @@ -4162,7 +4193,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "0dda6269-4ce3-5b98-9cea-d40737a68500", + "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0", "name": "qa", "result": [("answer", "doc1,doc2,doc3,doc4")], }, @@ -4436,7 +4467,7 @@ async def test_branch_then() -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "d6e87693-41fb-58f5-8e0d-ee9ab46890b5", + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", "name": "prepare", "input": {"my_key": "value", "market": "DE"}, "triggers": ["start:prepare"], @@ -4447,7 +4478,7 @@ async def test_branch_then() -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "d6e87693-41fb-58f5-8e0d-ee9ab46890b5", + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", "name": "prepare", "result": [("my_key", " prepared")], }, @@ -4481,7 +4512,7 @@ async def test_branch_then() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "b1826010-0028-5aa7-abd2-ed24984614ea", + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", "name": "tool_two_slow", "input": {"my_key": "value prepared", "market": "DE"}, "triggers": ["branch:prepare:condition:tool_two_slow"], @@ -4492,7 +4523,7 @@ async def test_branch_then() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "b1826010-0028-5aa7-abd2-ed24984614ea", + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", "name": "tool_two_slow", "result": [("my_key", " slow")], }, @@ -4529,7 +4560,7 @@ async def test_branch_then() -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "a22dbd2d-f136-57f0-a86a-bc2c234ffcb1", + "id": "ceada3c5-5f25-59e4-9ea5-544599ce1d2f", "name": "finish", "input": {"my_key": "value prepared slow", "market": "DE"}, "triggers": ["branch:prepare:condition:then"], @@ -4540,7 +4571,7 @@ async def test_branch_then() -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "a22dbd2d-f136-57f0-a86a-bc2c234ffcb1", + "id": "ceada3c5-5f25-59e4-9ea5-544599ce1d2f", "name": "finish", "result": [("my_key", " finished")], }, @@ -4852,6 +4883,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: return {"docs": ["doc1", "doc2"]} async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) return {"docs": ["doc3", "doc4"]} async def qa(data: State) -> State: @@ -4882,10 +4914,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -4903,10 +4933,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: ) ] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] @@ -4942,6 +4970,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( return {"docs": ["doc1", "doc2"]} async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) return {"docs": ["doc3", "doc4"]} async def qa(data: State) -> State: @@ -4976,10 +5005,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -4997,10 +5024,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ) ] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] @@ -5043,6 +5068,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( return {"docs": ["doc1", "doc2"]} async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) return {"docs": ["doc3", "doc4"]} async def qa(data: State) -> State: @@ -5084,10 +5110,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -5105,10 +5129,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ) ] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] @@ -5136,12 +5158,14 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: return {"query": f'query: {data["query"]}'} async def analyzer_one(data: State) -> State: + await asyncio.sleep(0.1) return {"query": f'analyzed: {data["query"]}'} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} async def retriever_two(data: State) -> State: + await asyncio.sleep(0.2) return {"docs": ["doc3", "doc4"]} async def qa(data: State) -> State: @@ -5176,11 +5200,9 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - "qa": {"answer": ""}, - }, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] @@ -5198,11 +5220,9 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: ) ] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - "qa": {"answer": ""}, - }, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] @@ -5236,6 +5256,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: return {"docs": ["doc1", "doc2"]} async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) return {"docs": ["doc3", "doc4"]} async def qa(data: State) -> State: @@ -5277,21 +5298,17 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, { "analyzer_one": { "query": "analyzed: query: analyzed: query: what is weather in sf" - }, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, - { - "retriever_one": {"docs": ["doc1", "doc2"]}, + } }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, ] @@ -5324,6 +5341,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> N return {"docs": ["doc1", "doc2"]} async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) return {"docs": ["doc3", "doc4"]} async def qa(data: State) -> State: @@ -5364,21 +5382,17 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> N assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ {"rewrite_query": {"query": "query: what is weather in sf"}}, - { - "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, { "analyzer_one": { "query": "analyzed: query: analyzed: query: what is weather in sf" - }, - "retriever_two": {"docs": ["doc3", "doc4"]}, - }, - { - "retriever_one": {"docs": ["doc1", "doc2"]}, + } }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, ]