From 508193272a3849ed478c025543959a55ffa482e2 Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Thu, 12 Mar 2026 15:15:49 -0700 Subject: [PATCH] anyof --- .../langgraph/advanced_graph/__init__.py | 12 +++ .../langgraph/advanced_graph/state.py | 97 ++++++++++++++++++- .../tests/advanced-graph/test_sub_agents.py | 48 ++++++--- 3 files changed, 142 insertions(+), 15 deletions(-) diff --git a/libs/langgraph/langgraph/advanced_graph/__init__.py b/libs/langgraph/langgraph/advanced_graph/__init__.py index 95bb5f03f..17a28ba1e 100644 --- a/libs/langgraph/langgraph/advanced_graph/__init__.py +++ b/libs/langgraph/langgraph/advanced_graph/__init__.py @@ -1,13 +1,25 @@ from langgraph.advanced_graph.state import ( AdvancedStateGraph, + AnyOfCondition, + ChannelCondition, CompiledGraphEngine, + TimerCondition, + any_of, + channel_condition, publish_to_channel, + timer_condition, wait_for, ) __all__ = ( "AdvancedStateGraph", + "AnyOfCondition", + "ChannelCondition", "CompiledGraphEngine", + "TimerCondition", + "any_of", + "channel_condition", "publish_to_channel", + "timer_condition", "wait_for", ) diff --git a/libs/langgraph/langgraph/advanced_graph/state.py b/libs/langgraph/langgraph/advanced_graph/state.py index 2e1b8487a..64e1c44b4 100644 --- a/libs/langgraph/langgraph/advanced_graph/state.py +++ b/libs/langgraph/langgraph/advanced_graph/state.py @@ -5,6 +5,7 @@ import contextvars import inspect from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from datetime import timedelta from typing import Any, Generic, TypeVar, cast from langgraph.types import Command, Send @@ -22,6 +23,25 @@ class _ChannelSpec: maxsize: int +@dataclass(frozen=True) +class ChannelCondition: + channel: str + n: int = 1 + + +@dataclass(frozen=True) +class TimerCondition: + seconds: float + + +@dataclass(frozen=True) +class AnyOfCondition: + conditions: tuple[WaitCondition, ...] + + +WaitCondition = ChannelCondition | TimerCondition | AnyOfCondition + + class AdvancedStateGraph(Generic[StateT]): """Experimental in-memory graph engine with async channels.""" @@ -199,7 +219,24 @@ class _GraphEngineRun: queue = self._get_channel(channel) queue.put_nowait(value) - async def wait_for(self, channel: str, n: int = 1) -> Any: + async def wait_for(self, target: str | WaitCondition, n: int = 1) -> Any: + if isinstance(target, str): + return await self._wait_for_channel_values(target, n=n) + if isinstance(target, ChannelCondition): + value = await self._wait_for_channel_values(target.channel, n=target.n) + return { + "condition": "channel", + "channel": target.channel, + "value": value, + } + if isinstance(target, TimerCondition): + await asyncio.sleep(target.seconds) + return {"condition": "timer", "seconds": target.seconds} + if isinstance(target, AnyOfCondition): + return await self._wait_for_any_of(target) + raise ValueError(f"Unsupported wait condition type: {type(target)!r}") + + async def _wait_for_channel_values(self, channel: str, n: int) -> Any: if n < 1: raise ValueError("wait_for count `n` must be >= 1") queue = self._get_channel(channel) @@ -210,6 +247,21 @@ class _GraphEngineRun: values.append(await queue.get()) return values + async def _wait_for_any_of(self, condition: AnyOfCondition) -> Any: + if not condition.conditions: + raise ValueError("any_of() requires at least one condition") + + tasks = [ + asyncio.create_task(self.wait_for(inner_condition, n=1)) + for inner_condition in condition.conditions + ] + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + first = done.pop() + return first.result() + def _get_channel(self, channel: str) -> asyncio.Queue[Any]: if channel not in self._channels: raise ValueError(f"Unknown channel `{channel}`") @@ -321,11 +373,11 @@ def _normalize_goto(goto: Any, *, default_arg: Any) -> list[Send]: return [] -async def wait_for(channel: str, n: int = 1) -> Any: +async def wait_for(target: str | WaitCondition, n: int = 1) -> Any: run = _CURRENT_RUN.get() if run is None: raise RuntimeError("wait_for() can only be used inside advanced_graph nodes") - return await run.wait_for(channel, n=n) + return await run.wait_for(target, n=n) def publish_to_channel(channel: str, value: Any) -> None: @@ -337,6 +389,45 @@ def publish_to_channel(channel: str, value: Any) -> None: run.publish_nowait(channel, value) +def channel_condition(channel: str, n: int = 1) -> ChannelCondition: + if n < 1: + raise ValueError("channel_condition `n` must be >= 1") + return ChannelCondition(channel=channel, n=n) + + +def timer_condition( + timeout: float | timedelta | None = None, + *, + seconds: float | None = None, + minutes: float | None = None, +) -> TimerCondition: + if timeout is not None and (seconds is not None or minutes is not None): + raise ValueError( + "Provide either `timeout` or named `seconds`/`minutes`, not both" + ) + + if isinstance(timeout, timedelta): + resolved_seconds = timeout.total_seconds() + elif isinstance(timeout, (int, float)): + resolved_seconds = float(timeout) + else: + resolved_seconds = 0.0 + if seconds is not None: + resolved_seconds += float(seconds) + if minutes is not None: + resolved_seconds += float(minutes) * 60.0 + + if resolved_seconds <= 0: + raise ValueError("timer_condition must be greater than 0 seconds") + return TimerCondition(seconds=resolved_seconds) + + +def any_of(*conditions: WaitCondition) -> AnyOfCondition: + if not conditions: + raise ValueError("any_of() requires at least one condition") + return AnyOfCondition(conditions=tuple(conditions)) + + def _infer_node_name(node: Callable[..., Any]) -> str: node_name = getattr(node, "__name__", "") if not node_name or node_name == "": diff --git a/libs/langgraph/tests/advanced-graph/test_sub_agents.py b/libs/langgraph/tests/advanced-graph/test_sub_agents.py index 8a7f8a6b7..31ded9793 100644 --- a/libs/langgraph/tests/advanced-graph/test_sub_agents.py +++ b/libs/langgraph/tests/advanced-graph/test_sub_agents.py @@ -5,7 +5,14 @@ from typing import Any, Literal import pytest from typing_extensions import TypedDict -from langgraph.advanced_graph import AdvancedStateGraph, publish_to_channel, wait_for +from langgraph.advanced_graph import ( + AdvancedStateGraph, + any_of, + channel_condition, + publish_to_channel, + timer_condition, + wait_for, +) from langgraph.constants import END, START from langgraph.graph import StateGraph from langgraph.types import Command, Send @@ -90,9 +97,26 @@ async def test_async_sub_graph() -> None: return Command(goto=sends) async def wait_node(state: MainAgentState) -> Command: - # Lightweight interrupt: only this node blocks on inbox. - msg = await wait_for("inbox") - state["output"].append(f"{msg['type']}: {msg['payload']}") + # Lightweight interrupt: only this node blocks for the next relevant signal. + event = await wait_for( + any_of( + channel_condition("tool_completion_channel"), + channel_condition("subagent_completion_channel"), + channel_condition("user_input_channel"), + timer_condition(minutes=1), + ) + ) + if event["condition"] == "channel": + channel = event["channel"] + payload = event["value"] + if channel == "tool_completion_channel": + state["output"].append(f"tool: {payload}") + elif channel == "subagent_completion_channel": + state["output"].append(f"sub_agent: {payload}") + elif channel == "user_input_channel": + state["output"].append(f"user_input: {payload}") + else: + state["output"].append("timer: no updates yet") # Loop back to planner with updated output. return Command(goto=Send("llm_node", state)) @@ -101,8 +125,8 @@ async def test_async_sub_graph() -> None: # Fire-and-forget style completion: publish result to inbox and exit. # (i.e., just complete without explicitly going to a next node) publish_to_channel( - "inbox", - {"type": "tool", "payload": f"tool completed for: {tool_input}"}, + "tool_completion_channel", + f"tool completed for: {tool_input}", ) async def sub_agent_node(sub_agent_input: str) -> None: @@ -112,8 +136,8 @@ async def test_async_sub_graph() -> None: ) # Same pattern as tool node: publish result and complete current node. publish_to_channel( - "inbox", - {"type": "sub_agent", "payload": sub_agent_output["output"]}, + "subagent_completion_channel", + sub_agent_output["output"], ) async def order_food_node(complete_message: str) -> dict[str, str]: @@ -121,7 +145,9 @@ async def test_async_sub_graph() -> None: advanced_flow = AdvancedStateGraph(MainAgentState) # Default behavior is an unbounded async channel (maxsize=None). - advanced_flow.add_async_channel("inbox", dict) + advanced_flow.add_async_channel("tool_completion_channel", str) + advanced_flow.add_async_channel("subagent_completion_channel", str) + advanced_flow.add_async_channel("user_input_channel", str) advanced_flow.add_entry_node(llm_node) advanced_flow.node(wait_node) advanced_flow.node(tool_node) @@ -151,9 +177,7 @@ async def test_async_sub_graph() -> None: # External input can be injected while graph execution is in progress. await asyncio.sleep(0.01) - await main_agent.apublish_to_channel( - "inbox", {"type": "user_input", "payload": "No spicy food please"} - ) + await main_agent.apublish_to_channel("user_input_channel", "No spicy food please") result = await started assert result == {