diff --git a/saf-python-sdk/python/saf_python_sdk/langgraph_rust_core.cpython-313-darwin.so b/saf-python-sdk/python/saf_python_sdk/langgraph_rust_core.cpython-313-darwin.so new file mode 100755 index 000000000..9e3fd37d3 Binary files /dev/null and b/saf-python-sdk/python/saf_python_sdk/langgraph_rust_core.cpython-313-darwin.so differ 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 f0675e52f..01bf37ac0 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 @@ -15,14 +15,7 @@ from saf_python_sdk.advanced_graph import ( 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 +from langgraph.graph import END, START, StateGraph RUNS = 100 MIDDLE_COUNT = 10 @@ -95,71 +88,65 @@ def build_advanced_sequential() -> Any: return graph.compile() -def _build_stategraph_suites() -> list[tuple[str, Any]]: - if not HAS_STATEGRAPH: - return [] +def build_stategraph_parallel() -> Any: + graph = StateGraph(dict) - def build_stategraph_parallel() -> Any: - graph = StateGraph(dict) + async def start_node(state: dict[str, Any]) -> None: + _ = state - async def start_node(state: dict[str, Any]) -> None: + 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) - async def end_node(state: dict[str, Any]) -> dict[str, Any]: - out = dict(state) - out["done"] = True - return out + graph.add_node(f"middle_{i}", middle_node) + graph.add_node("end_node", end_node) - graph.add_node("start_node", start_node) - for i in range(MIDDLE_COUNT): + 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() - 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) - def build_stategraph_sequential() -> Any: - graph = StateGraph(dict) + async def start_node(state: dict[str, Any]) -> None: + _ = state - async def start_node(state: dict[str, Any]) -> None: + 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) - async def end_node(state: dict[str, Any]) -> dict[str, Any]: - out = dict(state) - out["done"] = True - return out + graph.add_node(f"middle_{i}", middle_node) + graph.add_node("end_node", end_node) - 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()), - ] + 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() async def run_benchmark(name: str, compiled: Any) -> float: @@ -176,14 +163,13 @@ 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()), ] - 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") diff --git a/saf-python-sdk/tests/advanced-graph/test_sub_agents.py b/saf-python-sdk/tests/advanced-graph/test_sub_agents.py index 0d49cafcb..73cc0ea57 100644 --- a/saf-python-sdk/tests/advanced-graph/test_sub_agents.py +++ b/saf-python-sdk/tests/advanced-graph/test_sub_agents.py @@ -1,4 +1,5 @@ import asyncio +import json from dataclasses import dataclass from typing import Any, Literal @@ -12,6 +13,8 @@ from saf_python_sdk.advanced_graph import ( channel_condition, timer_condition, ) +from langgraph.constants import END, START +from langgraph.graph import StateGraph from saf_python_sdk.types import Command, Send pytestmark = pytest.mark.anyio @@ -49,22 +52,30 @@ class MockLLM: return response -class FakeSubAgent: - async def ainvoke(self, state: SubAgentState) -> SubAgentState: - await asyncio.sleep(5) - return {"input": state["input"], "output": f"research sub agent completed for: {state['input']}"} - - def build_sub_agent() -> Any: - return FakeSubAgent() + # 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. + await asyncio.sleep(5) + return {"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_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, @@ -75,10 +86,12 @@ 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"), @@ -96,22 +109,28 @@ 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"], @@ -125,9 +144,11 @@ 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) @@ -144,12 +165,17 @@ 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")], ] @@ -157,6 +183,7 @@ 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() @@ -187,3 +214,5 @@ async def test_async_sub_graph() -> None: assert first_sub_idx < second_sub_idx < order_food_idx assert llm._idx == len(llm.responses) + print(json.dumps(result, ensure_ascii=False, indent=2)) + diff --git a/saf-python-sdk/uv.lock b/saf-python-sdk/uv.lock new file mode 100644 index 000000000..ea0ad8d69 --- /dev/null +++ b/saf-python-sdk/uv.lock @@ -0,0 +1,8 @@ +version = 1 +revision = 3 +requires-python = ">=3.10" + +[[package]] +name = "saf-python-sdk" +version = "0.1.1" +source = { editable = "." }