Merge pull request #1902 from langchain-ai/dqbd/debug-parent-config

fix(debug): send parent_config in debug checkpoint events
This commit is contained in:
Nuno Campos
2024-09-29 18:30:24 -07:00
committed by GitHub
4 changed files with 131 additions and 0 deletions
+3
View File
@@ -52,6 +52,7 @@ class CheckpointPayload(TypedDict):
metadata: CheckpointMetadata
values: dict[str, Any]
next: list[str]
parent_config: Optional[RunnableConfig]
tasks: list[CheckpointTask]
@@ -136,6 +137,7 @@ def map_debug_checkpoint(
checkpoint: Checkpoint,
tasks: Iterable[PregelExecutableTask],
pending_writes: list[PendingWrite],
parent_config: Optional[RunnableConfig],
) -> Iterator[DebugOutputCheckpoint]:
"""Produce "checkpoint" events for stream_mode=debug."""
yield {
@@ -144,6 +146,7 @@ def map_debug_checkpoint(
"step": step,
"payload": {
"config": config,
"parent_config": parent_config,
"values": read_channels(channels, stream_channels),
"metadata": metadata,
"next": [t.name for t in tasks],
+15
View File
@@ -176,6 +176,7 @@ class PregelLoop:
checkpoint_metadata: CheckpointMetadata
checkpoint_pending_writes: List[PendingWrite]
checkpoint_previous_versions: dict[str, Union[str, float, int]]
prev_checkpoint_config: Optional[RunnableConfig]
step: int
stop: int
@@ -249,6 +250,13 @@ class PregelLoop:
if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
else ()
)
self.prev_checkpoint_config = (
self.checkpoint_config
if self.checkpoint_config
and CONF in self.checkpoint_config
and CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF]
else None
)
def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
"""Put writes for a task, to be read by the next tick."""
@@ -386,6 +394,7 @@ class PregelLoop:
self.checkpoint,
self.tasks.values(),
self.checkpoint_pending_writes,
self.prev_checkpoint_config,
)
# if no more tasks, we're done
@@ -535,6 +544,12 @@ class PregelLoop:
self.checkpoint = create_checkpoint(self.checkpoint, self.channels, self.step)
# bail if no checkpointer
if self._checkpointer_put_after_previous is not None:
self.prev_checkpoint_config = (
self.checkpoint_config
if CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF]
and self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID]
else None
)
self.checkpoint_metadata = metadata
self.checkpoint_config = {
**self.checkpoint_config,
+45
View File
@@ -7019,6 +7019,7 @@ def test_branch_then(
"step": -1,
"writes": {"__start__": {"my_key": "value", "market": "DE"}},
},
"parent_config": None,
"next": ["__start__"],
"tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}],
},
@@ -7049,6 +7050,17 @@ def test_branch_then(
"step": 0,
"writes": None,
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "10"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "10",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": ["prepare"],
"tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}],
},
@@ -7102,6 +7114,17 @@ def test_branch_then(
"step": 1,
"writes": {"prepare": {"my_key": " prepared"}},
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "10"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "10",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": ["tool_two_slow"],
"tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}],
},
@@ -7155,6 +7178,17 @@ def test_branch_then(
"step": 2,
"writes": {"tool_two_slow": {"my_key": " slow"}},
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "10"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "10",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": ["finish"],
"tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}],
},
@@ -7208,6 +7242,17 @@ def test_branch_then(
"step": 3,
"writes": {"finish": {"my_key": " finished"}},
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "10"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "10",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": [],
"tasks": [],
},
+68
View File
@@ -5651,6 +5651,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
"step": -1,
"writes": {"__start__": {"my_key": "value", "market": "DE"}},
},
"parent_config": None,
"next": ["__start__"],
"tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}],
},
@@ -5681,6 +5682,17 @@ async def test_branch_then(checkpointer_name: str) -> None:
"step": 0,
"writes": None,
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "10"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "10",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": ["prepare"],
"tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}],
},
@@ -5734,6 +5746,17 @@ async def test_branch_then(checkpointer_name: str) -> None:
"step": 1,
"writes": {"prepare": {"my_key": " prepared"}},
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "10"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "10",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": ["tool_two_slow"],
"tasks": [
{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}
@@ -5789,6 +5812,17 @@ async def test_branch_then(checkpointer_name: str) -> None:
"step": 2,
"writes": {"tool_two_slow": {"my_key": " slow"}},
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "10"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "10",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": ["finish"],
"tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}],
},
@@ -5842,6 +5876,17 @@ async def test_branch_then(checkpointer_name: str) -> None:
"step": 3,
"writes": {"finish": {"my_key": " finished"}},
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "10"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "10",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": [],
"tasks": [],
},
@@ -5888,6 +5933,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
"step": -1,
"writes": {"__start__": {"my_key": "value", "market": "DE"}},
},
"parent_config": None,
"next": ["__start__"],
"tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}],
},
@@ -5918,6 +5964,17 @@ async def test_branch_then(checkpointer_name: str) -> None:
"step": 0,
"writes": None,
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "11"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "11",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": ["prepare"],
"tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}],
},
@@ -5971,6 +6028,17 @@ async def test_branch_then(checkpointer_name: str) -> None:
"step": 1,
"writes": {"prepare": {"my_key": " prepared"}},
},
"parent_config": {
"tags": [],
"metadata": {"thread_id": "11"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "11",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"next": ["tool_two_slow"],
"tasks": [
{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}