Flows

import asyncio
from pydantic import BaseModel
from augments.adk import Flow, Runner, flow_start, flow_listen, flow_router
 
class ResearchState(BaseModel):
    topic: str = ""
    research: str = ""
 
class ResearchFlow(Flow[ResearchState]):
    state_factory = ResearchState  # explicit, never inferred from generic
 
    @flow_start
    async def kickoff(self) -> None:
        self.state.topic = "decentralized AI"
 
    @flow_listen(kickoff)
    async def research(self) -> None:
        # Build prompt from self.state and call an Agent inline:
        # result = await Runner.arun(agent, f"Research {self.state.topic}")
        # self.state.research = str(result.final_output)
        self.state.research = "synthetic finding"
 
    @flow_router(research)
    async def classify(self) -> str:
        return "high_risk" if "danger" in self.state.research else "low_risk"
 
    @flow_listen("high_risk")
    async def escalate(self) -> None: ...
 
    @flow_listen("low_risk")
    async def approve(self) -> None: ...
 
flow = ResearchFlow()
result = asyncio.run(Runner.arun_flow(flow))
print(result.final_state.research, result.status)

class MyFlow(Flow[State]):
    state_factory = State
 
    @flow_start
    async def a(self) -> None: ...
 
    @flow_start
    async def b(self) -> None: ...
 
    # OR gate — fires ONCE on first arrival; gate is consumed.
    @flow_listen(a | b)
    async def either(self) -> None: ...
 
    # AND gate — fires ONCE after BOTH complete.
    @flow_listen(a & b)
    async def both(self) -> None: ...
 
    # Chaining flattens left-associatively:
    @flow_listen(a | b | "external_label")
    async def any_of_three(self) -> None: ...

class State(BaseModel):
    counter: int = 0
 
# Option 1: state_factory class attribute (works without args)
class A(Flow[State]):
    state_factory = State
    @flow_start
    async def kickoff(self): ...
 
flow_a = A()  # state_factory() builds a fresh State()
 
# Option 2: explicit initial_state
class B(Flow[State]):
    @flow_start
    async def kickoff(self): ...
 
flow_b = B(initial_state=State(counter=42))

>>> flow
ResearchFlow(flow_id='flow-a1b2c3d4', steps=3, routers=1, state=ResearchState)
>>> result
FlowRunResult(flow_id='flow-a1b2c3d4', status='completed', steps=3, final_state=...)

>>> ResearchFlow.research.role       # "start" | "listen" | "router"
'listen'
>>> ResearchFlow.research.triggers   # trigger names declared by the decorator
('kickoff',)

# Async
result = await Runner.arun_flow(flow, config=FlowConfig(max_steps=200))
 
# Sync wrapper
result = Runner.run_flow(flow, config=FlowConfig(max_steps=200))
 
# Streaming events
streaming = Runner.arun_flow_streamed(flow)
async for event in streaming.stream_events():
    print(event.type, getattr(event, "step_name", ""))
print(streaming.status, streaming.final_state)
 
# Resume from checkpoint
result = await Runner.arun_flow_from_checkpoint(flow, checkpoint)

from augments.adk import FlowCheckpoint
 
# Capture (typically from a hook or after partial run)
checkpoint = FlowCheckpoint(
    flow_id=flow.flow_id,
    completed_steps=("step1", "step2"),
    pending_steps=("step3",),
    and_gate_arrivals={},
    consumed_gates=(),
    state_data=flow.state.model_dump_json(),
)
blob = checkpoint.to_json()  # store anywhere — DB, file, queue
 
# Resume in another process:
loaded = FlowCheckpoint.from_json(blob)  # raises on schema mismatch
state_back = MyState.model_validate_json(loaded.state_data)
flow_resumed = MyFlow(initial_state=state_back)
result = await Runner.arun_flow_from_checkpoint(flow_resumed, loaded)

from augments.adk.flows import Flow, FlowStepContext, flow_listen, flow_start
 
def needs_review(ctx: FlowStepContext[CartState]) -> bool:
    return ctx.flow_state.amount >= 1000
 
