This commit is contained in:
Nuno Campos
2025-04-15 16:52:52 -07:00
parent 48c08421fa
commit 5cca153b72
2 changed files with 12 additions and 13 deletions
+7 -7
View File
@@ -286,7 +286,7 @@ class PregelLoop(LoopProtocol):
)
else:
self.checkpoint_config = self.config
if thread_id := config[CONF].get(CONFIG_KEY_THREAD_ID):
if thread_id := self.checkpoint_config[CONF].get(CONFIG_KEY_THREAD_ID):
if not isinstance(thread_id, str):
self.checkpoint_config = patch_configurable(
self.checkpoint_config,
@@ -1050,16 +1050,16 @@ class SyncPregelLoop(PregelLoop, ContextManager):
saved = None
if saved is None:
saved = CheckpointTuple(
self.config, empty_checkpoint(), {"step": -2}, None, []
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
)
elif self._migrate_checkpoint is not None:
self._migrate_checkpoint(saved.checkpoint)
self.checkpoint_config = {
**self.config,
**self.checkpoint_config,
**saved.config,
CONF: {
CONFIG_KEY_CHECKPOINT_NS: "",
**self.config.get(CONF, {}),
**self.checkpoint_config.get(CONF, {}),
**saved.config.get(CONF, {}),
},
}
@@ -1200,16 +1200,16 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
saved = None
if saved is None:
saved = CheckpointTuple(
self.config, empty_checkpoint(), {"step": -2}, None, []
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
)
elif self._migrate_checkpoint is not None:
self._migrate_checkpoint(saved.checkpoint)
self.checkpoint_config = {
**self.config,
**self.checkpoint_config,
**saved.config,
CONF: {
CONFIG_KEY_CHECKPOINT_NS: "",
**self.config.get(CONF, {}),
**self.checkpoint_config.get(CONF, {}),
**saved.config.get(CONF, {}),
},
}
+5 -6
View File
@@ -6169,7 +6169,7 @@ def test_falsy_return_from_task(
"checkpoint_ns": "",
"thread_id": AnyStr(),
},
"metadata": configurable["configurable"],
"metadata": {},
"recursion_limit": 25,
"tags": [],
},
@@ -6177,7 +6177,6 @@ def test_falsy_return_from_task(
"parents": {},
"source": "input",
"step": -1,
"thread_id": AnyStr(),
"writes": {
"__start__": {
"a": 5,
@@ -6268,6 +6267,7 @@ def test_falsy_return_from_task(
"type": "task_result",
},
]
print(type(configurable["configurable"]["thread_id"]))
assert [
c
for c in graph.stream(Command(resume="123"), configurable, stream_mode="debug")
@@ -6281,7 +6281,7 @@ def test_falsy_return_from_task(
"checkpoint_ns": "",
"thread_id": AnyStr(),
},
"metadata": configurable["configurable"],
"metadata": {},
"recursion_limit": 25,
"tags": [],
},
@@ -6376,7 +6376,7 @@ def test_falsy_return_from_task(
"checkpoint_ns": "",
"thread_id": AnyStr(),
},
"metadata": configurable["configurable"],
"metadata": {},
"recursion_limit": 25,
"tags": [],
},
@@ -6384,7 +6384,6 @@ def test_falsy_return_from_task(
"parents": {},
"source": "loop",
"step": 0,
"thread_id": AnyStr(),
"writes": {
"falsy_task": False,
"graph": None,
@@ -6398,7 +6397,7 @@ def test_falsy_return_from_task(
"checkpoint_ns": "",
"thread_id": AnyStr(),
},
"metadata": configurable["configurable"],
"metadata": {},
"recursion_limit": 25,
"tags": [],
},