Parallel fan-out and gather

adk-parallel-gather-lesson/
  requirements.txt
  .env.example
  release-3.5.0.md
  release-3.6.0.md
  site-3.5.0/
  harness.py
  run.py
  release_prep.py
  prep-run.json
  test_release_prep.py
  test_runner.py

## 3.5.0
- Build pages in parallel; set the worker count with `--jobs`. See /guide/parallel-builds/.
- BREAKING: `docuctl build` no longer reads `docuctl.yaml`; rename it to `docuctl.toml`. See /reference/config/.
- Fix the search index missing pages whose title contains a colon. See /reference/search/.
sitemap: site-3.5.0/sitemap.txt

"""Prepare a release page from three independent branches that run at the same time."""
import asyncio
import re
from pathlib import Path

from google.adk import Agent, Event, Workflow
from google.adk.workflow import JoinNode, RetryConfig, node
from google.genai import types
from pydantic import BaseModel

HERE = Path(__file__).parent
PAGE = re.compile(r"See (/\S*/)\.")


class Release(BaseModel):
    version: str
    entries: list[str]
    sitemap: str


class Highlights(BaseModel):
    bullets: list[str]


class UpgradeSteps(BaseModel):
    steps: list[str]


class LinkReport(BaseModel):
    missing: list[str]


def load_release(node_input: str):
    """Function node: parse the release file the run starts from."""
    lines = node_input.strip().splitlines()
    entries = [line[2:] for line in lines if line.startswith("- ")]
    sitemap = next(line[8:].strip() for line in lines if line.startswith("sitemap:"))
    return Event(output=Release(version=lines[0][3:], entries=entries, sitemap=sitemap),
                 state={"version": lines[0][3:]})


def check_links(node_input: Release):
    """Function node: every page an entry points to must be in the built site's sitemap.
    It raises if the site has not been built, because then nothing can be checked."""
    sitemap = HERE / node_input.sitemap
    if not sitemap.is_file():
        raise FileNotFoundError(f"{node_input.sitemap} not found: build the site first")
    built = set(sitemap.read_text().split())
    wanted = [page for entry in node_input.entries for page in PAGE.findall(entry)]
    return LinkReport(missing=[page for page in wanted if page not in built])


def assemble(node_input: dict, version: str):
    """Function node after the join: node_input maps each branch name to its output."""
    lines = [f"# docuctl {version} release notes", ""]
    lines += [f"- {b}" for b in node_input["write_highlights"]["bullets"]]
    lines += ["", "## Upgrade", *(f"1. {s}" for s in node_input["write_upgrade_steps"]["steps"])]
    links = node_input.get("check_links")
    if links is None:
        lines += ["", "Link check: not run (see the error above)."]
    elif links["missing"]:
        lines += ["", "Link check: missing " + ", ".join(links["missing"])]
    else:
        lines += ["", "Link check: every linked page exists."]
    return Event(message="\n".join(lines))


def make_agents(model):
    model = model or "gemini-3.8-flash"  # or LiteLlm(model="ollama_chat/gemma4:e2b")
    retry = RetryConfig(max_attempts=3, initial_delay=1.0)
    highlights = Agent(name="write_highlights", model=model,
                       generate_content_config=types.GenerateContentConfig(
                           max_output_tokens=1024,
                           thinking_config=types.ThinkingConfig(thinking_level=types.ThinkingLevel.LOW)), input_schema=Release,
                       output_schema=Highlights,
                       instruction="Write one short user-facing bullet per entry. No marketing.")
    upgrade = Agent(name="write_upgrade_steps", model=model,
                       generate_content_config=types.GenerateContentConfig(
                           max_output_tokens=1024,
                           thinking_config=types.ThinkingConfig(thinking_level=types.ThinkingLevel.LOW)), input_schema=Release,
                    output_schema=UpgradeSteps,
                    instruction=("Write the steps a user must take to upgrade, one per entry "
                                 "marked BREAKING. Return an empty list if there are none."))
    return node(highlights, retry_config=retry), node(upgrade, retry_config=retry)


def build_graph(model=None, max_concurrency=None) -> Workflow:
    """Graph fan-out: three rows from load_release, gathered by a JoinNode. Fail-all."""
    highlights, upgrade = make_agents(model)
    gather = JoinNode(name="gather")
    return Workflow(
        name="release_prep",
        max_concurrency=max_concurrency,
        edges=[
            ("START", load_release),
            (load_release, highlights, gather),
            (load_release, upgrade, gather),
            (load_release, check_links, gather),
            (gather, assemble),
        ],
    )


