From eee91d21292052476988a5bbcf341e38d98f891c Mon Sep 17 00:00:00 2001 From: isaac hershenson Date: Thu, 8 Aug 2024 15:29:40 -0700 Subject: [PATCH 1/5] wip --- libs/langgraph/langgraph/pregel/__init__.py | 7 +++++++ libs/langgraph/langgraph/pregel/debug.py | 9 ++++++--- libs/langgraph/tests/any_int.py | 6 ++++++ libs/langgraph/tests/test_pregel.py | 21 +++++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 21 +++++++++++++++++++++ 5 files changed, 61 insertions(+), 3 deletions(-) create mode 100644 libs/langgraph/tests/any_int.py diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 8f9583964..e8d43d0c9 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -4,6 +4,7 @@ import asyncio import concurrent.futures import time from collections import deque +from datetime import datetime, timezone from functools import partial from inspect import signature from typing import ( @@ -913,6 +914,8 @@ class Pregel( interrupt_after=interrupt_after, manager=run_manager, ): + # start time of superstep + step_start_time = datetime.now(timezone.utc) # debug flag if self.debug: print_step_checkpoint( @@ -987,6 +990,7 @@ class Pregel( loop.step, [task], self.stream_channels_list, + step_start_time, ), ) else: @@ -1157,6 +1161,8 @@ class Pregel( interrupt_after=interrupt_after, manager=run_manager, ): + # start time of superstep + step_start_time = datetime.now(timezone.utc) # debug flag if self.debug: print_step_checkpoint( @@ -1233,6 +1239,7 @@ class Pregel( loop.step, [task], self.stream_channels_list, + step_start_time, ), ): yield chunk diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 9e152cf87..d4123ba2b 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -90,9 +90,10 @@ def map_debug_task_results( step: int, tasks: list[PregelExecutableTask], stream_channels_list: Sequence[str], + step_start_time: datetime, ) -> Iterator[DebugOutputTaskResult]: - ts = datetime.now(timezone.utc).isoformat() - for name, _, _, writes, config, _, _, _ in tasks: + ts = datetime.now(timezone.utc) + for name, _, _, writes, config, _, _, task_id in tasks: if config is not None and TAG_HIDDEN in config.get("tags", []): continue @@ -102,11 +103,13 @@ def map_debug_task_results( yield { "type": "task_result", - "timestamp": ts, + "timestamp": ts.isoformat(), "step": step, "payload": { "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))), + "task_id": task_id, "name": name, + "node_exec_ms": int((ts - step_start_time).total_seconds() * 1000), "result": [w for w in writes if w[0] in stream_channels_list], }, } diff --git a/libs/langgraph/tests/any_int.py b/libs/langgraph/tests/any_int.py new file mode 100644 index 000000000..2fb2dba55 --- /dev/null +++ b/libs/langgraph/tests/any_int.py @@ -0,0 +1,6 @@ +class AnyInt(int): + def __init__(self) -> None: + super().__init__() + + def __eq__(self, other: object) -> bool: + return isinstance(other, int) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 71f2f8c59..388e34507 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -61,6 +61,7 @@ from langgraph.prebuilt.chat_agent_executor import ( from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy +from tests.any_int import AnyInt from tests.any_str import AnyStr from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -1010,6 +1011,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "payload": { "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923", "name": "one", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("inbox", 3)], }, }, @@ -1020,6 +1023,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "payload": { "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e", "name": "two", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("output", 13)], }, }, @@ -1041,6 +1046,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "payload": { "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62", "name": "two", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("output", 4)], }, }, @@ -5881,6 +5888,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf", "name": "rewrite_query", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("query", "query: what is weather in sf")], }, }, @@ -5935,6 +5944,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2", "name": "retriever_two", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("docs", ["doc3", "doc4"])], }, }, @@ -5952,6 +5963,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b", "name": "retriever_one", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("docs", ["doc1", "doc2"])], }, }, @@ -5991,6 +6004,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0", "name": "qa", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("answer", "doc1,doc2,doc3,doc4")], }, }, @@ -6269,6 +6284,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "payload": { "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", "name": "prepare", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("my_key", " prepared")], }, }, @@ -6317,6 +6334,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "payload": { "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", "name": "tool_two_slow", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("my_key", " slow")], }, }, @@ -6365,6 +6384,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "payload": { "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", "name": "finish", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("my_key", " finished")], }, }, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index e5c305dd8..e27cc2b58 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -58,6 +58,7 @@ from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy +from tests.any_int import AnyInt from tests.any_str import AnyStr from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -1117,6 +1118,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "payload": { "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923", "name": "one", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("inbox", 3)], }, }, @@ -1127,6 +1130,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "payload": { "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e", "name": "two", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("output", 13)], }, }, @@ -1148,6 +1153,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "payload": { "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62", "name": "two", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("output", 4)], }, }, @@ -4490,6 +4497,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf", "name": "rewrite_query", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("query", "query: what is weather in sf")], }, }, @@ -4544,6 +4553,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2", "name": "retriever_two", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("docs", ["doc3", "doc4"])], }, }, @@ -4561,6 +4572,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b", "name": "retriever_one", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("docs", ["doc1", "doc2"])], }, }, @@ -4600,6 +4613,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0", "name": "qa", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("answer", "doc1,doc2,doc3,doc4")], }, }, @@ -4904,6 +4919,8 @@ async def test_branch_then() -> None: "payload": { "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", "name": "prepare", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("my_key", " prepared")], }, }, @@ -4952,6 +4969,8 @@ async def test_branch_then() -> None: "payload": { "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", "name": "tool_two_slow", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("my_key", " slow")], }, }, @@ -5000,6 +5019,8 @@ async def test_branch_then() -> None: "payload": { "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", "name": "finish", + "task_id": AnyStr(), + "node_exec_ms": AnyInt(), "result": [("my_key", " finished")], }, }, From 01765fbda45c37cb1382fdbf32ce8b2a9cbc56cf Mon Sep 17 00:00:00 2001 From: isaac hershenson Date: Fri, 30 Aug 2024 17:39:45 -0700 Subject: [PATCH 2/5] x --- libs/langgraph/langgraph/pregel/debug.py | 7 +-- libs/langgraph/langgraph/pregel/loop.py | 2 +- libs/langgraph/tests/test_pregel.py | 58 +++++++++++---------- libs/langgraph/tests/test_pregel_async.py | 61 ++++++++++------------- 4 files changed, 56 insertions(+), 72 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 73f992364..9b20c3c3a 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -104,7 +104,6 @@ def map_debug_task_results( step: int, tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]], stream_keys: Union[str, Sequence[str]], - step_start_time: datetime, ) -> Iterator[DebugOutputTaskResult]: stream_channels_list = ( [stream_keys] if isinstance(stream_keys, str) else stream_keys @@ -123,12 +122,8 @@ def map_debug_task_results( "timestamp": ts.isoformat(), "step": step, "payload": { - "id": str( - uuid5(TASK_NAMESPACE, json.dumps((task.name, step, metadata))) - ), - "task_id": task.task_id, + "id": task.id, "name": task.name, - "node_exec_ms": int((ts - step_start_time).total_seconds() * 1000), "error": next((w[1] for w in writes if w[0] == ERROR), None), "result": [w for w in writes if w[0] in stream_channels_list], "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT], diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index b2eaac472..7c1d09203 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -241,7 +241,7 @@ class PregelLoop: self.stream.extend( (self.config["configurable"].get("checkpoint_ns", ""), "debug", v) for v in map_debug_task_results( - self.step, [(task, writes)], self.stream_keys + self.step, [(task, writes)], self.stream_keys, ) ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 7825af3ea..0a69d9385 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1084,10 +1084,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923", + "id": AnyStr(), "name": "one", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("inbox", 3)], "error": None, "interrupts": [], @@ -1098,10 +1096,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e", + "id": AnyStr(), "name": "two", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("output", 13)], "error": None, "interrupts": [], @@ -1123,10 +1119,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62", + "id": AnyStr(), "name": "two", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("output", 4)], "error": None, "interrupts": [], @@ -6383,9 +6377,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf", + "id": AnyStr(), "name": "rewrite_query", - "task_id": AnyStr(), "node_exec_ms": AnyInt(), "result": [("query", "query: what is weather in sf")], "error": None, @@ -6433,9 +6426,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2", + "id": AnyStr(), "name": "retriever_two", - "task_id": AnyStr(), "node_exec_ms": AnyInt(), "result": [("docs", ["doc3", "doc4"])], "error": None, @@ -6454,9 +6446,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b", + "id": AnyStr(), "name": "retriever_one", - "task_id": AnyStr(), "node_exec_ms": AnyInt(), "result": [("docs", ["doc1", "doc2"])], "error": None, @@ -6496,9 +6487,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0", + "id": AnyStr(), "name": "qa", - "task_id": AnyStr(), "node_exec_ms": AnyInt(), "result": [("answer", "doc1,doc2,doc3,doc4")], "error": None, @@ -6898,16 +6888,19 @@ def test_branch_then( "next": ["prepare"], "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, - { + }, + { "type": "task_result", "timestamp": AnyStr(), "step": 1, "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "id": AnyStr(), "name": "prepare", - "task_id": AnyStr(), "node_exec_ms": AnyInt(), "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, }, { "type": "task", @@ -6925,7 +6918,7 @@ def test_branch_then( "timestamp": AnyStr(), "step": 1, "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "id": AnyStr(), "name": "prepare", "result": [("my_key", " prepared")], "error": None, @@ -6960,17 +6953,19 @@ def test_branch_then( }, "next": ["tool_two_slow"], "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], + }, }, { "type": "task_result", "timestamp": AnyStr(), "step": 2, "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "id": AnyStr(), "name": "tool_two_slow", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("my_key", " slow")], + "error": None, + "interrupts": [], + }, }, { "type": "task", @@ -6988,7 +6983,7 @@ def test_branch_then( "timestamp": AnyStr(), "step": 2, "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "id": AnyStr(), "name": "tool_two_slow", "result": [("my_key", " slow")], "error": None, @@ -7024,16 +7019,19 @@ def test_branch_then( "next": ["finish"], "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, - { + }, + { "type": "task_result", "timestamp": AnyStr(), "step": 3, "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "id": AnyStr(), "name": "finish", - "task_id": AnyStr(), "node_exec_ms": AnyInt(), - "result": [("my_key", " finished")], + "result": [("my_key", " finished")], + "error": None, + "interrupts": [], + }, }, { "type": "task", @@ -7051,7 +7049,7 @@ def test_branch_then( "timestamp": AnyStr(), "step": 3, "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "id": AnyStr(), "name": "finish", "result": [("my_key", " finished")], "error": None, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 1c63c8565..8208e80f5 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1317,10 +1317,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923", + "id": AnyStr(), "name": "one", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("inbox", 3)], "error": None, "interrupts": [], @@ -1331,10 +1329,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 0, "payload": { - "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e", + "id": AnyStr(), "name": "two", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("output", 13)], "error": None, "interrupts": [], @@ -1356,10 +1352,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62", + "id": AnyStr(), "name": "two", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("output", 4)], "error": None, "interrupts": [], @@ -5006,10 +5000,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 1, "payload": { - "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf", + "id": AnyStr(), "name": "rewrite_query", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("query", "query: what is weather in sf")], "error": None, "interrupts": [], @@ -5056,10 +5048,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2", + "id": AnyStr(), "name": "retriever_two", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("docs", ["doc3", "doc4"])], "error": None, "interrupts": [], @@ -5077,10 +5067,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 2, "payload": { - "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b", + "id": AnyStr(), "name": "retriever_one", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("docs", ["doc1", "doc2"])], "error": None, "interrupts": [], @@ -5119,10 +5107,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "timestamp": AnyStr(), "step": 3, "payload": { - "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0", + "id": AnyStr(), "name": "qa", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("answer", "doc1,doc2,doc3,doc4")], "error": None, "interrupts": [], @@ -5447,17 +5433,19 @@ async def test_branch_then( }, "next": ["prepare"], "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + }, }, { "type": "task_result", "timestamp": AnyStr(), "step": 1, "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "id": AnyStr(), "name": "prepare", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), - "result": [("my_key", " prepared")], + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, }, { "type": "task", @@ -5475,7 +5463,7 @@ async def test_branch_then( "timestamp": AnyStr(), "step": 1, "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "id": AnyStr(), "name": "prepare", "result": [("my_key", " prepared")], "error": None, @@ -5510,17 +5498,19 @@ async def test_branch_then( }, "next": ["tool_two_slow"], "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], + }, }, { "type": "task_result", "timestamp": AnyStr(), "step": 2, "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "id": AnyStr(), "name": "tool_two_slow", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("my_key", " slow")], + "error": None, + "interrupts": [], + }, }, { "type": "task", @@ -5538,7 +5528,7 @@ async def test_branch_then( "timestamp": AnyStr(), "step": 2, "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "id": AnyStr(), "name": "tool_two_slow", "result": [("my_key", " slow")], "error": None, @@ -5573,17 +5563,18 @@ async def test_branch_then( }, "next": ["finish"], "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], + }, }, { "type": "task_result", "timestamp": AnyStr(), "step": 3, "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "id": AnyStr(), "name": "finish", - "task_id": AnyStr(), - "node_exec_ms": AnyInt(), "result": [("my_key", " finished")], + "error": None, + "interrupts": [], }, }, { @@ -5602,7 +5593,7 @@ async def test_branch_then( "timestamp": AnyStr(), "step": 3, "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "id": AnyStr(), "name": "finish", "result": [("my_key", " finished")], "error": None, @@ -5730,7 +5721,7 @@ async def test_branch_then( "timestamp": AnyStr(), "step": 1, "payload": { - "id": "1a591be4-f85c-558f-8d00-1ccac0d1877f", + "id": AnyStr(), "name": "prepare", "result": [("my_key", " prepared")], "error": None, From d14477c891271392ffb788b2390306409250c70a Mon Sep 17 00:00:00 2001 From: isaac hershenson Date: Fri, 30 Aug 2024 18:22:22 -0700 Subject: [PATCH 3/5] x --- libs/langgraph/tests/test_pregel.py | 57 +++-------------------- libs/langgraph/tests/test_pregel_async.py | 37 --------------- 2 files changed, 7 insertions(+), 87 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0a69d9385..74a16b010 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -69,7 +69,6 @@ from langgraph.pregel import ( StateSnapshot, ) from langgraph.pregel.retry import RetryPolicy -from tests.any_int import AnyInt from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence @@ -6379,7 +6378,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": AnyStr(), "name": "rewrite_query", - "node_exec_ms": AnyInt(), "result": [("query", "query: what is weather in sf")], "error": None, "interrupts": [], @@ -6428,7 +6426,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": AnyStr(), "name": "retriever_two", - "node_exec_ms": AnyInt(), "result": [("docs", ["doc3", "doc4"])], "error": None, "interrupts": [], @@ -6448,7 +6445,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": AnyStr(), "name": "retriever_one", - "node_exec_ms": AnyInt(), "result": [("docs", ["doc1", "doc2"])], "error": None, "interrupts": [], @@ -6489,7 +6485,6 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "payload": { "id": AnyStr(), "name": "qa", - "node_exec_ms": AnyInt(), "result": [("answer", "doc1,doc2,doc3,doc4")], "error": None, "interrupts": [], @@ -6880,27 +6875,14 @@ def test_branch_then( "market": "DE", }, "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "node_exec_ms": AnyInt(), - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + }, }, { "type": "task", @@ -6955,18 +6937,6 @@ def test_branch_then( "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], }, }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], - }, - }, { "type": "task", "timestamp": AnyStr(), @@ -7020,19 +6990,6 @@ def test_branch_then( "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "node_exec_ms": AnyInt(), - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], - }, - }, { "type": "task", "timestamp": AnyStr(), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 8208e80f5..70d90bb62 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -68,7 +68,6 @@ from langgraph.pregel import ( StateSnapshot, ) from langgraph.pregel.retry import RetryPolicy -from tests.any_int import AnyInt from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence @@ -5435,18 +5434,6 @@ async def test_branch_then( "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, { "type": "task", "timestamp": AnyStr(), @@ -5500,18 +5487,6 @@ async def test_branch_then( "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], }, }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], - }, - }, { "type": "task", "timestamp": AnyStr(), @@ -5565,18 +5540,6 @@ async def test_branch_then( "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], - }, - }, { "type": "task", "timestamp": AnyStr(), From ee5801511b7f5a0b099305914e9b4a5a8fca518f Mon Sep 17 00:00:00 2001 From: isaac hershenson Date: Fri, 30 Aug 2024 18:25:02 -0700 Subject: [PATCH 4/5] fmt --- libs/langgraph/langgraph/pregel/__init__.py | 5 ----- libs/langgraph/langgraph/pregel/loop.py | 4 +++- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 7cd946430..01f0ad91d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -4,7 +4,6 @@ import asyncio import concurrent.futures import time from collections import deque -from datetime import datetime, timezone from functools import partial from typing import ( Any, @@ -1232,8 +1231,6 @@ class Pregel( interrupt_after=interrupt_after, manager=run_manager, ): - # start time of superstep - step_start_time = datetime.now(timezone.utc) # debug flag if debug: print_step_checkpoint( @@ -1485,8 +1482,6 @@ class Pregel( interrupt_after=interrupt_after, manager=run_manager, ): - # start time of superstep - step_start_time = datetime.now(timezone.utc) # debug flag if debug: print_step_checkpoint( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7c1d09203..3e55d0df8 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -241,7 +241,9 @@ class PregelLoop: self.stream.extend( (self.config["configurable"].get("checkpoint_ns", ""), "debug", v) for v in map_debug_task_results( - self.step, [(task, writes)], self.stream_keys, + self.step, + [(task, writes)], + self.stream_keys, ) ) From ec51760f5715affcf997a1f464f97de1f1079066 Mon Sep 17 00:00:00 2001 From: isaac hershenson Date: Tue, 3 Sep 2024 13:45:22 -0700 Subject: [PATCH 5/5] x --- libs/langgraph/tests/test_pregel_async.py | 88 +++++++++++------------ 1 file changed, 44 insertions(+), 44 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 3da3b0886..d8747bd0a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5491,17 +5491,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "triggers": ["start:prepare"], }, }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, }, { "type": "checkpoint", @@ -5546,17 +5546,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "triggers": ["branch:prepare:condition:tool_two_slow"], }, }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "tool_two_slow", + "result": [("my_key", " slow")], + "error": None, + "interrupts": [], + }, }, { "type": "checkpoint", @@ -5599,17 +5599,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "triggers": ["branch:prepare:condition::then"], }, }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "finish", + "result": [("my_key", " finished")], + "error": None, + "interrupts": [], + }, }, { "type": "checkpoint", @@ -5728,17 +5728,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "triggers": ["start:prepare"], }, }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, }, { "type": "checkpoint",