diff --git a/advanced-graph-examples/README.md b/advanced-graph-examples/README.md new file mode 100644 index 000000000..a2caddbdd --- /dev/null +++ b/advanced-graph-examples/README.md @@ -0,0 +1,14 @@ +# advanced-graph-examples + +Examples for the published `saf-python-sdk` package. + +## Quick start + +```bash +cd advanced-graph-examples +uv sync +uv run python examples/basic_run.py +``` + +> Note: current published wheel is for Python 3.13 on macOS arm64. + diff --git a/advanced-graph-examples/examples/basic_run.py b/advanced-graph-examples/examples/basic_run.py new file mode 100644 index 000000000..9e61a765a --- /dev/null +++ b/advanced-graph-examples/examples/basic_run.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import asyncio +from typing import TypedDict + +from saf_python_sdk import Command, Send +from saf_python_sdk.advanced_graph import AdvancedStateGraph + + +class MyState(TypedDict): + count: int + logs: list[str] + + +async def start_node(state: MyState) -> Command: + state["logs"].append("start") + return Command(update=state, goto=Send("finish_node", "hello")) + + +async def finish_node(input: str, state: MyState) -> Command: + state["logs"].append(f"finish:{input}") + state["count"] += 1 + return Command(update=state) + + +async def main() -> None: + graph = AdvancedStateGraph(MyState) + graph.add_entry_node(start_node) + graph.add_finish_node(finish_node) + result = await graph.compile().ainvoke({"count": 0, "logs": []}) + print(result) + + +if __name__ == "__main__": + asyncio.run(main()) + diff --git a/advanced-graph-examples/pyproject.toml b/advanced-graph-examples/pyproject.toml new file mode 100644 index 000000000..2cf6e68f6 --- /dev/null +++ b/advanced-graph-examples/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "advanced-graph-examples" +version = "0.1.0" +requires-python = ">=3.13,<3.14" +dependencies = [ + "saf-python-sdk>=0.1.1", +] + +[tool.uv] +dev-dependencies = [] + diff --git a/advanced-graph-examples/uv.lock b/advanced-graph-examples/uv.lock new file mode 100644 index 000000000..d49931b02 --- /dev/null +++ b/advanced-graph-examples/uv.lock @@ -0,0 +1,25 @@ +version = 1 +revision = 3 +requires-python = "==3.13.*" + +[[package]] +name = "advanced-graph-examples" +version = "0.1.0" +source = { virtual = "." } +dependencies = [ + { name = "saf-python-sdk" }, +] + +[package.metadata] +requires-dist = [{ name = "saf-python-sdk", specifier = ">=0.1.1" }] + +[package.metadata.requires-dev] +dev = [] + +[[package]] +name = "saf-python-sdk" +version = "0.1.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/56/73bcc01ff0fbf67fdfd22c2b06f31a7331fd2f6fe27f682da296027891f7/saf_python_sdk-0.1.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:920eaf3451d52b62a20ffbad5d9e881c0645f89f5efedf540f5b5d9b65768005", size = 502816, upload-time = "2026-03-16T23:09:46.081Z" }, +] diff --git a/libs/langgraph/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py b/libs/langgraph/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py deleted file mode 100644 index ff4373193..000000000 --- a/libs/langgraph/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py +++ /dev/null @@ -1,306 +0,0 @@ -from __future__ import annotations - -import asyncio -import os -import time -from typing import Any - -# Configure advanced-graph runtime pools for this benchmark run. -os.environ["LANGGRAPH_RUN_POOL_SIZE"] = "10" -os.environ["LANGGRAPH_NODE_POOL_SIZE"] = "1000" - -from langgraph.advanced_graph import ( - AdvancedStateGraph, - channel_condition, - timer_condition, -) -from langgraph.graph import END, START, StateGraph -from langgraph.types import Command, Send - -RUNS = 100 -MIDDLE_COUNT = 10 -SLEEP_SECONDS = 2.0 -BLOCKING_SECONDS = 0.1 -STATE_BYTES = 10 * 1024 - - -def make_initial_state() -> dict[str, Any]: - return {"payload": "x" * STATE_BYTES, "done": False} - - -def build_advanced_parallel() -> Any: - graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict) - done_channel = "__bench_done_channel" - graph.add_async_channel(done_channel, str) - - async def start_node(state: dict[str, Any]) -> Command: - _ = state - sends = [Send(f"middle_{i}", None) for i in range(MIDDLE_COUNT)] - sends.append(Send("end_node", None)) - 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)) - out = dict(state) - out["done"] = True - return out - - graph.add_entry_node(start_node) - for i in range(MIDDLE_COUNT): - async def middle_node(ctx: Any, state: dict[str, Any], idx: int = i) -> None: - _ = idx - _ = state - await ctx.wait_for(timer_condition(seconds=SLEEP_SECONDS)) - ctx.publish_to_channel(done_channel, "done") - - graph.add_node(f"middle_{i}", middle_node) - graph.add_finish_node(end_node) - return graph.compile() - - -def build_advanced_sequential() -> Any: - graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict) - - async def start_node(state: dict[str, Any]) -> Command: - _ = state - return Command(goto=Send("middle_0", None)) - - async def end_node(state: dict[str, Any]) -> dict[str, Any]: - out = dict(state) - out["done"] = True - return out - - graph.add_entry_node(start_node) - def make_middle(target: str): - async def middle_node(ctx: Any, state: dict[str, Any]) -> Command: - _ = state - await ctx.wait_for(timer_condition(seconds=SLEEP_SECONDS)) - return Command(goto=Send(target, None)) - - return middle_node - - for i in range(MIDDLE_COUNT): - next_name = "end_node" if i == MIDDLE_COUNT - 1 else f"middle_{i+1}" - graph.add_node(f"middle_{i}", make_middle(next_name)) - graph.add_finish_node(end_node) - return graph.compile() - - -def build_stategraph_parallel() -> Any: - graph = StateGraph(dict) - - async def start_node(state: dict[str, Any]) -> None: - _ = state - - async def end_node(state: dict[str, Any]) -> dict[str, Any]: - out = dict(state) - out["done"] = True - return out - - graph.add_node("start_node", start_node) - for i in range(MIDDLE_COUNT): - async def middle_node(state: dict[str, Any], idx: int = i) -> None: - _ = idx - _ = state - await asyncio.sleep(SLEEP_SECONDS) - - graph.add_node(f"middle_{i}", middle_node) - graph.add_node("end_node", end_node) - - graph.add_edge(START, "start_node") - for i in range(MIDDLE_COUNT): - graph.add_edge("start_node", f"middle_{i}") - graph.add_edge(f"middle_{i}", "end_node") - graph.add_edge("end_node", END) - return graph.compile() - - -def build_stategraph_sequential() -> Any: - graph = StateGraph(dict) - - async def start_node(state: dict[str, Any]) -> None: - _ = state - - async def end_node(state: dict[str, Any]) -> dict[str, Any]: - out = dict(state) - out["done"] = True - return out - - graph.add_node("start_node", start_node) - for i in range(MIDDLE_COUNT): - async def middle_node(state: dict[str, Any], idx: int = i) -> None: - _ = idx - _ = state - await asyncio.sleep(SLEEP_SECONDS) - - graph.add_node(f"middle_{i}", middle_node) - graph.add_node("end_node", end_node) - - graph.add_edge(START, "start_node") - graph.add_edge("start_node", "middle_0") - for i in range(MIDDLE_COUNT - 1): - graph.add_edge(f"middle_{i}", f"middle_{i+1}") - graph.add_edge(f"middle_{MIDDLE_COUNT - 1}", "end_node") - graph.add_edge("end_node", END) - return graph.compile() - - -def build_advanced_parallel_blocking() -> Any: - graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict) - done_channel = "__bench_done_channel_blocking" - graph.add_async_channel(done_channel, str) - - async def start_node(state: dict[str, Any]) -> Command: - _ = state - sends = [Send(f"middle_blocking_{i}", None) for i in range(MIDDLE_COUNT)] - sends.append(Send("end_node_blocking", None)) - 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)) - out = dict(state) - out["done"] = True - return out - - graph.add_entry_node(start_node) - for i in range(MIDDLE_COUNT): - async def middle_blocking( - ctx: Any, state: dict[str, Any], idx: int = i - ) -> None: - _ = idx - _ = state - time.sleep(BLOCKING_SECONDS) - ctx.publish_to_channel(done_channel, "done") - - graph.add_node(f"middle_blocking_{i}", middle_blocking) - graph.add_finish_node(end_node_blocking) - return graph.compile() - - -def build_advanced_sequential_blocking() -> Any: - graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict) - - async def start_node(state: dict[str, Any]) -> Command: - _ = state - return Command(goto=Send("middle_blocking_seq_0", None)) - - async def end_node_blocking_seq(state: dict[str, Any]) -> dict[str, Any]: - out = dict(state) - out["done"] = True - return out - - graph.add_entry_node(start_node) - - def make_middle(target: str): - async def middle_blocking_seq(state: dict[str, Any]) -> Command: - _ = state - time.sleep(BLOCKING_SECONDS) - return Command(goto=Send(target, None)) - - return middle_blocking_seq - - for i in range(MIDDLE_COUNT): - next_name = ( - "end_node_blocking_seq" - if i == MIDDLE_COUNT - 1 - else f"middle_blocking_seq_{i+1}" - ) - graph.add_node(f"middle_blocking_seq_{i}", make_middle(next_name)) - graph.add_finish_node(end_node_blocking_seq) - return graph.compile() - - -def build_stategraph_parallel_blocking() -> Any: - graph = StateGraph(dict) - - async def start_node(state: dict[str, Any]) -> None: - _ = state - - async def end_node(state: dict[str, Any]) -> dict[str, Any]: - out = dict(state) - out["done"] = True - return out - - graph.add_node("start_node", start_node) - for i in range(MIDDLE_COUNT): - async def middle_blocking(state: dict[str, Any], idx: int = i) -> None: - _ = idx - _ = state - time.sleep(BLOCKING_SECONDS) - - graph.add_node(f"middle_blocking_{i}", middle_blocking) - graph.add_node("end_node", end_node) - - graph.add_edge(START, "start_node") - for i in range(MIDDLE_COUNT): - graph.add_edge("start_node", f"middle_blocking_{i}") - graph.add_edge(f"middle_blocking_{i}", "end_node") - graph.add_edge("end_node", END) - return graph.compile() - - -def build_stategraph_sequential_blocking() -> Any: - graph = StateGraph(dict) - - async def start_node(state: dict[str, Any]) -> None: - _ = state - - async def end_node(state: dict[str, Any]) -> dict[str, Any]: - out = dict(state) - out["done"] = True - return out - - graph.add_node("start_node", start_node) - for i in range(MIDDLE_COUNT): - async def middle_blocking_seq(state: dict[str, Any], idx: int = i) -> None: - _ = idx - _ = state - time.sleep(BLOCKING_SECONDS) - - graph.add_node(f"middle_blocking_seq_{i}", middle_blocking_seq) - graph.add_node("end_node", end_node) - - graph.add_edge(START, "start_node") - graph.add_edge("start_node", "middle_blocking_seq_0") - for i in range(MIDDLE_COUNT - 1): - graph.add_edge( - f"middle_blocking_seq_{i}", - f"middle_blocking_seq_{i+1}", - ) - graph.add_edge(f"middle_blocking_seq_{MIDDLE_COUNT - 1}", "end_node") - graph.add_edge("end_node", END) - return graph.compile() - - -async def run_benchmark(name: str, compiled: Any) -> float: - started = time.perf_counter() - tasks = [asyncio.create_task(compiled.ainvoke(make_initial_state())) for _ in range(RUNS)] - results = await asyncio.gather(*tasks) - elapsed = time.perf_counter() - started - if not all(item.get("done") is True for item in results): - raise RuntimeError(f"{name} produced unfinished runs") - return elapsed - - -async def main() -> None: - suites = [ - ("advanced-graph-parallel", build_advanced_parallel()), - ("advanced-graph-sequential", build_advanced_sequential()), - ("state-graph-parallel", build_stategraph_parallel()), - ("state-graph-sequential", build_stategraph_sequential()), - ("advanced-graph-parallel-blocking", build_advanced_parallel_blocking()), - ("advanced-graph-sequential-blocking", build_advanced_sequential_blocking()), - ("state-graph-parallel-blocking", build_stategraph_parallel_blocking()), - ("state-graph-sequential-blocking", build_stategraph_sequential_blocking()), - ] - print( - f"runs={RUNS}, middle_nodes={MIDDLE_COUNT}, sleep={SLEEP_SECONDS}s, " - f"blocking_sleep={BLOCKING_SECONDS}s, state_bytes={STATE_BYTES}" - ) - for name, compiled in suites: - elapsed = await run_benchmark(name, compiled) - print(f"{name}: {elapsed:.3f}s") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/saf-python-sdk/Cargo.lock b/saf-python-sdk/Cargo.lock new file mode 100644 index 000000000..082228bd7 --- /dev/null +++ b/saf-python-sdk/Cargo.lock @@ -0,0 +1,338 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + +[[package]] +name = "bitflags" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "itoa" +version = "1.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" + +[[package]] +name = "langgraph_rust_core" +version = "0.1.0" +dependencies = [ + "libc", + "parking_lot", + "pyo3", + "serde", + "serde_json", + "tokio", +] + +[[package]] +name = "libc" +version = "0.2.183" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "memchr" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "smallvec" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + +[[package]] +name = "tokio" +version = "1.50.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d" +dependencies = [ + "pin-project-lite", + "tokio-macros", +] + +[[package]] +name = "tokio-macros" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/saf-python-sdk/Cargo.toml b/saf-python-sdk/Cargo.toml new file mode 100644 index 000000000..d506f921d --- /dev/null +++ b/saf-python-sdk/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "langgraph_rust_core" +version = "0.1.0" +edition = "2021" + +[lib] +name = "langgraph_rust_core" +path = "../rust-core/src/lib.rs" +crate-type = ["cdylib", "rlib"] + +[features] +default = ["python-bindings"] +python-bindings = ["dep:pyo3"] + +[dependencies] +pyo3 = { version = "0.23.5", features = ["extension-module"], optional = true } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +parking_lot = "0.12" +libc = "0.2" +tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time"] } + diff --git a/saf-python-sdk/Makefile b/saf-python-sdk/Makefile new file mode 100644 index 000000000..ff0099649 --- /dev/null +++ b/saf-python-sdk/Makefile @@ -0,0 +1,10 @@ +PYPI_REPOSITORY ?= pypi +PYPI_TOKEN ?= + +.PHONY: publish-to-pypi-saf-python-sdk +publish-to-pypi-saf-python-sdk: + @test -n "$(PYPI_TOKEN)" || (echo "PYPI_TOKEN is required"; exit 1) + cd "$(CURDIR)" && \ + MATURIN_PYPI_TOKEN="$(PYPI_TOKEN)" \ + uvx maturin publish --repository $(PYPI_REPOSITORY) --non-interactive --skip-existing --no-sdist + diff --git a/saf-python-sdk/README.md b/saf-python-sdk/README.md new file mode 100644 index 000000000..45ada8435 --- /dev/null +++ b/saf-python-sdk/README.md @@ -0,0 +1,16 @@ +# saf-python-sdk + +Standalone Python SDK for the `advanced_graph` runtime backed by the Rust engine. + +This package intentionally contains only: + +- `saf_python_sdk.advanced_graph` (Python API) +- `langgraph_rust_core` (Rust execution engine via PyO3) + +It does not package the original `langgraph` `stategraph` stack. + +## Moved assets + +- Original advanced graph design doc: `evolve-extend-langgraph.md` +- Original advanced graph tests: `tests/advanced-graph/` + diff --git a/libs/langgraph/langgraph/advanced_graph/evolve-extend-langgraph.md b/saf-python-sdk/evolve-extend-langgraph.md similarity index 98% rename from libs/langgraph/langgraph/advanced_graph/evolve-extend-langgraph.md rename to saf-python-sdk/evolve-extend-langgraph.md index 435457575..bc1984280 100644 --- a/libs/langgraph/langgraph/advanced_graph/evolve-extend-langgraph.md +++ b/saf-python-sdk/evolve-extend-langgraph.md @@ -32,11 +32,10 @@ Sub-agents also cannot simply be modeled as subgraphs, because subgraphs today e The closest workaround today is double texting, but it has a fundamental flaw: when a new audio input arrives, the previous one is interrupted and canceled rather than being allowed to gracefully complete. The workflow code itself should have the control to decide whether to stop running. -LangGraph is, at its core, a general-purpose workflow engine. Although we focus primarily on agent development, none of the primitives it offers are exclusive to agents or dedicated solely to agentic use cases. Conversely, there is nothing that a general-purpose workflow engine provides that we can safely assume agent development will _never_ need. +LangGraph is, at its core, a general-purpose workflow engine. Although we focus primarily on agent development, none of the primitives it offers are exclusive to agents or dedicated solely to agentic use cases. Conversely, there is nothing that a general-purpose workflow engine provides that we can safely assume agent development will _never_ need. The difference is probably only priority. For example, durable timer where a step can sleep for hours, days or months before resuming. Traditional workflow engines — those built for general microservice orchestration(which doesn't need streaming) -- they may need durable timers. In the agent development world today, most agents are still relatively simple. There are not yet many scenarios that require a step to wait for hours or days before proceeding. - ## Deriving What's Needed from First Principles Before jumping to solutions, it is worth stepping back and asking a fundamental question: what is an orchestration engine, and what do users expect it to provide? @@ -47,7 +46,7 @@ Starting from the simple. A developer could write a simple `main` function — a But if that machine crashes, you probably do not want the process to start over from scratch. You want it to resume from the last step that completed successfully. And if a step fails, you might want it to retry automatically before giving up. -LangGraph handles this case very well. +LangGraph handles this case very well. There is an important constraint worth calling out explicitly: LangGraph requires the developer to organize their code into **nodes**, which serve as the boundaries at which checkpoints can be taken. This is a constraint shared by every workflow engine — it is simply not feasible to persist a checkpoint after every single line of arbitrary code. @@ -59,9 +58,9 @@ This is precisely why LangGraph's superstep restriction feels awkward in practic Multiple threads and processes do, however, need to coordinate with each other. In concurrent programming, channels are an essential primitive precisely because they provide a safe, structured way for threads to communicate and synchronize without relying on shared mutable memory — avoiding data races and deadlocks. In some cases, threads may use locking for coordination, but the preferred approach is message passing through channels. -NOTE: "channel" is overloaded term here as it's also an internal term within current LangGraph pregel algorithm. +NOTE: "channel" is overloaded term here as it's also an internal term within current LangGraph pregel algorithm. -LangGraph already has a mechanism that is closely related: `interrupt`. A run can be interrupted, and then another run can resume it. If we look at this through the lens of channels, `interrupt` is essentially a **channel with size 0** — a synchronous rendezvous point where one side blocks until the other side is ready. +LangGraph already has a mechanism that is closely related: `interrupt`. A run can be interrupted, and then another run can resume it. If we look at this through the lens of channels, `interrupt` is essentially a **channel with size 0** — a synchronous rendezvous point where one side blocks until the other side is ready. The natural extension: @@ -97,7 +96,6 @@ Branch 2: a → b2 → b22 → ... Each branch advances at its own pace. `b1` finishing triggers `b11` immediately, without waiting for `b2`. - No API change is needed from the user's perspective — the graph definition stays the same. The change is in the execution semantics: the engine no longer forces all parallel nodes to synchronize at each step boundary. Each branch is checkpointed independently, so if `b1 → b11` completes while `b2` is still running, `b11`'s result is already persisted. This is necessary for the next one -- Light-weight Interrupt: Only Block the Current Node. Because we want to let other nodes continue to run while a node is waiting on something. @@ -118,13 +116,12 @@ The `wait_for` call takes a channel name and optionally a count `N`, meaning "wa See [test_sub_agents.py](../libs/langgraph/tests/advanced-graph/test_sub_agents.py) - ### P2: likely needed #### subGraph redesign #### durable timers #### more flexiable waiting conditions on interrupts #### locking on state fields - ### P3: future needed or nice to have -#### RPC \ No newline at end of file +#### RPC + diff --git a/saf-python-sdk/pyproject.toml b/saf-python-sdk/pyproject.toml new file mode 100644 index 000000000..f9b43392f --- /dev/null +++ b/saf-python-sdk/pyproject.toml @@ -0,0 +1,26 @@ +[build-system] +requires = ["maturin>=1.7,<2"] +build-backend = "maturin" + +[project] +name = "saf-python-sdk" +version = "0.1.1" +description = "Standalone advanced graph runtime powered by Rust engine" +readme = "README.md" +requires-python = ">=3.10" +license = "MIT" +authors = [{ name = "LangGraph Contributors" }] +classifiers = [ + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Rust", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", +] + +[tool.maturin] +python-source = "python" +module-name = "saf_python_sdk.langgraph_rust_core" +bindings = "pyo3" +features = ["python-bindings"] + diff --git a/saf-python-sdk/python/saf_python_sdk/__init__.py b/saf-python-sdk/python/saf_python_sdk/__init__.py new file mode 100644 index 000000000..0ce193fd1 --- /dev/null +++ b/saf-python-sdk/python/saf_python_sdk/__init__.py @@ -0,0 +1,4 @@ +from .types import Command, Send + +__all__ = ["Command", "Send"] + diff --git a/libs/langgraph/langgraph/advanced_graph/__init__.py b/saf-python-sdk/python/saf_python_sdk/advanced_graph/__init__.py similarity index 87% rename from libs/langgraph/langgraph/advanced_graph/__init__.py rename to saf-python-sdk/python/saf_python_sdk/advanced_graph/__init__.py index 851d6bc7a..d0cd77021 100644 --- a/libs/langgraph/langgraph/advanced_graph/__init__.py +++ b/saf-python-sdk/python/saf_python_sdk/advanced_graph/__init__.py @@ -1,4 +1,4 @@ -from langgraph.advanced_graph.state import ( +from .state import ( AdvancedStateGraph, AnyOfCondition, ChannelCondition, @@ -11,15 +11,16 @@ from langgraph.advanced_graph.state import ( timer_condition, ) -__all__ = ( +__all__ = [ "AdvancedStateGraph", - "AnyOfCondition", - "ChannelCondition", - "Context", "CompiledGraphEngine", + "Context", "GraphRunHandler", + "ChannelCondition", "TimerCondition", - "any_of", + "AnyOfCondition", "channel_condition", "timer_condition", -) + "any_of", +] + diff --git a/libs/langgraph/langgraph/advanced_graph/state.py b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py similarity index 98% rename from libs/langgraph/langgraph/advanced_graph/state.py rename to saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py index 98a1a0027..c8a6ed958 100644 --- a/libs/langgraph/langgraph/advanced_graph/state.py +++ b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py @@ -2,18 +2,18 @@ from __future__ import annotations import atexit import asyncio -import os import inspect +import os import threading -from concurrent.futures import ThreadPoolExecutor from collections.abc import Callable, Coroutine, Sequence +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from datetime import timedelta from typing import Any, Generic, TypeVar, cast -from langgraph_rust_core import PyRustEngine # type: ignore[import-untyped] +from saf_python_sdk.langgraph_rust_core import PyRustEngine # type: ignore[import-untyped] -from langgraph.types import Command, Send +from saf_python_sdk.types import Command, Send StateT = TypeVar("StateT") @@ -53,7 +53,7 @@ def _advanced_graph_executor() -> ThreadPoolExecutor: worker_count = max(worker_count, 1) _EXECUTOR = ThreadPoolExecutor( max_workers=worker_count, - thread_name_prefix="langgraph-advanced-py", + thread_name_prefix="saf-advanced-py", ) atexit.register(_shutdown_advanced_graph_executor) return _EXECUTOR @@ -359,6 +359,7 @@ class _GraphEngineRun: self._local.worker_loop = loop return loop.run_until_complete(awaitable) + def _normalize_result_to_sends(result: Any, *, default_input: Any) -> list[Send]: if result is None: return [] @@ -542,3 +543,4 @@ def _invoke_node(node: Callable[..., Any], ctx: Context, node_input: Any, state: return node(node_input, state) return node(ctx, node_input, state) + diff --git a/saf-python-sdk/python/saf_python_sdk/py.typed b/saf-python-sdk/python/saf_python_sdk/py.typed new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/saf-python-sdk/python/saf_python_sdk/py.typed @@ -0,0 +1 @@ + diff --git a/saf-python-sdk/python/saf_python_sdk/types.py b/saf-python-sdk/python/saf_python_sdk/types.py new file mode 100644 index 000000000..92b3f2bcd --- /dev/null +++ b/saf-python-sdk/python/saf_python_sdk/types.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Generic, TypeVar + +N = TypeVar("N") + + +@dataclass +class Send(Generic[N]): + node: N + arg: Any = None + + +@dataclass +class Command: + update: Any = None + goto: Any = field(default=None) + 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 new file mode 100644 index 000000000..f0675e52f --- /dev/null +++ b/saf-python-sdk/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +import asyncio +import os +import time +from typing import Any + +# Configure advanced-graph runtime pools for this benchmark run. +os.environ["LANGGRAPH_RUN_POOL_SIZE"] = "10" +os.environ["LANGGRAPH_NODE_POOL_SIZE"] = "1000" + +from saf_python_sdk.advanced_graph import ( + AdvancedStateGraph, + channel_condition, + timer_condition, +) +from saf_python_sdk.types import Command, Send + +try: + from langgraph.graph import END, START, StateGraph # type: ignore + + HAS_STATEGRAPH = True +except Exception: + HAS_STATEGRAPH = False + END = START = StateGraph = None # type: ignore + +RUNS = 100 +MIDDLE_COUNT = 10 +SLEEP_SECONDS = 2.0 +BLOCKING_SECONDS = 0.1 +STATE_BYTES = 10 * 1024 + + +def make_initial_state() -> dict[str, Any]: + return {"payload": "x" * STATE_BYTES, "done": False} + + +def build_advanced_parallel() -> Any: + graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict) + done_channel = "__bench_done_channel" + graph.add_async_channel(done_channel, str) + + async def start_node(state: dict[str, Any]) -> Command: + _ = state + sends = [Send(f"middle_{i}", None) for i in range(MIDDLE_COUNT)] + sends.append(Send("end_node", None)) + 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)) + out = dict(state) + out["done"] = True + return out + + graph.add_entry_node(start_node) + for i in range(MIDDLE_COUNT): + + async def middle_node(ctx: Any, state: dict[str, Any], idx: int = i) -> None: + _ = idx + _ = state + await ctx.wait_for(timer_condition(seconds=SLEEP_SECONDS)) + ctx.publish_to_channel(done_channel, "done") + + graph.add_node(f"middle_{i}", middle_node) + graph.add_finish_node(end_node) + return graph.compile() + + +def build_advanced_sequential() -> Any: + graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict) + + async def start_node(state: dict[str, Any]) -> Command: + _ = state + return Command(goto=Send("middle_0", None)) + + async def end_node(state: dict[str, Any]) -> dict[str, Any]: + out = dict(state) + out["done"] = True + return out + + graph.add_entry_node(start_node) + + def make_middle(target: str): + async def middle_node(ctx: Any, state: dict[str, Any]) -> Command: + _ = state + await ctx.wait_for(timer_condition(seconds=SLEEP_SECONDS)) + return Command(goto=Send(target, None)) + + return middle_node + + for i in range(MIDDLE_COUNT): + next_name = "end_node" if i == MIDDLE_COUNT - 1 else f"middle_{i+1}" + graph.add_node(f"middle_{i}", make_middle(next_name)) + graph.add_finish_node(end_node) + return graph.compile() + + +def _build_stategraph_suites() -> list[tuple[str, Any]]: + if not HAS_STATEGRAPH: + return [] + + def build_stategraph_parallel() -> Any: + graph = StateGraph(dict) + + async def start_node(state: dict[str, Any]) -> None: + _ = state + + async def end_node(state: dict[str, Any]) -> dict[str, Any]: + out = dict(state) + out["done"] = True + return out + + graph.add_node("start_node", start_node) + for i in range(MIDDLE_COUNT): + + async def middle_node(state: dict[str, Any], idx: int = i) -> None: + _ = idx + _ = state + await asyncio.sleep(SLEEP_SECONDS) + + graph.add_node(f"middle_{i}", middle_node) + graph.add_node("end_node", end_node) + graph.add_edge(START, "start_node") + for i in range(MIDDLE_COUNT): + graph.add_edge("start_node", f"middle_{i}") + graph.add_edge(f"middle_{i}", "end_node") + graph.add_edge("end_node", END) + return graph.compile() + + def build_stategraph_sequential() -> Any: + graph = StateGraph(dict) + + async def start_node(state: dict[str, Any]) -> None: + _ = state + + async def end_node(state: dict[str, Any]) -> dict[str, Any]: + out = dict(state) + out["done"] = True + return out + + graph.add_node("start_node", start_node) + for i in range(MIDDLE_COUNT): + + async def middle_node(state: dict[str, Any], idx: int = i) -> None: + _ = idx + _ = state + await asyncio.sleep(SLEEP_SECONDS) + + graph.add_node(f"middle_{i}", middle_node) + graph.add_node("end_node", end_node) + graph.add_edge(START, "start_node") + graph.add_edge("start_node", "middle_0") + for i in range(MIDDLE_COUNT - 1): + graph.add_edge(f"middle_{i}", f"middle_{i+1}") + graph.add_edge(f"middle_{MIDDLE_COUNT - 1}", "end_node") + graph.add_edge("end_node", END) + return graph.compile() + + return [ + ("state-graph-parallel", build_stategraph_parallel()), + ("state-graph-sequential", build_stategraph_sequential()), + ] + + +async def run_benchmark(name: str, compiled: Any) -> float: + started = time.perf_counter() + tasks = [asyncio.create_task(compiled.ainvoke(make_initial_state())) for _ in range(RUNS)] + results = await asyncio.gather(*tasks) + elapsed = time.perf_counter() - started + if not all(item.get("done") is True for item in results): + raise RuntimeError(f"{name} produced unfinished runs") + return elapsed + + +async def main() -> None: + suites = [ + ("advanced-graph-parallel", build_advanced_parallel()), + ("advanced-graph-sequential", build_advanced_sequential()), + ] + suites.extend(_build_stategraph_suites()) + print( + f"runs={RUNS}, middle_nodes={MIDDLE_COUNT}, sleep={SLEEP_SECONDS}s, " + f"blocking_sleep={BLOCKING_SECONDS}s, state_bytes={STATE_BYTES}" + ) + if not HAS_STATEGRAPH: + print("stategraph benchmarks skipped (langgraph not installed)") + for name, compiled in suites: + elapsed = await run_benchmark(name, compiled) + print(f"{name}: {elapsed:.3f}s") + + +if __name__ == "__main__": + asyncio.run(main()) + diff --git a/libs/langgraph/tests/advanced-graph/test_primitives.py b/saf-python-sdk/tests/advanced-graph/test_primitives.py similarity index 94% rename from libs/langgraph/tests/advanced-graph/test_primitives.py rename to saf-python-sdk/tests/advanced-graph/test_primitives.py index 2248a3102..ee794e775 100644 --- a/libs/langgraph/tests/advanced-graph/test_primitives.py +++ b/saf-python-sdk/tests/advanced-graph/test_primitives.py @@ -1,8 +1,8 @@ import pytest from typing_extensions import TypedDict -from langgraph.advanced_graph import AdvancedStateGraph, Context -from langgraph.types import Command, Send +from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context +from saf_python_sdk.types import Command, Send pytestmark = pytest.mark.anyio @@ -64,3 +64,4 @@ async def test_run_ends_without_finish_node() -> None: assert result["counter"] == 8 assert result["done"] == "stopped" assert result["logs"] == ["start", "middle:from_start"] + diff --git a/libs/langgraph/tests/advanced-graph/test_run_pool_size.py b/saf-python-sdk/tests/advanced-graph/test_run_pool_size.py similarity index 90% rename from libs/langgraph/tests/advanced-graph/test_run_pool_size.py rename to saf-python-sdk/tests/advanced-graph/test_run_pool_size.py index db79c0046..df6030a3c 100644 --- a/libs/langgraph/tests/advanced-graph/test_run_pool_size.py +++ b/saf-python-sdk/tests/advanced-graph/test_run_pool_size.py @@ -8,8 +8,8 @@ def test_run_pool_size_one_still_allows_parallel_runs() -> None: import asyncio import time from typing_extensions import TypedDict -from langgraph.advanced_graph import AdvancedStateGraph, Context, timer_condition -from langgraph.types import Command, Send +from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context, timer_condition +from saf_python_sdk.types import Command, Send class RunState(TypedDict): @@ -53,3 +53,4 @@ asyncio.run(main()) ) elapsed = float(completed.stdout.strip().splitlines()[-1]) assert elapsed < 0.35, completed.stdout + diff --git a/libs/langgraph/tests/advanced-graph/test_sub_agents.py b/saf-python-sdk/tests/advanced-graph/test_sub_agents.py similarity index 74% rename from libs/langgraph/tests/advanced-graph/test_sub_agents.py rename to saf-python-sdk/tests/advanced-graph/test_sub_agents.py index 2613edec3..0d49cafcb 100644 --- a/libs/langgraph/tests/advanced-graph/test_sub_agents.py +++ b/saf-python-sdk/tests/advanced-graph/test_sub_agents.py @@ -5,16 +5,14 @@ from typing import Any, Literal import pytest from typing_extensions import TypedDict -from langgraph.advanced_graph import ( +from saf_python_sdk.advanced_graph import ( AdvancedStateGraph, Context, any_of, channel_condition, timer_condition, ) -from langgraph.constants import END, START -from langgraph.graph import StateGraph -from langgraph.types import Command, Send +from saf_python_sdk.types import Command, Send pytestmark = pytest.mark.anyio @@ -51,30 +49,22 @@ class MockLLM: return response -def build_sub_agent() -> Any: - # Sub-agent uses the regular/simple StateGraph API. - sub_agent = StateGraph(SubAgentState) - - async def research_node(state: SubAgentState) -> dict[str, str]: - # Intentionally slower than timer_condition(seconds=1) to validate timer path. +class FakeSubAgent: + async def ainvoke(self, state: SubAgentState) -> SubAgentState: await asyncio.sleep(5) - return {"output": f"research sub agent completed for: {state['input']}"} + return {"input": state["input"], "output": f"research sub agent completed for: {state['input']}"} - sub_agent.add_node("research_node", research_node) - sub_agent.add_edge(START, "research_node") - sub_agent.add_edge("research_node", END) - return sub_agent.compile() + +def build_sub_agent() -> Any: + return FakeSubAgent() def build_main_agent(planner: MockLLM, sub_agent: Any) -> Any: async def llm_node(state: MainAgentState) -> Command: - # Planner decides whether to call a tool, spawn a sub-agent, or finish. decisions = await planner.ainvoke(state) sends: list[Send] = [] for decision in decisions: if decision.type == "end": - # NOTE: this can be simplified further in the future with a dedicated - # complete primitive, instead of routing to a finish node manually. return Command( goto=Send( order_food_node, @@ -85,12 +75,10 @@ def build_main_agent(planner: MockLLM, sub_agent: Any) -> Any: sends.append(Send("sub_agent_node", decision.sub_agent)) if decision.type == "tool" and decision.tool: sends.append(Send("tool_node", decision.tool)) - # Keep the main loop responsive: wait for one inbound message and continue. sends.append(Send("wait_node", None)) return Command(goto=sends) async def wait_node(ctx: Context, state: MainAgentState) -> Command: - # Lightweight interrupt: only this node blocks for the next relevant signal. event = await ctx.wait_for( any_of( channel_condition("tool_completion_channel"), @@ -108,28 +96,22 @@ def build_main_agent(planner: MockLLM, sub_agent: Any) -> Any: state["output"].append(f"sub_agent: {payload}") elif channel == "user_input_channel": state["output"].append(f"user_input: {payload}") - # State changed -> ask planner what to do next. return Command(update=state, goto=Send("llm_node", None)) else: state["output"].append("timer: no updates yet") - # No meaningful state change -> keep waiting without calling planner. return Command(update=state, goto=Send("wait_node", None)) async def tool_node(ctx: Context, tool_input: str) -> None: await asyncio.sleep(0.1) - # Fire-and-forget style completion: publish result to inbox and exit. - # (i.e., just complete without explicitly going to a next node) ctx.publish_to_channel( "tool_completion_channel", f"tool completed for: {tool_input}", ) async def sub_agent_node(ctx: Context, sub_agent_input: str) -> None: - # Sub-agent remains a regular StateGraph, compiled independently. sub_agent_output = await sub_agent.ainvoke( {"input": sub_agent_input, "output": ""} ) - # Same pattern as tool node: publish result and complete current node. ctx.publish_to_channel( "subagent_completion_channel", sub_agent_output["output"], @@ -143,11 +125,9 @@ def build_main_agent(planner: MockLLM, sub_agent: Any) -> Any: } advanced_flow = AdvancedStateGraph(MainAgentState) - # Default behavior is an unbounded async channel like Rust channel 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) - # nodes are the same as in the regular StateGraph API advanced_flow.add_entry_node(llm_node) advanced_flow.add_node(wait_node) advanced_flow.add_node(tool_node) @@ -164,17 +144,12 @@ async def test_async_sub_graph() -> None: llm.responses = [ [ - # First planner pass triggers one slow sub-agent. Decision(type="sub_agent", sub_agent="research lunch options"), Decision(type="tool", tool="slack_tool"), ], - # After user input. [], - # After tool completion. [], - # After first sub-agent completion, planner decides to run second research. [Decision(type="sub_agent", sub_agent="find vegetarian fallback")], - # After second sub-agent completion, planner decides to end. [Decision(type="end", complete="order submitted")], ] @@ -182,7 +157,6 @@ async def test_async_sub_graph() -> None: {"input": "help me get something for lunch", "output": [], "done": None} ) - # External input can be injected while graph execution is in progress. await asyncio.sleep(0.01) await handler.apublish_to_channel("user_input_channel", "No spicy food please") result = await handler.aresult() @@ -212,6 +186,4 @@ async def test_async_sub_graph() -> None: order_food_idx = output.index("order_food: order submitted") assert first_sub_idx < second_sub_idx < order_food_idx assert llm._idx == len(llm.responses) - import json - print(json.dumps(result, ensure_ascii=False, indent=2)) diff --git a/libs/langgraph/tests/advanced-graph/test_update_elision.py b/saf-python-sdk/tests/advanced-graph/test_update_elision.py similarity index 93% rename from libs/langgraph/tests/advanced-graph/test_update_elision.py rename to saf-python-sdk/tests/advanced-graph/test_update_elision.py index a3597d366..453f88ecc 100644 --- a/libs/langgraph/tests/advanced-graph/test_update_elision.py +++ b/saf-python-sdk/tests/advanced-graph/test_update_elision.py @@ -1,12 +1,12 @@ import asyncio from dataclasses import dataclass + +import pytest from pydantic import BaseModel from typing_extensions import TypedDict -import pytest - -from langgraph.advanced_graph import AdvancedStateGraph, CompiledGraphEngine -from langgraph.types import Command, Send +from saf_python_sdk.advanced_graph import AdvancedStateGraph, CompiledGraphEngine +from saf_python_sdk.types import Command, Send pytestmark = pytest.mark.anyio @@ -46,7 +46,6 @@ def _initial_state() -> UpdateElisionState: ) - async def test_noop_slow_update_does_not_override_fast_update() -> None: graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState) @@ -65,7 +64,6 @@ async def test_noop_slow_update_does_not_override_fast_update() -> None: async def slow_node(state: UpdateElisionState) -> UpdateElisionState: await asyncio.sleep(0.1) - # Returns the same values as the initial snapshot. return state graph.add_entry_node(start_node) @@ -101,7 +99,6 @@ async def test_changed_slow_update_overrides_fast_update() -> None: async def slow_node(state: UpdateElisionState) -> UpdateElisionState: await asyncio.sleep(0.1) - # Slow node makes real changes for all field types. state.x = 2 state.dc.value = 2 state.model.value = 2 @@ -124,3 +121,4 @@ async def test_changed_slow_update_overrides_fast_update() -> None: assert result.td == {"flag": False, "n": 2} assert result.obj == {"n": 2} assert result.items == [0, 1, 2] +