Update managed value usage in local_read (#4854)

This commit is contained in:
Nuno Campos
2025-05-28 14:05:14 -07:00
committed by GitHub
2 changed files with 44 additions and 19 deletions
@@ -72,6 +72,7 @@ from langgraph.errors import (
from langgraph.managed.base import ManagedValueSpec
from langgraph.pregel.algo import (
PregelTaskWrites,
_scratchpad,
apply_writes,
local_read,
prepare_next_tasks,
@@ -1782,6 +1783,15 @@ class Pregel(PregelProtocol):
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: partial(
local_read,
_scratchpad(
None,
[],
task_id,
"",
None,
step,
step + 2,
),
channels,
managed,
task,
@@ -2188,6 +2198,15 @@ class Pregel(PregelProtocol):
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: partial(
local_read,
_scratchpad(
None,
[],
task_id,
"",
None,
step,
step + 2,
),
channels,
managed,
task,
+25 -19
View File
@@ -172,6 +172,7 @@ def should_interrupt(
def local_read(
scratchpad: PregelScratchpad,
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
task: WritesProtocol,
@@ -208,7 +209,7 @@ def local_read(
else:
values = read_channels(channels, select)
if managed_keys:
values.update({k: managed[k]() for k in managed_keys})
values.update({k: managed[k].get(scratchpad) for k in managed_keys})
return values
@@ -571,6 +572,15 @@ def prepare_single_task(
)
else:
cache_key = None
scratchpad = _scratchpad(
config[CONF].get(CONFIG_KEY_SCRATCHPAD),
pending_writes,
task_id,
xxh3_128_hexdigest(task_checkpoint_ns.encode()),
config[CONF].get(CONFIG_KEY_RESUME_MAP),
step,
stop,
)
return PregelExecutableTask(
name,
call.input,
@@ -587,6 +597,7 @@ def prepare_single_task(
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: partial(
local_read,
scratchpad,
channels,
managed,
PregelTaskWrites(task_path, name, writes, triggers),
@@ -601,15 +612,7 @@ def prepare_single_task(
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_SCRATCHPAD: _scratchpad(
config[CONF].get(CONFIG_KEY_SCRATCHPAD),
pending_writes,
task_id,
xxh3_128_hexdigest(task_checkpoint_ns.encode()),
config[CONF].get(CONFIG_KEY_RESUME_MAP),
step,
stop,
),
CONFIG_KEY_SCRATCHPAD: scratchpad,
},
),
triggers,
@@ -695,6 +698,15 @@ def prepare_single_task(
)
else:
cache_key = None
scratchpad = _scratchpad(
config[CONF].get(CONFIG_KEY_SCRATCHPAD),
pending_writes,
task_id,
xxh3_128_hexdigest(task_checkpoint_ns.encode()),
config[CONF].get(CONFIG_KEY_RESUME_MAP),
step,
stop,
)
return PregelExecutableTask(
packet.node,
packet.arg,
@@ -712,6 +724,7 @@ def prepare_single_task(
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: partial(
local_read,
scratchpad,
channels,
managed,
PregelTaskWrites(task_path, packet.node, writes, triggers),
@@ -726,15 +739,7 @@ def prepare_single_task(
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_SCRATCHPAD: _scratchpad(
config[CONF].get(CONFIG_KEY_SCRATCHPAD),
pending_writes,
task_id,
xxh3_128_hexdigest(task_checkpoint_ns.encode()),
config[CONF].get(CONFIG_KEY_RESUME_MAP),
step,
stop,
),
CONFIG_KEY_SCRATCHPAD: scratchpad,
CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get(
PREVIOUS, None
),
@@ -860,6 +865,7 @@ def prepare_single_task(
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: partial(
local_read,
scratchpad,
channels,
managed,
PregelTaskWrites(