class RefundFlow(Flow[CartState]):
    @flow_start
    async def intake(self) -> None: ...
 
    @flow_listen("intake", requires_approval=needs_review)
    async def refund(self) -> None: ...
 
result = await Runner.arun_flow(flow)
if result.requires_action:
    checkpoint = result.checkpoint
    # Approve — by step name (or pass the FlowDeferredStep object itself):
    checkpoint.approve(
        "refund",
        approver_id="alice@example.com",
        approver_role="ops_lead",
        reason="amount within manual-review band",
    )
    # Or reject (message is routed to the @flow_listen(FLOW_ERROR_TRIGGER) handler):
    checkpoint.reject("refund", "fraud signal exceeded threshold")
    resumed = await Runner.arun_flow_from_checkpoint(flow_again, checkpoint)

from augments.adk.flows import FlowApprovalPolicy
 
class RefundFlow(Flow[CartState]):
    @flow_listen(
        "intake",
        requires_approval=needs_review,
        approval_policy=FlowApprovalPolicy(quorum=2, allowed_roles=("ops_lead",)),
    )
    async def refund(self) -> None: ...

from augments.adk.flows import (
    Flow, FlowStepCachePolicy, FlowStepGuardrails,
    FlowStepGuardrailVerdict, FlowStepRateLimit,
    flow_listen, flow_start,
)
 
class Pipeline(Flow[PipelineState]):
    @flow_start(rate_limit=FlowStepRateLimit(rpm=120))
    async def intake(self) -> None: ...
 
    @flow_listen(
        "intake",
        guardrails=FlowStepGuardrails(pre=(amount_under_threshold,)),
    )
    async def risk_check(self) -> None: ...
 
    @flow_listen(
        "risk_check",
        cache=FlowStepCachePolicy(cache_key_fn=lambda ctx: tier_of(ctx.flow_state)),
    )
    async def classify(self) -> None: ...

from augments.adk.flows import arun_flow_agent
 
class CartFlow(Flow[CartState]):
    @flow_start
    async def review(self) -> None:
        result = await arun_flow_agent(self, agent, "approve $5000", defer_key="purchase")
        # If the agent's tool deferred, arun_flow_agent raises
        # FlowAgentDeferred and the executor halts before this line.
        self.state.outcome = result.final_output

ds = result.deferred_steps[0]
resumed_state = RunState.from_dict(json.loads(ds.agent_run_state))
# resumed_state.approve(call) / .reject(call, message) — same shape as plain tool HITL
resumed = await Runner.arun_flow_from_checkpoint(
    flow_again,
    result.checkpoint,
    agent_resolutions={ds.defer_key: json.dumps(resumed_state.to_dict())},
)

results = await Runner.arun_flow_for_each(
    lambda state: SentimentFlow(state),
    [HeadlineState(headline=h) for h in headlines],
    concurrency=3,   # cap parallel runs; default 1 = sequential
)
for r in results:
    if r.status == "completed":
        print(r.final_state.label)
    else:
        print("failed:", r.error)

from augments.adk.flows import SqliteFlowWorkerBackend
from augments.adk.run.runner import Runner
 
backend = SqliteFlowWorkerBackend(path="/var/state/flows.db")
result = await Runner.arun_flow_distributed(
    flow, backend, worker_id="hostA-pid1234"
)

# Flow inside Graph (via FlowExecutable adapter)
from augments.adk.graphs import GraphBuilder
graph = (
    GraphBuilder.new("g")
    .node("flow_node", my_flow)  # auto-wraps via to_executable()
    .node("agent_node", my_agent)
    .edge("flow_node", "agent_node")
    .entry("flow_node")
    .terminal("agent_node")
    .compile()
)
 
# Graph / Swarm / Agent inside a Flow step body
class MyFlow(Flow[State]):
    state_factory = State
 
    @flow_start
    async def kickoff(self) -> None:
        # Run a Graph inside the step
        graph_result = await Runner.arun_graph(some_graph)
        self.state.intermediate = str(graph_result.final_output)