diff --git a/contributing/samples/workflows/fan_out_preempt/README.md b/contributing/samples/workflows/fan_out_preempt/README.md new file mode 100644 index 00000000000..e455bbc42ce --- /dev/null +++ b/contributing/samples/workflows/fan_out_preempt/README.md @@ -0,0 +1,72 @@ +# Parallel Fan-Out with Mid-Stream Preemption + +## Overview + +Read several data sources **in parallel**, and **stop reading any source the +moment it becomes apparent it's irrelevant** — don't stream 100 pages of +analysis when the model realizes on "page 3" that the source doesn't apply. + +Each source is analyzed by its own streaming agent wrapped in a +`StreamingRouterNode`. A `monitor` predicate watches that branch's token +stream; the instant the source declares itself irrelevant, the node commits an +`{"relevant": false}` output and **cancels the rest of that generation** via +ADK's cooperative `aclosing` cancellation. The other branches are untouched. + +A `JoinNode` fans the results back in and a synthesizer drops the irrelevant +ones. + +## Why it's actually parallel (and fast) + +- The workflow scheduler launches every fan-out branch as its own + `asyncio.create_task` and awaits them together, so the reads truly overlap. +- `enable_uvloop()` puts them on a libuv event loop. Install with + `pip install "google-adk[uvloop]"` (or set `ADK_UVLOOP=1`). +- Preemption saves wall-clock by killing the tail of an irrelevant read + instead of paying for tokens nobody uses. + +## How the preemption works + +```python +def monitor(view: StreamView) -> StreamDecision | None: + if view.text.lstrip().upper().startswith("IRRELEVANT"): + return StreamDecision(output={"source": source, "relevant": False}) + return None # keep reading + +StreamingRouterNode(name=f"reader_{source}", agent=reader, monitor=monitor, + timeout=60) +``` + +- Returning a `StreamDecision` commits the output and (by default) cancels the + remaining generation for that branch. +- Returning `None` keeps streaming; a relevant read finishes normally and its + final text becomes the branch's output. +- `timeout=60` is a hard cap so a stuck/slow source can never hold up the join. + +## Graph + +```mermaid +graph TD + START --> stash_query + stash_query --> reader_sharepoint + stash_query --> reader_havian + stash_query --> reader_wiki + stash_query --> reader_crm + stash_query --> reader_docs + reader_sharepoint --> join_sources + reader_havian --> join_sources + reader_wiki --> join_sources + reader_crm --> join_sources + reader_docs --> join_sources + join_sources --> synthesize +``` + +## Notes and limits + +- This cancels the **LLM analysis stream**. If the expensive work is the data + **fetch** itself (a network/tool call), that tool must be async and + cooperatively cancellable for the cancel to unwind it; consider a cheap + relevance pre-scan before the deep read. +- `JoinNode` waits for **all** branches. Preemption just makes the irrelevant + ones return sooner. To proceed the moment the relevant subset is in (dropping + irrelevant branches from the join entirely), you'd need a custom any/first-N + join. diff --git a/contributing/samples/workflows/fan_out_preempt/agent.py b/contributing/samples/workflows/fan_out_preempt/agent.py new file mode 100644 index 00000000000..e2e3e72e02d --- /dev/null +++ b/contributing/samples/workflows/fan_out_preempt/agent.py @@ -0,0 +1,103 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Parallel fan-out with per-branch mid-stream preemption + libuv. + +Reads five data sources concurrently. Each source is analyzed by its own +streaming agent. The instant an agent's stream reveals the source is +irrelevant to the query, that branch is preempted -- the rest of its +generation is cancelled -- while the other branches keep running. A JoinNode +fans the results back in and a synthesizer drops the irrelevant ones. + +Each branch runs as its own asyncio task, so the reads truly overlap; +``enable_uvloop()`` puts them on a faster libuv loop. +""" + +from typing import Any + +from google.adk import Agent +from google.adk import enable_uvloop +from google.adk import Event +from google.adk import Workflow +from google.adk.workflow import JoinNode +from google.adk.workflow import StreamDecision +from google.adk.workflow import StreamingRouterNode +from google.adk.workflow import StreamView + +enable_uvloop() + +SOURCES = ('sharepoint', 'havian', 'wiki', 'crm', 'docs') + +# Sentinel the reader emits as its first line when a source does not apply. +IRRELEVANT = 'IRRELEVANT' + + +def stash_query(node_input: str): + """Puts the user query in state so every reader can template it in.""" + yield Event(state={'query': node_input}) + + +def _make_reader(source: str) -> StreamingRouterNode: + reader = Agent( + name=f'read_{source}', + instruction=( + f'You are reading the "{source}" source to answer: {{query}}.\n' + 'If this source is clearly irrelevant to the query, your FIRST' + f' line must be exactly "{IRRELEVANT}". Otherwise, extract only the' + ' facts from this source that help answer the query.' + ), + output_key=f'{source}_result', + ) + + def monitor(view: StreamView) -> StreamDecision | None: + # As soon as the model declares irrelevance, stop reading this source. + if view.text.lstrip().upper().startswith(IRRELEVANT): + return StreamDecision(output={'source': source, 'relevant': False}) + # Otherwise keep streaming; a relevant read finishes normally and its + # final text becomes this branch's output. + return None + + return StreamingRouterNode( + name=f'reader_{source}', + agent=reader, + monitor=monitor, + # Bound the deep read so a stuck/slow source can never hold up the join. + timeout=60, + ) + + +readers = tuple(_make_reader(source) for source in SOURCES) +join_sources = JoinNode(name='join_sources') + + +async def synthesize(node_input: dict[str, Any]): + """Fan-in: drop the branches that preempted as irrelevant, then answer.""" + relevant = { + name: result + for name, result in node_input.items() + if not (isinstance(result, dict) and result.get('relevant') is False) + } + skipped = sorted(set(node_input) - set(relevant)) + yield Event( + message=( + f'Answer synthesized from {sorted(relevant)}.\n' + f'Skipped (irrelevant, preempted mid-read): {skipped}.' + ), + ) + + +root_agent = Workflow( + name='root_agent', + edges=[('START', stash_query, readers, join_sources, synthesize)], +) diff --git a/contributing/samples/workflows/search_fanout_first_answer/README.md b/contributing/samples/workflows/search_fanout_first_answer/README.md new file mode 100644 index 00000000000..a59b18acbe3 --- /dev/null +++ b/contributing/samples/workflows/search_fanout_first_answer/README.md @@ -0,0 +1,75 @@ +# Search Fan-Out: Answer From the First Source That Has It + +## Overview + +The enterprise-search shape (Recall@k): a keyword search returns the top-k +candidate sources — SharePoint, Confluence, CRM, a drive, email — and you don't +know which one holds the answer. So you read them **in parallel** and answer the +moment *any* one of them does, without waiting on (or paying to finish) the rest. + +This sample combines two mechanisms: + +1. **Per-branch preemption** — each source is read by a `StreamingRouterNode` + whose `monitor` stops that read as soon as the model can say "answer is here" + or "not here". A single irrelevant source never streams to the end. +2. **Cross-branch first-answer-wins** — `FirstMatchNode` races the branches and, + the instant one returns `found=True`, **cancels the still-running siblings** + (tearing down their in-flight model calls) and returns that answer. + +`FirstMatchNode` is the complement of `JoinNode`: `JoinNode` waits for **all** +predecessors; `FirstMatchNode` returns the **first** matching one and cancels the +losers. + +## How the race works + +```python +first_answer = FirstMatchNode( + name="first_answer", + nodes=[reader_sharepoint, reader_confluence, reader_crm, ...], + match=lambda r: isinstance(r, dict) and r.get("found"), + no_match_output={"found": False, "answer": "Not found in any source."}, +) +``` + +- Every branch is launched concurrently and handed the same input. +- The first result the `match` predicate accepts wins; the rest are cancelled and + awaited (so their reads are actually torn down before the graph advances). +- A branch that *fails* is logged and skipped, so one flaky source can't deny an + answer another source can give. +- If no branch matches, the node yields `no_match_output`. + +## Why the cancellation is real + +Cancelling a branch propagates cooperatively: the `FirstMatchNode` cancels the +branch's `asyncio` task → the dynamic scheduler's `await` on the child unwinds → +`CancelledError` reaches the `StreamingRouterNode`'s `Aclosing` block → the SSE +model call is closed. The loser stops **decoding** immediately. + +`enable_uvloop()` puts the concurrent reads on a libuv loop. Install with +`pip install "google-adk[uvloop]"` (or set `ADK_UVLOOP=1`). + +## What it saves — and what it doesn't + +- **Saves:** the losers' output generation and the wall-clock of waiting on the + slowest branch. You answer at ~the speed of the fastest source that has it. +- **Does not save:** the input already read. A branch that started still paid to + *prefill* its source. Preemption/racing is a **decode-side** win. +- **To also cut prefill:** pass sources in rank order and set `max_parallel` + (e.g. `2`) so an early win short-circuits before lower-ranked sources are ever + read. With `max_parallel=1` it degrades to a cheap sequential gate. + +## Graph + +```mermaid +graph TD + START --> stash_query + stash_query --> first_answer + first_answer --> respond + subgraph first_answer [FirstMatchNode: race, first-wins, cancel losers] + reader_sharepoint + reader_confluence + reader_crm + reader_gdrive + reader_email + end +``` diff --git a/contributing/samples/workflows/search_fanout_first_answer/agent.py b/contributing/samples/workflows/search_fanout_first_answer/agent.py new file mode 100644 index 00000000000..ffed64bc17e --- /dev/null +++ b/contributing/samples/workflows/search_fanout_first_answer/agent.py @@ -0,0 +1,130 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Search fan-out that answers from the first source that has it (Recall@k). + +The classic enterprise-search shape: a keyword search returns the top-k candidate +sources (SharePoint, Confluence, CRM, a drive, email...). You don't know which +one holds the answer, so you read them **in parallel** -- and you want to answer +the moment *any* one of them does, without waiting on (or paying to finish) the +rest. + +Two mechanisms combine here: + + 1. Per-branch preemption -- each source is read by a ``StreamingRouterNode`` + whose ``monitor`` stops that read as soon as the model can say "answer is + here" or "not here", so a single irrelevant source never streams to the end. + 2. Cross-branch first-answer-wins -- ``FirstMatchNode`` races the branches and, + the instant one returns ``found=True``, cancels the still-running siblings + (tearing down their in-flight model calls) and returns that answer. + +``enable_uvloop()`` puts the concurrent reads on a libuv loop. + +Honest limit: branches that already started still paid to *read* their source +(prefill). This saves the losers' *generation* and the wall-clock of waiting on +them. To also avoid reading low-ranked sources, pass them in rank order and set +``max_parallel`` so a win short-circuits before the tail is ever started. +""" + +from google.adk import Agent +from google.adk import enable_uvloop +from google.adk import Event +from google.adk import Workflow +from google.adk.workflow import FirstMatchNode +from google.adk.workflow import StreamDecision +from google.adk.workflow import StreamingRouterNode +from google.adk.workflow import StreamView + +enable_uvloop() + +# Top-k candidates as if returned by a keyword search, in rank order. +SOURCES = ('sharepoint', 'confluence', 'crm', 'gdrive', 'email') + +_FOUND = 'FOUND:' +_NOT_HERE = 'NOTHERE' + + +def stash_query(node_input: str): + """Puts the user query in state so every reader can template it in.""" + yield Event(state={'query': node_input}) + + +def _make_reader(source: str) -> StreamingRouterNode: + reader = Agent( + name=f'read_{source}', + instruction=( + f'You are searching the "{source}" source to answer: {{query}}.\n' + 'Read only as far as you must. As soon as you can tell, output ONE' + ' line and nothing after it:\n' + f' "{_FOUND} " if this source answers the question,' + ' or\n' + f' "{_NOT_HERE}" if it clearly does not.' + ), + output_key=f'{source}_answer', + ) + + def monitor(view: StreamView) -> StreamDecision | None: + upper = view.text.upper() + idx = upper.find(_FOUND) + if idx != -1: + line, sep, _ = view.text[idx + len(_FOUND) :].partition('\n') + if sep or line.strip(): + # Answer found: commit it and preempt this branch's generation. + return StreamDecision( + output={'found': True, 'source': source, 'answer': line.strip()} + ) + if _NOT_HERE in upper: + # Source is irrelevant: stop reading this one early. + return StreamDecision(output={'found': False, 'source': source}) + return None + + return StreamingRouterNode( + name=f'reader_{source}', + agent=reader, + monitor=monitor, + # Hard cap so a stuck/slow source can never hold up the race. + timeout=60, + ) + + +# The complement of JoinNode: race the readers, return the first that answers, +# and cancel the losers mid-read. +first_answer = FirstMatchNode( + name='first_answer', + nodes=[_make_reader(source) for source in SOURCES], + match=lambda r: isinstance(r, dict) and r.get('found'), + no_match_output={ + 'found': False, + 'answer': 'Not found in any of the top sources.', + }, + # Read all k at once. For rank-ordered prefill savings, set e.g. + # max_parallel=2 so a win short-circuits before lower ranks are read. +) + + +async def respond(node_input: dict): + if node_input and node_input.get('found'): + yield Event( + message=( + f"Answer (from {node_input['source']}): {node_input['answer']}" + ) + ) + else: + yield Event(message='No source contained the answer.') + + +root_agent = Workflow( + name='root_agent', + edges=[('START', stash_query, first_answer, respond)], +) diff --git a/contributing/samples/workflows/speculative_tool/README.md b/contributing/samples/workflows/speculative_tool/README.md new file mode 100644 index 00000000000..f7495b30ae0 --- /dev/null +++ b/contributing/samples/workflows/speculative_tool/README.md @@ -0,0 +1,58 @@ +# Speculative Tool Dispatch + +## Overview + +Don't wait for the model to finish asking — start the work the moment you can +guess what it's asking for, then verify. + +`SpeculativeRouterNode` is the aggressive counterpart to `StreamingRouterNode`: + +| | `StreamingRouterNode` | `SpeculativeRouterNode` | +|---|---|---| +| Strategy | **conservative** — wait for a committed decision, then cancel the tail | **aggressive** — act on a *partial* call, then verify | +| On the args | doesn't touch them | **repairs** truncated JSON and dispatches early | +| Risk | ~none | can mis-guess → cancel + re-run (must be idempotent) | + +## How it works + +The model emits a directive line: + +``` +TOOL_CALL: {"name": "read_file", "arguments": {"path": "src/main.c"}} +``` + +As it streams, the node: + +1. **Extracts + repairs.** The default extractor finds `TOOL_CALL:`, takes the + JSON after it, and — if it's still truncated (`{"path": "src/ma`) — runs it + through `repair_json` to get a parseable, best-effort payload. +2. **Dispatches early.** The first payload that passes `should_speculate` is used + to run the `target` node **immediately**, overlapping with generation. +3. **Verifies.** When the finalized call arrives it compares (via `same`): + - **hit** → keep the speculative result (already done or nearly so); + - **miss** → cancel the speculative run (cooperatively tearing down its + in-flight work) and re-run the target with the correct payload. + +The node's output is the *verified* target output. + +## Safety: idempotent targets only + +Speculation means the target may run on a wrong guess and be cancelled. Only use +it for **read-only / idempotent** work — file/DB reads, search, retrieval. Never +speculate a side-effecting action (sending mail, writing files, charging a card). + +Knobs: + +- `should_speculate(payload) -> bool` — gate early dispatch (e.g. only once a + path is long enough to be worth guessing). +- `same(a, b) -> bool` — how hit/miss is decided (here: same resolved `path`). +- `extract(text) -> payload | None` — plug in a different protocol, or add a + parameter-prediction step (e.g. complete a partial path against the repo). +- `timeout` — bounds the speculative read so a bad guess can't hang the turn. + +## Provenance + +This mirrors a libuv-based agent runtime that repairs partial tool-call JSON and +fires the call before the stream closes, then reconciles at end-of-stream — ported +to ADK as a first-class, cancellable graph node built on `ctx.run_node`. +``` diff --git a/contributing/samples/workflows/speculative_tool/agent.py b/contributing/samples/workflows/speculative_tool/agent.py new file mode 100644 index 00000000000..c23a4099124 --- /dev/null +++ b/contributing/samples/workflows/speculative_tool/agent.py @@ -0,0 +1,79 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Speculative tool dispatch: run the read before the model finishes asking. + +The model is told to answer by emitting a single directive line: + + TOOL_CALL: {"name": "read_file", "arguments": {"path": ""}} + +``SpeculativeRouterNode`` watches that line stream in. As soon as the (still +incomplete) JSON can be repaired into a plausible call, it dispatches the +``read_file`` target **immediately** -- overlapping the read with the rest of +generation -- then verifies against the finalized call: keep the result on a +match, cancel and re-read on a mismatch. + +The target here is a read (idempotent, safe to cancel), which is the only kind of +work speculation is appropriate for. ``enable_uvloop()`` runs the overlapped work +on a libuv loop. +""" + +import pathlib + +from google.adk import Agent +from google.adk import enable_uvloop +from google.adk import Event +from google.adk.workflow import SpeculativeRouterNode + +enable_uvloop() + + +def _path_of(payload: dict) -> str: + return (payload or {}).get('arguments', {}).get('path', '') + + +def read_file(node_input: dict): + """The speculatively-dispatched target: read a file (idempotent).""" + path = _path_of(node_input) + try: + text = pathlib.Path(path).read_text() + preview = text[:500] + yield Event(output={'path': path, 'ok': True, 'preview': preview}) + except OSError as e: + yield Event(output={'path': path, 'ok': False, 'error': str(e)}) + + +reader = Agent( + name='planner', + instruction=( + 'The user will ask about a file. Respond with EXACTLY one line and' + ' nothing else:\n' + ' TOOL_CALL: {"name": "read_file", "arguments": {"path": ""}}\n' + 'Use the most likely repository-relative path for what they asked.' + ), +) + + +root_agent = SpeculativeRouterNode( + name='speculative_read', + agent=reader, + target=read_file, + # Only speculate once the path looks substantial enough to be worth a guess; + # short prefixes are too likely to be revised. + should_speculate=lambda payload: len(_path_of(payload)) >= 6, + # Verify hit/miss on the resolved path. + same=lambda a, b: _path_of(a) == _path_of(b), + # Bound the speculative read so a bad guess can't hang the turn. + timeout=30, +) diff --git a/contributing/samples/workflows/streaming_route/README.md b/contributing/samples/workflows/streaming_route/README.md new file mode 100644 index 00000000000..6965ebc737f --- /dev/null +++ b/contributing/samples/workflows/streaming_route/README.md @@ -0,0 +1,75 @@ +# Streaming (Preemptive) Routing + uvloop Sample + +## Overview + +This sample demonstrates two speed-oriented features: + +1. **Mid-stream preemptive graph advancement** via `StreamingRouterNode`. A + classifier agent streams its answer token-by-token. The moment the routing + decision is present in the stream, the node commits the route and cancels + the rest of the generation — the workflow advances mid-stream instead of + waiting for the model to finish the turn. +1. **libuv event loop** via `enable_uvloop()`, which puts the whole process on + a faster asyncio runtime. + +## How it works + +`StreamingRouterNode` runs a wrapped agent in SSE streaming mode and hands +every streamed delta to a `monitor` predicate. When the monitor returns a +`StreamDecision`, the node: + +- commits that decision's `route` / `output` to the context, and +- (by default) closes the model stream, cancelling the remaining generation. + +Because the node returns promptly, the scheduler advances the graph on the +committed route. This is deterministic: the graph moves only once the decision +is unambiguously present in the stream, so no branch has to be revised later. + +```python +def route_when_category_streams(view: StreamView) -> StreamDecision | None: + text = view.text.lower() + for category in ("billing", "technical", "sales"): + if category in text: + return StreamDecision(route=category) + return None + +intent_router = StreamingRouterNode( + name="intent_router", + agent=classifier, + monitor=route_when_category_streams, +) +``` + +## Enabling uvloop + +`enable_uvloop()` installs the libuv event-loop policy process-wide. It is a +no-op (with a log line) when uvloop is not installed: + +```bash +pip install "google-adk[uvloop]" +``` + +You can also enable it without touching code by setting `ADK_UVLOOP=1`, which +the synchronous `Runner.run` path honours automatically. + +> **Note:** uvloop only accelerates work that actually awaits on an asyncio +> loop. Sync clients offloaded to a thread pool see no benefit until they are +> moved onto the loop (async client + `asyncio.gather`). libuv is the last +> 10%, not a 10x on its own. + +## Sample Inputs + +- `My invoice charged me twice this month.` → `billing` +- `The app crashes when I click export.` → `technical` +- `Do you offer volume discounts?` → `sales` + +## Graph + +```mermaid +graph TD + START --> process_input + process_input --> intent_router + intent_router -->|billing| billing_agent + intent_router -->|technical| technical_agent + intent_router -->|sales| sales_agent +``` diff --git a/contributing/samples/workflows/streaming_route/agent.py b/contributing/samples/workflows/streaming_route/agent.py new file mode 100644 index 00000000000..19c9763ad8e --- /dev/null +++ b/contributing/samples/workflows/streaming_route/agent.py @@ -0,0 +1,104 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mid-stream preemptive routing + libuv (uvloop). + +The classifier only needs to emit a single category word. A plain routing +node would wait for the model to finish the turn before advancing the graph. +``StreamingRouterNode`` instead watches the streamed tokens and, the instant +the category word appears, commits the route and cancels the rest of the +generation — the graph advances mid-stream. + +``enable_uvloop()`` swaps the process onto a libuv event loop for a faster +asyncio runtime. It is a no-op (with a log line) when uvloop is not installed; +install it with ``pip install "google-adk[uvloop]"``. +""" + +from google.adk import Agent +from google.adk import enable_uvloop +from google.adk import Event +from google.adk import Workflow +from google.adk.workflow import StreamDecision +from google.adk.workflow import StreamingRouterNode +from google.adk.workflow import StreamView + +# Put the whole process on libuv. Call once, before anything runs. +enable_uvloop() + +CATEGORIES = ('billing', 'technical', 'sales') + + +def process_input(node_input: str): + """Stashes the raw user message in state for downstream agents.""" + yield Event(state={'input': node_input}) + + +classifier = Agent( + name='classifier', + instruction=( + "Classify the user's request into exactly one category. Reply with" + ' ONLY that single lowercase word and nothing else: billing,' + ' technical, or sales.\n\nRequest: {input}' + ), +) + + +def route_when_category_streams(view: StreamView) -> StreamDecision | None: + """Advances the graph as soon as a category word appears in the stream. + + Returning a ``StreamDecision`` commits the route and (by default) cancels + the remainder of the model call. Returning ``None`` means "keep streaming". + """ + text = view.text.lower() + for category in CATEGORIES: + if category in text: + return StreamDecision(route=category) + return None + + +intent_router = StreamingRouterNode( + name='intent_router', + agent=classifier, + monitor=route_when_category_streams, +) + + +billing_agent = Agent( + name='billing_agent', + instruction='Help the user with their billing issue: {input}', +) +technical_agent = Agent( + name='technical_agent', + instruction='Help the user with their technical issue: {input}', +) +sales_agent = Agent( + name='sales_agent', + instruction='Help the user with their sales question: {input}', +) + + +root_agent = Workflow( + name='root_agent', + edges=[ + ('START', process_input, intent_router), + ( + intent_router, + { + 'billing': billing_agent, + 'technical': technical_agent, + 'sales': sales_agent, + }, + ), + ], +) diff --git a/docs/guides/workflow/streaming_preemption/index.md b/docs/guides/workflow/streaming_preemption/index.md new file mode 100644 index 00000000000..15c0c973a9a --- /dev/null +++ b/docs/guides/workflow/streaming_preemption/index.md @@ -0,0 +1,161 @@ +# Streaming preemption (`StreamingRouterNode`) + +`StreamingRouterNode` advances a workflow graph **mid-stream**: it runs a wrapped +agent in SSE mode, watches the model's tokens as they arrive, and the moment a +caller-supplied monitor is confident, it commits a route/output and **cancels the +rest of the generation**. This turns "wait for the whole turn, then move on" into +"move on the instant the answer is knowable". + +## Introduction + +The stock LlmAgent-as-node wrapper only commits a node's output — the thing that +fires downstream triggers — on the *final*, non-partial event. The graph +therefore always advances at turn granularity: the model finishes generating, +then the scheduler moves on. + +`StreamingRouterNode` closes that gap: + +1. It forces the wrapped agent into `StreamingMode.SSE`. +2. It hands every streamed delta to a `monitor` callback. +3. When the monitor returns a `StreamDecision`, the node commits that decision's + `route`/`output` and (by default) **closes the model stream**. + +Closing the generator propagates `GeneratorExit` down ADK's `aclosing` chain, +which cancels the in-flight model call cooperatively — the same mechanism the +runtime already uses for node timeouts and interrupts. Because the node's +`run()` returns promptly, the scheduler advances the graph immediately instead +of paying for the tail of a turn the model has already effectively decided. + +This is deterministic, mid-stream advancement: the graph moves as soon as the +decision is unambiguously present in the stream, and no wasted output tokens are +paid for. + +## Get started + +Fan out to several documents in parallel and, for each, stop generating the +moment a verdict has streamed in: + +```python +from typing import Optional +from google.adk import Agent, Event, Workflow +from google.adk.workflow import ( + JoinNode, + StreamDecision, + StreamingRouterNode, + StreamView, +) + + +def verdict_monitor(view: StreamView) -> Optional[StreamDecision]: + """Preempt as soon as a `VERDICT:` line has fully streamed in.""" + idx = view.text.upper().find("VERDICT:") + if idx == -1: + return None + line, sep, _ = view.text[idx + len("VERDICT:") :].partition("\n") + if not sep: # verdict line still streaming — keep reading + return None + verdict = line.strip() + if verdict.upper().startswith("IRRELEVANT"): + return StreamDecision(output={"relevant": False}) + if verdict.upper().startswith("RELEVANT"): + return StreamDecision(output={"relevant": True, "verdict": verdict}) + return None + + +def make_reader(i: int, document: str) -> StreamingRouterNode: + prompt = ( + f"PAPER:\n{document}\n\n---\n" + "Is this a computer-science AI/ML paper?\n" + "FIRST output a line 'VERDICT: RELEVANT - ' or " + "'VERDICT: IRRELEVANT'.\nTHEN write a long summary." + ) + return StreamingRouterNode( + name=f"reader_{i}", + # A callable instruction bypasses {var} templating, so raw braces in the + # document are sent verbatim. + agent=Agent(name=f"reader_{i}", model="gemini-3.5-flash-lite", + instruction=lambda _ctx, _p=prompt: _p), + monitor=verdict_monitor, + ) +``` + +Wire the readers into a fan-out/fan-in graph with a `JoinNode`; the scheduler +runs them concurrently and each one abandons its generation as soon as its +verdict lands. + +## Benefits (measured) + +The integration test +`tests/integration/test_streaming_router_preemption_timing.py` reads five whole +arXiv papers in parallel and asks "is this an AI paper?", comparing: + +- **A** — read + answer, stream to completion. +- **B** — read + answer + SSE preemption (cut once the verdict streams in). + +A representative run on real Vertex `gemini-3.5-flash-lite` (whole documents, no +chunking, ~199k shared input tokens): + +| Metric | A (full) | B (preempt) | +| --- | --- | --- | +| Wall clock | 6.63s | 1.88s (**3.5x faster**) | +| Output tokens | 3,614 | 128 (**~28x fewer**) | +| Cost @ $0.30/$2.50 per 1M | $0.0686 | $0.0599 (**~13% cheaper**) | +| Cost w/ context caching (input @ $0.03/1M) | $0.0150 | $0.0063 (**~58% cheaper**) | + +### Reading the numbers + +- **Preemption saves *generation*, not *reading*.** Each document is one whole + prompt in one call, so the model must prefill the entire input before it emits + any token. Preemption cancels the *output* stream, which happens strictly after + prefill — the input is already paid for. To also save reading you must not send + the whole document in one call (incremental input), which is a different design. +- At standard pricing the per-query cost is **input-bound** (199k input vs a few + thousand output tokens), so preemption's dollar impact is modest (~13%). +- With **context caching** the input read is 10x cheaper, output becomes the + dominant cost, and preemption's ~28x output cut drives ~58% total savings. + Context caching (amortize reading) and preemption (cut generation) are + complementary. + +## How it works + +```python +async with Aclosing(self.agent.run_async(ic)) as run_iter: + async for event in run_iter: + if event.partial: + # accumulate streamed text, hand it to the monitor + decision = await self._invoke_monitor(StreamView(...)) + if decision is not None: + self._apply_decision(ctx, decision) + if decision.stop: + return # GeneratorExit -> aclosing cancels the model call +``` + +Key fields on `StreamingRouterNode`: + +- `agent`: the (tool-free classifier/router) agent to stream. +- `monitor`: `Callable[[StreamView], Optional[StreamDecision]]`, sync or async. +- `forward_partials` (default `True`): re-yield partials as user-visible + messages (typewriter effect) in addition to driving the monitor. +- `include_thoughts` (default `False`): include model `thought` parts in + `StreamView.text`. + +`StreamDecision(route=..., output=..., stop=True)` commits a routing value and/or +an output; `stop=True` (default) cancels the remaining generation, `stop=False` +lets it run to completion while the decision stands. + +## Related: `enable_uvloop()` + +For network-bound agent workloads you can install the libuv event loop with a +one-line switch at your entrypoint: + +```python +import google.adk + +google.adk.enable_uvloop() # process-wide; call once before Runner.run +``` + +Deployments can opt in without touching code via `ADK_UVLOOP=1`, which the sync +`Runner.run` path honours. uvloop only accelerates code that actually awaits on +the loop (async client + `asyncio.gather`); it is the last 10%, not a 10x on its +own. Install with `pip install "google-adk[uvloop]"` (no Windows wheels). +``` diff --git a/pyproject.toml b/pyproject.toml index a500bd1ffe1..f10169130c6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -275,6 +275,9 @@ optional-dependencies.toolbox = [ "toolbox-adk>=1,<2" ] optional-dependencies.tools = [ "google-api-python-client>=2.157,<3", ] +optional-dependencies.uvloop = [ + "uvloop>=0.21; platform_system != 'Windows'", # libuv event loop; no Windows wheels. +] urls.changelog = "https://github.com/google/adk-python/blob/main/CHANGELOG.md" urls.documentation = "https://google.github.io/adk-docs/" urls.homepage = "https://google.github.io/adk-docs/" diff --git a/src/google/adk/__init__.py b/src/google/adk/__init__.py index 83cd21d150c..c37f3bb4385 100644 --- a/src/google/adk/__init__.py +++ b/src/google/adk/__init__.py @@ -24,6 +24,7 @@ from .agents.llm_agent import Agent from .events.event import Event from .runners import Runner + from .utils.event_loop import enable_uvloop from .workflow import Workflow __version__ = version.__version__ @@ -33,7 +34,8 @@ 'Event': '.events.event', 'Runner': '.runners', 'Workflow': '.workflow', + 'enable_uvloop': '.utils.event_loop', } -__all__ = ['Agent', 'Context', 'Event', 'Runner', 'Workflow'] +__all__ = ['Agent', 'Context', 'Event', 'Runner', 'Workflow', 'enable_uvloop'] __getattr__, __dir__ = _lazy.accessors(globals(), _LAZY_MEMBERS) diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index cb69b845510..4c1692f8948 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -1083,6 +1083,13 @@ async def _invoke_run_async() -> None: def _asyncio_thread_main() -> None: try: + # Honour ADK_UVLOOP=1 so deployments can run this sync shim on libuv + # without changing application code. asyncio.run() below then picks up + # the installed uvloop policy. Explicit enable_uvloop() at the + # entrypoint has the same effect and takes precedence. + from .utils.event_loop import maybe_enable_uvloop_from_env + + maybe_enable_uvloop_from_env() asyncio.run(_invoke_run_async()) finally: event_queue.put(None) diff --git a/src/google/adk/utils/event_loop.py b/src/google/adk/utils/event_loop.py new file mode 100644 index 00000000000..8990e0075a3 --- /dev/null +++ b/src/google/adk/utils/event_loop.py @@ -0,0 +1,131 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Event-loop selection helpers. + +ADK never pins an event loop or installs a loop policy of its own; every +``asyncio.run`` inside the framework (the sync ``Runner.run`` shim, the CLI, +tests) honours whatever policy is active in the process. That makes swapping +in `uvloop `_ (a libuv-backed loop that +is meaningfully faster for network-bound workloads) a one-line change at your +program's entrypoint:: + + import google.adk + + google.adk.enable_uvloop() + +Call it once, before the first ``asyncio.run``/``Runner.run``. It sets the +process-wide event-loop policy, so any subsequently created loop — including +the one the sync ``Runner`` spins up on its background thread — uses libuv. + +Caveat worth internalising: uvloop only accelerates code that actually awaits +on an asyncio loop. Work offloaded to ``ThreadPoolExecutor`` with a *sync* +client sees no benefit until it is moved onto the loop (async client + +``asyncio.gather``). libuv is the last 10%, not a 10x on its own. +""" + +from __future__ import annotations + +import asyncio +import logging + +from .env_utils import is_env_enabled + +logger = logging.getLogger('google_adk.' + __name__) + +_UVLOOP_ENV_VAR = 'ADK_UVLOOP' +"""When truthy (``1``/``true``), the sync ``Runner.run`` path auto-enables uvloop.""" + +# Tracks whether we have already installed the uvloop policy so repeated calls +# (e.g. one at the entrypoint and one auto-triggered by the env var) are cheap +# no-ops instead of re-installing the policy on every invocation. +_uvloop_installed = False + + +def is_uvloop_available() -> bool: + """Returns True if the ``uvloop`` package can be imported.""" + try: + import uvloop # noqa: F401 pylint: disable=g-import-not-at-top,unused-import + + return True + except ImportError: + return False + + +def enable_uvloop(*, strict: bool = False) -> bool: + """Installs the uvloop (libuv) event-loop policy process-wide. + + Idempotent: calling it more than once is a cheap no-op after the first + successful install. + + Args: + strict: If True, raise ``RuntimeError`` when uvloop is not installed + instead of silently falling back to the default asyncio loop. Use + this when the speedup is a hard requirement and a silent fallback + would mask a misconfigured deployment. + + Returns: + True if uvloop is now the active policy, False if it was unavailable + and ``strict`` is False. + + Raises: + RuntimeError: If ``strict`` is True and uvloop cannot be imported. + """ + global _uvloop_installed + if _uvloop_installed: + return True + + try: + import uvloop # pylint: disable=g-import-not-at-top + except ImportError as e: + message = ( + 'uvloop is not installed. Install it with `pip install' + ' "google-adk[uvloop]"` (uvloop does not support Windows).' + ) + if strict: + raise RuntimeError(message) from e + logger.info('%s Falling back to the default asyncio event loop.', message) + return False + + uvloop.install() + _uvloop_installed = True + logger.info('uvloop (libuv) event-loop policy installed.') + return True + + +def is_uvloop_active() -> bool: + """Returns True if the running (or default-policy) loop is a uvloop loop. + + Checks the currently running loop when called from inside a coroutine, and + otherwise inspects the active event-loop policy's loop factory. + """ + try: + loop = asyncio.get_running_loop() + return type(loop).__module__.startswith('uvloop') + except RuntimeError: + # No running loop; fall back to inspecting the installed policy. + policy = asyncio.get_event_loop_policy() + return type(policy).__module__.startswith('uvloop') + + +def maybe_enable_uvloop_from_env() -> bool: + """Enables uvloop iff the ``ADK_UVLOOP`` env var is truthy. + + Lets deployments opt into libuv without touching application code. Called + internally by the sync ``Runner.run`` path. Returns True if uvloop is + active after the call. + """ + if is_env_enabled(_UVLOOP_ENV_VAR): + return enable_uvloop() + return _uvloop_installed diff --git a/src/google/adk/workflow/__init__.py b/src/google/adk/workflow/__init__.py index b18156f281b..20202c52c11 100644 --- a/src/google/adk/workflow/__init__.py +++ b/src/google/adk/workflow/__init__.py @@ -22,6 +22,7 @@ from ._base_node import BaseNode from ._base_node import START from ._errors import NodeTimeoutError + from ._first_match_node import FirstMatchNode from ._function_node import FunctionNode from ._graph import DEFAULT_ROUTE from ._graph import Edge @@ -29,33 +30,53 @@ from ._node import Node from ._node import node from ._retry_config import RetryConfig + from ._speculative_router import make_marker_extractor + from ._speculative_router import repair_json + from ._speculative_router import SpeculativeRouterNode + from ._streaming_router import StreamDecision + from ._streaming_router import StreamingRouterNode + from ._streaming_router import StreamView from ._workflow import Workflow _LAZY_MEMBERS: dict[str, str] = { 'BaseNode': '._base_node', 'DEFAULT_ROUTE': '._graph', 'Edge': '._graph', + 'FirstMatchNode': '._first_match_node', 'FunctionNode': '._function_node', 'JoinNode': '._join_node', 'Node': '._node', 'NodeTimeoutError': '._errors', 'RetryConfig': '._retry_config', 'START': '._base_node', + 'SpeculativeRouterNode': '._speculative_router', + 'StreamDecision': '._streaming_router', + 'StreamingRouterNode': '._streaming_router', + 'StreamView': '._streaming_router', 'Workflow': '._workflow', + 'make_marker_extractor': '._speculative_router', 'node': '._node', + 'repair_json': '._speculative_router', } __all__ = [ 'BaseNode', 'DEFAULT_ROUTE', 'Edge', + 'FirstMatchNode', 'FunctionNode', 'JoinNode', 'Node', 'NodeTimeoutError', 'RetryConfig', 'START', + 'SpeculativeRouterNode', + 'StreamDecision', + 'StreamingRouterNode', + 'StreamView', 'Workflow', + 'make_marker_extractor', 'node', + 'repair_json', ] __getattr__, __dir__ = _lazy.accessors(globals(), _LAZY_MEMBERS) diff --git a/src/google/adk/workflow/_first_match_node.py b/src/google/adk/workflow/_first_match_node.py new file mode 100644 index 00000000000..c243682c8bc --- /dev/null +++ b/src/google/adk/workflow/_first_match_node.py @@ -0,0 +1,174 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""First-match fan-out: race several branches, cancel the losers. + +``JoinNode`` is fan-in that waits for *all* predecessors. ``FirstMatchNode`` is +its opposite: it fans out over a set of branch nodes, runs them concurrently, +and the instant one returns a result the ``match`` predicate accepts, it +*cancels the still-running siblings* and yields the winner. + +This is the "search Recall@k" pattern: retrieve k candidate sources, read them +in parallel, and answer the moment any one of them contains the answer -- you +neither wait for the slow branches nor keep paying to generate their output. +Cancellation propagates cooperatively down ``ctx.run_node`` into each branch's +in-flight work (e.g. a ``StreamingRouterNode``'s SSE model call is torn down via +its ``aclosing`` chain), so the losers stop *decoding* immediately. + +What it does *not* do: it cannot un-send input already prefilled. Branches that +have started reading their source still paid that read; ``FirstMatchNode`` saves +the losers' generation and wall-clock, not their prefill. Pair it with a cheap +relevance gate (or rank-ordered ``max_parallel``) if you also need to avoid +reading low-ranked sources at all. + +Example:: + + FirstMatchNode( + name='first_source_with_answer', + nodes=[read_sharepoint, read_wiki, read_crm, read_docs, read_drive], + match=lambda r: isinstance(r, dict) and r.get('found'), + ) + +Each branch is handed the same ``node_input`` (broadcast); encapsulate each +source inside its own branch node. If no branch matches, the node yields +``no_match_output`` (``None`` by default). +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator +from collections.abc import Callable +import logging +from typing import Any + +from pydantic import ConfigDict +from pydantic import Field +from pydantic import PrivateAttr +from typing_extensions import override + +from ..agents.context import Context +from ._base_node import BaseNode +from ._graph import NodeLike +from ._retry_config import RetryConfig +from .utils._workflow_graph_utils import build_node + +logger = logging.getLogger('google_adk.' + __name__) + + +def _default_match(result: Any) -> bool: + """Accepts any non-``None`` result as a match.""" + return result is not None + + +class FirstMatchNode(BaseNode): + """Races branch nodes and returns the first matching result, cancelling the rest. + + Attributes: + max_parallel: Maximum branches to run at once. ``None`` runs them all + concurrently. Set this (with branches supplied in priority order) to read + higher-ranked sources first and avoid ever starting lower-ranked ones once + an earlier branch wins. + no_match_output: The node's output when no branch produces a matching + result. Defaults to ``None``. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + max_parallel: int | None = Field(default=None) + no_match_output: Any = Field(default=None) + + _nodes: list[BaseNode] = PrivateAttr() + _match: Callable[[Any], bool] = PrivateAttr() + + def __init__( + self, + *, + name: str, + nodes: list[NodeLike], + match: Callable[[Any], bool] | None = None, + max_parallel: int | None = None, + no_match_output: Any = None, + retry_config: RetryConfig | None = None, + timeout: float | None = None, + ): + if not nodes: + raise ValueError('FirstMatchNode requires at least one branch node.') + if max_parallel is not None and max_parallel < 1: + raise ValueError('max_parallel must be >= 1.') + built = [build_node(n) for n in nodes] + super().__init__( + name=name, + rerun_on_resume=True, + retry_config=retry_config, + timeout=timeout, + max_parallel=max_parallel, + no_match_output=no_match_output, + ) + self._nodes = built + self._match = match or _default_match + + async def _run_one(self, ctx: Context, node: BaseNode, node_input: Any) -> Any: + return await ctx.run_node(node, node_input=node_input, use_sub_branch=True) + + @override + async def _run_impl( + self, + *, + ctx: Context, + node_input: Any, + ) -> AsyncGenerator[Any, None]: + pending: set[asyncio.Task[Any]] = set() + remaining = list(self._nodes) + winner: Any = self.no_match_output + found = False + + def _launch_next() -> None: + while remaining and ( + self.max_parallel is None or len(pending) < self.max_parallel + ): + node = remaining.pop(0) + pending.add(asyncio.create_task(self._run_one(ctx, node, node_input))) + + try: + _launch_next() + while pending and not found: + done, pending = await asyncio.wait( + pending, return_when=asyncio.FIRST_COMPLETED + ) + for task in done: + # A failed branch must not sink the whole race; log and move on so a + # single flaky source can't deny an answer another source can give. + if task.cancelled(): + continue + exc = task.exception() + if exc is not None: + logger.warning('FirstMatchNode %s: branch failed: %s', self.name, exc) + continue + result = task.result() + if self._match(result): + winner = result + found = True + break + if not found: + _launch_next() + finally: + # Cancel every still-running loser and wait for their teardown so the + # in-flight reads are actually torn down before we advance the graph. + for task in pending: + task.cancel() + if pending: + await asyncio.wait(pending) + + yield winner diff --git a/src/google/adk/workflow/_speculative_router.py b/src/google/adk/workflow/_speculative_router.py new file mode 100644 index 00000000000..c67238d9dfb --- /dev/null +++ b/src/google/adk/workflow/_speculative_router.py @@ -0,0 +1,399 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Speculative mid-stream dispatch: act on a partial call, verify, roll back. + +Where :class:`StreamingRouterNode` is *conservative* — it waits for a committed +decision in the stream and then cancels the tail — ``SpeculativeRouterNode`` is +*aggressive*. As the model streams a structured call (e.g. a ``TOOL_CALL: {...}`` +directive), this node repairs the still-incomplete JSON, dispatches a downstream +target node **immediately** with that best-effort payload, and lets the model +keep generating. When the finalized call arrives it verifies: + + * **hit** — the finalized payload matches the speculated one → keep the + speculative result (its work overlapped generation, so it's already done or + nearly so), and + * **miss** — they differ → cancel the speculative run (cooperatively tearing + down its in-flight work) and re-dispatch the target with the correct + payload. + +This trades peak latency for speculation risk: a wrong guess wastes a run and +must be safe to cancel. Use it only for **idempotent / side-effect-free** +targets (reads, searches, retrieval) — never for a target that, say, sends an +email or charges a card. This is the ADK analogue of a libuv agent runtime that +fires repaired tool calls before the stream closes. + +The default extractor looks for a marker (``TOOL_CALL:`` by default), takes the +JSON that follows, and — if it is truncated — runs it through :func:`repair_json` +before parsing. Supply your own ``extract`` to plug in a different protocol or a +parameter-prediction step (e.g. completing a partial file path). +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator +from collections.abc import Callable +import json +import logging +from typing import Any +from typing import Optional +from typing import TYPE_CHECKING + +from pydantic import ConfigDict +from pydantic import Field +from pydantic import PrivateAttr +from typing_extensions import override + +from ..agents.llm_agent import LlmAgent +from ..events.event import Event +from ..utils.context_utils import Aclosing +from ._base_node import BaseNode +from ._graph import NodeLike +from ._retry_config import RetryConfig +from ._streaming_router import _model_text +from ._streaming_router import build_sse_invocation_context +from .utils._workflow_graph_utils import build_node + +if TYPE_CHECKING: + from ..agents.context import Context + +logger = logging.getLogger('google_adk.' + __name__) + +_UNSET = object() + + +def repair_json(fragment: str) -> str: + """Completes a truncated JSON fragment into a parseable string (best effort). + + A stack-based state machine (ported from the ``syncrig`` C repairer): it closes + an open string, finishes a truncated ``true``/``false``/``null`` literal, drops + a dangling comma, fills a dangling key/colon with ``null``, and appends the + ``}``/``]`` needed to balance every open object/array. Given + ``'{"path": "src/ma'`` it returns ``'{"path": "src/ma"}'``. + """ + _OBJ, _ARR = 0, 1 + out: list[str] = [] + stack: list[int] = [] + state: list[int] = [] # per-frame object state: 0 key,1 colon,2 value,3 comma + in_string = False + escaped = False + + for c in fragment: + out.append(c) + if in_string: + if escaped: + escaped = False + elif c == '\\': + escaped = True + elif c == '"': + in_string = False + if stack and stack[-1] == _OBJ: + if state[-1] == 0: + state[-1] = 1 + elif state[-1] == 2: + state[-1] = 3 + continue + if c == '"': + in_string = True + elif c == '{': + stack.append(_OBJ) + state.append(0) + elif c == '[': + stack.append(_ARR) + state.append(-1) + elif c == '}': + if stack and stack[-1] == _OBJ: + stack.pop() + state.pop() + if stack and stack[-1] == _OBJ and state[-1] == 2: + state[-1] = 3 + elif c == ']': + if stack and stack[-1] == _ARR: + stack.pop() + state.pop() + if stack and stack[-1] == _OBJ and state[-1] == 2: + state[-1] = 3 + elif c == ':': + if stack and stack[-1] == _OBJ: + state[-1] = 2 + elif c == ',': + if stack and stack[-1] == _OBJ: + state[-1] = 0 + + if in_string: + if escaped: + out.pop() # dangling backslash + out.append('"') + if stack and stack[-1] == _OBJ: + if state[-1] == 0: + state[-1] = 1 + elif state[-1] == 2: + state[-1] = 3 + + while out and out[-1] in ' \n\r\t': + out.pop() + + if out and out[-1] == ',': + out.pop() + while out and out[-1] in ' \n\r\t': + out.pop() + if stack and stack[-1] == _OBJ: + state[-1] = 3 + + if out and (out[-1].isalnum() or out[-1] in '.-'): + s = ''.join(out) + j = len(s) + while j > 0 and s[j - 1].isalpha(): + j -= 1 + frag = s[j:] + completed = False + for kw in ('true', 'false', 'null'): + if frag and frag != kw and kw.startswith(frag): + out.extend(kw[len(frag) :]) + completed = True + break + if stack and stack[-1] == _OBJ and state[-1] == 2: + state[-1] = 3 + del completed + + if stack and stack[-1] == _OBJ: + if state[-1] == 1: + out.append(':null') + state[-1] = 3 + elif state[-1] == 2: + out.append('null') + state[-1] = 3 + + while stack: + out.append('}' if stack.pop() == _OBJ else ']') + + return ''.join(out) + + +def make_marker_extractor(marker: str = 'TOOL_CALL:') -> Callable[[str], Any]: + """Builds an extractor that pulls the (possibly-partial) JSON after ``marker``. + + Returns the parsed object, or ``None`` if the marker/JSON has not appeared yet. + A complete object is parsed as-is (trailing text ignored); a truncated one is + run through :func:`repair_json` first. + """ + + def extract(text: str) -> Any: + idx = text.find(marker) + if idx == -1: + return None + rest = text[idx + len(marker) :] + brace = rest.find('{') + if brace == -1: + return None + fragment = rest[brace:] + try: + obj, _ = json.JSONDecoder().raw_decode(fragment) + return obj + except json.JSONDecodeError: + pass + try: + return json.loads(repair_json(fragment)) + except json.JSONDecodeError: + return None + + return extract + + +class SpeculativeRouterNode(BaseNode): + """Speculatively dispatches a target node from a partial streamed call. + + Wrap a streaming agent and a ``target`` node. As the agent streams, ``extract`` + turns the accumulated text into a payload (repairing truncated JSON); the first + time it yields something ``should_speculate`` accepts, the target is dispatched + with that payload while generation continues. When the finalized payload + arrives it is compared with ``same``: on a match the speculative result is kept; + otherwise the speculative run is cancelled and the target is re-run with the + finalized payload. The node's output is the target's (verified) output. + + If ``combine`` is supplied, the node's output is instead + ``combine(agent_full_text, target_output)`` — this makes the agent's complete + streamed text (e.g. a planner's rationale) a *required* returned deliverable + rather than a discarded tail, so a non-speculative baseline that must also emit + that text gets no "just plan and stop early" shortcut. + + The target must be safe to run speculatively and to cancel — use read-only / + idempotent work only. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + agent: LlmAgent = Field(...) + """The streaming agent whose output is parsed for a call to dispatch.""" + + marker: str = 'TOOL_CALL:' + """Marker the default extractor searches for before the JSON payload.""" + + forward_partials: bool = False + """If True, re-yield the agent's partial events as user-visible messages.""" + + include_thoughts: bool = False + """If True, model ``thought`` parts are included in the accumulated text.""" + + emit_speculation_events: bool = True + """If True, yield lightweight hit/miss/dispatch events for observability.""" + + rerun_on_resume: bool = True + + _target: BaseNode = PrivateAttr() + _extract: Callable[[str], Any] = PrivateAttr() + _should_speculate: Callable[[Any], bool] = PrivateAttr() + _same: Callable[[Any, Any], bool] = PrivateAttr() + _combine: Optional[Callable[[str, Any], Any]] = PrivateAttr(default=None) + + def __init__( + self, + *, + name: str, + agent: LlmAgent, + target: NodeLike, + extract: Optional[Callable[[str], Any]] = None, + should_speculate: Optional[Callable[[Any], bool]] = None, + same: Optional[Callable[[Any, Any], bool]] = None, + combine: Optional[Callable[[str, Any], Any]] = None, + marker: str = 'TOOL_CALL:', + forward_partials: bool = False, + include_thoughts: bool = False, + emit_speculation_events: bool = True, + retry_config: RetryConfig | None = None, + timeout: float | None = None, + ): + super().__init__( + name=name, + agent=agent, + marker=marker, + forward_partials=forward_partials, + include_thoughts=include_thoughts, + emit_speculation_events=emit_speculation_events, + retry_config=retry_config, + timeout=timeout, + ) + self._target = build_node(target) + self._extract = extract or make_marker_extractor(marker) + self._should_speculate = should_speculate or (lambda _payload: True) + self._same = same or (lambda a, b: a == b) + self._combine = combine + + async def _run_target(self, ctx: Context, payload: Any) -> Any: + return await ctx.run_node( + self._target, node_input=payload, use_sub_branch=True + ) + + def _info(self, kind: str, payload: Any) -> Event: + return Event(author=self.name, message=f'[speculation:{kind}] {payload}') + + def _commit(self, ctx: Context, result: Any, plan_text: str = '') -> None: + # When ``combine`` is set the node's deliverable is a function of *both* the + # agent's full text (e.g. a planner's rationale that the caller requires) and + # the target's result — so the streamed text is a returned artifact, not a + # throwaway tail that only exists to make speculation look good. + out = self._combine(plan_text, result) if self._combine else result + ctx.output = out + output_key = getattr(self.agent, 'output_key', None) + if output_key and out is not None: + ctx.actions.state_delta[output_key] = out + + @override + async def _run_impl( + self, *, ctx: Context, node_input: Any + ) -> AsyncGenerator[Any, None]: + ic = build_sse_invocation_context(self.agent, ctx, node_input) + + accumulated: list[str] = [] + final_text = '' + spec_task: Optional[asyncio.Task[Any]] = None + spec_payload: Any = _UNSET + all_tasks: list[asyncio.Task[Any]] = [] + + try: + async with Aclosing(self.agent.run_async(ic)) as run_iter: + async for event in run_iter: + if event.partial: + delta = _model_text(event, include_thoughts=self.include_thoughts) + if delta: + accumulated.append(delta) + if self.forward_partials: + yield event + payload = self._extract(''.join(accumulated)) + if ( + payload is not None + and self._should_speculate(payload) + and not ( + spec_payload is not _UNSET + and self._same(payload, spec_payload) + ) + ): + # (Re)dispatch: the stream refined the payload, so abandon the + # prior guess and speculate on the newer one. + if spec_task is not None and not spec_task.done(): + spec_task.cancel() + spec_payload = payload + spec_task = asyncio.create_task(self._run_target(ctx, payload)) + all_tasks.append(spec_task) + if self.emit_speculation_events: + yield self._info('dispatch', payload) + continue + + # Non-partial (aggregated / final) event: capture the full text. + text = _model_text(event, include_thoughts=self.include_thoughts) + if text: + final_text = text + + plan_text = final_text or ''.join(accumulated) + final_payload = self._extract(plan_text) + + if final_payload is None: + # Nothing actionable ever materialized; drop any speculation. + return + + if spec_task is not None and self._same(final_payload, spec_payload): + # HIT: the speculated payload was right — keep its (overlapped) result. + try: + result = await spec_task + except asyncio.CancelledError: + raise + except Exception as e: # pylint: disable=broad-except + logger.warning( + 'SpeculativeRouterNode %s: speculative run failed (%s);' + ' re-running.', + self.name, + e, + ) + result = await self._run_target(ctx, final_payload) + else: + if self.emit_speculation_events: + yield self._info('hit', final_payload) + self._commit(ctx, result, plan_text) + return + + # MISS (or never speculated): cancel the wrong guess, run the real one. + if spec_task is not None: + spec_task.cancel() + await asyncio.gather(spec_task, return_exceptions=True) + if self.emit_speculation_events: + yield self._info('miss', final_payload) + result = await self._run_target(ctx, final_payload) + self._commit(ctx, result, plan_text) + finally: + leftovers = [t for t in all_tasks if not t.done()] + for t in leftovers: + t.cancel() + if leftovers: + await asyncio.gather(*leftovers, return_exceptions=True) diff --git a/src/google/adk/workflow/_streaming_router.py b/src/google/adk/workflow/_streaming_router.py new file mode 100644 index 00000000000..5d4843a70ad --- /dev/null +++ b/src/google/adk/workflow/_streaming_router.py @@ -0,0 +1,295 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mid-stream preemptive graph advancement. + +The stock LlmAgent-as-node wrapper only commits a node's output — the thing +that fires downstream triggers — on the *final*, non-partial event. The graph +therefore always advances at turn granularity: the model finishes generating, +then the scheduler moves on. + +``StreamingRouterNode`` closes that gap. It runs a wrapped agent in SSE mode +and hands every streamed delta to a caller-supplied ``monitor``. The moment the +monitor can decide — e.g. it has seen the routing token, a confident answer +prefix, or the classification JSON — it returns a :class:`StreamDecision`. The +node then: + + 1. commits that decision's ``route`` / ``output`` to the context, and + 2. (by default) closes the model stream, cancelling the rest of the + generation. + +Closing the generator propagates ``GeneratorExit`` down ADK's ``aclosing`` +chain, which cancels the in-flight model call cooperatively — the same +mechanism the runtime already uses for node timeouts and interrupts. Because +the node's ``run()`` then returns promptly, the workflow scheduler advances the +graph immediately instead of waiting for the tail of a turn the model has +already effectively decided. + +This is deterministic, mid-stream advancement: the graph moves as soon as the +decision is unambiguously present in the stream, and no wasted tokens are paid +for. It intentionally does *not* speculatively dispatch a branch the model +might later revise; advancing only on a committed decision keeps the +correctness story simple. Keep the wrapped agent tool-free (a classifier / +router persona); function calls are streamed through but never trigger a +decision on their own. +""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from collections.abc import Awaitable +from collections.abc import Callable +from dataclasses import dataclass +import inspect +from typing import Any +from typing import Optional +from typing import TYPE_CHECKING +from typing import Union + +from pydantic import Field +from pydantic import model_validator +from typing_extensions import override + +from ..agents._streaming_mode import StreamingMode +from ..agents.llm_agent import LlmAgent +from ..agents.run_config import RunConfig +from ..events.event import Event +from ..utils.context_utils import Aclosing +from ._base_node import BaseNode + +if TYPE_CHECKING: + from ..agents.context import Context + from ._graph import RouteValue + + +@dataclass(frozen=True) +class StreamView: + """A read-only snapshot of the model stream handed to the monitor. + + Attributes: + text: All model text accumulated so far in this turn (thought parts + excluded unless ``include_thoughts`` is set on the node). + delta: The text carried by the current partial event (may be empty for + non-text partials such as streaming function-call arguments). + event: The raw partial :class:`Event` currently being processed. + """ + + text: str + delta: str + event: Event + + +class StreamDecision: + """The monitor's verdict for advancing the graph mid-stream. + + Return an instance from a ``monitor`` to commit a routing decision and/or an + output before the wrapped agent finishes its turn. Returning ``None`` means + "not decided yet — keep streaming". + + At least one of ``route`` or ``output`` must be provided. + + Attributes: + route: Routing value for conditional edges (a single value or a list). + Read by the workflow scheduler to pick downstream edges. + output: The node's output value. Also written to the agent's + ``output_key`` (when set) as a state delta. + stop: If True (default), the model stream is closed and the remaining + generation cancelled the moment this decision is returned — the + "preempt". If False, generation continues to completion while the + decision stands (useful when you still want the full text persisted). + """ + + def __init__( + self, + *, + route: Optional[Union[RouteValue, list[RouteValue]]] = None, + output: Any = None, + stop: bool = True, + ) -> None: + if route is None and output is None: + raise ValueError( + 'StreamDecision requires at least one of `route` or `output`.' + ) + self.route = route + self.output = output + self.stop = stop + + +# A monitor inspects each streamed delta and returns a decision (or None to +# keep waiting). It may be sync or async. +StreamMonitorCallback = Callable[ + [StreamView], + Union[Optional[StreamDecision], Awaitable[Optional[StreamDecision]]], +] + + +def _model_text(event: Event, *, include_thoughts: bool) -> str: + """Concatenates the model text on an event, skipping thoughts by default.""" + if not event.content or not event.content.parts: + return '' + return ''.join( + part.text + for part in event.content.parts + if part.text and (include_thoughts or not part.thought) + ) + + +def build_sse_invocation_context(agent: LlmAgent, ctx: Context, node_input: Any): + """Prepares an InvocationContext that streams ``agent`` via SSE. + + Shared by the streaming nodes (:class:`StreamingRouterNode` and the + speculative variant) so mid-stream monitoring always sees progressive + partial events regardless of the caller's default streaming mode. + """ + from ._llm_agent_wrapper import prepare_llm_agent_context + from ._llm_agent_wrapper import prepare_llm_agent_input + + if agent.mode is None: + agent.mode = 'single_turn' + # As a single-turn node, default to not replaying prior turns unless the + # author opted in — matching run_llm_agent_as_node. + if ( + agent.mode == 'single_turn' + and 'include_contents' not in agent.model_fields_set + ): + agent.include_contents = 'none' + + agent_ctx = prepare_llm_agent_context(agent, ctx) + prepare_llm_agent_input(agent, agent_ctx, node_input) + + ic = agent_ctx.get_invocation_context() + run_config = (ic.run_config or RunConfig()).model_copy( + update={'streaming_mode': StreamingMode.SSE} + ) + update: dict[str, Any] = {'agent': agent, 'run_config': run_config} + iso = getattr(agent_ctx, 'isolation_scope', None) + if agent.mode in ('task', 'single_turn') and iso: + update['isolation_scope'] = iso + return ic.model_copy(update=update) + + +class StreamingRouterNode(BaseNode): + """Advances the workflow graph mid-stream based on the model's own output. + + Wrap a classifier / router agent and supply a ``monitor`` predicate. The node + streams the agent in SSE mode and, as soon as the monitor returns a + :class:`StreamDecision`, commits the route/output and (by default) cancels the + rest of the generation so the graph advances immediately. + + Example:: + + def route_on_label(view: StreamView) -> StreamDecision | None: + low = view.text.lower() + if 'billing' in low: + return StreamDecision(route='billing') + if 'technical' in low: + return StreamDecision(route='technical') + return None + + router = StreamingRouterNode( + name='intent_router', + agent=LlmAgent(name='classifier', model='gemini-2.5-flash', + instruction='Reply with the single word intent.'), + monitor=route_on_label, + ) + """ + + agent: LlmAgent = Field(...) + """The agent to stream. Should be a tool-free classifier / router persona.""" + + monitor: StreamMonitorCallback = Field(...) + """Predicate called on every streamed delta; returns a decision or None.""" + + forward_partials: bool = True + """If True, streamed partial events are re-yielded as user-visible messages + (typewriter effect) in addition to driving the monitor.""" + + include_thoughts: bool = False + """If True, model ``thought`` parts are included in ``StreamView.text``.""" + + # Dynamic scheduling / resume support. Mirrors the LlmAgent-as-node default. + rerun_on_resume: bool = True + + @model_validator(mode='after') + def _validate_monitor(self) -> StreamingRouterNode: + if not callable(self.monitor): + raise ValueError('`monitor` must be callable.') + return self + + def _apply_decision(self, ctx: Context, decision: StreamDecision) -> None: + """Commits a decision's output and/or route onto the context.""" + if decision.output is not None: + ctx.output = decision.output + output_key = getattr(self.agent, 'output_key', None) + if output_key: + ctx.actions.state_delta[output_key] = decision.output + if decision.route is not None: + ctx.route = decision.route + + async def _invoke_monitor(self, view: StreamView) -> Optional[StreamDecision]: + result = self.monitor(view) + if inspect.isawaitable(result): + return await result + return result + + def _build_streaming_ic(self, ctx: Context, node_input: Any) -> Any: + """Prepares an InvocationContext that streams the wrapped agent via SSE.""" + return build_sse_invocation_context(self.agent, ctx, node_input) + + @override + async def _run_impl( + self, *, ctx: Context, node_input: Any + ) -> AsyncGenerator[Any, None]: + ic = self._build_streaming_ic(ctx, node_input) + + decided = False + accumulated: list[str] = [] + + async with Aclosing(self.agent.run_async(ic)) as run_iter: + async for event in run_iter: + if event.partial: + delta = _model_text(event, include_thoughts=self.include_thoughts) + if delta: + accumulated.append(delta) + if self.forward_partials: + yield event + if not decided: + decision = await self._invoke_monitor( + StreamView(text=''.join(accumulated), delta=delta, event=event) + ) + if decision is not None: + self._apply_decision(ctx, decision) + decided = True + if decision.stop: + # Returning closes the stream (GeneratorExit -> aclosing), + # cancelling the rest of the model call, and lets the + # scheduler advance on the committed route/output. + return + continue + + # Non-partial (aggregated / final) event. + if decided: + # A decision already owns this node's output/route. Stream the + # event for visibility but strip any output it carries to avoid a + # double-set on the context. + if event.output is not None: + event = event.model_copy(update={'output': None}) + yield event + continue + + # No early decision: fall back to standard output extraction. + from ._llm_agent_wrapper import process_llm_agent_output + + process_llm_agent_output(self.agent, ctx, event) + yield event diff --git a/tests/integration/test_speculative_router_arxiv.py b/tests/integration/test_speculative_router_arxiv.py new file mode 100644 index 00000000000..00debb51c12 --- /dev/null +++ b/tests/integration/test_speculative_router_arxiv.py @@ -0,0 +1,198 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Real-LLM: SpeculativeRouterNode overlaps a downstream handler with generation. + +Contrived to isolate the speculative win, on a whole real arXiv paper -- NO +chunking, no truncation, one streaming call. + +The model reads the paper and is told to FIRST emit a directive line: + + TOOL_CALL: {"name": "lookup", "arguments": {"topic": ""}} + +...THEN write a long analysis. A downstream ``lookup`` handler (here: a fixed +2s sleep, standing in for a retrieval/DB call) depends on that directive. + + A. sequential -- wait for the whole generation, then run the handler. + B. speculative -- the instant the (repaired, partial) directive streams in, + dispatch the handler so it runs *while* the model is still + writing the analysis; verify against the finalized directive. + +B must be faster by ~the handler's duration, because that work overlaps the long +tail of generation instead of following it. + + ADK_TEST_MODEL=gemini-3.5-flash-lite \\ + uv run pytest -s -p no:cacheprovider \\ + tests/integration/test_speculative_router_arxiv.py + +Requires Vertex (GOOGLE_CLOUD_PROJECT via ADC), ``pypdf`` and network access to +arXiv; skips otherwise. +""" + +import asyncio +import os +import pathlib +import tempfile +import time +from typing import Any +import urllib.error +import urllib.request + +from dotenv import load_dotenv +from google.adk import Agent +from google.adk import Event +from google.adk import Workflow +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.workflow import SpeculativeRouterNode +from google.genai import types +import pytest + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] +load_dotenv(_REPO_ROOT / '.env', override=False) + +pytestmark = pytest.mark.skipif( + not os.environ.get('GOOGLE_CLOUD_PROJECT'), + reason='Real-LLM speculative test requires Vertex (GOOGLE_CLOUD_PROJECT).', +) + +MODEL = os.environ.get('ADK_TEST_MODEL', 'gemini-3.5-flash-lite') + +# The relevant paper; whole document, no chunking. +_ARXIV_ID = '1706.03762' # Attention Is All You Need + +# Stand-in for a real downstream dependency (retrieval / DB / API call). Fixed so +# the overlap saving is deterministic and measurable. +_HANDLER_SECONDS = 2.0 + +# The directive line the model must emit first. Kept out of the f-string so the +# literal JSON braces are not doubled. +_DIRECTIVE = ( + 'FIRST output EXACTLY one line, nothing else on it:\n' + 'TOOL_CALL: {"name": "lookup", "arguments": {"topic":' + ' ""}}\n' + 'THEN write a long, detailed, multi-paragraph summary of the paper (at' + ' least 300 words).' +) + + +def _download_paper_text(arxiv_id: str) -> str: + from pypdf import PdfReader + + cache = pathlib.Path(tempfile.gettempdir()) / 'adk_arxiv_papers' + cache.mkdir(parents=True, exist_ok=True) + pdf_path = cache / (arxiv_id.replace('/', '_') + '.pdf') + if not pdf_path.exists(): + url = 'https://arxiv.org/pdf/' + arxiv_id + req = urllib.request.Request(url, headers={'User-Agent': 'adk-test/1.0'}) + pdf_path.write_bytes(urllib.request.urlopen(req, timeout=60).read()) # noqa: S310 + reader = PdfReader(str(pdf_path)) + return ''.join((page.extract_text() or '') for page in reader.pages) + + +def _prompt(doc: str) -> str: + return ( + f'PAPER:\n{doc}\n\n----------------------------------------\n' + 'You just read the paper above.\n\n' + _DIRECTIVE + ) + + +def _topic(payload: Any) -> str: + if isinstance(payload, dict): + return payload.get('arguments', {}).get('topic', '') + return '' + + +def _build(doc: str, *, speculative: bool) -> tuple[Workflow, dict[str, Any]]: + stats: dict[str, Any] = {'dispatches': 0} + + async def lookup(node_input: Any): + # The downstream dependency: expensive work keyed on the directive's topic. + stats['dispatches'] += 1 + try: + await asyncio.sleep(_HANDLER_SECONDS) + except asyncio.CancelledError: + raise + yield Event(output={'topic': _topic(node_input), 'handled': True}) + + agent = Agent( + name='reader', + model=MODEL, + # Callable instruction bypasses {var} templating so raw paper braces are + # sent verbatim -- no escaping, no truncation. + instruction=lambda _ctx, _p=_prompt(doc): _p, + ) + node = SpeculativeRouterNode( + name='speculative_read', + agent=agent, + target=lookup, + # A: never speculate (sequential). B: speculate (overlap). + should_speculate=(lambda p: bool(_topic(p))) if speculative else ( + lambda p: False + ), + same=lambda a, b: _topic(a) == _topic(b), + timeout=120, + ) + return Workflow(name='spec_wf', edges=[('START', node)]), stats + + +async def _run(wf: Workflow) -> tuple[float, Any]: + ss = InMemorySessionService() + runner = Runner(app_name=wf.name, node=wf, session_service=ss) + session = await ss.create_session(app_name=wf.name, user_id='u') + msg = types.Content(parts=[types.Part(text='go')], role='user') + output = None + start = time.perf_counter() + async for event in runner.run_async( + user_id='u', session_id=session.id, new_message=msg + ): + if isinstance(event, Event) and event.output is not None: + output = event.output + return time.perf_counter() - start, output + + +@pytest.mark.asyncio +@pytest.mark.parametrize('llm_backend', ['VERTEX'], indirect=True) +async def test_speculative_dispatch_overlaps_generation(llm_backend): + doc = _download_paper_text(_ARXIV_ID) + + a_wf, _a_stats = _build(doc, speculative=False) + a_time, a_out = await _run(a_wf) + + b_wf, b_stats = _build(doc, speculative=True) + b_time, b_out = await _run(b_wf) + + saved = a_time - b_time + print( + f'\n[speculative dispatch] model={MODEL} doc_chars={len(doc)} ' + '(whole paper, no chunking)\n' + f' handler work (overlappable): {_HANDLER_SECONDS:6.2f}s\n' + f' A sequential (gen, then run): {a_time:6.2f}s topic={_topic_out(a_out)!r}\n' + f' B speculative (run overlaps): {b_time:6.2f}s topic={_topic_out(b_out)!r}\n' + f' wall-clock saved by overlap: {saved:6.2f}s' + f' (dispatches={b_stats["dispatches"]})\n' + ) + + # Both strategies produce the handled result. + assert isinstance(a_out, dict) and a_out.get('handled') + assert isinstance(b_out, dict) and b_out.get('handled') + # Speculation dispatched the handler early (at least once). + assert b_stats['dispatches'] >= 1 + # Overlap hides most of the handler's cost behind generation. + assert b_time < a_time + assert saved > _HANDLER_SECONDS * 0.5 + + +def _topic_out(out: Any) -> str: + return out.get('topic', '') if isinstance(out, dict) else '' diff --git a/tests/integration/test_speculative_router_chained_llm.py b/tests/integration/test_speculative_router_chained_llm.py new file mode 100644 index 00000000000..8de79912e61 --- /dev/null +++ b/tests/integration/test_speculative_router_chained_llm.py @@ -0,0 +1,243 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Real-LLM: two chained model calls, where speculation overlaps the second. + +This is the *real* use case for ``SpeculativeRouterNode`` -- no sleeps, no +simulated work. A DAG of two LLM calls: + + planner (LLM) --topic--> worker (LLM) + +The ``planner`` reads a whole arXiv paper (NO chunking) and is told to FIRST emit +a routing directive naming the topic to dig into: + + TOOL_CALL: {"name": "explain", "arguments": {"topic": ""}} + +...THEN write a long rationale. The ``worker`` is a second, real LLM call that +writes the detailed answer for that topic. It only depends on the *directive*, +which streams out at the very start -- long before the planner finishes its +rationale. + +Honest framing (no "just plan" escape) +-------------------------------------- +The planner's rationale is a **required deliverable**: the workflow returns +``{"plan": , "answer": }`` via the +node's ``combine`` hook, and the test asserts BOTH are present and substantial +for *both* strategies. So a baseline cannot cheat by telling the planner to emit +only the directive and stop -- it is contractually obligated to produce the whole +rationale, exactly like the speculative path. The only difference between A and B +is *when* the worker runs: + + A. sequential -- produce the whole rationale, THEN call the worker. + B. speculative -- the instant the directive streams in, start the worker so + its generation overlaps the (still-required) rationale tail; + verify against the finalized directive. + +Both emit identical deliverables; B just hides the worker's genuine multi-second +generation behind work the caller already demanded. The assertion ties the saving +to the worker's *measured* solo runtime -- not a magic constant -- so it reflects +real overlap. + + ADK_TEST_MODEL=gemini-3.5-flash-lite \\ + uv run pytest -s -p no:cacheprovider \\ + tests/integration/test_speculative_router_chained_llm.py + +Requires Vertex (GOOGLE_CLOUD_PROJECT via ADC), ``pypdf`` and arXiv; skips +otherwise. +""" + +import os +import pathlib +import tempfile +import time +from typing import Any +import urllib.error +import urllib.request + +from dotenv import load_dotenv +from google.adk import Agent +from google.adk import Event +from google.adk import Workflow +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.workflow import make_marker_extractor +from google.adk.workflow import SpeculativeRouterNode +from google.genai import types +import pytest + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] +load_dotenv(_REPO_ROOT / '.env', override=False) + +pytestmark = pytest.mark.skipif( + not os.environ.get('GOOGLE_CLOUD_PROJECT'), + reason='Real-LLM speculative test requires Vertex (GOOGLE_CLOUD_PROJECT).', +) + +MODEL = os.environ.get('ADK_TEST_MODEL', 'gemini-3.5-flash-lite') +_ARXIV_ID = '1706.03762' # Attention Is All You Need + +_DIRECTIVE = ( + 'FIRST output EXACTLY one line, nothing else on it:\n' + 'TOOL_CALL: {"name": "explain", "arguments": {"topic":' + ' ""}}\n' + 'THEN write a long, detailed rationale for your choice (at least 300' + ' words).' +) + +_WORKER_INSTRUCTION = ( + 'You are given the name of a technical method from a research paper. Write' + ' a detailed, ~200 word technical explanation of that method: what it is,' + ' how it works, and why it matters.' +) + + +def _download_paper_text(arxiv_id: str) -> str: + from pypdf import PdfReader + + cache = pathlib.Path(tempfile.gettempdir()) / 'adk_arxiv_papers' + cache.mkdir(parents=True, exist_ok=True) + pdf_path = cache / (arxiv_id.replace('/', '_') + '.pdf') + if not pdf_path.exists(): + url = 'https://arxiv.org/pdf/' + arxiv_id + req = urllib.request.Request(url, headers={'User-Agent': 'adk-test/1.0'}) + pdf_path.write_bytes(urllib.request.urlopen(req, timeout=60).read()) # noqa: S310 + reader = PdfReader(str(pdf_path)) + return ''.join((page.extract_text() or '') for page in reader.pages) + + +def _planner_prompt(doc: str) -> str: + return ( + f'PAPER:\n{doc}\n\n----------------------------------------\n' + 'You just read the paper above.\n\n' + _DIRECTIVE + ) + + +def _extract_topic(text: str) -> Any: + """Returns just the directive's ``topic`` string (or None), repairing partials.""" + obj = make_marker_extractor('TOOL_CALL:')(text) + if not isinstance(obj, dict): + return None + topic = obj.get('arguments', {}).get('topic') + return topic or None + + +def _worker_agent() -> Agent: + return Agent(name='worker', model=MODEL, instruction=_WORKER_INSTRUCTION) + + +def _build_chain(doc: str, *, speculative: bool) -> Workflow: + planner = Agent( + name='planner', + model=MODEL, + instruction=lambda _ctx, _p=_planner_prompt(doc): _p, + ) + node = SpeculativeRouterNode( + name='plan_then_explain', + agent=planner, + target=_worker_agent(), + extract=_extract_topic, + should_speculate=( + (lambda topic: bool(topic) and len(topic) >= 4) + if speculative + else (lambda _topic: False) + ), + same=lambda a, b: (a or '').strip().lower() == (b or '').strip().lower(), + # The planner's full rationale is a required deliverable: both strategies + # must return it alongside the worker's answer, so B has no "just plan and + # stop" shortcut -- it produces the exact same artifact, only faster. + combine=lambda plan, answer: {'plan': plan, 'answer': _text(answer)}, + timeout=120, + ) + return Workflow(name='chain', edges=[('START', node)]) + + +def _worker_only(topic: str) -> Workflow: + def give_topic(node_input: Any): + yield Event(output=topic) + + return Workflow(name='worker_only', edges=[('START', give_topic, _worker_agent())]) + + +async def _run(wf: Workflow) -> tuple[float, Any]: + ss = InMemorySessionService() + runner = Runner(app_name=wf.name, node=wf, session_service=ss) + session = await ss.create_session(app_name=wf.name, user_id='u') + msg = types.Content(parts=[types.Part(text='go')], role='user') + output = None + start = time.perf_counter() + async for event in runner.run_async( + user_id='u', session_id=session.id, new_message=msg + ): + if isinstance(event, Event) and event.output is not None: + output = event.output + return time.perf_counter() - start, output + + +def _text(out: Any) -> str: + if isinstance(out, str): + return out + if isinstance(out, types.Content) and out.parts: + return ''.join(p.text or '' for p in out.parts) + return str(out) if out is not None else '' + + +def _plan(out: Any) -> str: + """The planner's full rationale from the ``{plan, answer}`` deliverable.""" + return out.get('plan', '') if isinstance(out, dict) else '' + + +def _answer(out: Any) -> str: + """The worker's explanation from the ``{plan, answer}`` deliverable.""" + return out.get('answer', '') if isinstance(out, dict) else _text(out) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('llm_backend', ['VERTEX'], indirect=True) +async def test_speculative_chain_overlaps_worker_llm(llm_backend): + doc = _download_paper_text(_ARXIV_ID) + + # Measure the worker's solo runtime so the saving is tied to real work. + w_time, _w_out = await _run(_worker_only('the self-attention mechanism')) + + a_time, a_out = await _run(_build_chain(doc, speculative=False)) + b_time, b_out = await _run(_build_chain(doc, speculative=True)) + + saved = a_time - b_time + print( + f'\n[speculative chain] model={MODEL} doc_chars={len(doc)} ' + '(two chained LLM calls, whole paper, no chunking)\n' + f' worker LLM alone: {w_time:6.2f}s\n' + f' A sequential (plan -> work): {a_time:6.2f}s' + f' plan={len(_plan(a_out)):5d} chars, answer={len(_answer(a_out)):4d}' + ' chars\n' + f' B speculative (work overlaps):{b_time:6.2f}s' + f' plan={len(_plan(b_out)):5d} chars, answer={len(_answer(b_out)):4d}' + ' chars\n' + f' wall-clock saved by overlap: {saved:6.2f}s' + f' ({100 * saved / w_time:4.0f}% of the worker call hidden)\n' + ) + + # Honest deliverable: BOTH strategies must return the planner's full rationale + # AND the worker's answer. The rationale is required, so a baseline can't cheat + # by "just planning" and skipping the tail -- both produce the same artifact. + for out in (a_out, b_out): + assert len(_plan(out)) > 200, 'planner rationale is a required deliverable' + assert len(_answer(out)) > 80, 'worker answer is a required deliverable' + # Speculation overlaps the worker with the (still-required) planner tail -> + # faster while producing the identical two-part deliverable... + assert b_time < a_time + # ...by a meaningful fraction of the worker's *measured* real runtime (not a + # magic constant). Overlap can't hide more than the worker takes; requiring + # >=30% keeps this robust to run-to-run LLM variance. + assert saved > 0.30 * w_time diff --git a/tests/integration/test_streaming_router_preemption_timing.py b/tests/integration/test_streaming_router_preemption_timing.py new file mode 100644 index 00000000000..ffc43753e15 --- /dev/null +++ b/tests/integration/test_streaming_router_preemption_timing.py @@ -0,0 +1,427 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Real-LLM: fan out over N whole documents, keep the summaries you asked for, +and skip the ones you were going to throw away. + +NO chunking. Each document is handed to Gemini whole, in a single streaming API +call. Each reader must FIRST emit a ``VERDICT:`` line (RELEVANT / IRRELEVANT), +THEN write a long (>=300 word) summary of the paper. + +Honest framing (no "just emit the verdict" escape) +-------------------------------------------------- +The summary of a **RELEVANT** paper is a *required deliverable* -- the caller +consumes it -- so both strategies must produce it in full. The only legitimate +saving is refusing to summarize the papers you have already judged IRRELEVANT and +are about to discard. So the two strategies differ only in that: + + A. read + answer -- every reader streams its full summary, even + the four papers that turn out to be + irrelevant (4 summaries with no consumer). + B. read + answer + preemption -- the monitor cancels a reader the instant it + declares itself IRRELEVANT, but lets a + RELEVANT reader stream its summary to + completion. + +Both classify all five papers correctly AND both return the full summary of the +one relevant paper (asserted). B is faster/cheaper purely because it does not pay +to summarize the four documents it is discarding -- not because A was rigged to +manufacture output nobody wanted. + +Five real arXiv papers are read in parallel; only "Attention Is All You Need" is +a CS AI/ML paper. + + ADK_TEST_MODEL=gemini-3.5-flash-lite \\ + uv run pytest -s -p no:cacheprovider \\ + tests/integration/test_streaming_router_preemption_timing.py + +Requires Vertex (GOOGLE_CLOUD_PROJECT via ADC), ``pypdf`` and network access to +arXiv; skips otherwise. +""" + +import os +import pathlib +import tempfile +import time +from typing import Any +from typing import Optional +import urllib.error +import urllib.request + +from dotenv import load_dotenv +from google import genai +from google.adk import Agent +from google.adk import Event +from google.adk import Workflow +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.workflow import JoinNode +from google.adk.workflow import StreamDecision +from google.adk.workflow import StreamingRouterNode +from google.adk.workflow import StreamView +from google.genai import types +import pytest + +# Load the repo-root .env (Vertex project/location/model live there). +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] +load_dotenv(_REPO_ROOT / '.env', override=False) + +# Vertex backend (ADC auth, no API key). Skip unless a project is configured. +pytestmark = pytest.mark.skipif( + not os.environ.get('GOOGLE_CLOUD_PROJECT'), + reason='Real-LLM timing test requires Vertex (GOOGLE_CLOUD_PROJECT).', +) + +# gemini-3.5-flash-lite by default (verified on Vertex). Override w/ ADK_TEST_MODEL. +MODEL = os.environ.get('ADK_TEST_MODEL', 'gemini-3.5-flash-lite') + +# Published gemini-3.5-flash-lite pricing (USD per 1M tokens). Output is 8.3x +# input; the cached-input rate is 10x cheaper than a fresh read. +_PRICE_IN = 0.30 / 1e6 +_PRICE_OUT = 2.50 / 1e6 +_PRICE_IN_CACHED = 0.03 / 1e6 + +QUERY = ( + 'Is this paper about artificial intelligence, machine learning, or neural' + ' networks (i.e. a computer-science AI paper)?' +) + +# The relevant paper is first; the other four are real but not AI papers. +_PAPERS: list[tuple[str, str, bool]] = [ + ('1706.03762', 'Attention Is All You Need', True), + ('1602.03837', 'Observation of Gravitational Waves (GW150914)', False), + ('1207.7214', 'Observation of the Higgs boson (ATLAS)', False), + ('astro-ph/9805201', 'Accelerating universe / dark energy', False), + ('math/0211159', 'The entropy formula for the Ricci flow', False), +] + +_VERDICT_MARKER = 'VERDICT:' +_CACHE_DIR = pathlib.Path(tempfile.gettempdir()) / 'adk_arxiv_papers' + +_CLIENT: Optional[genai.Client] = None + + +def _client() -> genai.Client: + """A process-wide Vertex genai client, used only for token counting.""" + global _CLIENT + if _CLIENT is None: + _CLIENT = genai.Client( + vertexai=True, + project=os.environ['GOOGLE_CLOUD_PROJECT'], + location=os.environ.get('GOOGLE_CLOUD_LOCATION', 'global'), + ) + return _CLIENT + + +def _count_tokens(text: str) -> int: + """Real token count for ``text`` via the model's tokenizer (0 if empty).""" + if not text.strip(): + return 0 + return _client().models.count_tokens(model=MODEL, contents=text).total_tokens + + +def _download_paper_text(arxiv_id: str) -> str: + """Downloads an arXiv PDF (cached) and returns its full extracted text.""" + from pypdf import PdfReader + + _CACHE_DIR.mkdir(parents=True, exist_ok=True) + pdf_path = _CACHE_DIR / (arxiv_id.replace('/', '_') + '.pdf') + if not pdf_path.exists(): + url = 'https://arxiv.org/pdf/' + arxiv_id + req = urllib.request.Request(url, headers={'User-Agent': 'adk-test/1.0'}) + data = urllib.request.urlopen(req, timeout=60).read() # noqa: S310 + pdf_path.write_bytes(data) + reader = PdfReader(str(pdf_path)) + return ''.join((page.extract_text() or '') for page in reader.pages) + + +def _load_documents() -> list[str]: + """Fetches every paper as one whole-document string; skips if unavailable.""" + pytest.importorskip('pypdf', reason='PDF extraction needs pypdf.') + docs: list[str] = [] + for arxiv_id, label, _ in _PAPERS: + try: + text = _download_paper_text(arxiv_id) + except (urllib.error.URLError, TimeoutError) as e: + pytest.skip(f'Could not fetch arXiv:{arxiv_id} ({label}): {e}') + docs.append(text) # whole document, no truncation + return docs + + +def _verdict_after_marker(text: str) -> Optional[str]: + """Returns the text after ``VERDICT:`` once that line has fully streamed.""" + idx = text.upper().find(_VERDICT_MARKER) + if idx == -1: + return None + after = text[idx + len(_VERDICT_MARKER) :] + line, sep, _ = after.partition('\n') + if not sep: + return None # verdict line still streaming + return line.strip() + + +def _summary_body(text: str) -> str: + """The summary the reader wrote *after* its ``VERDICT:`` line (or '').""" + idx = text.upper().find(_VERDICT_MARKER) + if idx == -1: + return text.strip() # no verdict line; treat the whole thing as body + after = text[idx + len(_VERDICT_MARKER) :] + _, sep, rest = after.partition('\n') + return rest.strip() if sep else '' + + +def _word_count(text: str) -> int: + return len(text.split()) + + +def _decide(view: StreamView) -> Optional[StreamDecision]: + """Preemption decision: cut ONLY the papers we don't want summarized. + + This is the honest crux. The summary is a *required deliverable for a RELEVANT + paper*, so we must NOT preempt those -- we let the relevant reader stream its + full summary to completion. We only cut a reader once it has declared itself + IRRELEVANT, because for those papers the summary is genuinely unwanted: we have + everything we need (the verdict) and refuse to pay to summarize a document we + are about to discard. + """ + verdict = _verdict_after_marker(view.text) + if verdict is None: + return None + if verdict.upper().startswith('IRRELEVANT'): + return StreamDecision(output={'relevant': False, 'verdict': 'IRRELEVANT'}) + # RELEVANT (or anything else): do NOT preempt -- keep generating the summary, + # which is a deliverable the caller actually consumes. + return None + + +def _reader_prompt(doc: str) -> str: + """The whole-document, single-call prompt handed to one reader.""" + return ( + # Whole document first, then the question -- one prompt, one call. + f'PAPER:\n{doc}\n\n----------------------------------------\nYou' + ' just read the paper above. Judge ONLY from it.\n\nQUESTION:' + f' {QUERY}\n\nFIRST output a line beginning "VERDICT:" that is' + ' exactly one of:\n VERDICT: IRRELEVANT (NOT a' + ' computer-science AI/ML/neural-networks paper -- e.g. physics,' + ' astronomy, or pure mathematics, even if it uses the word' + ' "model"), or\n VERDICT: RELEVANT - (a' + ' computer-science paper about AI, machine learning, or neural' + ' networks -- name its subject).\nTHEN write a long, detailed,' + ' multi-paragraph summary of the paper (at least 300 words).' + ) + + +def _build_workflow( + docs: list[str], + *, + preempt: bool, + sink: dict[str, Any], + gen: dict[int, str], +) -> Workflow: + readers = [] + for i, doc in enumerate(docs): + prompt = _reader_prompt(doc) + reader = Agent( + name=f'reader_{i}', + model=MODEL, + # A callable (provider) instruction bypasses {var} state-injection, so + # raw LaTeX/math braces in the papers are sent verbatim -- no escaping, + # no truncation. + instruction=lambda _ctx, _p=prompt: _p, + ) + + def monitor(view: StreamView, _i: int = i) -> Optional[StreamDecision]: + # Capture the latest streamed output text so we can count the tokens the + # model actually generated under each strategy. + gen[_i] = view.text + return _decide(view) if preempt else None + + readers.append( + StreamingRouterNode( + name=f'reader_{i}', + agent=reader, + monitor=monitor, + forward_partials=False, + timeout=180, + ) + ) + + join = JoinNode(name='join_sources') + + async def collect(node_input: dict[str, Any]): + sink['fan_in'] = node_input + yield Event(message='done') + + return Workflow( + name='parallel_read_answer', + edges=[('START', tuple(readers), join, collect)], + ) + + +async def _run(wf: Workflow) -> float: + ss = InMemorySessionService() + runner = Runner(app_name=wf.name, node=wf, session_service=ss) + session = await ss.create_session(app_name=wf.name, user_id='u') + msg = types.Content(parts=[types.Part(text='go')], role='user') + start = time.perf_counter() + async for _ in runner.run_async( + user_id='u', session_id=session.id, new_message=msg + ): + pass + return time.perf_counter() - start + + +def _status_and_verdict(value: Any) -> tuple[str, str]: + """Normalises a reader's output (dict from B, raw text from A) to a verdict.""" + if isinstance(value, dict): + return ( + 'RELEVANT' if value.get('relevant') else 'IRRELEVANT', + str(value.get('verdict', '')), + ) + verdict = _verdict_after_marker(str(value)) or '' + status = ( + 'RELEVANT' if verdict.upper().startswith('RELEVANT') else 'IRRELEVANT' + ) + return status, verdict + + +def _value_for_index(fan_in: dict[str, Any], index: int) -> Any: + for key, value in fan_in.items(): + digits = ''.join(ch for ch in key if ch.isdigit()) + if digits and int(digits) == index: + return value + raise KeyError(f'No fan-in entry for reader index {index}: {list(fan_in)}') + + +def _generated_text( + index: int, fan_in: dict[str, Any], gen: dict[int, str] +) -> str: + """The text a reader actually generated: full answer (A) or up-to-cut (B).""" + value = _value_for_index(fan_in, index) + if isinstance(value, str) and value.strip(): + return value # A: the whole streamed answer is the node output + return gen.get(index, '') # B: streamed text captured up to preemption + + +def _output_tokens(fan_in: dict[str, Any], gen: dict[int, str]) -> list[int]: + return [ + _count_tokens(_generated_text(i, fan_in, gen)) + for i in range(len(_PAPERS)) + ] + + +def _print_qa( + title: str, fan_in: dict[str, Any], out_tokens: list[int] +) -> None: + print(f'\n===== {title} =====') + print(f'Q: {QUERY}\n') + for i, (arxiv_id, label, _relevant) in enumerate(_PAPERS): + status, verdict = _status_and_verdict(_value_for_index(fan_in, i)) + mark = 'RELEVANT ' if status == 'RELEVANT' else 'irrelevant' + print(f' [{mark}] {out_tokens[i]:5d} out-tok {label} (arXiv:{arxiv_id})') + print(f' -> {verdict or status}') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('llm_backend', ['VERTEX'], indirect=True) +async def test_sse_preemption_beats_full_generation(llm_backend): + docs = _load_documents() + + # A: read + answer, stream to completion. + a_sink: dict[str, Any] = {} + a_gen: dict[int, str] = {} + a_time = await _run( + _build_workflow(docs, preempt=False, sink=a_sink, gen=a_gen) + ) + + # B: read + answer, SSE + preemption (cut once the verdict streams in). + b_sink: dict[str, Any] = {} + b_gen: dict[int, str] = {} + b_time = await _run( + _build_workflow(docs, preempt=True, sink=b_sink, gen=b_gen) + ) + + # Real token counts. Input is identical for A and B (same whole-doc prompts); + # the difference is entirely in generated output tokens. + input_tokens = sum(_count_tokens(_reader_prompt(doc)) for doc in docs) + a_out = _output_tokens(a_sink['fan_in'], a_gen) + b_out = _output_tokens(b_sink['fan_in'], b_gen) + a_out_total, b_out_total = sum(a_out), sum(b_out) + + speedup = a_time / b_time if b_time else float('inf') + tok_ratio = a_out_total / b_out_total if b_out_total else float('inf') + + # Cost (gemini-3.5-flash-lite pricing). Input is identical for A and B; the + # only difference is output tokens. With context caching the input read is 10x + # cheaper, so preemption's output savings dominate the total. + a_cost = input_tokens * _PRICE_IN + a_out_total * _PRICE_OUT + b_cost = input_tokens * _PRICE_IN + b_out_total * _PRICE_OUT + a_cost_cached = input_tokens * _PRICE_IN_CACHED + a_out_total * _PRICE_OUT + b_cost_cached = input_tokens * _PRICE_IN_CACHED + b_out_total * _PRICE_OUT + save = 100 * (1 - b_cost / a_cost) if a_cost else 0.0 + save_cached = ( + 100 * (1 - b_cost_cached / a_cost_cached) if a_cost_cached else 0 + ) + + print( + f'\n[SSE preemption] model={MODEL} docs={len(_PAPERS)} (whole docs, no' + ' chunking)\n' + f' input tokens (both): {input_tokens:6d}\n' + f' A read+answer: {a_time:6.2f}s ' + f'{a_out_total:6d} out-tok\n' + f' B read+answer+preemption: {b_time:6.2f}s ' + f'{b_out_total:6d} out-tok\n' + f' speedup: {speedup:5.2f}x\n' + f' output tokens saved: {a_out_total - b_out_total:6d}' + f' (B uses {tok_ratio:4.2f}x fewer)\n' + ' cost @ $0.30/$2.50 per 1M (in/out):\n' + f' A ${a_cost:.4f} B ${b_cost:.4f} (B saves {save:4.1f}%)\n' + ' cost w/ context caching (in @ $0.03/1M):\n' + f' A ${a_cost_cached:.4f} B ${b_cost_cached:.4f}' + f' (B saves {save_cached:4.1f}%)\n' + ) + _print_qa('A: read + answer (stream to completion)', a_sink['fan_in'], a_out) + _print_qa('B: read + answer + SSE preemption', b_sink['fan_in'], b_out) + + # The RELEVANT paper's summary is a required deliverable -- in BOTH strategies + # the reader must have streamed a full (>=300 word) summary, not just a verdict. + # This is what kills the weasel: B cannot win by skipping wanted output, only by + # skipping the summaries of papers it judged irrelevant. + a_rel_summary = _summary_body(_generated_text(0, a_sink['fan_in'], a_gen)) + b_rel_summary = _summary_body(_generated_text(0, b_sink['fan_in'], b_gen)) + print( + ' relevant-paper summary kept (a required deliverable):\n' + f' A {_word_count(a_rel_summary):4d} words ' + f'B {_word_count(b_rel_summary):4d} words\n' + ) + + # Both strategies must classify all five papers correctly. + for fan_in in (a_sink['fan_in'], b_sink['fan_in']): + assert _status_and_verdict(_value_for_index(fan_in, 0))[0] == 'RELEVANT' + for idx in range(1, len(_PAPERS)): + assert _status_and_verdict(_value_for_index(fan_in, idx))[0] == ( + 'IRRELEVANT' + ) + + # Both must actually deliver the relevant paper's full summary (>=250 words -- + # a small margin under the requested 300 for LLM variance). B is NOT allowed to + # drop the deliverable the caller wanted. + assert _word_count(a_rel_summary) >= 250, 'A must summarize the relevant paper' + assert _word_count(b_rel_summary) >= 250, 'B must keep the relevant summary' + + # Preemption only skips the FOUR irrelevant summaries the caller discards -> B + # generates meaningfully fewer output tokens... + assert b_out_total < a_out_total + # ...and it is faster. + assert b_time < a_time diff --git a/tests/integration/test_streaming_router_tsla_10k.py b/tests/integration/test_streaming_router_tsla_10k.py new file mode 100644 index 00000000000..dce7f09622e --- /dev/null +++ b/tests/integration/test_streaming_router_tsla_10k.py @@ -0,0 +1,300 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Real-LLM demo: StreamingRouterNode reads a whole Tesla 10-K and preempts. + +This is the arXiv timing test's big-document sibling. Instead of five small +papers it hands Gemini a single, enormous real-world filing -- the latest Tesla +annual report (Form 10-K), pulled live from SEC EDGAR (~100k+ input tokens) -- +in one whole-document streaming call. NO chunking. + +The point it makes on a huge doc: the input is paid for once (prefill), so the +two strategies differ only in how much they *generate*: + + A. read + answer -- stream the model's answer to completion. + B. read + answer + preemption -- the StreamingRouterNode monitor watches the + SSE tokens and cancels the stream the moment + the "VERDICT:" line has streamed in, so the + long analysis that follows is never decoded. + +B is dramatically faster (no long decode) and, once the static filing is context +-cached, dramatically cheaper -- the output cut becomes the whole bill. + + ADK_TEST_MODEL=gemini-3.5-flash-lite \\ + uv run pytest -s -p no:cacheprovider \\ + tests/integration/test_streaming_router_tsla_10k.py + +Requires Vertex (GOOGLE_CLOUD_PROJECT via ADC) and network access to SEC EDGAR; +skips otherwise. +""" + +import os +import pathlib +import re +import time +from typing import Any +from typing import Optional +import urllib.error +import urllib.request + +from dotenv import load_dotenv +from google import genai +from google.adk import Agent +from google.adk import Event +from google.adk import Workflow +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.workflow import StreamDecision +from google.adk.workflow import StreamingRouterNode +from google.adk.workflow import StreamView +from google.genai import types +import pytest + +# Load the repo-root .env (Vertex project/location/model live there). +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] +load_dotenv(_REPO_ROOT / '.env', override=False) + +# Vertex backend (ADC auth, no API key). Skip unless a project is configured. +pytestmark = pytest.mark.skipif( + not os.environ.get('GOOGLE_CLOUD_PROJECT'), + reason='Real-LLM timing test requires Vertex (GOOGLE_CLOUD_PROJECT).', +) + +# gemini-3.5-flash-lite by default (verified on Vertex). Override w/ ADK_TEST_MODEL. +MODEL = os.environ.get('ADK_TEST_MODEL', 'gemini-3.5-flash-lite') + +# Published gemini-3.5-flash-lite pricing (USD per 1M tokens). Output is 8.3x +# input; the cached-input rate is 10x cheaper than a fresh read. +_PRICE_IN = 0.30 / 1e6 +_PRICE_OUT = 2.50 / 1e6 +_PRICE_IN_CACHED = 0.03 / 1e6 + +# Tesla, Inc. central index key on SEC EDGAR. +_TSLA_CIK = '0001318605' +# SEC requires a descriptive User-Agent with contact info or it returns 403. +_SEC_UA = {'User-Agent': 'google-adk integration-test contact@example.com'} + +QUERY = ( + 'What are the three most significant risk factors Tesla identifies in this' + ' filing?' +) +_VERDICT_MARKER = 'VERDICT:' + +_CLIENT: Optional[genai.Client] = None + + +def _client() -> genai.Client: + """A process-wide Vertex genai client, used only for token counting.""" + global _CLIENT + if _CLIENT is None: + _CLIENT = genai.Client( + vertexai=True, + project=os.environ['GOOGLE_CLOUD_PROJECT'], + location=os.environ.get('GOOGLE_CLOUD_LOCATION', 'global'), + ) + return _CLIENT + + +def _count_tokens(text: str) -> int: + """Real token count for ``text`` via the model's tokenizer (0 if empty).""" + if not text.strip(): + return 0 + return _client().models.count_tokens(model=MODEL, contents=text).total_tokens + + +def _get(url: str) -> bytes: + req = urllib.request.Request(url, headers=_SEC_UA) + return urllib.request.urlopen(req, timeout=60).read() # noqa: S310 + + +def _html_to_text(html: str) -> str: + """Crudely strips a 10-K .htm down to readable prose (no external deps).""" + html = re.sub(r'(?is)<(script|style).*?', ' ', html) + html = re.sub(r'(?is)', '\n', html) + html = re.sub(r'(?is)', '\n', html) + text = re.sub(r'(?is)<[^>]+>', ' ', html) + import html as _htmllib + + text = _htmllib.unescape(text) + text = re.sub(r'[ \t]+', ' ', text) + return re.sub(r'\n\s*\n+', '\n\n', text).strip() + + +def _load_tsla_10k() -> str: + """Fetches the latest Tesla 10-K primary document as whole text (no cache).""" + import json + + try: + subs = json.loads(_get(f'https://data.sec.gov/submissions/CIK{_TSLA_CIK}.json')) + recent = subs['filings']['recent'] + idx = next(i for i, f in enumerate(recent['form']) if f == '10-K') + accession = recent['accessionNumber'][idx].replace('-', '') + document = recent['primaryDocument'][idx] + url = ( + f'https://www.sec.gov/Archives/edgar/data/{int(_TSLA_CIK)}/' + f'{accession}/{document}' + ) + return _html_to_text(_get(url).decode('utf-8', 'ignore')) + except (urllib.error.URLError, TimeoutError, StopIteration, KeyError) as e: + pytest.skip(f'Could not fetch Tesla 10-K from SEC EDGAR: {e}') + + +def _verdict_after_marker(text: str) -> Optional[str]: + """Returns the text after ``VERDICT:`` once that line has fully streamed.""" + idx = text.upper().find(_VERDICT_MARKER) + if idx == -1: + return None + line, sep, _ = text[idx + len(_VERDICT_MARKER) :].partition('\n') + return line.strip() if sep else None + + +def _reader_prompt(doc: str) -> str: + """The whole-document, single-call prompt handed to the reader.""" + return ( + # Whole filing first, then the question -- one prompt, one call. + f'FILING (Tesla annual report / Form 10-K):\n{doc}\n\n' + '----------------------------------------\n' + 'You just read the filing above. Answer ONLY from it.\n\n' + f'QUESTION: {QUERY}\n\n' + 'FIRST output a line beginning "VERDICT:" that names the top three risks' + ' in one sentence.\nTHEN write a detailed multi-paragraph analysis (at' + ' least 400 words).' + ) + + +def _build_workflow( + doc: str, + *, + preempt: bool, + sink: dict[str, Any], + gen: dict[str, str], +) -> Workflow: + prompt = _reader_prompt(doc) + # A callable (provider) instruction bypasses {var} state-injection, so raw + # braces in the filing are sent verbatim -- no escaping, no truncation. + reader = Agent( + name='reader', model=MODEL, instruction=lambda _ctx, _p=prompt: _p + ) + + def monitor(view: StreamView) -> Optional[StreamDecision]: + # Capture the latest streamed text so we can count generated tokens. + gen['text'] = view.text + if not preempt: + return None + verdict = _verdict_after_marker(view.text) + return StreamDecision(output={'answer': verdict}) if verdict else None + + node = StreamingRouterNode( + name='reader', + agent=reader, + monitor=monitor, + forward_partials=False, + timeout=300, + ) + + async def collect(node_input: Any): + sink['out'] = node_input + yield Event(message='done') + + return Workflow(name='tsla_10k', edges=[('START', node, collect)]) + + +async def _run(wf: Workflow) -> float: + ss = InMemorySessionService() + runner = Runner(app_name=wf.name, node=wf, session_service=ss) + session = await ss.create_session(app_name=wf.name, user_id='u') + msg = types.Content(parts=[types.Part(text='go')], role='user') + start = time.perf_counter() + async for _ in runner.run_async( + user_id='u', session_id=session.id, new_message=msg + ): + pass + return time.perf_counter() - start + + +def _generated_text(sink: dict[str, Any], gen: dict[str, str]) -> str: + """The text the reader actually generated: full answer (A) or up-to-cut (B).""" + value = sink.get('out') + if isinstance(value, str) and value.strip(): + return value # A: the whole streamed answer is the node output + return gen.get('text', '') # B: streamed text captured up to preemption + + +def _verdict_of(sink: dict[str, Any], gen: dict[str, str]) -> str: + value = sink.get('out') + if isinstance(value, dict): + return str(value.get('answer') or '') + return _verdict_after_marker(_generated_text(sink, gen)) or '' + + +@pytest.mark.asyncio +@pytest.mark.parametrize('llm_backend', ['VERTEX'], indirect=True) +async def test_sse_preemption_on_tsla_10k(llm_backend): + doc = _load_tsla_10k() + + # A: read + answer, stream to completion. + a_sink: dict[str, Any] = {} + a_gen: dict[str, str] = {} + a_time = await _run(_build_workflow(doc, preempt=False, sink=a_sink, gen=a_gen)) + + # B: read + answer, SSE + preemption (cut once the verdict streams in). + b_sink: dict[str, Any] = {} + b_gen: dict[str, str] = {} + b_time = await _run(_build_workflow(doc, preempt=True, sink=b_sink, gen=b_gen)) + + # Real token counts. Input is identical for A and B (same whole-doc prompt); + # the difference is entirely in generated output tokens. + input_tokens = _count_tokens(_reader_prompt(doc)) + a_out = _count_tokens(_generated_text(a_sink, a_gen)) + b_out = _count_tokens(_generated_text(b_sink, b_gen)) + + speedup = a_time / b_time if b_time else float('inf') + tok_ratio = a_out / b_out if b_out else float('inf') + + # Cost (gemini-3.5-flash-lite pricing). Input is identical for A and B; only + # output differs. With context caching the huge filing is 10x cheaper to read, + # so preemption's output savings dominate the total. + a_cost = input_tokens * _PRICE_IN + a_out * _PRICE_OUT + b_cost = input_tokens * _PRICE_IN + b_out * _PRICE_OUT + a_cost_cached = input_tokens * _PRICE_IN_CACHED + a_out * _PRICE_OUT + b_cost_cached = input_tokens * _PRICE_IN_CACHED + b_out * _PRICE_OUT + save = 100 * (1 - b_cost / a_cost) if a_cost else 0.0 + save_cached = 100 * (1 - b_cost_cached / a_cost_cached) if a_cost_cached else 0 + + print( + f'\n[SSE preemption / TSLA 10-K] model={MODEL} doc_chars={len(doc)}' + ' (whole filing, no chunking)\n' + f' input tokens (both): {input_tokens:6d}\n' + f' A read+answer: {a_time:6.2f}s {a_out:6d} out-tok\n' + f' B read+answer+preemption: {b_time:6.2f}s {b_out:6d} out-tok\n' + f' speedup: {speedup:5.2f}x\n' + f' output tokens saved: {a_out - b_out:6d}' + f' (B uses {tok_ratio:4.2f}x fewer)\n' + ' cost @ $0.30/$2.50 per 1M (in/out):\n' + f' A ${a_cost:.4f} B ${b_cost:.4f} (B saves {save:4.1f}%)\n' + ' cost w/ context caching (in @ $0.03/1M):\n' + f' A ${a_cost_cached:.4f} B ${b_cost_cached:.4f}' + f' (B saves {save_cached:4.1f}%)\n' + ) + print(f'Q: {QUERY}') + print(f'A verdict: {_verdict_of(a_sink, a_gen)}') + print(f'B verdict: {_verdict_of(b_sink, b_gen)}') + + # The model must actually answer from the filing under both strategies. + assert _verdict_of(a_sink, a_gen) + assert _verdict_of(b_sink, b_gen) + # Preemption stops generating the long analysis -> fewer output tokens... + assert b_out < a_out + # ...and it is faster. + assert b_time < a_time diff --git a/tests/unittests/utils/test_event_loop.py b/tests/unittests/utils/test_event_loop.py new file mode 100644 index 00000000000..2b515c3aecc --- /dev/null +++ b/tests/unittests/utils/test_event_loop.py @@ -0,0 +1,84 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the uvloop event-loop helpers.""" + +import asyncio + +from google.adk.utils import event_loop +import pytest + + +@pytest.fixture(autouse=True) +def _reset_install_flag(monkeypatch): + # Isolate each test from the process-wide install flag and restore the + # global event-loop policy so installing uvloop here does not leak into + # the rest of the test session. + monkeypatch.setattr(event_loop, '_uvloop_installed', False) + original_policy = asyncio.get_event_loop_policy() + try: + yield + finally: + asyncio.set_event_loop_policy(original_policy) + + +def test_enable_uvloop_installs_policy_when_available(): + if not event_loop.is_uvloop_available(): + pytest.skip('uvloop not installed in this environment.') + + assert event_loop.enable_uvloop() is True + + async def _loop_module() -> str: + return type(asyncio.get_running_loop()).__module__ + + assert asyncio.run(_loop_module()).startswith('uvloop') + assert event_loop.is_uvloop_active() is True + + +def test_enable_uvloop_is_idempotent(): + if not event_loop.is_uvloop_available(): + pytest.skip('uvloop not installed in this environment.') + + assert event_loop.enable_uvloop() is True + # Second call short-circuits on the install flag and stays True. + assert event_loop.enable_uvloop() is True + + +def test_enable_uvloop_strict_raises_when_unavailable(monkeypatch): + # Simulate uvloop being absent regardless of what is installed. + import builtins + + real_import = builtins.__import__ + + def _fake_import(name, *args, **kwargs): + if name == 'uvloop': + raise ImportError('no uvloop') + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, '__import__', _fake_import) + + assert event_loop.enable_uvloop() is False + with pytest.raises(RuntimeError): + event_loop.enable_uvloop(strict=True) + + +def test_maybe_enable_uvloop_from_env(monkeypatch): + if not event_loop.is_uvloop_available(): + pytest.skip('uvloop not installed in this environment.') + + monkeypatch.delenv('ADK_UVLOOP', raising=False) + assert event_loop.maybe_enable_uvloop_from_env() is False + + monkeypatch.setenv('ADK_UVLOOP', '1') + assert event_loop.maybe_enable_uvloop_from_env() is True diff --git a/tests/unittests/workflow/test_first_match_node.py b/tests/unittests/workflow/test_first_match_node.py new file mode 100644 index 00000000000..4b0d5927f4a --- /dev/null +++ b/tests/unittests/workflow/test_first_match_node.py @@ -0,0 +1,156 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for FirstMatchNode (race fan-out; cancel losers on first match).""" + +import asyncio +import time +from typing import Any +from typing import AsyncGenerator + +from google.adk.agents.context import Context +from google.adk.workflow import FirstMatchNode +from google.adk.workflow._base_node import BaseNode +from google.adk.workflow._workflow import Workflow +from pydantic import ConfigDict +import pytest +from typing_extensions import override + +from .workflow_testing_utils import get_outputs +from .workflow_testing_utils import run_workflow + +# Branch lifecycle events, recorded as (state, name). Module-global so it +# survives the scheduler's ``model_copy`` of branch nodes. +_EVENTS: list[tuple[str, str]] = [] + + +class _SleepBranch(BaseNode): + """A branch that sleeps, then yields ``result`` -- unless cancelled first.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + delay: float = 0.0 + result: Any = None + + @override + async def _run_impl( + self, *, ctx: Context, node_input: Any + ) -> AsyncGenerator[Any, None]: + try: + await asyncio.sleep(self.delay) + except asyncio.CancelledError: + _EVENTS.append(('cancelled', self.name)) + raise + _EVENTS.append(('done', self.name)) + yield self.result + + +def _found(result: Any) -> bool: + return isinstance(result, dict) and bool(result.get('found')) + + +@pytest.mark.asyncio +async def test_first_match_cancels_slower_branches(): + _EVENTS.clear() + fast = _SleepBranch( + name='fast', delay=0.05, result={'found': True, 'src': 'fast'} + ) + slow1 = _SleepBranch(name='slow1', delay=5.0, result={'found': True}) + slow2 = _SleepBranch(name='slow2', delay=5.0, result={'found': True}) + race = FirstMatchNode( + name='race', nodes=[fast, slow1, slow2], match=_found + ) + wf = Workflow(name='w_first', edges=[('START', race)]) + + start = time.perf_counter() + events, _, _ = await run_workflow(wf) + elapsed = time.perf_counter() - start + + # The fast branch's result is the node output. + assert {'found': True, 'src': 'fast'} in get_outputs(events) + # We returned in ~fast time, not waiting on the 5s losers. + assert elapsed < 2.0 + # The winner finished; both losers were cancelled mid-flight (never "done"). + assert ('done', 'fast') in _EVENTS + assert ('cancelled', 'slow1') in _EVENTS + assert ('cancelled', 'slow2') in _EVENTS + assert ('done', 'slow1') not in _EVENTS + assert ('done', 'slow2') not in _EVENTS + + +@pytest.mark.asyncio +async def test_no_branch_matches_yields_default_and_runs_all(): + _EVENTS.clear() + a = _SleepBranch(name='a', delay=0.02, result={'found': False}) + b = _SleepBranch(name='b', delay=0.04, result={'found': False}) + race = FirstMatchNode( + name='race', + nodes=[a, b], + match=_found, + no_match_output={'found': False, 'reason': 'nobody had it'}, + ) + wf = Workflow(name='w_none', edges=[('START', race)]) + + events, _, _ = await run_workflow(wf) + + assert {'found': False, 'reason': 'nobody had it'} in get_outputs(events) + # With no early winner, every branch is allowed to finish. + assert ('done', 'a') in _EVENTS + assert ('done', 'b') in _EVENTS + + +@pytest.mark.asyncio +async def test_max_parallel_one_never_starts_later_branches_after_win(): + """Rank-ordered gate: a win short-circuits before lower-ranked reads start.""" + _EVENTS.clear() + first = _SleepBranch(name='first', delay=0.02, result={'found': True}) + second = _SleepBranch(name='second', delay=0.02, result={'found': True}) + race = FirstMatchNode( + name='race', nodes=[first, second], match=_found, max_parallel=1 + ) + wf = Workflow(name='w_serial', edges=[('START', race)]) + + events, _, _ = await run_workflow(wf) + + assert {'found': True} in get_outputs(events) + assert ('done', 'first') in _EVENTS + # The second branch was never launched -- not even started, so not cancelled. + assert ('done', 'second') not in _EVENTS + assert ('cancelled', 'second') not in _EVENTS + + +@pytest.mark.asyncio +async def test_failing_branch_does_not_sink_the_race(): + _EVENTS.clear() + + class _Boom(BaseNode): + model_config = ConfigDict(arbitrary_types_allowed=True) + + @override + async def _run_impl( + self, *, ctx: Context, node_input: Any + ) -> AsyncGenerator[Any, None]: + raise RuntimeError('branch blew up') + yield # pragma: no cover + + boom = _Boom(name='boom') + good = _SleepBranch(name='good', delay=0.05, result={'found': True}) + race = FirstMatchNode(name='race', nodes=[boom, good], match=_found) + wf = Workflow(name='w_flaky', edges=[('START', race)]) + + events, _, _ = await run_workflow(wf) + + # A single failing source must not deny the answer another source can give. + assert {'found': True} in get_outputs(events) + assert ('done', 'good') in _EVENTS diff --git a/tests/unittests/workflow/test_speculative_router.py b/tests/unittests/workflow/test_speculative_router.py new file mode 100644 index 00000000000..4481fd40985 --- /dev/null +++ b/tests/unittests/workflow/test_speculative_router.py @@ -0,0 +1,264 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for SpeculativeRouterNode and the JSON repairer.""" + +import asyncio +import json +from typing import Any +from typing import AsyncGenerator + +from google.adk.agents.context import Context +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.llm_agent import LlmAgent +from google.adk.events.event import Event +from google.adk.workflow import repair_json +from google.adk.workflow import SpeculativeRouterNode +from google.adk.workflow._base_node import BaseNode +from google.adk.workflow._workflow import Workflow +from google.genai import types +from pydantic import ConfigDict +from pydantic import Field +import pytest +from typing_extensions import override + +from .workflow_testing_utils import get_outputs +from .workflow_testing_utils import run_workflow + +# Target lifecycle, recorded as (state, path). Module-global to survive the +# scheduler's model_copy of the target node. +_EVENTS: list[tuple[str, Any]] = [] + + +def _partial(text: str) -> Event: + return Event( + author='scripted', + content=types.Content(role='model', parts=[types.Part(text=text)]), + partial=True, + ) + + +def _final(text: str) -> Event: + return Event( + author='scripted', + content=types.Content(role='model', parts=[types.Part(text=text)]), + partial=False, + ) + + +class _ScriptedAgent(LlmAgent): + """An LlmAgent whose stream is a fixed script of events.""" + + script: list[Event] = Field(default_factory=list) + + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + for event in self.script: + yield event + + +def _path(payload: Any) -> Any: + if isinstance(payload, dict): + return payload.get('arguments', {}).get('path') + return None + + +class _CaptureTarget(BaseNode): + """Records the payload it ran with; sleeps so speculation can overlap.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + delay: float = 0.0 + + @override + async def _run_impl( + self, *, ctx: Context, node_input: Any + ) -> AsyncGenerator[Any, None]: + p = _path(node_input) + _EVENTS.append(('start', p)) + try: + await asyncio.sleep(self.delay) + except asyncio.CancelledError: + _EVENTS.append(('cancelled', p)) + raise + _EVENTS.append(('done', p)) + yield {'read': p} + + +# --- repair_json -------------------------------------------------------------- + + +@pytest.mark.parametrize( + 'fragment,expected', + [ + ('{"a": "b', {'a': 'b'}), + ('{"a": [1, 2', {'a': [1, 2]}), + ('{"name":"x","arguments":{"path":"src/ma', { + 'name': 'x', + 'arguments': {'path': 'src/ma'}, + }), + ('{"ok": tru', {'ok': True}), + ('{"a":1,', {'a': 1}), + ('{"a":1, "b":', {'a': 1, 'b': None}), + ], +) +def test_repair_json_completes_truncated_fragments(fragment, expected): + assert json.loads(repair_json(fragment)) == expected + + +def test_repair_json_leaves_complete_json_parseable(): + assert json.loads(repair_json('{"a": 1, "b": [2, 3]}')) == { + 'a': 1, + 'b': [2, 3], + } + + +# --- SpeculativeRouterNode ---------------------------------------------------- + + +def _spec_node(agent: _ScriptedAgent, target: _CaptureTarget): + return SpeculativeRouterNode(name='spec', agent=agent, target=target) + + +@pytest.mark.asyncio +async def test_speculation_hit_keeps_result_and_runs_target_once(): + _EVENTS.clear() + # Partial already carries the full path; final only closes the braces. + agent = _ScriptedAgent( + name='m', + script=[ + _partial('TOOL_CALL: {"name":"read_file","arguments":{"path":"src/main.c"'), + _final('TOOL_CALL: {"name":"read_file","arguments":{"path":"src/main.c"}}'), + ], + ) + node = _spec_node(agent, _CaptureTarget(name='reader', delay=0.05)) + wf = Workflow(name='w_hit', edges=[('START', node)]) + + events, _, _ = await run_workflow(wf) + + assert {'read': 'src/main.c'} in get_outputs(events) + starts = [p for state, p in _EVENTS if state == 'start'] + # Speculation was correct -> the target ran exactly once (no re-run). + assert starts == ['src/main.c'] + assert ('done', 'src/main.c') in _EVENTS + assert ('cancelled', 'src/main.c') not in _EVENTS + + +@pytest.mark.asyncio +async def test_speculation_miss_cancels_and_reruns_with_final_payload(): + _EVENTS.clear() + # Partial repairs to the wrong (truncated) path; final has the real one. + agent = _ScriptedAgent( + name='m', + script=[ + _partial('TOOL_CALL: {"name":"read_file","arguments":{"path":"src/ma'), + _final('TOOL_CALL: {"name":"read_file","arguments":{"path":"src/main.c"}}'), + ], + ) + # Long delay so the speculative run is still in-flight when the final arrives. + node = _spec_node(agent, _CaptureTarget(name='reader', delay=5.0)) + wf = Workflow(name='w_miss', edges=[('START', node)]) + + events, _, _ = await run_workflow(wf) + + # The verified (final) payload wins. + assert {'read': 'src/main.c'} in get_outputs(events) + # The wrong speculative guess was started then cancelled... + assert ('start', 'src/ma') in _EVENTS + assert ('cancelled', 'src/ma') in _EVENTS + assert ('done', 'src/ma') not in _EVENTS + # ...and the correct payload was run to completion. + assert ('done', 'src/main.c') in _EVENTS + + +@pytest.mark.asyncio +async def test_no_partial_call_runs_target_once_on_final(): + _EVENTS.clear() + agent = _ScriptedAgent( + name='m', + script=[ + _partial('thinking about it...'), + _final('TOOL_CALL: {"name":"read_file","arguments":{"path":"x"}}'), + ], + ) + node = _spec_node(agent, _CaptureTarget(name='reader', delay=0.01)) + wf = Workflow(name='w_none', edges=[('START', node)]) + + events, _, _ = await run_workflow(wf) + + assert {'read': 'x'} in get_outputs(events) + starts = [p for state, p in _EVENTS if state == 'start'] + assert starts == ['x'] # never speculated; ran once on the final call + + +@pytest.mark.asyncio +async def test_combine_returns_agent_text_and_target_result(): + _EVENTS.clear() + # The agent emits a directive plus a trailing rationale that the caller wants. + rationale = ' because this file holds the entrypoint.' + agent = _ScriptedAgent( + name='m', + script=[ + _partial('TOOL_CALL: {"name":"read_file","arguments":{"path":"src/main.c"'), + _final( + 'TOOL_CALL: {"name":"read_file","arguments":{"path":"src/main.c"}}' + + rationale + ), + ], + ) + node = SpeculativeRouterNode( + name='spec', + agent=agent, + target=_CaptureTarget(name='reader', delay=0.02), + combine=lambda plan, result: {'plan': plan, 'result': result}, + ) + wf = Workflow(name='w_combine', edges=[('START', node)]) + + events, _, _ = await run_workflow(wf) + + outputs = get_outputs(events) + combined = next(o for o in outputs if isinstance(o, dict) and 'plan' in o) + # The target's verified result is carried through... + assert combined['result'] == {'read': 'src/main.c'} + # ...alongside the agent's FULL text (directive + required rationale), proving + # the streamed tail is a returned deliverable, not a discarded artifact. + assert rationale.strip() in combined['plan'] + + +@pytest.mark.asyncio +async def test_should_speculate_gate_suppresses_early_dispatch(): + _EVENTS.clear() + agent = _ScriptedAgent( + name='m', + script=[ + _partial('TOOL_CALL: {"name":"read_file","arguments":{"path":"s'), + _final('TOOL_CALL: {"name":"read_file","arguments":{"path":"src/main.c"}}'), + ], + ) + # Only speculate once the path looks long enough to be worth a guess. + node = SpeculativeRouterNode( + name='spec', + agent=agent, + target=_CaptureTarget(name='reader', delay=0.01), + should_speculate=lambda p: len(_path(p) or '') >= 6, + ) + wf = Workflow(name='w_gate', edges=[('START', node)]) + + events, _, _ = await run_workflow(wf) + + assert {'read': 'src/main.c'} in get_outputs(events) + # The short 's' guess was gated out, so the target only ran on the final. + starts = [p for state, p in _EVENTS if state == 'start'] + assert starts == ['src/main.c'] diff --git a/tests/unittests/workflow/test_streaming_router.py b/tests/unittests/workflow/test_streaming_router.py new file mode 100644 index 00000000000..977181feae6 --- /dev/null +++ b/tests/unittests/workflow/test_streaming_router.py @@ -0,0 +1,371 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for StreamingRouterNode (mid-stream preemptive graph advancement).""" + +import asyncio +from typing import Any +from typing import AsyncGenerator +from unittest.mock import AsyncMock +from unittest.mock import MagicMock + +from google.adk.agents.base_agent import BaseAgent +from google.adk.agents.context import Context +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.llm_agent import LlmAgent +from google.adk.events.event import Event +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.sessions.session import Session +from google.adk.workflow import Edge +from google.adk.workflow import JoinNode +from google.adk.workflow import START +from google.adk.workflow import StreamDecision +from google.adk.workflow import StreamingRouterNode +from google.adk.workflow import StreamView +from google.adk.workflow._node_runner import NodeRunner +from google.adk.workflow._workflow import Workflow +from google.genai import types +from pydantic import Field +import pytest + +from .workflow_testing_utils import run_workflow +from .workflow_testing_utils import simplify_events_with_node +from .workflow_testing_utils import TestingNode as _RoutingNode + + +def _partial(text: str) -> Event: + return Event( + author='scripted', + content=types.Content(role='model', parts=[types.Part(text=text)]), + partial=True, + ) + + +def _final(text: str) -> Event: + return Event( + author='scripted', + content=types.Content(role='model', parts=[types.Part(text=text)]), + partial=False, + ) + + +class _ScriptedAgent(LlmAgent): + """An LlmAgent whose stream is a fixed script of events. + + ``consumed`` counts how many scripted events were actually pulled — a + preempting router closes the stream early, so a preempted run consumes + fewer events than the script contains. + """ + + script: list[Event] = Field(default_factory=list) + consumed: list[Event] = Field(default_factory=list) + + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + for event in self.script: + self.consumed.append(event) + yield event + + +async def _run_router( + router: StreamingRouterNode, node_input: str = 'hi' +) -> tuple[Context, list[Event]]: + """Drives a router node via NodeRunner in isolation. + + ``_enqueue_event`` is mocked to collect events without blocking; the real + method blocks non-partial events on the Runner main loop, which is absent + in this unit-level harness. Returns the child context and the events the + node emitted (state/artifact deltas are flushed onto these events and + cleared from ``ctx.actions``). + """ + session = Session(id='s', app_name='a', user_id='u') + ic = InvocationContext( + invocation_id='inv', + agent=MagicMock(spec=BaseAgent), + session=session, + session_service=InMemorySessionService(), + ) + collected: list[Event] = [] + + async def _enqueue(event: Event) -> None: + collected.append(event) + + object.__setattr__(ic, '_enqueue_event', AsyncMock(side_effect=_enqueue)) + parent_ctx = Context(invocation_context=ic, node_path='', run_id='1') + child = await NodeRunner(node=router, parent_ctx=parent_ctx, run_id='1').run( + node_input=node_input + ) + return child, collected + + +def _billing_monitor(view: StreamView): + low = view.text.lower() + if 'billing' in low: + return StreamDecision(route='billing') + if 'technical' in low: + return StreamDecision(route='technical') + return None + + +@pytest.mark.asyncio +async def test_preempts_and_routes_midstream(): + agent = _ScriptedAgent( + name='classifier', + script=[ + _partial('Bil'), + _partial('ling'), + _partial(' department, definitely'), + _final('Billing department, definitely'), + ], + ) + router = StreamingRouterNode( + name='intent_router', agent=agent, monitor=_billing_monitor + ) + + child, _ = await _run_router(router) + + assert child.route == 'billing' + # The decision fired after the second chunk completed "Billing"; the + # remaining two events must never have been pulled from the stream. + assert len(agent.consumed) == 2 + + +@pytest.mark.asyncio +async def test_commits_output_and_state_delta_midstream(): + agent = _ScriptedAgent( + name='classifier', + output_key='intent', + script=[_partial('ans'), _partial('wer=42'), _final('answer=42')], + ) + + def monitor(view: StreamView): + if '=' in view.text: + return StreamDecision(output=view.text) + return None + + router = StreamingRouterNode( + name='intent_router', agent=agent, monitor=monitor + ) + + child, events = await _run_router(router) + + assert child.output == 'answer=42' + # The output_key delta is flushed onto an emitted event. + state_deltas = [ + e.actions.state_delta for e in events if e.actions.state_delta + ] + assert {'intent': 'answer=42'} in state_deltas + assert len(agent.consumed) == 2 + + +@pytest.mark.asyncio +async def test_no_decision_falls_back_to_final_output(): + agent = _ScriptedAgent( + name='classifier', + script=[_partial('Hel'), _partial('lo'), _final('Hello world')], + ) + # Monitor never decides. + router = StreamingRouterNode( + name='intent_router', agent=agent, monitor=lambda view: None + ) + + child, _ = await _run_router(router) + + assert child.route is None + assert child.output == 'Hello world' + # No preemption: the whole script is consumed. + assert len(agent.consumed) == 3 + + +@pytest.mark.asyncio +async def test_async_monitor_supported(): + agent = _ScriptedAgent( + name='classifier', + script=[_partial('techni'), _partial('cal'), _final('technical')], + ) + + async def monitor(view: StreamView): + await asyncio.sleep(0) + if 'technical' in view.text.lower(): + return StreamDecision(route='technical') + return None + + router = StreamingRouterNode( + name='intent_router', agent=agent, monitor=monitor + ) + + child, _ = await _run_router(router) + + assert child.route == 'technical' + assert len(agent.consumed) == 2 + + +@pytest.mark.asyncio +async def test_stop_false_continues_streaming_without_double_output(): + agent = _ScriptedAgent( + name='classifier', + script=[_partial('go'), _partial(' now'), _final('go now')], + ) + + def monitor(view: StreamView): + if 'go' in view.text: + return StreamDecision(route='fast', stop=False) + return None + + router = StreamingRouterNode( + name='intent_router', agent=agent, monitor=monitor + ) + + # Must not raise "Output already set": the final event's output is + # suppressed because the decision already owns routing. + child, _ = await _run_router(router) + + assert child.route == 'fast' + # stop=False lets generation run to completion. + assert len(agent.consumed) == 3 + + +@pytest.mark.asyncio +async def test_forward_partials_false_suppresses_partial_messages(): + agent = _ScriptedAgent( + name='classifier', + script=[_partial('a'), _partial('b'), _final('ab')], + ) + seen: list[bool] = [] + + def monitor(view: StreamView): + seen.append(view.event.partial) + return None + + router = StreamingRouterNode( + name='intent_router', + agent=agent, + monitor=monitor, + forward_partials=False, + ) + + child, events = await _run_router(router) + + # Monitor still saw the partials even though they were not forwarded. + assert seen == [True, True] + assert child.output == 'ab' + # No partial (streaming-message) events were emitted downstream. + assert not any(e.partial for e in events) + + +@pytest.mark.asyncio +async def test_streaming_router_in_workflow_advances_graph(): + agent = _ScriptedAgent( + name='classifier', + script=[ + _partial('Bil'), + _partial('ling'), + _partial(' and more text the model never gets to finish'), + _final('Billing and more text the model never gets to finish'), + ], + ) + router = StreamingRouterNode( + name='intent_router', agent=agent, monitor=_billing_monitor + ) + billing = _RoutingNode(name='billing_node', output='handled-billing') + technical = _RoutingNode(name='technical_node', output='handled-technical') + + wf = Workflow( + name='support_wf', + edges=[ + Edge(from_node=START, to_node=router), + Edge(from_node=router, to_node=billing, route='billing'), + Edge(from_node=router, to_node=technical, route='technical'), + ], + ) + + events, _, _ = await run_workflow(wf, message='my invoice is wrong') + simplified = simplify_events_with_node(events) + + authors = [author for author, _ in simplified] + assert any('billing_node' in a for a in authors) + assert not any('technical_node' in a for a in authors) + # Preemption held even through the full workflow run. + assert len(agent.consumed) == 2 + + +def test_stream_decision_requires_route_or_output(): + with pytest.raises(ValueError): + StreamDecision() + + +def _relevance_monitor(view: StreamView): + """Stops reading a source the instant it declares itself irrelevant.""" + if view.text.lstrip().upper().startswith('IRRELEVANT'): + return StreamDecision(output={'relevant': False}) + return None + + +@pytest.mark.asyncio +async def test_fan_out_preempts_only_irrelevant_branch(): + """Parallel readers: the irrelevant branch cancels, the others run on.""" + irrelevant = _ScriptedAgent( + name='src_a', + script=[ + _partial('IRRELE'), + _partial('VANT'), + _partial(' the model keeps talking but nobody is listening'), + _final('IRRELEVANT ...'), + ], + ) + relevant_1 = _ScriptedAgent( + name='src_b', + script=[_partial('fact'), _partial(' one'), _final('fact one')], + ) + relevant_2 = _ScriptedAgent( + name='src_c', + script=[_partial('fact'), _partial(' two'), _final('fact two')], + ) + + reader_a = StreamingRouterNode( + name='reader_a', agent=irrelevant, monitor=_relevance_monitor + ) + reader_b = StreamingRouterNode( + name='reader_b', agent=relevant_1, monitor=_relevance_monitor + ) + reader_c = StreamingRouterNode( + name='reader_c', agent=relevant_2, monitor=_relevance_monitor + ) + + join = JoinNode(name='join_sources') + captured: dict[str, Any] = {} + + async def synthesize(node_input: dict[str, Any]): + captured['fan_in'] = node_input + yield Event(message='synthesized') + + wf = Workflow( + name='fan_out_wf', + edges=[('START', (reader_a, reader_b, reader_c), join, synthesize)], + ) + + await run_workflow(wf, message='the query') + + # Only the irrelevant branch was preempted; it consumed 2 of its 4 events. + assert len(irrelevant.consumed) == 2 + # The relevant branches ran their full streams uninterrupted. + assert len(relevant_1.consumed) == 3 + assert len(relevant_2.consumed) == 3 + + # The join fanned all three branches back in, keyed by node name. + fan_in = captured['fan_in'] + assert fan_in['reader_a'] == {'relevant': False} + assert fan_in['reader_b'] == 'fact one' + assert fan_in['reader_c'] == 'fact two'