diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 7bbce9641..97b13ed30 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -40,7 +40,7 @@ MIGRATIONS = [ channel TEXT NOT NULL, version TEXT NOT NULL, type TEXT NOT NULL, - blob BYTEA NOT NULL, + blob BYTEA, PRIMARY KEY (thread_id, checkpoint_ns, channel, version) );""", """CREATE TABLE IF NOT EXISTS checkpoint_writes ( @@ -54,6 +54,7 @@ MIGRATIONS = [ blob BYTEA NOT NULL, PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) );""", + "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;", ] SELECT_SQL = """ @@ -140,6 +141,7 @@ class BasePostgresSaver(BaseCheckpointSaver): return { k.decode(): self.serde.loads_typed((t.decode(), v)) for k, t, v in blob_values + if t.decode() != "empty" } def _dump_blobs( @@ -162,10 +164,13 @@ class BasePostgresSaver(BaseCheckpointSaver): checkpoint_ns, k, ver, - *self.serde.dumps_typed(values[k]), + *( + self.serde.dumps_typed(values[k]) + if k in values + else ("empty", None) + ), ) for k, ver in versions.items() - if k in values ] def _load_writes( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 55e6ad976..728effc0e 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -36,6 +36,7 @@ from syrupy import SnapshotAssertion from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context +from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue @@ -9378,3 +9379,29 @@ def test_xray_lance(snapshot: SnapshotAssertion): # View assert graph.get_graph().to_json() == snapshot assert graph.get_graph(xray=1).to_json() == snapshot + + +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_channel_values(request: pytest.FixtureRequest, checkpointer_name: str) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + config = {"configurable": {"thread_id": "1"}} + chain = Channel.subscribe_to("input") | Channel.write_to("output") + app = Pregel( + nodes={ + "one": chain, + }, + channels={ + "ephemeral": EphemeralValue(Any), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels=["input", "ephemeral"], + output_channels="output", + checkpointer=checkpointer, + ) + app.invoke({"input": 1, "ephemeral": "meow"}, config) + assert checkpointer.get(config)["channel_values"] == {"input": 1, "output": 1}