Hierarchical task decomposition

adk-task-decomposition-lesson/
  requirements.txt
  .env.example
  docs/
  harness.py
  run.py
  upgrade_guide.py
  historical_upgrade_guide.py
  historical-docs/
  guide-run.json
  test_upgrade_guide.py
  test_runner.py

# Changelog

## 3.6.0
- Add `docuctl preview --port` to choose the preview server port.
- BREAKING: Remove the `--legacy-nav` flag. Navigation now comes only from `nav.toml`.

## 3.5.0
- Build pages in parallel; set the worker count with `--jobs`.
- BREAKING: `docuctl build` no longer reads `docuctl.yaml`; rename it to `docuctl.toml`.

"""Write an upgrade guide with a three-level agent tree: plan, write sections, research docs."""
import re
from pathlib import Path

from google.adk import Agent, Event, Workflow
from google.genai import types
from google.adk.plugins import BasePlugin
from google.adk.tools import AgentTool

DOCS = Path(__file__).with_name("docs")
SOURCE = re.compile(r"^Source: (\S+\.md)$", re.M)


def search_docs(query: str) -> list[dict]:
    """Find docs pages that contain every word of the query. Returns path and title."""
    words = query.lower().split()
    hits = []
    for page in sorted(DOCS.rglob("*.md")):
        text = page.read_text()
        if all(word in text.lower() for word in words):
            hits.append({"path": str(page.relative_to(DOCS)), "title": text.splitlines()[0][2:]})
    return hits


def read_page(path: str) -> dict:
    """Return the text of one docs page, by the path that search_docs returned."""
    page = (DOCS / path).resolve()
    if not page.is_relative_to(DOCS.resolve()) or not page.is_file():
        return {"error": f"no docs page at {path}"}
    return {"path": path, "text": page.read_text()}


# One entry per level, top first. The depth limit cuts the tree: the last level
# that is built gets the docs tools instead of a delegation tool.
LEVELS = [
    ("plan_guide",
     "You write upgrade guides for docuctl, a docs-site CLI. Split the goal into one "
     "section per breaking change, and ask write_section to write each one. Then return "
     "the sections in version order under one '# ' title. Do not write sections yourself."),
    ("write_section",
     "Write one section of an upgrade guide: a '## ' heading, the steps a user takes, and a "
     "last line 'Source: <path>' naming the docs page the steps come from. Get every fact "
     "from research_docs; never use your own knowledge of docuctl."),
    ("research_docs",
     "Answer the question from the docuctl docs only. Use search_docs with one or two "
     "words, then read_page. Quote exact commands and the page path."),
]
MAX_DEPTH = len(LEVELS)


def build_tree(model, max_depth: int = MAX_DEPTH, level: int = 0) -> Agent:
    name, instruction = LEVELS[level]
    if level + 1 < min(max_depth, len(LEVELS)):
        tools = [AgentTool(build_tree(model, max_depth, level + 1))]
    else:
        tools = [search_docs, read_page]  # a leaf: it can only read, not delegate
        if name == "write_section":
            instruction = instruction.replace(
                "Get every fact from research_docs;",
                "Use search_docs with one or two words, then read_page. "
                "Get every fact from those pages;")
    return Agent(name=name, model=model,
                       generate_content_config=types.GenerateContentConfig(
                           max_output_tokens=1024,
                           thinking_config=types.ThinkingConfig(thinking_level=types.ThinkingLevel.LOW)), instruction=instruction, tools=tools,
                 description=instruction.split(".")[0])


def check_sources(node_input: str):
    """Function node: every section must cite a docs page that exists."""
    sections = re.split(r"(?m)^## [^\n]+\n", node_input)[1:]
    if not sections:
        raise ValueError("guide needs at least one section")
    cited = []
    for section in sections:
        lines = section.strip().splitlines()
        sources = SOURCE.findall(section)
        if len(sources) != 1 or not SOURCE.fullmatch(lines[-1]):
            raise ValueError("each section must end with one Source: <path> line")
        cited.append(sources[0])
    root = DOCS.resolve()
    missing = [path for path in cited
               if not (page := (root / path).resolve()).is_relative_to(root)
               or not page.is_file()]
    if missing:
        raise ValueError(f"missing pages: {missing}")
    return Event(message=node_input)


class DelegationLog(BasePlugin):
    """Records every agent start and tool call in the tree, with the agent's depth.
    AgentTool passes the runner's plugins to the nested runs (include_plugins=True)."""

    def __init__(self, max_model_calls: int | None = None):
        super().__init__(name="delegation_log")
        self.rows: list[dict] = []
        self.model_calls = 0
        self.max_model_calls = max_model_calls

    async def before_agent_callback(self, *, agent, callback_context):
        depth = [name for name, _ in LEVELS].index(agent.name) + 1
        self.rows.append({"depth": depth, "agent": agent.name})

    async def before_tool_callback(self, *, tool, tool_args, tool_context):
        self.rows.append({"agent": tool_context.agent_name, "calls": tool.name, "args": tool_args})

    async def before_model_callback(self, *, callback_context, llm_request):
        # One budget for the whole tree: each nested run counts its own max_llm_calls.
        self.model_calls += 1
        if self.max_model_calls is not None and self.model_calls > self.max_model_calls:
            raise RuntimeError(f"tree-wide model call budget of {self.max_model_calls} exceeded")


def build_workflow(model=None, max_depth: int = MAX_DEPTH) -> Workflow:
    model = model or "gemini-3.8-flash"  # or LiteLlm(model="ollama_chat/gemma4:e2b")
    return Workflow(name="upgrade_guide",
                    edges=[("START", build_tree(model, max_depth), check_sources)])

"""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 upgrade_guide import build_workflow, DelegationLog, check_sources
from historical_upgrade_guide import build_workflow as historical_workflow

async def compare(*, live=False):
    record = json.loads((HERE / "guide-run.json").read_text())
    saved = next(r for r in record["runs"] if (r["max_depth"], r["trial"]) == (2, 1))
    for budget in (64, 1):
        model = None if live else replay(saved, record)
        log = DelegationLog(max_model_calls=budget)
        try:
            builder = build_workflow if live else historical_workflow
            rows = await invoke(builder(model, max_depth=2), (record["input"].replace("shipctl", "docuctl") if live else record["input"]),
                                max_llm_calls=64, plugins=[log])
        except RuntimeError as error:
            if "tree-wide model call budget" not in str(error):
                raise
            print(f"\nbudget: {budget}\n{type(error).__name__}: {error}")
        else:
            display(f"depth: 2 | budget: {budget}", rows)
    # A file-existence check is independent of model generation.
    try:
        check_sources("# Guide\n\n## Upgrade\nRun docuctl migrate.\nSource: missing.md")
    except ValueError as error:
        print(f"\nmissing citation\n{type(error).__name__}: {error}")

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"])
    messages = [row["message"] for row in rows if row.get("message")]
    if messages:
        print(messages[-1])


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 the historical shipctl workflow and saved replies")
    asyncio.run(compare(live=args.live))


if __name__ == "__main__":
    main()