batchContext = new HashMap<>();
114 | String result = batchFlow.run(batchContext);
115 |
116 | assertEquals("batch_complete", result);
117 | assertTrue((Boolean)batchContext.get("postBatchCalled"));
118 | assertEquals("SimpleLogNode executed with multiplier: 2", batchContext.get("last_message_from_batch_2"));
119 | assertEquals("SimpleLogNode executed with multiplier: 4", batchContext.get("last_message_from_batch_4"));
120 | }
121 | }
--------------------------------------------------------------------------------
/pom.xml:
--------------------------------------------------------------------------------
1 |
2 |
5 | 4.0.0
6 |
7 |
8 | io.github.the-pocket
9 | PocketFlow
10 | 1.0.0
11 | jar
12 |
13 |
14 | PocketFlow
15 | Pocket Flow: A minimalist LLM framework. Let Agents build Agents!
16 | https://github.com/The-Pocket/PocketFlow-Java
17 |
18 |
19 |
20 |
21 | MIT License
22 | http://www.opensource.org/licenses/mit-license.php
23 | repo
24 |
25 |
26 |
27 |
28 |
29 |
30 | zachary62
31 | Zachary Huang
32 | https://github.com/zachary62
33 |
34 |
35 |
36 |
37 |
38 | scm:git:git://github.com/The-Pocket/PocketFlow-Java.git
39 | scm:git:ssh://git@github.com/The-Pocket/PocketFlow-Java.git
40 | https://github.com/The-Pocket/PocketFlow-Java/tree/main
41 | HEAD
42 |
43 |
44 |
45 |
46 | UTF-8
47 | 11
48 | 11
49 | 5.9.2
50 |
51 |
52 |
53 |
54 |
55 |
56 |
57 | org.junit.jupiter
58 | junit-jupiter-api
59 | ${junit.jupiter.version}
60 | test
61 |
62 |
63 | org.junit.jupiter
64 | junit-jupiter-engine
65 | ${junit.jupiter.version}
66 | test
67 |
68 |
69 |
70 |
71 |
72 |
73 | ossrh
74 | https://s01.oss.sonatype.org/content/repositories/snapshots
75 |
76 |
77 | ossrh
78 | https://s01.oss.sonatype.org/service/local/staging/deploy/maven2/
79 |
80 |
81 |
82 |
83 |
84 |
85 |
86 |
87 | org.apache.maven.plugins
88 | maven-compiler-plugin
89 | 3.11.0
90 |
91 | ${maven.compiler.source}
92 | ${maven.compiler.target}
93 |
94 |
95 |
96 |
97 | org.apache.maven.plugins
98 | maven-surefire-plugin
99 | 3.0.0
100 |
101 |
102 |
103 | org.apache.maven.plugins
104 | maven-source-plugin
105 | 3.2.1
106 |
107 |
108 | attach-sources
109 |
110 | jar-no-fork
111 |
112 |
113 |
114 |
115 |
116 |
117 | org.apache.maven.plugins
118 | maven-javadoc-plugin
119 | 3.5.0
120 |
121 | ${maven.compiler.source}
122 |
123 |
124 |
125 | attach-javadocs
126 |
127 | jar
128 |
129 |
130 |
131 |
132 |
133 |
134 |
135 | org.apache.maven.plugins
136 | maven-gpg-plugin
137 | 3.1.0
138 |
139 |
140 | sign-artifacts
141 | verify
142 |
143 | sign
144 |
145 |
146 |
147 |
148 |
149 |
150 | --pinentry-mode
151 | loopback
152 |
153 |
154 |
155 |
156 |
157 |
158 |
159 |
160 | org.sonatype.central
161 | central-publishing-maven-plugin
162 | 0.7.0
163 | true
164 |
165 |
166 | central
167 |
168 | true
169 |
170 |
171 |
172 |
173 |
174 |
175 |
176 |
177 |
--------------------------------------------------------------------------------
/src/main/java/io/github/the_pocket/PocketFlow.java:
--------------------------------------------------------------------------------
1 | package io.github.the_pocket;
2 |
3 | import java.util.ArrayList;
4 | import java.util.Collections;
5 | import java.util.HashMap;
6 | import java.util.List;
7 | import java.util.Map;
8 | import java.util.Objects;
9 | import java.util.concurrent.TimeUnit;
10 |
11 | /**
12 | * PocketFlow - A simple, synchronous workflow library for Java in a single file.
13 | */
14 | public final class PocketFlow {
15 |
16 | private PocketFlow() {}
17 |
18 | private static void logWarn(String message) {
19 | System.err.println("WARN: PocketFlow - " + message);
20 | }
21 |
22 | public static class PocketFlowException extends RuntimeException {
23 | public PocketFlowException(String message) { super(message); }
24 | public PocketFlowException(String message, Throwable cause) { super(message, cause); }
25 | }
26 |
27 | /**
28 | * Base class for all nodes in the workflow.
29 | *
30 | * @param Type of the result from the prep phase.
31 | * @param Type of the result from the exec phase.
32 | * The return type of post/run is always String (or null for default action).
33 | */
34 | public static abstract class BaseNode {
35 | protected Map params = new HashMap<>();
36 | protected final Map> successors = new HashMap<>();
37 | public static final String DEFAULT_ACTION = "default";
38 |
39 | public BaseNode setParams(Map params) {
40 | this.params = params != null ? new HashMap<>(params) : new HashMap<>();
41 | return this;
42 | }
43 |
44 | public BaseNode next(BaseNode node) {
45 | return next(node, DEFAULT_ACTION);
46 | }
47 |
48 | public BaseNode next(BaseNode node, String action) {
49 | Objects.requireNonNull(node, "Successor node cannot be null");
50 | Objects.requireNonNull(action, "Action cannot be null");
51 | if (this.successors.containsKey(action)) {
52 | logWarn("Overwriting successor for action '" + action + "' in node " + this.getClass().getSimpleName());
53 | }
54 | this.successors.put(action, node);
55 | return node;
56 | }
57 |
58 | public P prep(Map sharedContext) { return null; }
59 | public abstract E exec(P prepResult);
60 | /** Post method MUST return a String action, or null for the default action. */
61 | public String post(Map sharedContext, P prepResult, E execResult) { return null; }
62 |
63 | protected E internalExec(P prepResult) { return exec(prepResult); }
64 |
65 | protected String internalRun(Map sharedContext) {
66 | P prepRes = prep(sharedContext);
67 | E execRes = internalExec(prepRes);
68 | return post(sharedContext, prepRes, execRes);
69 | }
70 |
71 | public String run(Map sharedContext) {
72 | if (!successors.isEmpty()) {
73 | logWarn("Node " + getClass().getSimpleName() + " has successors, but run() was called. Successors won't be executed. Use Flow.");
74 | }
75 | return internalRun(sharedContext);
76 | }
77 |
78 | protected BaseNode, ?> getNextNode(String action) {
79 | String actionKey = (action != null) ? action : DEFAULT_ACTION;
80 | BaseNode, ?> nextNode = successors.get(actionKey);
81 | if (nextNode == null && !successors.isEmpty() && !successors.containsKey(actionKey)) {
82 | logWarn("Flow might end: Action '" + actionKey + "' not found in successors "
83 | + successors.keySet() + " of node " + this.getClass().getSimpleName());
84 | }
85 | return nextNode;
86 | }
87 | }
88 |
89 | /**
90 | * A synchronous node with built-in retry capabilities.
91 | */
92 | public static abstract class Node extends BaseNode
{
93 | protected final int maxRetries;
94 | protected final long waitMillis;
95 | protected int currentRetry = 0;
96 |
97 | public Node() { this(1, 0); }
98 | public Node(int maxRetries, long waitMillis) {
99 | if (maxRetries < 1) throw new IllegalArgumentException("maxRetries must be at least 1");
100 | if (waitMillis < 0) throw new IllegalArgumentException("waitMillis cannot be negative");
101 | this.maxRetries = maxRetries;
102 | this.waitMillis = waitMillis;
103 | }
104 |
105 | public E execFallback(P prepResult, Exception lastException) throws Exception { throw lastException; }
106 |
107 | @Override
108 | protected E internalExec(P prepResult) {
109 | Exception lastException = null;
110 | for (currentRetry = 0; currentRetry < maxRetries; currentRetry++) {
111 | try { return exec(prepResult); }
112 | catch (Exception e) {
113 | lastException = e;
114 | if (currentRetry < maxRetries - 1 && waitMillis > 0) {
115 | try { TimeUnit.MILLISECONDS.sleep(waitMillis); }
116 | catch (InterruptedException ie) { Thread.currentThread().interrupt(); throw new PocketFlowException("Thread interrupted during retry wait", ie); }
117 | }
118 | }
119 | }
120 | try {
121 | if (lastException == null) throw new PocketFlowException("Execution failed, but no exception was captured.");
122 | return execFallback(prepResult, lastException);
123 | } catch (Exception fallbackException) {
124 | if (lastException != null && fallbackException != lastException) fallbackException.addSuppressed(lastException);
125 | if (fallbackException instanceof RuntimeException) throw (RuntimeException) fallbackException;
126 | else throw new PocketFlowException("Fallback execution failed", fallbackException);
127 | }
128 | }
129 | }
130 |
131 | /**
132 | * A synchronous node that processes a list of items individually.
133 | */
134 | public static abstract class BatchNode extends Node, List> {
135 | public BatchNode() { super(); }
136 | public BatchNode(int maxRetries, long waitMillis) { super(maxRetries, waitMillis); }
137 |
138 | public abstract OUT_ITEM execItem(IN_ITEM item);
139 | public OUT_ITEM execItemFallback(IN_ITEM item, Exception lastException) throws Exception { throw lastException; }
140 |
141 | @Override
142 | protected List internalExec(List batchPrepResult) {
143 | if (batchPrepResult == null || batchPrepResult.isEmpty()) return Collections.emptyList();
144 | List results = new ArrayList<>(batchPrepResult.size());
145 | for (IN_ITEM item : batchPrepResult) {
146 | Exception lastItemException = null;
147 | OUT_ITEM itemResult = null;
148 | boolean itemSuccess = false;
149 | for (currentRetry = 0; currentRetry < maxRetries; currentRetry++) {
150 | try { itemResult = execItem(item); itemSuccess = true; break; }
151 | catch (Exception e) {
152 | lastItemException = e;
153 | if (currentRetry < maxRetries - 1 && waitMillis > 0) {
154 | try { TimeUnit.MILLISECONDS.sleep(waitMillis); }
155 | catch (InterruptedException ie) { Thread.currentThread().interrupt(); throw new PocketFlowException("Interrupted batch retry wait: " + item, ie); }
156 | }
157 | }
158 | }
159 | if (!itemSuccess) {
160 | try {
161 | if (lastItemException == null) throw new PocketFlowException("Item exec failed without exception: " + item);
162 | itemResult = execItemFallback(item, lastItemException);
163 | } catch (Exception fallbackException) {
164 | if (lastItemException != null && fallbackException != lastItemException) fallbackException.addSuppressed(lastItemException);
165 | if (fallbackException instanceof RuntimeException) throw (RuntimeException) fallbackException;
166 | else throw new PocketFlowException("Item fallback failed: " + item, fallbackException);
167 | }
168 | }
169 | results.add(itemResult);
170 | }
171 | return results;
172 | }
173 | }
174 |
175 | /**
176 | * Orchestrates the execution of a sequence of connected nodes.
177 | */
178 | public static class Flow extends BaseNode {
179 | protected BaseNode, ?> startNode;
180 |
181 | public Flow() { this(null); }
182 | public Flow(BaseNode, ?> startNode) { this.start(startNode); }
183 |
184 | public BaseNode start(BaseNode startNode) {
185 | this.startNode = Objects.requireNonNull(startNode, "Start node cannot be null");
186 | return startNode;
187 | }
188 |
189 | @Override public final String exec(Void prepResult) {
190 | throw new UnsupportedOperationException("Flow.exec() is internal and should not be called directly.");
191 | }
192 |
193 |
194 | @SuppressWarnings({"unchecked", "rawtypes"}) // Raw types needed for successors map
195 | protected String orchestrate(Map sharedContext, Map initialParams) {
196 | if (startNode == null) { logWarn("Flow started with no start node."); return null; }
197 | BaseNode, ?> currentNode = this.startNode;
198 | String lastAction = null;
199 | Map currentParams = new HashMap<>(this.params);
200 | if (initialParams != null) { currentParams.putAll(initialParams); }
201 | while (currentNode != null) {
202 | currentNode.setParams(currentParams);
203 | lastAction = (String) ((BaseNode)currentNode).internalRun(sharedContext);
204 | currentNode = currentNode.getNextNode(lastAction);
205 | }
206 | return lastAction;
207 | }
208 |
209 | @Override
210 | protected String internalRun(Map sharedContext) {
211 | Void prepRes = prep(sharedContext);
212 | String orchRes = orchestrate(sharedContext, null);
213 | return post(sharedContext, prepRes, orchRes);
214 | }
215 |
216 | /** Post method for the Flow itself. Default returns the last action from orchestration. */
217 | @Override public String post(Map sharedContext, Void prepResult, String execResult) {
218 | return execResult;
219 | }
220 | }
221 |
222 | /**
223 | * A flow that runs its entire sequence for each parameter set from `prepBatch`.
224 | */
225 | public static abstract class BatchFlow extends Flow {
226 | public BatchFlow() { super(); }
227 | public BatchFlow(BaseNode, ?> startNode) { super(startNode); }
228 |
229 | public abstract List