From 88dfc0a6bb3af2a4ad46fb96fe2b1a3f4ad7caee Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Wed, 18 Mar 2026 10:18:10 -0700 Subject: [PATCH] minmax --- langgraph-go/advancedgraph/types.go | 12 ++-- langgraph-go/stategraph/state_graph.go | 6 +- langgraph-go/tests/sub_agents_test.go | 6 +- langgraph-go/tests/test_primitives_test.go | 69 +++++++++++++++++++ rust-core/src/engine.rs | 45 ++++++++---- rust-core/src/lib_py.rs | 32 +++++---- .../saf_python_sdk/advanced_graph/state.py | 41 ++++++++--- .../benchmark_stategraph_vs_advancedgraph.py | 4 +- .../tests/advanced-graph/test_primitives.py | 27 +++++++- 9 files changed, 191 insertions(+), 51 deletions(-) diff --git a/langgraph-go/advancedgraph/types.go b/langgraph-go/advancedgraph/types.go index 8be5be54a..6c2d9064b 100644 --- a/langgraph-go/advancedgraph/types.go +++ b/langgraph-go/advancedgraph/types.go @@ -11,18 +11,20 @@ type WaitCondition interface { type ChannelCondition struct { Channel string - N int + Min int + Max int } func (c ChannelCondition) toAny() map[string]any { - n := c.N - if n <= 0 { - n = 1 + min := c.Min + if min <= 0 { + min = 1 } return map[string]any{ "kind": "channel", "channel": c.Channel, - "n": n, + "min": min, + "max": c.Max, } } diff --git a/langgraph-go/stategraph/state_graph.go b/langgraph-go/stategraph/state_graph.go index ca1516367..7d56c3f80 100644 --- a/langgraph-go/stategraph/state_graph.go +++ b/langgraph-go/stategraph/state_graph.go @@ -26,7 +26,7 @@ func (c *Context) Interrupt(name string) (any, error) { } event, err := c.inner.WaitFor(ag.AnyOf(ag.ChannelCondition{ Channel: internalInterruptChannel, - N: 1, + Min: 1, })) if err != nil { if waitReq, ok := ag.AsErrWaitRequested(err); ok { @@ -172,7 +172,7 @@ func (g *BasicStateGraph[StateT]) Compile() *CompiledBasicStateGraph[StateT] { if len(cond.Conditions) == 0 { cond = ag.AnyOf(ag.ChannelCondition{ Channel: internalInterruptChannel, - N: 1, + Min: 1, }) } return ag.Command{}, ag.ErrWaitRequested{Condition: cond} @@ -202,7 +202,7 @@ func (g *BasicStateGraph[StateT]) Compile() *CompiledBasicStateGraph[StateT] { needed := len(levels[nextStep-1]) _, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{ Channel: internalBarrierChannel, - N: needed, + Min: needed, })) if err != nil { return ag.Command{}, err diff --git a/langgraph-go/tests/sub_agents_test.go b/langgraph-go/tests/sub_agents_test.go index 07db4c642..3516be3e5 100644 --- a/langgraph-go/tests/sub_agents_test.go +++ b/langgraph-go/tests/sub_agents_test.go @@ -64,9 +64,9 @@ func (w *lunchWorkflow) llmNode(ctx *ag.Context, _ any, _ lunchState) (ag.Comman func (w *lunchWorkflow) waitNode(ctx *ag.Context, _ any, state lunchState) (ag.Command, error) { event, err := ctx.WaitFor( ag.AnyOf( - ag.ChannelCondition{Channel: "tool_completion_channel", N: 1}, - ag.ChannelCondition{Channel: "subagent_completion_channel", N: 1}, - ag.ChannelCondition{Channel: "user_input_channel", N: 1}, + ag.ChannelCondition{Channel: "tool_completion_channel"}, + ag.ChannelCondition{Channel: "subagent_completion_channel"}, + ag.ChannelCondition{Channel: "user_input_channel"}, ag.TimerCondition{Seconds: 1}, ), ) diff --git a/langgraph-go/tests/test_primitives_test.go b/langgraph-go/tests/test_primitives_test.go index 1aa6f4486..b6319e468 100644 --- a/langgraph-go/tests/test_primitives_test.go +++ b/langgraph-go/tests/test_primitives_test.go @@ -1,6 +1,7 @@ package tests import ( + "encoding/json" "fmt" "testing" @@ -122,3 +123,71 @@ func TestRunEndsWithoutFinishNode(t *testing.T) { t.Fatalf("unexpected logs: %#v", result.Logs) } } + +func (w *primitiveWorkflow) startWaitBatchNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) { + if err := ctx.PublishToChannel("events", "a"); err != nil { + return ag.Command{}, err + } + if err := ctx.PublishToChannel("events", "b"); err != nil { + return ag.Command{}, err + } + if err := ctx.PublishToChannel("events", "c"); err != nil { + return ag.Command{}, err + } + return ag.Command{ + Update: state, + Goto: []ag.Send{ + {Node: w.waitBatchNode, NodeInput: nil}, + }, + }, nil +} + +func (w *primitiveWorkflow) waitBatchNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) { + event, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{ + Channel: "events", + Min: 2, + Max: 4, + })) + if err != nil { + return ag.Command{}, err + } + var values []string + if err := json.Unmarshal(event.Value, &values); err != nil { + return ag.Command{}, err + } + state.Count = len(values) + state.Logs = values + state.Done = "ok" + return ag.Command{Update: state}, nil +} + +func TestChannelWaitRespectsMaxM(t *testing.T) { + workflow := &primitiveWorkflow{} + graph := ag.NewAdvancedStateGraph[primitiveState]() + graph.AddAsyncChannel("events") + graph.AddEntryNode(workflow.startWaitBatchNode) + graph.AddFinishNode(workflow.waitBatchNode) + + handler, err := graph.Compile().Start(nil, primitiveState{ + Count: 0, + Logs: []string{}, + Done: "", + }) + if err != nil { + t.Fatalf("start failed: %v", err) + } + + result, err := handler.WaitForResult() + if err != nil { + t.Fatalf("result failed: %v", err) + } + if result.Count != 3 { + t.Fatalf("unexpected count: %v", result.Count) + } + if len(result.Logs) != 3 || result.Logs[0] != "a" || result.Logs[1] != "b" || result.Logs[2] != "c" { + t.Fatalf("unexpected logs: %#v", result.Logs) + } + if result.Done != "ok" { + t.Fatalf("unexpected done: %v", result.Done) + } +} diff --git a/rust-core/src/engine.rs b/rust-core/src/engine.rs index 789086fc0..9caf57235 100644 --- a/rust-core/src/engine.rs +++ b/rust-core/src/engine.rs @@ -16,7 +16,12 @@ use tokio::sync::{mpsc as tokio_mpsc, Mutex as AsyncMutex, Notify}; #[serde(tag = "kind")] pub enum WaitCondition { #[serde(rename = "channel")] - Channel { channel: String, n: usize }, + Channel { + channel: String, + min: usize, + #[serde(default)] + max: usize, + }, #[serde(rename = "timer")] Timer { seconds: f64 }, } @@ -379,12 +384,15 @@ impl Engine { pub async fn wait_for_async(&self, cond: &WaitCondition) -> Result { debug_log(&format!("Engine::wait_for_async(cond={cond:?})")); match cond { - WaitCondition::Channel { channel, n } => { - if *n < 1 { - return Err("channel condition n must be >= 1".to_string()); + WaitCondition::Channel { channel, min, max } => { + if *min < 1 { + return Err("channel condition min must be >= 1".to_string()); + } + if *max != 0 && *max < *min { + return Err("channel condition max must be 0 or >= min".to_string()); } loop { - if let Some(event) = self.try_take_channel_event(channel, *n)? { + if let Some(event) = self.try_take_channel_event(channel, *min, *max)? { return Ok(event); } self.channel_notify.notified().await; @@ -425,11 +433,14 @@ impl Engine { loop { for cond in &any_of.conditions { - if let WaitCondition::Channel { channel, n } = cond { - if *n < 1 { - return Err("channel condition n must be >= 1".to_string()); + if let WaitCondition::Channel { channel, min, max } = cond { + if *min < 1 { + return Err("channel condition min must be >= 1".to_string()); } - if let Some(event) = self.try_take_channel_event(channel, *n)? { + if *max != 0 && *max < *min { + return Err("channel condition max must be 0 or >= min".to_string()); + } + if let Some(event) = self.try_take_channel_event(channel, *min, *max)? { return Ok(event); } } @@ -462,15 +473,21 @@ impl Engine { run_loop_block_on(self.wait_for_any_of_async(any_of)) } - fn try_take_channel_event(&self, channel: &str, n: usize) -> Result, String> { + fn try_take_channel_event( + &self, + channel: &str, + min: usize, + max: usize, + ) -> Result, String> { let mut channels = self.channels.lock().expect("channels mutex poisoned"); let queue = channels .get_mut(channel) .ok_or_else(|| format!("Unknown channel `{channel}`"))?; - if queue.len() < n { + if queue.len() < min { return Ok(None); } - if n == 1 { + let take_count = if max == 0 { min } else { queue.len().min(max) }; + if take_count == 1 { if let Some(value) = queue.pop_front() { return Ok(Some(WaitEvent::Channel { channel: channel.to_string(), @@ -479,8 +496,8 @@ impl Engine { } return Ok(None); } - let mut values = Vec::with_capacity(n); - for _ in 0..n { + let mut values = Vec::with_capacity(take_count); + for _ in 0..take_count { if let Some(v) = queue.pop_front() { values.push(v); } diff --git a/rust-core/src/lib_py.rs b/rust-core/src/lib_py.rs index dd57bcce7..c91d6eadc 100644 --- a/rust-core/src/lib_py.rs +++ b/rust-core/src/lib_py.rs @@ -45,19 +45,18 @@ impl PyRustEngine { .map_err(PyValueError::new_err) } - fn wait_any_of_json(&self, any_of_json: &str) -> PyResult { - let any_of: AnyOfCondition = serde_json::from_str(any_of_json) - .map_err(|e| PyValueError::new_err(format!("Invalid any_of JSON: {e}")))?; - let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of)) - .map_err(PyValueError::new_err)?; - serde_json::to_string(&event) - .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}"))) - } - - fn wait_channel(&self, py: Python<'_>, channel: &str, n: usize) -> PyResult> { + #[pyo3(signature = (channel, min, max=None))] + fn wait_channel( + &self, + py: Python<'_>, + channel: &str, + min: usize, + max: Option, + ) -> PyResult> { let cond = WaitCondition::Channel { channel: channel.to_string(), - n, + min, + max: max.unwrap_or(0), }; let event = run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?; @@ -77,10 +76,19 @@ impl PyRustEngine { fn wait_any_of_obj(&self, py: Python<'_>, any_of_payload: Py) -> PyResult> { let payload_json = py_obj_to_json_string(py, &any_of_payload.bind(py))?; - let event_json = self.wait_any_of_json(&payload_json)?; + let event_json = self._wait_any_of_json(&payload_json)?; json_string_to_py_obj(py, &event_json) } + fn _wait_any_of_json(&self, any_of_json: &str) -> PyResult { + let any_of: AnyOfCondition = serde_json::from_str(any_of_json) + .map_err(|e| PyValueError::new_err(format!("Invalid any_of JSON: {e}")))?; + let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of)) + .map_err(PyValueError::new_err)?; + serde_json::to_string(&event) + .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}"))) + } + fn wait_condition_json(&self, cond_json: &str) -> PyResult { let cond: WaitCondition = serde_json::from_str(cond_json) .map_err(|e| PyValueError::new_err(format!("Invalid condition JSON: {e}")))?; diff --git a/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py index d3cb448fe..4071551d7 100644 --- a/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py +++ b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py @@ -26,7 +26,8 @@ class _ChannelSpec: @dataclass(frozen=True) class ChannelCondition: channel: str - n: int = 1 + min: int = 1 + max: int = 0 @dataclass(frozen=True) @@ -288,7 +289,9 @@ class _GraphEngineRun: async def wait_for(self, target: WaitCondition | AnyOfCondition) -> Any: if isinstance(target, ChannelCondition): - value = await self._wait_for_channel_values(target.channel, n=target.n) + value = await self._wait_for_channel_values( + target.channel, min=target.min, max=target.max + ) return { "condition": "channel", "channel": target.channel, @@ -305,15 +308,22 @@ class _GraphEngineRun: 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") + async def _wait_for_channel_values( + self, channel: str, min: int, max: int = 0 + ) -> Any: + if min < 1: + raise ValueError("wait_for count `min` must be >= 1") + if max < 0: + raise ValueError("wait_for max count `max` must be >= 0") + if max != 0 and max < min: + raise ValueError("wait_for max count `max` must be 0 or >= min") loop = asyncio.get_running_loop() event = await loop.run_in_executor( _advanced_graph_executor(), self._rust_engine.wait_channel, channel, - n, + min, + max, ) return event["value"] @@ -437,10 +447,14 @@ def _normalize_goto(goto: Any, *, default_input: Any) -> list[Send]: return [] -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 channel_condition(channel: str, min: int = 1, max: int = 0) -> ChannelCondition: + if min < 1: + raise ValueError("channel_condition `min` must be >= 1") + if max < 0: + raise ValueError("channel_condition `max` must be >= 0") + if max != 0 and max < min: + raise ValueError("channel_condition `max` must be 0 or >= min") + return ChannelCondition(channel=channel, min=min, max=max) def timer_condition( @@ -478,7 +492,12 @@ def any_of(*conditions: WaitCondition) -> AnyOfCondition: def _condition_to_rust(condition: WaitCondition) -> dict[str, Any]: if isinstance(condition, ChannelCondition): - return {"kind": "channel", "channel": condition.channel, "n": condition.n} + return { + "kind": "channel", + "channel": condition.channel, + "min": condition.min, + "max": condition.max, + } if isinstance(condition, TimerCondition): return {"kind": "timer", "seconds": condition.seconds} raise TypeError(f"Unsupported condition type: {type(condition)!r}") diff --git a/saf-python-sdk/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py b/saf-python-sdk/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py index 86cbdaf00..f1eb2435d 100644 --- a/saf-python-sdk/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py +++ b/saf-python-sdk/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py @@ -40,7 +40,7 @@ def build_advanced_parallel() -> Any: return Command(goto=sends) async def end_node(ctx: Any, state: dict[str, Any]) -> dict[str, Any]: - await ctx.wait_for(channel_condition(done_channel, n=MIDDLE_COUNT)) + await ctx.wait_for(channel_condition(done_channel, min=MIDDLE_COUNT)) out = dict(state) out["done"] = True return out @@ -161,7 +161,7 @@ def build_advanced_parallel_blocking() -> Any: return Command(goto=sends) async def end_node_blocking(ctx: Any, state: dict[str, Any]) -> dict[str, Any]: - await ctx.wait_for(channel_condition(done_channel, n=MIDDLE_COUNT)) + await ctx.wait_for(channel_condition(done_channel, min=MIDDLE_COUNT)) out = dict(state) out["done"] = True return out diff --git a/saf-python-sdk/tests/advanced-graph/test_primitives.py b/saf-python-sdk/tests/advanced-graph/test_primitives.py index ee794e775..9d0d61ee4 100644 --- a/saf-python-sdk/tests/advanced-graph/test_primitives.py +++ b/saf-python-sdk/tests/advanced-graph/test_primitives.py @@ -1,7 +1,7 @@ import pytest from typing_extensions import TypedDict -from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context +from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context, channel_condition from saf_python_sdk.types import Command, Send pytestmark = pytest.mark.anyio @@ -65,3 +65,28 @@ async def test_run_ends_without_finish_node() -> None: assert result["done"] == "stopped" assert result["logs"] == ["start", "middle:from_start"] + +async def test_channel_wait_respects_max_m() -> None: + graph = AdvancedStateGraph(PrimitiveState) + graph.add_async_channel("events", list[str]) + + async def start_node(ctx: Context, state: PrimitiveState) -> Command: + ctx.publish_to_channel("events", "a") + ctx.publish_to_channel("events", "b") + ctx.publish_to_channel("events", "c") + return Command(update=state, goto=Send("wait_node", None)) + + async def wait_node(ctx: Context, _input: None, state: PrimitiveState) -> dict[str, object]: + event = await ctx.wait_for(channel_condition("events", min=2, max=4)) + values = event["value"] + assert isinstance(values, list) + return {"counter": len(values), "logs": values, "done": "ok"} + + graph.add_entry_node(start_node) + graph.add_finish_node(wait_node) + + result = await graph.compile().ainvoke({"counter": 0, "logs": [], "done": None}) + assert result["counter"] == 3 + assert result["logs"] == ["a", "b", "c"] + assert result["done"] == "ok" +