Add feature flag (default off) so we can merge this before releasing

- Add additional ci job to test with FF on
This commit is contained in:
Nuno Campos
2024-11-11 15:38:48 -08:00
parent ea64ac5c07
commit d0567dc7be
8 changed files with 191 additions and 52 deletions
+6
View File
@@ -19,9 +19,13 @@ jobs:
- "3.13"
core-version:
- "latest"
ff-send-v2:
- "false"
include:
- python-version: "3.11"
core-version: ">=0.2.42,<0.3.0"
- python-version: "3.11"
ff-send-v2: "true"
defaults:
run:
@@ -52,6 +56,8 @@ jobs:
- name: Run tests
shell: bash
env:
LANGGRAPH_FF_SEND_V2: ${{ matrix.ff-send-v2 }}
run: |
make test
+1 -1
View File
@@ -49,7 +49,7 @@ test:
exit $$EXIT_CODE
test_watch:
make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \
make start-postgres && poetry run ptw . -- --ff -vv -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \
EXIT_CODE=$$?; \
make stop-postgres; \
exit $$EXIT_CODE
+3
View File
@@ -1,4 +1,5 @@
import sys
from os import getenv
from types import MappingProxyType
from typing import Any, Literal, Mapping, cast
@@ -81,6 +82,8 @@ NS_END = sys.intern(":")
# for checkpoint_ns, for each level, separates the namespace from the task_id
CONF = cast(Literal["configurable"], sys.intern("configurable"))
# key for the configurable dict in RunnableConfig
FF_SEND_V2 = getenv("LANGGRAPH_FF_SEND_V2", "false").lower() == "true"
# temporary flag to enable new Send semantics
RESERVED = {
TAG_HIDDEN,
+1 -1
View File
@@ -173,7 +173,7 @@ def local_write(
"""Function injected under CONFIG_KEY_SEND in task config, to write to channels.
Validates writes and forwards them to `commit` function."""
for chan, value in writes:
if chan == PUSH:
if chan in (PUSH, TASKS):
if not isinstance(value, Send):
raise InvalidUpdateError(f"Expected Send, got {value}")
if value.node not in process_keys:
+6 -2
View File
@@ -14,7 +14,7 @@ from typing import (
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONF, CONFIG_KEY_SEND, PUSH, TASKS, Send
from langgraph.constants import CONF, CONFIG_KEY_SEND, FF_SEND_V2, PUSH, TASKS, Send
from langgraph.errors import InvalidUpdateError
from langgraph.utils.runnable import RunnableCallable
@@ -119,7 +119,11 @@ class ChannelWrite(RunnableCallable):
if w.value is PASSTHROUGH:
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
# split packets and entries
sends = [(PUSH, packet) for packet in writes if isinstance(packet, Send)]
sends = [
(PUSH if FF_SEND_V2 else TASKS, packet)
for packet in writes
if isinstance(packet, Send)
]
entries = [write for write in writes if isinstance(write, ChannelWriteEntry)]
# process entries into values
values = [
-1
View File
@@ -327,7 +327,6 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
ALL_CHECKPOINTERS_SYNC = [
"memory",
"sqlite",
"duckdb",
"postgres",
"postgres_pipe",
"postgres_pool",
+87 -25
View File
@@ -54,7 +54,14 @@ from langgraph.checkpoint.base import (
CheckpointTuple,
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
from langgraph.constants import (
CONFIG_KEY_NODE_FINISHED,
ERROR,
FF_SEND_V2,
PULL,
PUSH,
START,
)
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
from langgraph.graph import END, Graph, GraphCommand, StateGraph
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
@@ -1781,17 +1788,31 @@ def test_concurrent_emit_sends() -> None:
builder.add_conditional_edges("1.1", send_for_profit)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert graph.invoke(["0"]) == [
"0",
"1",
"1.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
"3.1",
]
assert graph.invoke(["0"]) == (
[
"0",
"1",
"1.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"1.1",
"3.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
]
)
def test_send_sequences() -> None:
@@ -1831,16 +1852,31 @@ def test_send_sequences() -> None:
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert graph.invoke(["0"]) == [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
"3.1",
]
assert (
graph.invoke(["0"])
== [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"3.1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
)
@pytest.mark.repeat(20)
@@ -1848,8 +1884,8 @@ def test_send_sequences() -> None:
def test_send_dedupe_on_resume(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
if checkpointer_name == "duckdb":
pytest.skip("DuckDB isn't returning the right history")
if not FF_SEND_V2:
pytest.skip("Send deduplication is only available in Send V2")
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class InterruptOnce:
@@ -2298,6 +2334,9 @@ def test_send_react_interrupt(
}
assert foo_called == 0
if not FF_SEND_V2:
return
# get state should show the pending task
state = graph.get_state(thread1)
assert state == StateSnapshot(
@@ -2741,6 +2780,9 @@ def test_send_react_interrupt_control(
}
assert foo_called == 1
if not FF_SEND_V2:
return
# interrupt-update-resume flow
foo_called = 0
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"])
@@ -5866,6 +5908,9 @@ def test_state_graph_packets(
{"__interrupt__": ()},
]
if not FF_SEND_V2:
return
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
@@ -12264,6 +12309,21 @@ def test_send_to_nested_graphs(
# check state
outer_state = graph.get_state(config)
if not FF_SEND_V2:
# update state of dogs joke graph
graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"})
# continue past interrupt
assert sorted(
graph.stream(None, config=config),
key=lambda d: d["generate_joke"]["jokes"][0],
) == [
{"generate_joke": {"jokes": ["Joke about cats - hohoho"]}},
{"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}},
]
return
assert outer_state == StateSnapshot(
values={"subjects": ["cats", "dogs"], "jokes": []},
tasks=(
@@ -12409,7 +12469,9 @@ def test_send_to_nested_graphs(
tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),),
)
# update state of dogs joke graph
graph.update_state(outer_state.tasks[2].state, {"subject": "turtles - hohoho"})
graph.update_state(
outer_state.tasks[2 if FF_SEND_V2 else 1].state, {"subject": "turtles - hohoho"}
)
# continue past interrupt
assert sorted(
+87 -22
View File
@@ -51,7 +51,14 @@ from langgraph.checkpoint.base import (
CheckpointTuple,
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
from langgraph.constants import (
CONFIG_KEY_NODE_FINISHED,
ERROR,
FF_SEND_V2,
PULL,
PUSH,
START,
)
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
from langgraph.graph import END, Graph, GraphCommand, StateGraph
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
@@ -2022,17 +2029,31 @@ async def test_concurrent_emit_sends() -> None:
builder.add_conditional_edges("1.1", send_for_profit)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert await graph.ainvoke(["0"]) == [
"0",
"1",
"1.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
"3.1",
]
assert await graph.ainvoke(["0"]) == (
[
"0",
"1",
"1.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"1.1",
"3.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
]
)
@pytest.mark.repeat(10)
@@ -2073,16 +2094,34 @@ async def test_send_sequences(checkpointer_name: str) -> None:
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert await graph.ainvoke(["0"]) == [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
"3.1",
]
assert (
await graph.ainvoke(["0"])
== [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"3.1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
)
if not FF_SEND_V2:
return
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["3.1"])
@@ -2110,6 +2149,9 @@ async def test_send_sequences(checkpointer_name: str) -> None:
@pytest.mark.repeat(20)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
if not FF_SEND_V2:
pytest.skip("Send deduplication is only available in Send V2")
class InterruptOnce:
ticks: int = 0
@@ -2542,6 +2584,9 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
assert foo_called == 0
if not FF_SEND_V2:
return
# get state should show the pending task
state = await graph.aget_state(thread1)
assert state == StateSnapshot(
@@ -3004,6 +3049,9 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None:
}
assert foo_called == 0
if not FF_SEND_V2:
return
# get state should show the pending task
state = await graph.aget_state(thread1)
assert state == StateSnapshot(
@@ -5826,6 +5874,9 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
{"__interrupt__": ()},
]
if not FF_SEND_V2:
return
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"messages": [
@@ -11029,6 +11080,20 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
# check state
outer_state = await graph.aget_state(config)
if not FF_SEND_V2:
# update state of dogs joke graph
await graph.aupdate_state(
outer_state.tasks[1].state, {"subject": "turtles - hohoho"}
)
# continue past interrupt
assert await graph.ainvoke(None, config=config) == {
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
}
return
assert outer_state == StateSnapshot(
values={"subjects": ["cats", "dogs"], "jokes": []},
tasks=(