mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-24 16:42:24 +02:00
218 lines
8.2 KiB
Python
218 lines
8.2 KiB
Python
"""Verify that SEND values land in checkpoint_blobs (TASKS channel),
|
|
not just in checkpoint_writes.
|
|
|
|
Strategy:
|
|
1. Build a tiny graph: node "1" emits a Send to node "2" via
|
|
conditional edges. interrupt_before=["2"] freezes the loop after
|
|
super-step 1 completes (cp_1 persisted) but before "2" runs.
|
|
2. Run graph.invoke against PostgresSaver.
|
|
3. Read checkpoint_blobs raw via psycopg.
|
|
4. Find the row with channel = '__pregel_tasks' and decode the blob
|
|
using the saver's serde. Assert it contains a non-empty list of
|
|
Send objects.
|
|
5. As a control, also dump checkpoint_writes to confirm the *same*
|
|
value is independently present in the writes table.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import operator
|
|
from typing import Annotated
|
|
from typing_extensions import TypedDict
|
|
from uuid import uuid4
|
|
|
|
from psycopg import Connection
|
|
from psycopg.rows import dict_row
|
|
|
|
from langgraph.checkpoint.postgres import PostgresSaver
|
|
from langgraph.checkpoint.serde.types import TASKS
|
|
from langgraph.graph import START, StateGraph
|
|
from langgraph.types import Send
|
|
|
|
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5441/"
|
|
|
|
|
|
class State(TypedDict):
|
|
history: Annotated[list[str], operator.add]
|
|
|
|
|
|
def make_graph(checkpointer: PostgresSaver, *, with_interrupt: bool = True):
|
|
def node_one(state: State) -> State:
|
|
return {"history": ["1"]}
|
|
|
|
def node_two(state: dict) -> State:
|
|
return {"history": [f"2:{state['payload']}"]}
|
|
|
|
def fanout(state: State):
|
|
return [Send("two", {"payload": x}) for x in ("a", "b", "c")]
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("one", node_one)
|
|
builder.add_node("two", node_two)
|
|
builder.add_edge(START, "one")
|
|
builder.add_conditional_edges("one", fanout, ["two"])
|
|
return builder.compile(
|
|
checkpointer=checkpointer,
|
|
interrupt_before=["two"] if with_interrupt else [],
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
database = f"verify_{uuid4().hex[:12]}"
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
uri = DEFAULT_POSTGRES_URI + database
|
|
with Connection.connect(
|
|
uri, autocommit=True, prepare_threshold=0, row_factory=dict_row
|
|
) as conn:
|
|
saver = PostgresSaver(conn)
|
|
saver.setup()
|
|
|
|
graph = make_graph(saver)
|
|
thread_id = "t1"
|
|
cfg = {"configurable": {"thread_id": thread_id}}
|
|
result = graph.invoke({"history": []}, cfg)
|
|
print("=== invoke result (interrupted before 'two') ===")
|
|
print(result)
|
|
print()
|
|
|
|
# Inspect raw checkpoint_blobs.
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT thread_id, checkpoint_ns, channel, version, type,
|
|
octet_length(blob) AS nbytes, blob
|
|
FROM checkpoint_blobs
|
|
WHERE thread_id = %s
|
|
ORDER BY channel, version
|
|
""",
|
|
(thread_id,),
|
|
).fetchall()
|
|
|
|
print(f"=== checkpoint_blobs ({len(rows)} rows) ===")
|
|
for r in rows:
|
|
preview = (
|
|
"(empty)"
|
|
if r["blob"] is None
|
|
else f"{r['blob'][:60]!r}{'...' if len(r['blob']) > 60 else ''}"
|
|
)
|
|
print(
|
|
f" channel={r['channel']!r:30s} v={r['version']!s:18s} "
|
|
f"type={r['type']:8s} nbytes={r['nbytes']!s:>5} blob={preview}"
|
|
)
|
|
print()
|
|
|
|
tasks_rows = [r for r in rows if r["channel"] == TASKS]
|
|
assert tasks_rows, (
|
|
f"expected at least one row in checkpoint_blobs for {TASKS!r}, "
|
|
f"got channels={[r['channel'] for r in rows]}"
|
|
)
|
|
|
|
# Decode the latest TASKS blob using the saver's serde.
|
|
tasks_row = tasks_rows[-1]
|
|
decoded = saver.serde.loads_typed((tasks_row["type"], tasks_row["blob"]))
|
|
print(f"=== decoded {TASKS!r} blob ({tasks_row['type']}) ===")
|
|
print(f" type(decoded) = {type(decoded).__name__}")
|
|
print(f" len(decoded) = {len(decoded)}")
|
|
for i, v in enumerate(decoded):
|
|
print(f" [{i}] type={type(v).__name__} value={v!r}")
|
|
print()
|
|
|
|
assert isinstance(decoded, list), (
|
|
f"expected list of Sends, got {type(decoded).__name__}"
|
|
)
|
|
assert len(decoded) == 3, f"expected 3 Sends (a, b, c), got {len(decoded)}"
|
|
for v in decoded:
|
|
assert isinstance(v, Send), (
|
|
f"expected Send instance in blob, got {type(v).__name__}"
|
|
)
|
|
assert v.node == "two"
|
|
print("✅ ASSERTION PASSED: checkpoint_blobs contains the Sends.")
|
|
print()
|
|
|
|
# Control: confirm the same Sends are also in checkpoint_writes.
|
|
wrows = conn.execute(
|
|
"""
|
|
SELECT task_id, idx, channel, type, octet_length(blob) AS nbytes, blob
|
|
FROM checkpoint_writes
|
|
WHERE thread_id = %s AND channel = %s
|
|
ORDER BY task_id, idx
|
|
""",
|
|
(thread_id, TASKS),
|
|
).fetchall()
|
|
print(
|
|
f"=== checkpoint_writes (channel={TASKS!r}, {len(wrows)} rows) ==="
|
|
)
|
|
for r in wrows:
|
|
v = saver.serde.loads_typed((r["type"], r["blob"]))
|
|
print(
|
|
f" task={r['task_id']} idx={r['idx']} type={r['type']} "
|
|
f"nbytes={r['nbytes']} -> {v!r}"
|
|
)
|
|
|
|
assert len(wrows) == 3, (
|
|
f"expected 3 writes-table rows for {TASKS!r}, got {len(wrows)}"
|
|
)
|
|
print(
|
|
"✅ ASSERTION PASSED: checkpoint_writes also contains the Sends."
|
|
)
|
|
print()
|
|
print(
|
|
"Conclusion (cp_1): SEND is double-stored. blobs row carries "
|
|
"the authoritative Send list; writes rows carry the same "
|
|
"payload as a redundant per-task record."
|
|
)
|
|
|
|
# ----------------------------------------------------------------
|
|
# Now resume: let 'two' actually run, so the Sends get consumed.
|
|
# We expect the next checkpoint's TASKS blob to be cleared,
|
|
# proving the accumulate=False lifecycle.
|
|
# ----------------------------------------------------------------
|
|
graph2 = make_graph(saver, with_interrupt=False)
|
|
result2 = graph2.invoke(None, cfg)
|
|
print()
|
|
print("=== invoke result (resumed; 'two' fanout completed) ===")
|
|
print(result2)
|
|
print()
|
|
|
|
rows2 = conn.execute(
|
|
"""
|
|
SELECT channel, version, type, octet_length(blob) AS nbytes, blob
|
|
FROM checkpoint_blobs
|
|
WHERE thread_id = %s AND channel = %s
|
|
ORDER BY version
|
|
""",
|
|
(thread_id, TASKS),
|
|
).fetchall()
|
|
print(f"=== checkpoint_blobs rows for {TASKS!r} across all cps ===")
|
|
for r in rows2:
|
|
decoded_v = saver.serde.loads_typed((r["type"], r["blob"]))
|
|
print(
|
|
f" v={r['version']} type={r['type']} nbytes={r['nbytes']} "
|
|
f"-> {decoded_v!r}"
|
|
)
|
|
|
|
assert len(rows2) >= 2, (
|
|
"expected at least 2 versions of TASKS blob across the run"
|
|
)
|
|
latest = saver.serde.loads_typed(
|
|
(rows2[-1]["type"], rows2[-1]["blob"])
|
|
)
|
|
assert latest == [], (
|
|
f"expected latest TASKS blob to be cleared (==[]) "
|
|
f"after Sends were consumed, got {latest!r}"
|
|
)
|
|
print()
|
|
print(
|
|
"✅ ASSERTION PASSED: after Sends were consumed, the latest "
|
|
"TASKS blob is [] (Topic with accumulate=False clears itself "
|
|
"via update(EMPTY_SEQ) in apply_writes line 304-311)."
|
|
)
|
|
finally:
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|