feat: workflow improvements

Author: alexcheng1982

chain-workflow/src/main/java/com/javaaidev/agenticpatterns/chainworkflow/ChainStep.java

@@ -0,0 +1,18 @@
+package com.javaaidev.agenticpatterns.chainworkflow;
+
+import java.util.Map;
+import org.springframework.core.Ordered;
+
+public interface ChainStep<Request, Response> extends Ordered {
+
+  /**
+   * Call the current step
+   *
+   * @param request Task input
+   * @param context Shared context between different steps
+   * @param chain   The chain, see {@linkplain WorkflowChain}
+   * @return Task output
+   */
+  Response call(Request request, Map<String, Object> context,
+      WorkflowChain<Request, Response> chain);
+}

chain-workflow/src/main/java/com/javaaidev/agenticpatterns/chainworkflow/ChainStepAgent.java

@@ -1,12 +1,17 @@
 package com.javaaidev.agenticpatterns.chainworkflow;
 
+import com.javaaidev.agenticpatterns.taskexecution.AbstractTaskExecutionAgentBuilder;
 import com.javaaidev.agenticpatterns.taskexecution.TaskExecutionAgent;
 import io.micrometer.observation.ObservationRegistry;
 import java.lang.reflect.Type;
 import java.util.Map;
+import java.util.function.Consumer;
+import java.util.function.Function;
 import org.jspecify.annotations.Nullable;
 import org.springframework.ai.chat.client.ChatClient;
+import org.springframework.ai.chat.client.ChatClient.ChatClientRequestSpec;
 import org.springframework.core.Ordered;
+import org.springframework.util.Assert;
 
 /**
  * A step in the chain
@@ -16,7 +21,7 @@
  */
 public abstract class ChainStepAgent<Request, Response> extends
     TaskExecutionAgent<Request, Response> implements
