WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-24 00:22:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Pregel.stream passing all tests

Browse Source
This commit is contained in:
Nuno Campos
2024-07-18 15:49:58 -07:00
parent 60f7a7d593
commit b3ee728839
7 changed files with 272 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+34 -54
View File
@@ -849,38 +849,34 @@ class Pregel(
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
while loop.tick(
interrupt_before=interrupt_before, interrupt_after=interrupt_after
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
):
# print debug output
# debug flag
if self.debug:
print_step_checkpoint(
loop.checkpoint_metadata,
loop.channels,
self.stream_channels_list,
)
# emit debug output
if self.checkpointer and "debug" in stream_modes:
yield from _with_mode(
"debug",
isinstance(stream_mode, list),
map_debug_checkpoint(
loop.checkpoint_metadata["step"],
loop.config,
loop.channels,
self.stream_channels_asis,
loop.checkpoint_metadata,
),
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# debug flag
if debug:
print_step_tasks(loop.checkpoint_metadata["step"], loop.tasks)
print_step_tasks(loop.step, loop.tasks)
# TODO move to tick() ?
if "debug" in stream_modes:
yield from _with_mode(
"debug",
isinstance(stream_mode, list),
map_debug_tasks(
loop.checkpoint_metadata["step"], loop.tasks
),
map_debug_tasks(loop.step, loop.tasks),
)
# execute tasks, and wait for one to fail or all to finish.
@@ -931,7 +927,7 @@ class Pregel(
"debug",
isinstance(stream_mode, list),
map_debug_task_results(
loop.checkpoint_metadata["step"],
loop.step,
[task],
self.stream_channels_list,
),
@@ -941,38 +937,31 @@ class Pregel(
del fut, task
# panic on failure or timeout
_panic_or_proceed(done, inflight, loop.checkpoint_metadata["step"])
_panic_or_proceed(done, inflight, loop.step)
# don't keep futures around in memory longer than needed
del done, inflight, futures
# combine pending writes from all tasks
pending_writes = deque[tuple[str, Any]]()
for task in loop.tasks:
pending_writes.extend(task.writes)
# debug flag
if debug:
print_step_writes(
loop.checkpoint_metadata["step"],
pending_writes,
loop.step,
[w for t in loop.tasks for w in t.writes],
self.stream_channels_list,
)
# yield values output
if "values" in stream_modes:
yield from _with_mode(
"values",
isinstance(stream_mode, list),
map_output_values(
output_keys, pending_writes, loop.channels
),
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached"
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# set final channel values as run output
run_manager.on_chain_end(read_channels(loop.channels, output_keys))
except BaseException as e:
@@ -1142,6 +1131,7 @@ class Pregel(
)
def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]:
print(metadata)
nonlocal checkpoint, checkpoint_config, channels
if self.checkpointer is None:
@@ -1268,12 +1258,7 @@ class Pregel(
break
# before execution, check if we should interrupt
if should_interrupt(
checkpoint,
interrupt_before,
self.stream_channels_list,
next_tasks,
):
if should_interrupt(checkpoint, interrupt_before, next_tasks):
break
else:
checkpoint = next_checkpoint
@@ -1405,12 +1390,7 @@ class Pregel(
yield chunk
# after execution, check if we should interrupt
if should_interrupt(
checkpoint,
interrupt_after,
self.stream_channels_list,
next_tasks,
):
if should_interrupt(checkpoint, interrupt_after, next_tasks):
break
else:
raise GraphRecursionError(
libs/langgraph/langgraph/pregel/algo.py
-1
View File
@@ -44,7 +44,6 @@ from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescript
def should_interrupt(
checkpoint: Checkpoint,
interrupt_nodes: Union[All, Sequence[str]],
snapshot_channels: Sequence[str],
tasks: list[PregelExecutableTask],
) -> bool:
version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
libs/langgraph/langgraph/pregel/debug.py
+4
View File
@@ -70,6 +70,8 @@ def map_debug_tasks(
if config is not None and TAG_HIDDEN in config.get("tags", []):
continue
print("map_debug_tasks", json.dumps((name, step, config["metadata"])))
yield {
"type": "task",
"timestamp": ts,
@@ -95,6 +97,8 @@ def map_debug_task_results(
if config is not None and TAG_HIDDEN in config.get("tags", []):
continue
print("map_debug_tasks_r", json.dumps((name, step, config["metadata"])))
yield {
"type": "task_result",
"timestamp": ts,
libs/langgraph/langgraph/pregel/loop.py
+91 -58
View File
@@ -4,13 +4,17 @@ from types import TracebackType
from typing import (
TYPE_CHECKING,
Any,
Callable,
ContextManager,
List,
Literal,
Mapping,
Optional,
Sequence,
Tuple,
Type,
TypeVar,
Union,
)
from langchain_core.runnables import RunnableConfig
@@ -35,31 +39,39 @@ from langgraph.pregel.algo import (
prepare_next_tasks,
should_interrupt,
)
from langgraph.pregel.debug import map_debug_checkpoint
from langgraph.pregel.executor import BackgroundExecutor, Submit
from langgraph.pregel.io import map_input
from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
from langgraph.pregel.types import PregelExecutableTask
if TYPE_CHECKING:
from langgraph.pregel import Pregel
V = TypeVar("V")
INPUT_DONE = object()
class PregelLoop(ContextManager):
input: Optional[Any]
config: RunnableConfig
checkpoint: Checkpoint
checkpoint_metadata: CheckpointMetadata
checkpoint_pending_writes: Optional[List[PendingWrite]]
checkpointer: Optional[BaseCheckpointSaver]
get_next_version: Callable[[Optional[V]], V]
graph: "Pregel"
submit: Submit
channels: Mapping[str, BaseChannel]
managed: ManagedValueMapping
checkpoint: Checkpoint
checkpoint_config: RunnableConfig
checkpoint_metadata: CheckpointMetadata
checkpoint_pending_writes: Optional[List[PendingWrite]]
status: Literal[
"pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
]
tasks: Sequence[PregelExecutableTask]
stream: deque[Tuple[str, Any]]
def __init__(
self,
@@ -69,6 +81,7 @@ class PregelLoop(ContextManager):
checkpointer: Optional[BaseCheckpointSaver],
graph: "Pregel",
) -> None:
self.stream = deque()
self.stack = ExitStack()
self.input = input
self.config = config
@@ -84,7 +97,7 @@ class PregelLoop(ContextManager):
saved = (
self.checkpointer.get_tuple(self.config) if self.checkpointer else None
) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
self.config = {
self.checkpoint_config = {
**self.config,
**saved.config,
"configurable": {
@@ -106,6 +119,7 @@ class PregelLoop(ContextManager):
)
)
self.status = "pending"
self.step = self.checkpoint_metadata["step"] + 1
return self
@@ -121,10 +135,11 @@ class PregelLoop(ContextManager):
def tick(
self,
*,
output_keys: Union[str, Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Sequence[str]] = None,
) -> bool:
print("ticking", self.status, self.checkpoint_metadata["step"])
print("tick", self.status, self.step)
if self.status != "pending":
raise RuntimeError(f"Cannot tick when status is {self.status}")
if self.input is not INPUT_DONE:
@@ -132,42 +147,57 @@ class PregelLoop(ContextManager):
elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len(
self.tasks
):
# assign writes to tasks, apply them in order
grouped: dict[str, list[tuple[str, Any]]] = {}
for tid, k, v in self.checkpoint_pending_writes:
grouped.setdefault(tid, []).append((k, v))
writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])]
# all tasks have finished
apply_writes(
self.checkpoint,
self.channels,
((k, v) for _, k, v in self.checkpoint_pending_writes),
writes,
self.get_next_version,
)
# produce values output
self.stream.extend(
("values", v)
for v in map_output_values(output_keys, writes, self.channels)
)
# clear pending writes
self.checkpoint_pending_writes.clear()
# save checkpoint
self.put_checkpoint({"source": "loop", "writes": None}) # TODO
self.put_checkpoint(
{
"source": "loop",
"writes": single(
map_output_updates(output_keys, self.tasks)
if self.graph.stream_mode == "updates"
else map_output_values(output_keys, writes, self.channels)
),
}
)
# after execution, check if we should interrupt
if should_interrupt(
self.checkpoint,
interrupt_after,
self.graph.stream_channels_list,
self.tasks,
):
if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
self.status = "interrupt_after"
return False
else:
return False
# check if iteration limit is reached
if self.checkpoint_metadata["step"] >= self.config["recursion_limit"]:
if self.step > self.config["recursion_limit"]:
self.status = "out_of_steps"
return False
# prepare next tasks
prev_checkpoint = self.checkpoint
self.checkpoint, self.tasks = prepare_next_tasks(
self.checkpoint,
self.graph.nodes,
self.channels,
self.managed,
self.config,
self.checkpoint_metadata["step"],
self.step,
for_execution=True,
get_next_version=self.get_next_version,
)
@@ -180,21 +210,14 @@ class PregelLoop(ContextManager):
# TODO how to make this work for both
# - online case: we should schedule remaining tasks
# - offline case: we should just bail, as other tasks were scheduled before
# assign pending writes to tasks
# if self.checkpoint_pending_writes:
# # if there are pending writes from a previous loop, apply them
# for tid, k, v in self.checkpoint_pending_writes:
# if task := next((t for t in self.tasks if t.id == tid), None):
# task.writes.append((k, v))
# if
# if there are pending writes from a previous loop, apply them
if self.checkpoint_pending_writes:
for tid, k, v in self.checkpoint_pending_writes:
if task := next((t for t in self.tasks if t.id == tid), None):
task.writes.append((k, v))
# before execution, check if we should interrupt
if should_interrupt(
self.checkpoint,
interrupt_before,
self.graph.stream_channels_list,
self.tasks,
):
if should_interrupt(prev_checkpoint, interrupt_before, self.tasks):
self.status = "interrupt_before"
return False
@@ -210,9 +233,10 @@ class PregelLoop(ContextManager):
self.channels,
self.managed,
self.config,
-1,
self.step,
for_execution=True,
get_next_version=self.get_next_version,
# TODO missing run_manager
)
# apply input writes
apply_writes(
@@ -239,9 +263,9 @@ class PregelLoop(ContextManager):
self.submit(
self.checkpointer.put_writes,
{
**self.config,
**self.checkpoint_config,
"configurable": {
**self.config["configurable"],
**self.checkpoint_config["configurable"],
"thread_ts": self.checkpoint["id"],
},
},
@@ -253,30 +277,39 @@ class PregelLoop(ContextManager):
self,
metadata: CheckpointMetadata,
) -> None:
# increment step
self.checkpoint_metadata = {
**metadata,
"step": self.checkpoint_metadata["step"] + 1,
}
# assign step
metadata["step"] = self.step
# bail if no checkpointer
if self.checkpointer is None:
return
# create new checkpoint
self.checkpoint = create_checkpoint(
self.checkpoint, self.channels, self.checkpoint_metadata["step"]
)
# save it, without blocking
self.submit(
self.checkpointer.put,
self.config,
copy_checkpoint(self.checkpoint),
self.checkpoint_metadata,
)
# update checkpoint config
self.config = {
**self.config,
"configurable": {
**self.config["configurable"],
"thread_ts": self.checkpoint["id"],
},
}
if self.checkpointer is not None:
# create new checkpoint
self.checkpoint_metadata = metadata
self.checkpoint = create_checkpoint(
self.checkpoint, self.channels, self.step
)
# save it, without blocking
self.submit(
self.checkpointer.put,
self.checkpoint_config,
copy_checkpoint(self.checkpoint),
self.checkpoint_metadata,
)
self.checkpoint_config = {
**self.checkpoint_config,
"configurable": {
**self.checkpoint_config["configurable"],
"thread_ts": self.checkpoint["id"],
},
}
# produce debug output
self.stream.extend(
("debug", v)
for v in map_debug_checkpoint(
self.step,
self.checkpoint_config,
self.channels,
self.graph.stream_channels_asis,
self.checkpoint_metadata,
)
)
# increment step
self.step += 1
libs/langgraph/tests/test_prebuilt.py
+115 -3
View File
@@ -1,5 +1,8 @@
from collections import defaultdict
from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.sqlite import SqliteSaver
import pytest
from langchain_core.callbacks import (
CallbackManagerForLLMRun,
@@ -27,6 +30,8 @@ from langgraph.prebuilt import (
ValidationNode,
create_react_agent,
)
from tests.any_str import AnyStr
from tests.memory_assert import MemorySaverAssertImmutable
class FakeToolCallingModel(BaseChatModel):
@@ -56,14 +61,121 @@ class FakeToolCallingModel(BaseChatModel):
return self
def test_no_modifier():
@pytest.mark.parametrize(
"checkpointer",
[
MemorySaverAssertImmutable(),
None,
],
ids=[
"memory",
"none",
],
)
def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]):
model = FakeToolCallingModel()
agent = create_react_agent(model, [])
agent = create_react_agent(model, [], checkpointer=checkpointer)
inputs = [HumanMessage("hi?")]
response = agent.invoke({"messages": inputs}, debug=True)
thread = {"configurable": {"thread_id": "123"}}
response = agent.invoke({"messages": inputs}, thread, debug=True)
expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
assert response == expected_response
if checkpointer:
saved = checkpointer.get_tuple(thread)
assert saved is not None
assert saved.checkpoint == {
"v": 1,
"ts": AnyStr(),
"id": AnyStr(),
"channel_values": {
"messages": [
HumanMessage(content="hi?", id=AnyStr()),
AIMessage(content="hi?", id="0"),
],
"agent": "agent",
},
"channel_versions": {
"__start__": 2,
"messages": 3,
"start:agent": 3,
"agent": 3,
},
"versions_seen": defaultdict(
dict,
{
"__start__": {"__start__": 1},
"agent": {"start:agent": 2},
"tools": {},
},
),
"pending_sends": [],
}
assert saved.metadata == {
"source": "loop",
"writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
"step": 1,
}
assert saved.pending_writes == []
@pytest.mark.parametrize(
"checkpointer",
[
MemorySaverAssertImmutable(),
None,
],
ids=[
"memory",
"none",
],
)
async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]):
model = FakeToolCallingModel()
agent = create_react_agent(model, [], checkpointer=checkpointer)
inputs = [HumanMessage("hi?")]
thread = {"configurable": {"thread_id": "123"}}
response = await agent.ainvoke({"messages": inputs}, thread, debug=True)
expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
assert response == expected_response
if checkpointer:
saved = await checkpointer.aget_tuple(thread)
assert saved is not None
assert saved.checkpoint == {
"v": 1,
"ts": AnyStr(),
"id": AnyStr(),
"channel_values": {
"messages": [
HumanMessage(content="hi?", id=AnyStr()),
AIMessage(content="hi?", id="0"),
],
"agent": "agent",
},
"channel_versions": {
"__start__": 2,
"messages": 3,
"start:agent": 3,
"agent": 3,
},
"versions_seen": defaultdict(
dict,
{
"__start__": {"__start__": 1},
"agent": {"start:agent": 2},
"tools": {},
},
),
"pending_sends": [],
}
assert saved.metadata == {
"source": "loop",
"writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
"step": 1,
}
assert saved.pending_writes == []
def test_passing_two_modifiers():
model = FakeToolCallingModel()
libs/langgraph/tests/test_pregel.py
+16 -3
View File
@@ -521,7 +521,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert app.invoke(2) == 4
with pytest.raises(GraphRecursionError):
app.invoke(2, {"recursion_limit": 1})
app.invoke(2, {"recursion_limit": 1}, debug=1)
graph = Graph()
graph.add_node("add_one", add_one)
@@ -533,7 +533,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert gapp.invoke(2) == 4
for step, values in enumerate(gapp.stream(2), start=1):
for step, values in enumerate(gapp.stream(2, debug=1), start=1):
if step == 1:
assert values == {
"add_one": 3,
@@ -6166,6 +6166,18 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
"my_key": "value ⛰️",
"market": "DE",
}
assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [
{
"source": "loop",
"step": 0,
"writes": None,
},
{
"source": "input",
"step": -1,
"writes": {"my_key": "value ⛰️", "market": "DE"},
},
]
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_two_slow",),
@@ -6875,7 +6887,8 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
},
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
print("yo")
assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}},
]
libs/langgraph/tests/test_pregel_async.py
+12
View File
@@ -4735,6 +4735,18 @@ async def test_start_branch_then() -> None:
"my_key": "value",
"market": "DE",
}
assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
{
"source": "loop",
"step": 0,
"writes": None,
},
{
"source": "input",
"step": -1,
"writes": {"my_key": "value", "market": "DE"},
},
]
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value", "market": "DE"},
next=("tool_two_slow",),