Use of strong typing for Context and more...

[jmdevall]

Congratulations on your project! I find the simplicity very intriguing without compromising on power.

I’ve reworked the PocketFlow library with a focus on modularity, type safety, and decoupled components while keeping the core concepts intact. Below is an overview of key changes and enhancements compared to the original design:

• Alignments with Your Original Design: Your foundational ideas (e.g., workflows, nodes, exec/prep/post patterns) are preserved and improved in PFJ
• Strong typing instead of Map<String, Object> for Context
• Replaced raw maps with generics (C, I, O) for compile-time safety and type clarity.
• Preference Interfaces over inheritance: Node is an interface, and Step implements it without forced hierarchies (no BaseNode or subclass chains).
• Reliable for retry logic:
• A standalone class wraps functions (not tied to a class hierarchy), enabling retries with customizable delays.
• Flows as Steps:
  • The newFlow() method creates a Step that orchestrates sub-flows, enabling nesting.
  • Flows can act as Steps within other flows (e.g., newFlow(startNode).run()).
• Simplifies configuration with minimal boilerplate (e.g., Step.builder() thanks to Lombok).
• 100 lines of code (but I removed the batch classes)

package pocketflowj;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import lombok.Builder;

public class PFJ {
 static final String DEFAULT_ACTION = "default";
 private static void logWarn(String message) { System.err.println("WARN: PocketFlow - " + message); }
 static class PocketFlowException extends RuntimeException {
 public PocketFlowException(String message) { super(message); }
 public PocketFlowException(String message, Throwable cause) { super(message, cause); }
 }
 interface Node<C> { String run(C context); Node<C> getNextNode(String action); }
 interface ContextAdapter<C, I, O> { I prep(C context); String post(C context, I input, O output); }
 @FunctionalInterface interface TriFunction<A, B, C, R> { R apply(A a, B b, C c); }
 
 static class Reliable<I, O> implements Function<I, O> {
 private final int maxRetries;
 private final long waitMillis;
 private final Function<I, O> function;

 public Reliable(Function<I, O> function) { this(function, 1, 0); }
 public Reliable(Function<I, O> function, int maxRetries, long waitMillis) {
 if(maxRetries <1) throw new IllegalArgumentException("maxRetries must be at least 1");
 if(waitMillis <0) throw new IllegalArgumentException("waitMillis cannot be negative");
 this.function = function; this.maxRetries = maxRetries; this.waitMillis = waitMillis;
 }
 @Override
 public O apply(I input) {
 Exception lastException = null;
 for(int currentRetry=0; currentRetry < maxRetries; currentRetry++) {
 try { return function.apply(input); }
 catch(Exception e) { lastException = e; 
 if(currentRetry < maxRetries-1 && waitMillis>0) waitTheTime();
 }
 }
 if(lastException == null) throw new PocketFlowException("Execution failed, but no exception was captured.");
 if(lastException instanceof RuntimeException) throw (RuntimeException)lastException;
 throw new PocketFlowException("Final execution failed", lastException);
 }
 private void waitTheTime() {
 try { TimeUnit.MILLISECONDS.sleep(waitMillis); }
 catch(InterruptedException ie) { Thread.currentThread().interrupt();
 throw new PocketFlowException("Thread interrupted during retry wait", ie); }
 }
 }
 
 @Builder static class DefaultContextAdapter<C, I, O> implements ContextAdapter<C, I, O> {
 private Function<C, I> prep;
 private TriFunction<C, I, O, String> post;
 public I prep(C context) { return prep.apply(context); }
 public String post(C context, I input, O output) { return post.apply(context, input, output); }
 }
 
 @Builder static class Step<C, I, O> implements Node<C> {
 private Function<C, I> prep;
 private TriFunction<C, I, O, String> post;

 	private final ContextAdapter<C, I, O> adapter;
 private Function<I, O> exec;
 @Builder.Default private final int maxRetries = 1;
 @Builder.Default private final long waitMillis = 0;
 private final Map<String, Node<C>> successors = new HashMap<>();
 
 private Function<I, O> getWrappedFunction() { return maxRetries>1 ? new Reliable<>(exec, maxRetries, waitMillis) : exec; }
 private ContextAdapter<C, I, O> getWrappedAdapter(){ return this.adapter!=null?this.adapter:new DefaultContextAdapter<>(prep,post); }

 public String run(C context) {
 	ContextAdapter<C,I,O> a=getWrappedAdapter();
 I prepRes = a.prep(context);
 O execRes = getWrappedFunction().apply(prepRes);
 return a.post(context, prepRes, execRes);
 }
 
 public Node<C> getNextNode(String action) {
 String actionKey = Objects.requireNonNullElse(action, DEFAULT_ACTION);
 Node<C> nextNode = successors.get(actionKey);
 if(nextNode == null && !successors.isEmpty()) logWarn("Action '"+actionKey+"' not found in successors of "+getClass().getSimpleName());
 return nextNode;
 }
 public Node<C> on(String action, Node<C> next) { successors.put(action, next); return next; }
 public Node<C> defaultNext(Node<C> next) { return on(DEFAULT_ACTION, next); }
 }
 
