diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 36e375751..5d18b5262 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -20,6 +20,8 @@ START = sys.intern("__start__") """The first (maybe virtual) node in graph-style Pregel.""" END = sys.intern("__end__") """The last (maybe virtual) node in graph-style Pregel.""" +SELF = sys.intern("__self__") +"""The implicit branch that handles each node's Control values.""" # --- Reserved write keys --- INPUT = sys.intern("__input__") diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index e3d11c1eb..afe6604dd 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -27,6 +27,7 @@ from typing_extensions import Self from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.constants import ( + EMPTY_SEQ, END, NS_END, NS_SEP, @@ -47,6 +48,7 @@ logger = logging.getLogger(__name__) class NodeSpec(NamedTuple): runnable: Runnable metadata: Optional[dict[str, Any]] = None + ends: Optional[tuple[str, ...]] = EMPTY_SEQ class Branch(NamedTuple): @@ -123,7 +125,7 @@ class Branch(NamedTuple): result: Any, config: RunnableConfig, ) -> Union[Runnable, Any]: - if not isinstance(result, list): + if not isinstance(result, (list, tuple)): result = [result] if self.ends: destinations: Sequence[Union[Send, str]] = [ @@ -364,6 +366,9 @@ class Graph: for node in self.nodes: if node != start and node != branch.then: all_sources.add(node) + for name, spec in self.nodes.items(): + if spec.ends: + all_sources.add(name) # validate sources for source in all_sources: if source not in self.nodes and source != START: @@ -387,6 +392,9 @@ class Graph: for node in self.nodes: if node != start and node != branch.then: all_targets.add(node) + for name, spec in self.nodes.items(): + if spec.ends: + all_targets.update(spec.ends) # validate targets for node in self.nodes: if node not in all_targets: @@ -620,5 +628,9 @@ class CompiledGraph(Pregel): ) if branch.then is not None: add_edge(end, branch.then) + for key, n in self.builder.nodes.items(): + if n.ends: + for end in n.ends: + add_edge(key, end, conditional=True) return graph diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index ad9a3edab..b64e54b21 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -14,6 +14,7 @@ from typing import ( Type, Union, cast, + get_args, get_origin, get_type_hints, overload, @@ -32,7 +33,7 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.named_barrier_value import NamedBarrierValue -from langgraph.constants import NS_END, NS_SEP, TAG_HIDDEN +from langgraph.constants import EMPTY_SEQ, NS_END, NS_SEP, SELF, TAG_HIDDEN from langgraph.errors import ErrorCode, InvalidUpdateError, create_error_message from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send from langgraph.managed.base import ( @@ -46,10 +47,10 @@ from langgraph.managed.base import ( from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import All, Checkpointer, RetryPolicy +from langgraph.types import All, Checkpointer, Control, RetryPolicy from langgraph.utils.fields import get_field_default from langgraph.utils.pydantic import create_model -from langgraph.utils.runnable import coerce_to_runnable +from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable logger = logging.getLogger(__name__) @@ -71,6 +72,7 @@ class StateNodeSpec(NamedTuple): metadata: Optional[dict[str, Any]] input: Type[Any] retry_policy: Optional[RetryPolicy] + ends: Optional[tuple[str, ...]] = EMPTY_SEQ class StateGraph(Graph): @@ -338,8 +340,29 @@ class StateGraph(Graph): f"'{character}' is a reserved character and is not allowed in the node names." ) - if input is None: - input = _get_input_schema_from_type_hint(action) + ends = EMPTY_SEQ + try: + if isfunction(action) and ( + hints := get_type_hints(getattr(action, "__call__")) + or get_type_hints(action) + ): + if input is None: + first_parameter_name = next( + iter(inspect.signature(action).parameters.keys()) + ) + if input_hint := hints.get(first_parameter_name): + if isinstance(input_hint, type) and get_type_hints(input_hint): + input = input_hint + if ( + (rtn := hints.get("return")) + and get_origin(rtn) is Control + and (rargs := get_args(rtn)) + and get_origin(rargs[0]) is Literal + and (vals := get_args(rargs[0])) + ): + ends = vals + except (TypeError, StopIteration): + pass if input is not None: self._add_schema(input) self.nodes[cast(str, node)] = StateNodeSpec( @@ -347,6 +370,7 @@ class StateGraph(Graph): metadata, input=input or self.schema, retry_policy=retry, + ends=ends, ) return self @@ -468,6 +492,9 @@ class StateGraph(Graph): for key, node in self.nodes.items(): compiled.attach_node(key, node) + for key, node in self.nodes.items(): + compiled.attach_branch(key, SELF, CONTROL_BRANCH, with_reader=False) + for start, end in self.edges: compiled.attach_edge(start, end) @@ -518,11 +545,23 @@ class CompiledStateGraph(CompiledGraph): if is_writable_managed_value(v) ] + def _get_root(input: Any) -> Any: + if isinstance(input, Control): + return input.update_state + else: + return input + def _get_state_key(input: Union[None, dict, Any], *, key: str) -> Any: if input is None: return SKIP_WRITE elif isinstance(input, dict): + if all(k not in output_keys for k in input): + raise InvalidUpdateError( + f"Expected node {key} to update at least one of {output_keys}, got {input}" + ) return input.get(key, SKIP_WRITE) + elif isinstance(input, Control): + return _get_state_key(input.update_state, key=key) elif get_type_hints(type(input)): value = getattr(input, key, SKIP_WRITE) return value if value is not None else SKIP_WRITE @@ -535,7 +574,7 @@ class CompiledStateGraph(CompiledGraph): # state updaters write_entries = ( - [ChannelWriteEntry("__root__", skip_none=True)] + [ChannelWriteEntry("__root__", skip_none=True, mapper=_get_root)] if output_keys == ["__root__"] else [ ChannelWriteEntry(key, mapper=partial(_get_state_key, key=key)) @@ -578,7 +617,6 @@ class CompiledStateGraph(CompiledGraph): ChannelWrite( [ChannelWriteEntry(key, key)] + write_entries, tags=[TAG_HIDDEN], - require_at_least_one_of=output_keys, ), ], metadata=node.metadata, @@ -615,7 +653,9 @@ class CompiledStateGraph(CompiledGraph): [ChannelWriteEntry(channel_name, start)], tags=[TAG_HIDDEN] ) - def attach_branch(self, start: str, name: str, branch: Branch) -> None: + def attach_branch( + self, start: str, name: str, branch: Branch, *, with_reader: bool = True + ) -> None: def branch_writer( packets: Sequence[Union[str, Send]], config: RunnableConfig ) -> None: @@ -648,7 +688,8 @@ class CompiledStateGraph(CompiledGraph): else self.builder.schema ) self.nodes[start] |= branch.run( - branch_writer, _get_state_reader(self.builder, schema) + branch_writer, + _get_state_reader(self.builder, schema) if with_reader else None, ) # attach branch subscribers @@ -697,6 +738,42 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: return schema(**input) +def _control_branch(value: Any) -> Sequence[Union[str, Send]]: + if not isinstance(value, Control): + return EMPTY_SEQ + rtn: list[Union[str, Send]] = [] + if isinstance(value.trigger, str): + rtn.append(value.trigger) + else: + rtn.extend(value.trigger) + if isinstance(value.send, Send): + rtn.append(value.send) + else: + rtn.extend(value.send) + return rtn + + +async def _acontrol_branch(value: Any) -> None: + if not isinstance(value, Control): + return EMPTY_SEQ + rtn: list[Union[str, Send]] = [] + if isinstance(value.trigger, str): + rtn.append(value.trigger) + else: + rtn.extend(value.trigger) + if isinstance(value.send, Send): + rtn.append(value.send) + else: + rtn.extend(value.send) + return rtn + + +CONTROL_BRANCH_PATH = RunnableCallable( + _control_branch, _acontrol_branch, tags=[TAG_HIDDEN], trace=False, recurse=False +) +CONTROL_BRANCH = Branch(CONTROL_BRANCH_PATH, None) + + def _get_channels( schema: Type[dict], ) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]: @@ -823,21 +900,3 @@ def _get_schema( if k in channels and isinstance(channels[k], BaseChannel) }, ) - - -def _get_input_schema_from_type_hint( - action: Optional[RunnableLike], -) -> Optional[Type[Any]]: - if not isfunction(action) and not ismethod(getattr(action, "__call__", None)): - return None - action = cast(Callable, action) - - try: - hints = get_type_hints(getattr(action, "__call__")) or get_type_hints(action) - first_parameter_name = next(iter(inspect.signature(action).parameters.keys())) - input_hint = hints.get(first_parameter_name) - if isinstance(input_hint, type) and get_type_hints(input_hint): - return input_hint - except (TypeError, StopIteration): - pass - return None diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 279594dec..f7f330b83 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -4,11 +4,13 @@ from typing import ( TYPE_CHECKING, Any, Callable, + Generic, Literal, NamedTuple, Optional, Sequence, Type, + TypeVar, Union, cast, ) @@ -221,6 +223,22 @@ class Send: ) +N = TypeVar("N") + + +class Control(Generic[N]): + def __init__( + self, + *, + update_state: Optional[dict[str, Any]] = None, + trigger: Union[str, Sequence[str]] = (), + send: Union[Send, Sequence[Send]] = (), + ) -> None: + self.update_state = update_state + self.trigger = trigger + self.send = send + + StreamChunk = tuple[tuple[str, ...], str, Any] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 71b7f4a4f..13bc25236 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6458,7 +6458,7 @@ def test_root_graph( ), AIMessage(content="answer", id="ai2"), AIMessage( - content="an extra message", id="00000000-0000-4000-8000-000000000091" + content="an extra message", id="00000000-0000-4000-8000-000000000092" ), HumanMessage(content="what is weather in la"), ], diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 03b119d0d..fe4bee325 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -62,7 +62,7 @@ from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore -from langgraph.types import Interrupt, PregelTask, Send, StreamWriter +from langgraph.types import Control, Interrupt, PregelTask, Send, StreamWriter from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, @@ -1969,6 +1969,65 @@ async def test_max_concurrency() -> None: assert node2.currently == 0 +async def test_max_concurrency_control() -> None: + async def node1(state) -> Control[Literal["2"]]: + return Control(update_state=["1"], send=[Send("2", state)] * 100) + + node2_currently = 0 + node2_max_currently = 0 + + async def node2(state) -> Control[Literal["3"]]: + nonlocal node2_currently, node2_max_currently + node2_currently += 1 + if node2_currently > node2_max_currently: + node2_max_currently = node2_currently + await asyncio.sleep(0.1) + node2_currently -= 1 + + return Control(update_state=["2"], trigger="3") + + async def node3(state) -> Literal["3"]: + return ["3"] + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node("1", node1) + builder.add_node("2", node2) + builder.add_node("3", node3) + builder.add_edge(START, "1") + graph = builder.compile() + + assert ( + graph.get_graph().draw_mermaid() + == """%%{init: {'flowchart': {'curve': 'linear'}}}%% +graph TD; + __start__([

__start__

]):::first + 1(1) + 2(2) + 3([3]):::last + __start__ --> 1; + 1 -.-> 2; + 2 -.-> 3; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc +""" + ) + + assert await graph.ainvoke(["0"], debug=True) == ["0", "1", *(["2"] * 100), "3"] + assert node2_max_currently == 100 + assert node2_currently == 0 + node2_max_currently = 0 + + assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [ + "0", + "1", + *(["2"] * 100), + "3", + ] + assert node2_max_currently == 10 + assert node2_currently == 0 + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_checkpoint_three( mocker: MockerFixture, checkpointer_name: str