publish

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")

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))
+

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 = "." }