	static <C> Step<C, C, String> newFlow(Node<C> firstStep) {
		ContextAdapter<C, C, String> adapter = DefaultContextAdapter.<C, C, String>builder().prep(Function.identity())
				.post((c, i, o) -> o).build();
		Function<C, String> exec = context -> {
			String lastAction = null;
			for (Node<C> current = firstStep; current != null; current = current.getNextNode(lastAction)) {
				lastAction = current.run(context);
			}
			return lastAction;
		};
		return Step.<C, C, String>builder().adapter(adapter).exec(exec).build();
	}
}

AgentTest.java.txt

[zachary62]

Amazing! Here are some quick thoughts:

• Oh, I love the Interface-over-inheritance design!
• TriFunction feels 'a bit ugly,' but I see it's pragmatic.
• Steps are not hidden from the user right. Is exposing Step potentially confusing?
• What does the developer experience look like? Could you provide a simple example code that uses this library to implement an LLM app? Like an agent?
• Could you reintroduce the batch node? It's actually very handy (unless there is another easier way to achieve something similar)? It's fine if it's not 100 lines for now - we can reduce later.
• Also, I'd be happy to add you as a collaborator on this repo if you want to maintain this project. Let me know!

[jmdevall]

Hello! , for an example of usage, see the attached file "AgentTest.java.txt", apologies if dependencies are missing, but it's because I had it in a project I was tinkering with

I see batch as a simple Step that processes a list of elements, I don't think it should require something special, let me show how should be...

In your documentation:

A BatchNode extends Node but changes prep() and exec():

 prep(shared): returns an iterable (e.g., list, generator).
 exec(item): called once per item in that iterable.
 post(shared, prep_res, exec_res_list): after all items are processed, receives a list of results (exec_res_list) and returns an Action.

Example: Summarize a Large File

class MapSummaries(BatchNode):
 def prep(self, shared):
 # Suppose we have a big file; chunk it
 content = shared["data"]
 chunk_size = 10000
 chunks = [content[i:i+chunk_size] for i in range(0, len(content), chunk_size)]
 return chunks

 def exec(self, chunk):
 prompt = f"Summarize this chunk in 10 words: {chunk}"
 summary = call_llm(prompt)
 return summary

 def post(self, shared, prep_res, exec_res_list):
 combined = "\n".join(exec_res_list)
 shared["summary"] = combined
 return "default"

map_summaries = MapSummaries()
flow = Flow(start=map_summaries)
flow.run(shared)

It would translate to the following equivalent code (sorry if there are any errors, as I did it on the fly)...

package pocketflowj;

import java.util.ArrayList;
import java.util.List;

import org.junit.jupiter.api.Test;

import lombok.Data;
import lombok.var;
import pocketflowj.PFJ.Reliable;
import pocketflowj.PFJ.Step;

public class BatchExample {

 public static List<String> splitIntoChunks(String str, int chunkSize) {
 List<String> chunks = new ArrayList<>();
 int length = str.length();

 for (int i = 0; i < length; i += chunkSize) {
 chunks.add(str.substring(i, Math.min(i + chunkSize, length)));
 }

 return chunks;
 }

 @Data
 static class SharedContext{
 	String fullText;
 	String fullSummary;
 }
 
 private String summary(String text) {
 // haría el resumen de un "chunk"
 	return SetupUtil.callLlm("Summarize this chunk in 10 words: "+text);
 }
 
 private List<String> summaryChunks(List<String> chunks){
 	return chunks.stream().map(
 		new Reliable<>(this::summary, 3, 500))
 	.toList();
 }
	
	@Test
	public void testBatchExample() {
		
		System.out.println("test batch");
		
		var simpleBatchStep = Step.<SharedContext, List<String>, List<String>>builder()
			 .prep((shared)->{
			 	 String fullContent=shared.getFullText();
			 	 return splitIntoChunks(fullContent,10000);
			 })
			 .exec(this::summaryChunks)
			 .post((SharedContext sharedContext, List<String> prepResult, List<String> execResult) -> {
			 	 sharedContext.setFullSummary(String.join("",execResult));
			 	 return PFJ.DEFAULT_ACTION;
			 })
			 .build();
		
		//continue creating the initial context and run...
		
	}
}

[zachary62]

Thanks for sending the example! I'll review them when I have a moment.
If you're interested in maintaining the Java branch repo to finish the library and create examples, let me know.

[cmf]

One difference in the Batch interface proposed above is that exec() now receives and returns a collection of items. In the original Python implementation, exec() was still specified to process a single item which is invoked repeatedly (a mapping operation), whereas in this proposal the node has to do that looping itself. I guess it's an exercise for the reader whether that is a desirable change or not.

[jmdevall]

Yes, this is a solution that is more in line with functional programming; we could create a function that encapsulates the loop and the reliable, and it would remain almost the same:

public static <I,O> Function<List<I>, List<O>> newBatchLoop(Function<I, O> execItem, int maxRetries, long waitMillis){
 	return (List<I> list) -> list.stream().map(maxRetries>1 ? new Reliable<>(execItem, maxRetries, waitMillis) : execItem).toList();
}

and the construction of the Step....

.exec(PFJ.newBatchLoop(this::summary,3,500))