-    Ordered {
+    Ordered, ChainStep<Request, Response> {
 
   protected ChainStepAgent(ChatClient chatClient,
       @Nullable ObservationRegistry observationRegistry) {
@@ -29,14 +34,62 @@ protected ChainStepAgent(ChatClient chatClient,
     super(chatClient, responseType, observationRegistry);
   }
 
-  /**
-   * Call the current step
-   *
-   * @param request Task input
-   * @param context Shared context between different steps
-   * @param chain   The chain, see {@linkplain WorkflowChain}
-   * @return Task output
-   */
-  protected abstract Response call(Request request, Map<String, Object> context,
-      WorkflowChain<Request, Response> chain);
+  public ChainStepAgent(ChatClient chatClient,
+      String promptTemplate,
+      @Nullable Type responseType,
+      @Nullable Function<Request, Map<String, Object>> promptTemplateContextProvider,
+      @Nullable Consumer<ChatClientRequestSpec> chatClientRequestSpecUpdater,
+      @Nullable String name,
+      @Nullable ObservationRegistry observationRegistry) {
+    super(chatClient, promptTemplate, responseType, promptTemplateContextProvider,
+        chatClientRequestSpecUpdater, name, observationRegistry);
+  }
+
+  public static <Req, Res> Builder<Req, Res> builder() {
+    return new Builder<>();
+  }
+
+  public static class Builder<Request, Response> extends
+      AbstractTaskExecutionAgentBuilder<Request, Response, Builder<Request, Response>> {
+
+    private Function<Response, Request> nextRequestPreparer;
+    private int order;
+
+    public Builder<Request, Response> nextRequestPreparer(
+        Function<Response, Request> nextRequestPreparer) {
+      this.nextRequestPreparer = nextRequestPreparer;
+      return this;
+    }
+
+    public Builder<Request, Response> order(int order) {
+      this.order = order;
+      return this;
+    }
+
+    @Override
+    public ChainStepAgent<Request, Response> build() {
+      Assert.notNull(chatClient, "ChatClient cannot be null");
+      Assert.hasText(promptTemplate, "Prompt template cannot be empty");
+      Assert.notNull(nextRequestPreparer, "nextRequestPreparer cannot be null");
+      return new ChainStepAgent<>(chatClient,
+          promptTemplate,
+          responseType,
+          promptTemplateContextProvider,
+          chatClientRequestSpecUpdater,
+          name,
+          observationRegistry) {
+        @Override
+        public Response call(Request request, Map<String, Object> context,
+            WorkflowChain<Request, Response> chain) {
+          var response = this.call(request);
+          return chain.callNext(nextRequestPreparer.apply(response), response);
+        }
+
+        @Override
+        public int getOrder() {
+          return order;
+        }
+      };
+    }
+  }
 }

chain-workflow/src/main/java/com/javaaidev/agenticpatterns/chainworkflow/ChainWorkflow.java

@@ -0,0 +1,54 @@
+package com.javaaidev.agenticpatterns.chainworkflow;
+
+import com.javaaidev.agenticpatterns.core.AbstractAgenticWorkflow;
+import com.javaaidev.agenticpatterns.core.AbstractAgenticWorkflowBuilder;
+import io.micrometer.observation.ObservationRegistry;
+import java.util.ArrayList;
+import java.util.List;
+import org.jspecify.annotations.Nullable;
+import org.springframework.util.Assert;
+
+public class ChainWorkflow<Request, Response> extends AbstractAgenticWorkflow<Request, Response> {
+
+  private final List<ChainStep<Request, Response>> steps;
+
+  ChainWorkflow(List<ChainStep<Request, Response>> steps,
+      @Nullable String workflowName,
+      @Nullable ObservationRegistry observationRegistry) {
+    super(workflowName, observationRegistry);
+    this.steps = new ArrayList<>(steps);
+  }
+
+  @Override
+  protected Response doExecute(@Nullable Request request) {
+    var chain = new WorkflowChain<>(steps);
+    return chain.callNext(request, null);
+  }
+
+  public static <Req, Res> Builder<Req, Res> builder() {
+    return new Builder<>();
+  }
+
+  public static final class Builder<Request, Response> extends
+      AbstractAgenticWorkflowBuilder<Request, Response, Builder<Request, Response>> {
+
+    private final List<ChainStep<Request, Response>> steps = new ArrayList<>();
+
+    public Builder<Request, Response> addStep(ChainStep<Request, Response> step) {
+      Assert.notNull(step, "Chain step cannot be null");
+      steps.add(step);
+      return this;
+    }
+
+    public Builder<Request, Response> addStepAgent(
+        ChainStepAgent<Request, Response> agent) {
+      Assert.notNull(agent, "Chain step agent cannot be null");
+      steps.add(agent);
+      return this;
+    }
+
+    public ChainWorkflow<Request, Response> build() {
+      return new ChainWorkflow<>(steps, name, observationRegistry);
+    }
+  }
+}

chain-workflow/src/main/java/com/javaaidev/agenticpatterns/chainworkflow/ChainWorkflowAgent.java

@@ -18,14 +18,14 @@
  */
 public class WorkflowChain<Request, Response> {
 
-  private final Deque<ChainStepAgent<Request, Response>> agents;
+  private final Deque<ChainStep<Request, Response>> steps;
   private final Map<String, Object> context = new HashMap<>();
 
   private static final Logger LOGGER = LoggerFactory.getLogger(WorkflowChain.class);
 
-  public WorkflowChain(List<ChainStepAgent<Request, Response>> agents) {
-    this.agents = new ArrayDeque<>(agents.stream().sorted(OrderComparator.INSTANCE).toList());
-    LOGGER.info("Added {} agents to the chain", this.agents.size());
+  public WorkflowChain(List<ChainStep<Request, Response>> steps) {
+    this.steps = new ArrayDeque<>(steps.stream().sorted(OrderComparator.INSTANCE).toList());
+    LOGGER.info("Added {} agents to the chain", this.steps.size());
   }
 
   /**
@@ -36,10 +36,10 @@ public WorkflowChain(List<ChainStepAgent<Request, Response>> agents) {
    * @return Task output
    */
   public Response callNext(Request request, @Nullable Response lastResponse) {
-    if (agents.isEmpty()) {
+    if (steps.isEmpty()) {
       return lastResponse;
     }
-    var agent = agents.pop();
-    return agent.call(request, context, this);
+    var step = steps.pop();
+    return step.call(request, context, this);
   }
 }

chain-workflow/src/main/java/com/javaaidev/agenticpatterns/chainworkflow/WorkflowChain.java

@@ -10,19 +10,19 @@ public abstract class AbstractAgenticWorkflow<Request, Response> implements
     AgenticWorkflow<Request, Response> {
 
   @Nullable
-  protected String workflowName;
+  protected String name;
   @Nullable
   protected ObservationRegistry observationRegistry;
 
-  protected AbstractAgenticWorkflow(@Nullable String workflowName,
+  protected AbstractAgenticWorkflow(@Nullable String name,
       @Nullable ObservationRegistry observationRegistry) {
-    this.workflowName = workflowName;
+    this.name = name;
     this.observationRegistry = observationRegistry;
   }
 
   @Override
   public String getName() {
-    return this.workflowName != null ? this.workflowName : this.getClass().getSimpleName();
+    return this.name != null ? this.name : this.getClass().getSimpleName();
   }
 
   @Override
@@ -51,4 +51,5 @@ public Response execute(@Nullable Request request) {
   }
 
   protected abstract Response doExecute(@Nullable Request request);
+
 }

core/src/main/java/com/javaaidev/agenticpatterns/core/AbstractAgenticWorkflow.java

@@ -0,0 +1,29 @@
+package com.javaaidev.agenticpatterns.core;
+
+import io.micrometer.observation.ObservationRegistry;
+import org.jspecify.annotations.Nullable;
+
+public abstract class AbstractAgenticWorkflowBuilder<Request, Response, T extends AbstractAgenticWorkflowBuilder<Request, Response, T>> implements
+    AgenticWorkflow.Builder<Request, Response, T> {
+
+  @Nullable
+  protected String name;
+  @Nullable
+  protected ObservationRegistry observationRegistry;
+
+  @SuppressWarnings("unchecked")
+  protected T self() {
+    return (T) this;
+  }
+
+  public T name(String name) {
+    this.name = name;
+    return self();
+  }
+
+  public T observationRegistry(
+      ObservationRegistry observationRegistry) {
+    this.observationRegistry = observationRegistry;
+    return self();
+  }
+}

core/src/main/java/com/javaaidev/agenticpatterns/core/AbstractAgenticWorkflowBuilder.java

@@ -1,5 +1,6 @@
 package com.javaaidev.agenticpatterns.core;
 
+import io.micrometer.observation.ObservationRegistry;
 import org.jspecify.annotations.Nullable;
 
 public interface AgenticWorkflow<Request, Response> {
@@ -9,4 +10,21 @@ public interface AgenticWorkflow<Request, Response> {
   default String getName() {
     return this.getClass().getSimpleName();
   }
+
+  interface Builder<Request, Response, T extends AgenticWorkflow.Builder<Request, Response, T>> {
+
+    T name(String name);
+
+    T observationRegistry(ObservationRegistry observationRegistry);
+
+    AgenticWorkflow<Request, Response> build();
+  }
+
+  static <Request, Response> SingleStepWorkflowBuilder<Request, Response> single() {
+    return new SingleStepWorkflowBuilder<>();
+  }
+
+  static <Request, Response> CompositeWorkflowBuilder<Request, Response> composite() {
+    return new CompositeWorkflowBuilder<>();
+  }
 }

core/src/main/java/com/javaaidev/agenticpatterns/core/AgenticWorkflow.java

@@ -0,0 +1,35 @@
+package com.javaaidev.agenticpatterns.core;
+
+import java.util.ArrayDeque;
+import java.util.Deque;
+import org.jspecify.annotations.Nullable;
+import org.springframework.util.Assert;
+
+public class CompositeWorkflowBuilder<Request, Response> extends
+    AbstractAgenticWorkflowBuilder<Request, Response, CompositeWorkflowBuilder<Request, Response>> {
+
+  private final Deque<AgenticWorkflow> deque = new ArrayDeque<>();
+
+  public CompositeWorkflowBuilder<Request, Response> addNext(AgenticWorkflow workflow) {
+    Assert.notNull(workflow, "Workflow cannot be null");
+    deque.addLast(workflow);
+    return this;
+  }
+
+  @Override
+  public AgenticWorkflow<Request, Response> build() {
+    Assert.notEmpty(deque, "At least one workflow is required");
+    return new AbstractAgenticWorkflow<>(name, observationRegistry) {
+      @Override
+      @SuppressWarnings("unchecked")
+      protected Response doExecute(@Nullable Request request) {
+        Object result = request;
+        do {
+          var workflow = deque.pollFirst();
+          result = workflow.execute(result);
+        } while (!deque.isEmpty());
+        return (Response) result;
+      }
+    };
+  }
+}