Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

lib: Make Control object serializable #2347

Merged
merged 1 commit into from
Nov 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import dataclasses

Check notice on line 1 in libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py

View workflow job for this annotation

GitHub Actions / benchmark

Benchmark results

......................................... fanout_to_subgraph_10x: Mean +- std dev: 58.0 ms +- 1.2 ms ......................................... fanout_to_subgraph_10x_sync: Mean +- std dev: 49.0 ms +- 0.5 ms ......................................... fanout_to_subgraph_10x_checkpoint: Mean +- std dev: 82.9 ms +- 1.7 ms ......................................... fanout_to_subgraph_10x_checkpoint_sync: Mean +- std dev: 90.5 ms +- 0.9 ms ......................................... fanout_to_subgraph_100x: Mean +- std dev: 535 ms +- 13 ms ......................................... fanout_to_subgraph_100x_sync: Mean +- std dev: 475 ms +- 4 ms ......................................... fanout_to_subgraph_100x_checkpoint: Mean +- std dev: 874 ms +- 50 ms ......................................... fanout_to_subgraph_100x_checkpoint_sync: Mean +- std dev: 891 ms +- 17 ms ......................................... react_agent_10x: Mean +- std dev: 30.0 ms +- 0.5 ms ......................................... react_agent_10x_sync: Mean +- std dev: 22.1 ms +- 0.2 ms ......................................... react_agent_10x_checkpoint: Mean +- std dev: 45.9 ms +- 1.0 ms ......................................... react_agent_10x_checkpoint_sync: Mean +- std dev: 36.0 ms +- 0.5 ms ......................................... react_agent_100x: Mean +- std dev: 335 ms +- 6 ms ......................................... react_agent_100x_sync: Mean +- std dev: 269 ms +- 3 ms ......................................... react_agent_100x_checkpoint: Mean +- std dev: 938 ms +- 13 ms ......................................... react_agent_100x_checkpoint_sync: Mean +- std dev: 836 ms +- 9 ms ......................................... wide_state_25x300: Mean +- std dev: 23.5 ms +- 0.4 ms ......................................... wide_state_25x300_sync: Mean +- std dev: 15.0 ms +- 0.1 ms ......................................... wide_state_25x300_checkpoint: Mean +- std dev: 279 ms +- 3 ms ......................................... wide_state_25x300_checkpoint_sync: Mean +- std dev: 267 ms +- 3 ms ......................................... wide_state_15x600: Mean +- std dev: 27.5 ms +- 0.4 ms ......................................... wide_state_15x600_sync: Mean +- std dev: 17.3 ms +- 0.1 ms ......................................... wide_state_15x600_checkpoint: Mean +- std dev: 480 ms +- 3 ms ......................................... wide_state_15x600_checkpoint_sync: Mean +- std dev: 466 ms +- 7 ms ......................................... wide_state_9x1200: Mean +- std dev: 27.5 ms +- 0.4 ms ......................................... wide_state_9x1200_sync: Mean +- std dev: 17.3 ms +- 0.2 ms ......................................... wide_state_9x1200_checkpoint: Mean +- std dev: 313 ms +- 3 ms ......................................... wide_state_9x1200_checkpoint_sync: Mean +- std dev: 299 ms +- 4 ms

Check notice on line 1 in libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py

View workflow job for this annotation

GitHub Actions / benchmark

Comparison against main