def build_partial(model=None) -> Workflow:
    """Dynamic fan-out: the same branches, but the link check may fail without
    failing the release. Required branches still stop the run."""
    highlights, upgrade = make_agents(model)

    @node(rerun_on_resume=True)
    async def prepare(ctx, node_input: Release):
        branches = {"write_highlights": highlights, "write_upgrade_steps": upgrade,
                    "check_links": check_links}
        results = await asyncio.gather(
            *(ctx.run_node(branch, node_input) for branch in branches.values()),
            return_exceptions=True)
        outputs = dict(zip(branches, results))
        error = outputs["check_links"]  # DynamicNodeFailError wraps the child's exception
        for name in ("write_highlights", "write_upgrade_steps"):
            if isinstance(outputs[name], BaseException):
                raise outputs[name]  # a required branch: no release page without it
        if isinstance(error, BaseException):
            outputs["check_links"] = None
            cause = getattr(error, "error", error)
            yield Event(message=f"{error.error_node_path} failed: {type(cause).__name__}: {cause}")
        yield Event(output={k: v.model_dump() if hasattr(v, "model_dump") else v
                            for k, v in outputs.items()})

    return Workflow(name="release_prep_partial",
                    edges=[("START", load_release, prepare, assemble)])

"""Compare the workflow's cases; replay saved replies unless --live is explicit."""
import argparse
import asyncio
import copy
import json
from pathlib import Path

from dotenv import load_dotenv
import os
from google.adk.agents.run_config import RunConfig
from google.adk.apps import App
from google.adk.runners import InMemoryRunner
from google.genai import types
from harness import ReplayLlm, summarize

HERE = Path(__file__).resolve().parent
from release_prep import build_graph, build_partial

async def compare(*, live=False):
    record = json.loads((HERE / "prep-run.json").read_text())
    for kind, filename in (("graph", "release-3.5.0.md"), ("graph", "release-3.6.0.md"),
                           ("partial", "release-3.6.0.md")):
        saved = next(r for r in record["runs"] if (r["workflow"], r["release"]) == (kind, filename))
        model = None if live else replay(saved, record)
        if live:
            builder = build_graph if kind == "graph" else build_partial
        else:
            original = __import__("importlib").import_module("historical-shipctl-source.release_prep")
            builder = original.build_graph if kind == "graph" else original.build_partial
        root = builder(model)
        try:
            rows = await invoke(root, (HERE / ("" if live else "historical-shipctl-source") / filename).read_text(), max_llm_calls=6)
        except FileNotFoundError as error:
            print(f"\n{kind} | {filename}\n{type(error).__name__}: {error}")
        else:
            display(f"{kind} | {filename}", rows)

async def invoke(root, text, *, max_llm_calls, plugins=()):
    """Run the actual workflow and close its in-memory Runner, including on failure."""
    app = App(name="docs_assistant", root_agent=root, plugins=list(plugins))
    runner = InMemoryRunner(app=app)
    session = await runner.session_service.create_session(app_name=runner.app_name, user_id="dev")
    try:
        return [summarize(event) async for event in runner.run_async(
            user_id="dev", session_id=session.id,
            new_message=types.Content(role="user", parts=[types.Part(text=text)]),
            run_config=RunConfig(max_llm_calls=max_llm_calls))]
    finally:
        await runner.close()


def replay(saved, record):
    return ReplayLlm(model=record["environment"]["model"],
                     calls=copy.deepcopy(saved["model_calls"]))


def display(label, rows):
    print(f"\n{label}")
    for row in rows:
        if row.get("route"):
            print("route:", row["route"])
    # The parent may emit a diagnostic before the assembled page.
    # Keep both so a partial page's "see the error above" has its actual cause.
    for row in rows:
        if row.get("message"):
            print(row["message"])


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--live", action="store_true", help="request new Gemini replies instead of replaying saved replies")
    args = parser.parse_args()
    if args.live:
        load_dotenv(HERE / ".env", override=False)
        if not (os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")):
            parser.error("Set your Gemini API key in .env before using --live")
    print("Mode:", "live Gemini requests" if args.live else "offline replay of saved model replies")
    asyncio.run(compare(live=args.live))


if __name__ == "__main__":
    main()