[Feature/agent_framework] Actually make the WorkflowStepFactory a Fac…

…tory (#243) Actually make the WorkflowStepFactory a Factory Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Dec 18, 2023 · e99f83c · e99f83c
1 parent 860bc54
commit e99f83c
Showing 1 changed file with 16 additions and 15 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java
@@ -18,13 +18,14 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.Supplier;
 
 /**
  * Generates instances implementing {@link WorkflowStep}.
  */
 public class WorkflowStepFactory {
 
-    private final Map<String, WorkflowStep> stepMap = new HashMap<>();
+    private final Map<String, Supplier<WorkflowStep>> stepMap = new HashMap<>();
 
     /**
      * Instantiate this class.
@@ -42,21 +43,21 @@ public WorkflowStepFactory(
         MachineLearningNodeClient mlClient,
         FlowFrameworkIndicesHandler flowFrameworkIndicesHandler
     ) {
-        stepMap.put(NoOpStep.NAME, new NoOpStep());
-        stepMap.put(CreateIndexStep.NAME, new CreateIndexStep(clusterService, client, flowFrameworkIndicesHandler));
-        stepMap.put(CreateIngestPipelineStep.NAME, new CreateIngestPipelineStep(client, flowFrameworkIndicesHandler));
+        stepMap.put(NoOpStep.NAME, NoOpStep::new);
+        stepMap.put(CreateIndexStep.NAME, () -> new CreateIndexStep(clusterService, client, flowFrameworkIndicesHandler));
+        stepMap.put(CreateIngestPipelineStep.NAME, () -> new CreateIngestPipelineStep(client, flowFrameworkIndicesHandler));
         stepMap.put(
             RegisterLocalModelStep.NAME,
-            new RegisterLocalModelStep(settings, clusterService, mlClient, flowFrameworkIndicesHandler)
+            () -> new RegisterLocalModelStep(settings, clusterService, mlClient, flowFrameworkIndicesHandler)
         );
-        stepMap.put(RegisterRemoteModelStep.NAME, new RegisterRemoteModelStep(mlClient, flowFrameworkIndicesHandler));
-        stepMap.put(DeployModelStep.NAME, new DeployModelStep(mlClient));
-        stepMap.put(UndeployModelStep.NAME, new UndeployModelStep(mlClient));
-        stepMap.put(CreateConnectorStep.NAME, new CreateConnectorStep(mlClient, flowFrameworkIndicesHandler));
-        stepMap.put(DeleteConnectorStep.NAME, new DeleteConnectorStep(mlClient));
-        stepMap.put(ModelGroupStep.NAME, new ModelGroupStep(mlClient, flowFrameworkIndicesHandler));
-        stepMap.put(ToolStep.NAME, new ToolStep());
-        stepMap.put(RegisterAgentStep.NAME, new RegisterAgentStep(mlClient));
+        stepMap.put(RegisterRemoteModelStep.NAME, () -> new RegisterRemoteModelStep(mlClient, flowFrameworkIndicesHandler));
+        stepMap.put(DeployModelStep.NAME, () -> new DeployModelStep(mlClient));
+        stepMap.put(UndeployModelStep.NAME, () -> new UndeployModelStep(mlClient));
+        stepMap.put(CreateConnectorStep.NAME, () -> new CreateConnectorStep(mlClient, flowFrameworkIndicesHandler));
+        stepMap.put(DeleteConnectorStep.NAME, () -> new DeleteConnectorStep(mlClient));
+        stepMap.put(ModelGroupStep.NAME, () -> new ModelGroupStep(mlClient, flowFrameworkIndicesHandler));
+        stepMap.put(ToolStep.NAME, ToolStep::new);
+        stepMap.put(RegisterAgentStep.NAME, () -> new RegisterAgentStep(mlClient));
     }
 
     /**
@@ -66,7 +67,7 @@ public WorkflowStepFactory(
      */
     public WorkflowStep createStep(String type) {
         if (stepMap.containsKey(type)) {
-            return stepMap.get(type);
+            return stepMap.get(type).get();
         }
         throw new FlowFrameworkException("Workflow step type [" + type + "] is not implemented.", RestStatus.NOT_IMPLEMENTED);
     }
@@ -75,7 +76,7 @@ public WorkflowStep createStep(String type) {
      * Gets the step map
      * @return a read-only copy of the step map
      */
-    public Map<String, WorkflowStep> getStepMap() {
+    public Map<String, Supplier<WorkflowStep>> getStepMap() {
         return Map.copyOf(this.stepMap);
     }
 }