+-----------------------------------+---------+-----------------------+ | Benchmark | main | changes | +===================================+=========+=======================+ | react_agent_100x_checkpoint_sync | 846 ms | 836 ms: 1.01x faster | +-----------------------------------+---------+-----------------------+ | react_agent_10x_checkpoint | 46.3 ms | 45.9 ms: 1.01x faster | +-----------------------------------+---------+-----------------------+ | react_agent_100x | 337 ms | 335 ms: 1.01x faster | +-----------------------------------+---------+-----------------------+ | wide_state_9x1200_checkpoint_sync | 300 ms | 299 ms: 1.00x faster | +-----------------------------------+---------+-----------------------+ | fanout_to_subgraph_100x_sync | 477 ms | 475 ms: 1.00x faster | +-----------------------------------+---------+-----------------------+ | react_agent_100x_sync | 269 ms | 269 ms: 1.00x faster | +-----------------------------------+---------+-----------------------+ | wide_state_15x600_checkpoint | 478 ms | 480 ms: 1.00x slower | +-----------------------------------+---------+-----------------------+ | fanout_to_subgraph_10x_sync | 48.9 ms | 49.0 ms: 1.00x slower | +-----------------------------------+---------+-----------------------+ | wide_state_15x600_checkpoint_sync | 464 ms | 466 ms: 1.01x slower | +-----------------------------------+---------+-----------------------+ | react_agent_10x_checkpoint_sync | 35.8 ms | 36.0 ms: 1.01x slower | +-----------------------------------+---------+-----------------------+ | react_agent_100x_checkpoint | 929 ms | 938 ms: 1.01x slower | +-----------------------------------+---------+-----------------------+ | wide_state_25x300 | 23.3 ms | 23.5 ms: 1.01x slower | +-----------------------------------+---------+-----------------------+ | wide_state_15x600 | 27.2 ms | 27.5 ms: 1.01x slower | +-----------------------------------+---------+-----------------------+ | wide_state_9x1200 | 27.1 ms | 27.5 ms: 1.01x slower | +-----------------------------------+---------+-----------------------+ | wide_state_25x300_sync | 14.8 ms | 15.0 ms: 1.02x slower | +-----------------------------------+---------+-----------------------+ | wide_state_9x1200_sync | 17.0 ms | 17.3 ms: 1.02x slower | +-----------------------------------+---------+-----------------------+ | wide_state_15x600_sync | 17.0 ms | 17.3 ms: 1.02x slower | +-----------------------------------+---------+-----------------------+ | Geometric mean | (ref) | 1.00x slower | +-----------------------------------+---------+-----------------------+ Benchmark hidden because not significant (11): fanout_to_subgraph_10x_checkpoint, react_agent_10x, fanout_to_subgraph_10x_checkpoint_sync, fanout_to_subgraph_100x_checkpoint_sync, fanout_to_subgraph_10x, react_agent_10x_sync, wide_state_9x1200_checkpoint, wide_state_25x300_checkpoint, fanout_to_subgraph_100x, wide_state_25x300_checkpoint_sync, fanout_to_subgraph_100x_checkpoint
import decimal
import importlib
import json
Expand All @@ -25,7 +25,7 @@
from zoneinfo import ZoneInfo

from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.types import SendProtocol
from langgraph.checkpoint.serde.types import ControlProtocol, SendProtocol
from langgraph.store.base import Item

LC_REVIVER = Reviver()
Expand Down Expand Up @@ -402,6 +402,21 @@
(obj.__class__.__module__, obj.__class__.__name__, (obj.node, obj.arg)),
),
)
elif isinstance(obj, ControlProtocol):
return msgpack.ExtType(
EXT_CONSTRUCTOR_KW_ARGS,
_msgpack_enc(
(
obj.__class__.__module__,
obj.__class__.__name__,
{
"update_state": obj.update_state,
"trigger": obj.trigger,
"send": obj.send,
},
),
),
)
elif dataclasses.is_dataclass(obj):
# doesn't use dataclasses.asdict to avoid deepcopy and recursion
return msgpack.ExtType(
Expand Down
11 changes: 11 additions & 0 deletions libs/checkpoint/langgraph/checkpoint/serde/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
Protocol,
Sequence,
TypeVar,
Union,
runtime_checkable,
)

Expand Down Expand Up @@ -48,3 +49,13 @@ def __hash__(self) -> int: ...
def __repr__(self) -> str: ...

def __eq__(self, value: object) -> bool: ...


@runtime_checkable
class ControlProtocol(Protocol):
# Mirrors langgraph.constants.Control
update_state: Optional[dict[str, Any]]
trigger: Union[str, Sequence[str]]
send: Union[Any, Sequence[Any]]

def __repr__(self) -> str: ...
8 changes: 7 additions & 1 deletion libs/langgraph/langgraph/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,10 @@ def __eq__(self, value: object) -> bool:


class Control(Generic[N]):
"""A control object to update the graph's state, trigger nodes, and send messages."""

__slots__ = ("update_state", "trigger", "send")

def __init__(
self,
*,
Expand All @@ -240,7 +244,9 @@ def __init__(

def __repr__(self) -> str:
contents = ", ".join(
f"{key}={value!r}" for key, value in self.__dict__.items() if value
f"{key}={value!r}"
for key in self.__slots__
if (value := getattr(self, key))
)
return f"Control({contents})"

Expand Down
Loading