Skip to content

Commit 4721491

Browse files
damianmomotgooglecopybara-github
authored andcommitted
feat: add a builder for workflow graph edges
PiperOrigin-RevId: 1000728765
1 parent 0ee0d4b commit 4721491

4 files changed

Lines changed: 763 additions & 0 deletions

File tree

Lines changed: 368 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,368 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.workflow;
18+
19+
import static com.google.common.base.Preconditions.checkArgument;
20+
import static com.google.common.base.Preconditions.checkNotNull;
21+
import static com.google.common.base.Preconditions.checkState;
22+
import static com.google.common.collect.ImmutableList.toImmutableList;
23+
24+
import com.google.adk.annotations.Experimental;
25+
import com.google.common.collect.ImmutableList;
26+
import com.google.common.collect.ImmutableSet;
27+
import com.google.common.collect.Lists;
28+
import com.google.errorprone.annotations.CanIgnoreReturnValue;
29+
import java.util.ArrayList;
30+
import java.util.Collection;
31+
import java.util.List;
32+
import java.util.function.Consumer;
33+
34+
/**
35+
* A builder for the edges of a workflow graph, declared as chains of calls that start at {@link
36+
* #from} or {@link #fromStart}. Create one with {@code new EdgesBuilder()} and call {@link #build},
37+
* or let {@link Workflow.Builder#edges(Consumer)} do both. A node is identified by its instance, so
38+
* create each node once.
39+
*
40+
* <pre>{@code
41+
* Workflow support =
42+
* Workflow.builder()
43+
* .name("support_triage")
44+
* .edges(
45+
* g -> {
46+
* g.fromStart()
47+
* .chain(classify, category)
48+
* .route(
49+
* r -> {
50+
* r.on("billing").then(billing);
51+
* r.otherwise(summarize);
52+
* });
53+
* g.from(billing).then(summarize);
54+
* })
55+
* .build();
56+
* }</pre>
57+
*/
58+
@Experimental
59+
public final class EdgesBuilder {
60+
61+
private final List<Edge> edges = new ArrayList<>();
62+
63+
/** Starts a chain of edges at {@code node}: {@code g.from(a).then(b).then(c)}. */
64+
public EdgeChain from(Node node) {
65+
return new EdgeChain(this, checkNotNull(node));
66+
}
67+
68+
/**
69+
* Starts a chain at {@code nodes} (such as a computed list) so they can join without a fan-out,
70+
* like {@link #from(Node, Node, Node...)}.
71+
*/
72+
public NodeGroupChain from(Collection<? extends Node> nodes) {
73+
return new NodeGroupChain(this, ImmutableList.copyOf(nodes));
74+
}
75+
76+
/**
77+
* Starts a chain at {@code first}, {@code second} and each node in {@code rest} so they can join
78+
* without a fan-out: {@code g.from(a, b).joinTo(join)}.
79+
*/
80+
public NodeGroupChain from(Node first, Node second, Node... rest) {
81+
return from(Lists.asList(first, second, rest));
82+
}
83+
84+
/** Starts a chain of edges at {@link Start}: {@code g.fromStart().then(a)}. */
85+
public EdgeChain fromStart() {
86+
return from(Start.INSTANCE);
87+
}
88+
89+
/** Returns the edges added so far, in the order they were added. */
90+
public ImmutableList<Edge> build() {
91+
return ImmutableList.copyOf(edges);
92+
}
93+
94+
void add(Edge edge) {
95+
edges.add(edge);
96+
}
97+
98+
void addAll(List<Edge> newEdges) {
99+
edges.addAll(newEdges);
100+
}
101+
102+
/** Returns an edge from {@code from} to each of {@code targets}, failing if a target is null. */
103+
static ImmutableList<Edge> edgesTo(
104+
Node from, Collection<? extends Node> targets, ImmutableList<Route> routes) {
105+
return targets.stream().map(to -> new Edge(from, to, routes)).collect(toImmutableList());
106+
}
107+
108+
/**
109+
* A chain of edges positioned at a node, started by {@link EdgesBuilder#from(Node)}, {@link
110+
* EdgesBuilder#fromStart} or {@link NodeGroupChain#joinTo}. {@link #then(Node)} and {@link
111+
* #chain} add edges and continue from the last node they connect, a fan-out continues from its
112+
* targets as a group, and a routing map ends the chain.
113+
*/
114+
@Experimental
115+
public static final class EdgeChain {
116+
117+
private final EdgesBuilder builder;
118+
private final Node node;
119+
120+
EdgeChain(EdgesBuilder builder, Node node) {
121+
this.builder = builder;
122+
this.node = node;
123+
}
124+
125+
/** Connects the current node to {@code next} and continues from {@code next}. */
126+
@CanIgnoreReturnValue
127+
public EdgeChain then(Node next) {
128+
builder.add(new Edge(node, next));
129+
return new EdgeChain(builder, next);
130+
}
131+
132+
/**
133+
* Fans out by connecting the current node to each of {@code targets} (such as a computed list)
134+
* so they run concurrently, and continues from them as a group that {@link
135+
* NodeGroupChain#joinTo} can join. An empty collection adds no edges.
136+
*/
137+
@CanIgnoreReturnValue
138+
public NodeGroupChain then(Collection<? extends Node> targets) {
139+
ImmutableList<Node> nodes = ImmutableList.copyOf(targets);
140+
builder.addAll(edgesTo(node, nodes, ImmutableList.of()));
141+
return new NodeGroupChain(builder, nodes);
142+
}
143+
144+
/**
145+
* Fans out by connecting the current node to {@code first}, {@code second} and each node in
146+
* {@code rest}, like {@link #then(Collection)}: {@code g.from(a).then(b, c).joinTo(join)}.
147+
*/
148+
@CanIgnoreReturnValue
149+
public NodeGroupChain then(Node first, Node second, Node... rest) {
150+
return then(Lists.asList(first, second, rest));
151+
}
152+
153+
/**
154+
* Connects the current node to {@code first}, {@code second} and each node in {@code rest} in
155+
* order, and continues from the last one: {@code g.fromStart().chain(a, b)} is short for {@code
156+
* g.fromStart().then(a).then(b)}.
157+
*/
158+
@CanIgnoreReturnValue
159+
public EdgeChain chain(Node first, Node second, Node... rest) {
160+
// Copying rejects a null node before any edge is added.
161+
ImmutableList<Node> nodes = ImmutableList.copyOf(Lists.asList(first, second, rest));
162+
EdgeChain chain = this;
163+
for (Node next : nodes) {
164+
chain = chain.then(next);
165+
}
166+
return chain;
167+
}
168+
169+
/**
170+
* Adds the current node's routing map, which {@code routeMapDefinition} fills, and ends the
171+
* chain.
172+
*
173+
* @throws IllegalArgumentException if the routing map has no entries
174+
* @throws IllegalStateException if an entry started by {@link RouteMap#on} has no target
175+
*/
176+
public void route(Consumer<? super RouteMap> routeMapDefinition) {
177+
RouteMap routeMap = new RouteMap(node);
178+
routeMapDefinition.accept(routeMap);
179+
ImmutableList<Edge> routeEdges = routeMap.build();
180+
checkArgument(
181+
!routeEdges.isEmpty(),
182+
"The routing map of node '%s' needs at least one entry.",
183+
node.name());
184+
builder.addAll(routeEdges);
185+
}
186+
}
187+
188+
/**
189+
* A chain of edges positioned at a node group from a fan-out, {@link EdgesBuilder#from(Node,
190+
* Node, Node...)}, or {@link EdgesBuilder#from(Collection)}, which {@link #joinTo} connects to a
191+
* {@link JoinNode}.
192+
*/
193+
@Experimental
194+
public static final class NodeGroupChain {
195+
196+
private final EdgesBuilder builder;
197+
private final ImmutableList<Node> nodes;
198+
199+
NodeGroupChain(EdgesBuilder builder, ImmutableList<Node> nodes) {
200+
this.builder = builder;
201+
this.nodes = nodes;
202+
}
203+
204+
/**
205+
* Connects each node in the group to {@code join}, which runs once all of them complete, and
206+
* continues from {@code join}.
207+
*/
208+
@CanIgnoreReturnValue
209+
public EdgeChain joinTo(JoinNode join) {
210+
checkNotNull(join);
211+
for (Node node : nodes) {
212+
builder.add(new Edge(node, join));
213+
}
214+
return new EdgeChain(builder, join);
215+
}
216+
}
217+
218+
/**
219+
* A routing map configured in {@link EdgeChain#route} that maps each route the source node may
220+
* emit to the nodes it leads to, written {@code r.on("billing").then(billing)}. Mapping a route
221+
* to several nodes fans out on that route.
222+
*/
223+
@Experimental
224+
public static final class RouteMap {
225+
226+
private final Node from;
227+
private final ImmutableList.Builder<Edge> edges = ImmutableList.builder();
228+
private final List<RouteMapEntry> pendingEntries = new ArrayList<>();
229+
private boolean built = false;
230+
231+
RouteMap(Node from) {
232+
this.from = from;
233+
}
234+
235+
/**
236+
* Starts an entry for {@code route} as a {@link Route.Tag}; {@link RouteMapEntry#then}
237+
* completes it.
238+
*/
239+
public RouteMapEntry on(String route) {
240+
return on(new Route.Tag(route));
241+
}
242+
243+
/**
244+
* Starts an entry for {@code route} as a {@link Route.Flag}; {@link RouteMapEntry#then}
245+
* completes it.
246+
*/
247+
public RouteMapEntry on(boolean route) {
248+
return on(new Route.Flag(route));
249+
}
250+
251+
/**
252+
* Starts an entry for {@code route}, such as a {@link Route.Num}; {@link RouteMapEntry#then}
253+
* completes it.
254+
*/
255+
public RouteMapEntry on(Route route) {
256+
return start(ImmutableList.of(route));
257+
}
258+
259+
/**
260+
* Starts an entry whose edges fire once when the source node emits any of {@code routes}, built
261+
* with {@link #anyOf}; {@link RouteMapEntry#then} completes it.
262+
*/
263+
public RouteMapEntry on(RouteSet routes) {
264+
return start(routes.routes());
265+
}
266+
267+
/**
268+
* Leads to each target when no other route of the source node matches, including when it emits
269+
* none.
270+
*/
271+
public void otherwise(Node first, Node... rest) {
272+
add(ImmutableList.of(Route.Default.INSTANCE), Lists.asList(first, rest));
273+
}
274+
275+
/** Groups string routes as {@link Route.Tag}s for {@link #on(RouteSet)}. */
276+
public RouteSet anyOf(String first, String second, String... rest) {
277+
return new RouteSet(Lists.transform(Lists.asList(first, second, rest), Route.Tag::new));
278+
}
279+
280+
/** Groups routes, such as {@link Route.Num}s, for {@link #on(RouteSet)}. */
281+
public RouteSet anyOf(Route first, Route second, Route... rest) {
282+
return new RouteSet(Lists.asList(first, second, rest));
283+
}
284+
285+
/**
286+
* Returns the map's edges and closes the map, so an entry started or completed after {@link
287+
* EdgeChain#route} returns fails instead of being lost.
288+
*
289+
* @throws IllegalStateException if an entry started by {@code on} has no target
290+
*/
291+
ImmutableList<Edge> build() {
292+
built = true;
293+
checkState(
294+
pendingEntries.isEmpty(),
295+
"on(...) in the routing map of node '%s' has no target; complete it with then(node).",
296+
from.name());
297+
return edges.build();
298+
}
299+
300+
/** Adds an edge for {@code routes} to each of {@code targets} and completes {@code entry}. */
301+
void complete(RouteMapEntry entry, ImmutableList<Route> routes, List<Node> targets) {
302+
add(routes, targets);
303+
pendingEntries.remove(entry);
304+
}
305+
306+
private RouteMapEntry start(ImmutableList<Route> routes) {
307+
checkNotBuilt();
308+
RouteMapEntry entry = new RouteMapEntry(this, routes);
309+
pendingEntries.add(entry);
310+
return entry;
311+
}
312+
313+
private void add(ImmutableList<Route> routes, List<Node> targets) {
314+
checkNotBuilt();
315+
edges.addAll(edgesTo(from, targets, routes));
316+
}
317+
318+
private void checkNotBuilt() {
319+
checkState(
320+
!built,
321+
"The routing map of node '%s' was already built; add its entries inside route(...).",
322+
from.name());
323+
}
324+
}
325+
326+
/**
327+
* A routing-map entry started by {@link RouteMap#on} whose routes have no target yet; {@link
328+
* #then} completes it.
329+
*/
330+
@Experimental
331+
public static final class RouteMapEntry {
332+
333+
private final RouteMap routeMap;
334+
private final ImmutableList<Route> routes;
335+
336+
RouteMapEntry(RouteMap routeMap, ImmutableList<Route> routes) {
337+
this.routeMap = routeMap;
338+
this.routes = routes;
339+
}
340+
341+
/**
342+
* Leads to {@code first} and each node in {@code rest}, with one edge to each, when the source
343+
* node emits any of this entry's routes.
344+
*/
345+
public void then(Node first, Node... rest) {
346+
routeMap.complete(this, routes, Lists.asList(first, rest));
347+
}
348+
}
349+
350+
/**
351+
* A set of routes that select a single edge, which fires once when its source emits any of them;
352+
* separate edges to the same target would fire once each. {@link RouteMap#anyOf} builds one and
353+
* drops duplicate routes.
354+
*/
355+
@Experimental
356+
public static final class RouteSet {
357+
358+
private final ImmutableList<Route> routes;
359+
360+
RouteSet(List<? extends Route> routes) {
361+
this.routes = ImmutableSet.<Route>copyOf(routes).asList();
362+
}
363+
364+
ImmutableList<Route> routes() {
365+
return routes;
366+
}
367+
}
368+
}

‎core/src/main/java/com/google/adk/workflow/Workflow.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import io.reactivex.rxjava3.core.Flowable;
2929
import java.util.List;
3030
import java.util.Optional;
31+
import java.util.function.Consumer;
3132
import org.jspecify.annotations.Nullable;
3233

3334
/**
@@ -118,6 +119,18 @@ public Builder edges(Edge... edges) {
118119
return edges(ImmutableList.copyOf(edges));
119120
}
120121

122+
/**
123+
* Sets the edges that {@code graphDefinition} adds to a new {@link EdgesBuilder}.
124+
*
125+
* <p>For example: {@code .edges(g -> g.fromStart().chain(a, b))}.
126+
*/
127+
@CanIgnoreReturnValue
128+
public Builder edges(Consumer<? super EdgesBuilder> graphDefinition) {
129+
EdgesBuilder builder = new EdgesBuilder();
130+
graphDefinition.accept(builder);
131+
return edges(builder.build());
132+
}
133+
121134
@CanIgnoreReturnValue
122135
public Builder description(String description) {
123136
this.description = checkNotNull(description);

0 commit comments

Comments
 (0)