From 01765fbda45c37cb1382fdbf32ce8b2a9cbc56cf Mon Sep 17 00:00:00 2001 From: isaac hershenson Date: Fri, 30 Aug 2024 17:39:45 -0700 Subject: [PATCH] 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,