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

minmax

Browse Source
This commit is contained in:
Quanzheng Long
2026-03-18 10:18:10 -07:00
parent cb9ba2bc8d
commit 88dfc0a6bb
9 changed files with 191 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph-go/advancedgraph/types.go
+7 -5
View File
@@ -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,
}
}
langgraph-go/stategraph/state_graph.go
+3 -3
View File
@@ -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
langgraph-go/tests/sub_agents_test.go
+3 -3
View File
@@ -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},
),
)
langgraph-go/tests/test_primitives_test.go
+69
View File
@@ -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)
}
}
rust-core/src/engine.rs
+31 -14
View File
@@ -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<WaitEvent, String> {
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<Option<WaitEvent>, String> {
fn try_take_channel_event(
&self,
channel: &str,
min: usize,
max: usize,
) -> Result<Option<WaitEvent>, 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);
}
rust-core/src/lib_py.rs
+20 -12
View File
@@ -45,19 +45,18 @@ impl PyRustEngine {
.map_err(PyValueError::new_err)
}
fn wait_any_of_json(&self, any_of_json: &str) -> PyResult<String> {
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<Py<PyAny>> {
#[pyo3(signature = (channel, min, max=None))]
fn wait_channel(
&self,
py: Python<'_>,
channel: &str,
min: usize,
max: Option<usize>,
) -> PyResult<Py<PyAny>> {
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<PyAny>) -> PyResult<Py<PyAny>> {
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<String> {
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<String> {
let cond: WaitCondition = serde_json::from_str(cond_json)
.map_err(|e| PyValueError::new_err(format!("Invalid condition JSON: {e}")))?;
saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py
+30 -11
View File
@@ -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}")
saf-python-sdk/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py
+2 -2
View File
@@ -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
saf-python-sdk/tests/advanced-graph/test_primitives.py
+26 -1
View File
@@ -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"