From eee91d21292052476988a5bbcf341e38d98f891c Mon Sep 17 00:00:00 2001 From: isaac hershenson Date: Thu, 8 Aug 2024 15:29:40 -0700 Subject: [PATCH] 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")], }, },