from google.adk.workflow import Workflow

# Inner workflow
inner = Workflow(
    name="inner_pipeline",
    edges=[
        ('START', step_a),
        (step_a, step_b),
    ],
)

# Outer workflow using inner as a node
outer = Workflow(
    name="outer_pipeline",
    edges=[
        ('START', pre_process),
        (pre_process, inner),      # Nested workflow
        (inner, post_process),
    ],
)

from google.adk.workflow import RetryConfig
from google.adk.workflow import FunctionNode

retry = RetryConfig(
    max_attempts=5,         # Max attempts (default: 5). 0 or 1 = no retry
    initial_delay=1.0,      # Seconds before first retry (default: 1.0)
    max_delay=60.0,         # Max seconds between retries (default: 60.0)
    backoff_factor=2.0,     # Delay multiplier per attempt (default: 2.0)
    jitter=1.0,             # Randomness factor (default: 1.0, 0.0 = none)
    exceptions=None,        # Exception types to retry (None = all)
)

node = FunctionNode(
    flaky_api_call,
    name="api_call",
    retry_config=retry,
)

delay = initial_delay * (backoff_factor ^ attempt)
delay = min(delay, max_delay)
delay = delay * (1 + random(0, jitter))

def my_node(ctx: Context, node_input: str) -> str:
  if ctx.retry_count > 0:
    print(f"Retry attempt {ctx.retry_count}")
  return "result"

from google.adk.workflow import BaseNode
from google.adk.events.event import Event
from google.adk.agents.context import Context
from pydantic import ConfigDict, Field
from typing import Any, AsyncGenerator
from typing_extensions import override

class BatchProcessorNode(BaseNode):
  """Processes items in batches."""
  model_config = ConfigDict(arbitrary_types_allowed=True)

  name: str = Field(default="batch_processor")
  batch_size: int = Field(default=10)

  def __init__(self, *, name: str = "batch_processor", batch_size: int = 10):
    super().__init__()
    object.__setattr__(self, 'name', name)
    object.__setattr__(self, 'batch_size', batch_size)

  @override
  def get_name(self) -> str:
    return self.name

  @override
  async def run(
      self,
      *,
      ctx: Context,
      node_input: Any,
  ) -> AsyncGenerator[Any, None]:
    items = node_input if isinstance(node_input, list) else [node_input]
    results = []
    for i in range(0, len(items), self.batch_size):
      batch = items[i:i + self.batch_size]
      batch_result = await process_batch(batch)
      results.extend(batch_result)
    yield Event(output=results)

from google.adk.workflow import BaseNode

class CollectorNode(BaseNode):
  wait_for_output: bool = True  # Stay in WAITING until output is yielded

  async def run(self, *, ctx, node_input):
    # Store partial input, don't yield output yet
    collected = ctx.state.get("collected", [])
    collected.append(node_input)
    yield Event(state={"collected": collected})

    # Only yield output when we have enough
    if len(collected) >= 3:
      yield Event(output=collected)
      # Now node transitions to COMPLETED and triggers downstream

from google.adk.workflow._tool_node import _ToolNode as ToolNode
from google.adk.tools.function_tool import FunctionTool

def search(query: str) -> str:
  """Search for information."""
  return f"Results for: {query}"

tool = FunctionTool(search)
tool_node = ToolNode(tool, name="search_node")

agent = Workflow(
    name="with_tool",
    edges=[
        ('START', prepare_query),
        (prepare_query, tool_node),  # Input must be dict (tool args) or None
        (tool_node, process_results),
    ],
)

from google.adk.workflow._agent_node import AgentNode
from google.adk.agents.loop_agent import LoopAgent

loop = LoopAgent(
    name="refine_loop",
    sub_agents=[writer, reviewer],
    max_iterations=3,
)

loop_node = AgentNode(agent=loop, name="refinement")

agent = Workflow(
    name="with_loop",
    edges=[
        ('START', loop_node),
        (loop_node, final_step),
    ],
)

from google.adk.workflow import Edge
from google.adk.workflow._workflow_graph import WorkflowGraph

# Tuple syntax (most common)
edges = [
    ('START', node_a),                    # Simple edge
    (node_a, node_b, "route"),            # Routed edge
    (node_a, (node_b, node_c)),           # Fan-out
    ((node_b, node_c), join_node),        # Fan-in
]

# Sequence shorthand (tuple with 3+ elements creates chain)
edges = [('START', node_a, node_b, node_c)]
# Equivalent to: [('START', node_a), (node_a, node_b), (node_b, node_c)]

# Routing map (dict syntax)
edges = [
    (classifier, {"success": handler_a, "error": handler_b}),
]

# Edge objects (explicit)
edges = [
    Edge(START, node_a),
    Edge(node_a, node_b, route="success"),
]

# Edge.chain helper
edges = Edge.chain('START', node_a, node_b, node_c)
# Returns: [(START, node_a), (node_a, node_b), (node_b, node_c)]

# WorkflowGraph.from_edge_items
graph = WorkflowGraph.from_edge_items([
    ('START', node_a),
    (node_a, node_b),
])
agent = Workflow(name="my_workflow", graph=graph)