This commit is contained in:
isaac hershenson
2024-08-08 15:29:40 -07:00
parent dce73fde66
commit eee91d2129
5 changed files with 61 additions and 3 deletions
@@ -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
+6 -3
View File
@@ -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],
},
}
+6
View File
@@ -0,0 +1,6 @@
class AnyInt(int):
def __init__(self) -> None:
super().__init__()
def __eq__(self, other: object) -> bool:
return isinstance(other, int)
+21
View File
@@ -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")],
},
},
+21
View File
@@ -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")],
},
},