Fix a template field name bug in demo

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Sep 25, 2023 · 14c3338 · 14c3338
1 parent 88f0183
commit 14c3338
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 15 deletions.
diff --git a/src/main/java/demo/CreateIndexWorkflowStep.java b/src/main/java/demo/CreateIndexWorkflowStep.java
@@ -43,20 +43,16 @@ public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
         // https://github.com/opensearch-project/opensearch-ai-flow-framework/issues/42
         CompletableFuture.runAsync(() -> {
             String inputIndex = null;
-            boolean first = true;
             for (WorkflowData wfData : data) {
                 logger.debug(
                     "{} sent params: {}, content: {}",
-                    first ? "Initialization" : "Previous step",
+                    inputIndex == null ? "Initialization" : "Previous step",
                     wfData.getParams(),
                     wfData.getContent()
                 );
-                if (first) {
-                    Map<String, String> params = data.get(0).getParams();
-                    if (params.containsKey("index")) {
-                        inputIndex = params.get("index");
-                    }
-                    first = false;
+                if (inputIndex == null) {
+                    inputIndex = wfData.getParams()
+                        .getOrDefault("index_name", (String) wfData.getContent().getOrDefault("index_name", "NOT FOUND"));
                 }
             }
             // do some work, simulating a REST API call
@@ -68,7 +64,7 @@ public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
             future.complete(new WorkflowData() {
                 @Override
                 public Map<String, Object> getContent() {
-                    return Map.of("index", response.index());
+                    return Map.of("index_created", response.index());
                 }
             });
         });

diff --git a/src/main/java/org/opensearch/flowframework/template/ProcessNode.java b/src/main/java/org/opensearch/flowframework/template/ProcessNode.java
@@ -129,9 +129,9 @@ public CompletableFuture<WorkflowData> execute() {
             }
             CompletableFuture<WorkflowData> stepFuture = this.workflowStep.execute(input);
             try {
-                stepFuture.join();
+                stepFuture.orTimeout(15, TimeUnit.SECONDS).join();
+                logger.info(">>> Finished {}.", this.id);
                 future.complete(stepFuture.get());
-                logger.debug("<<< Completed {}", this.id);
             } catch (InterruptedException | ExecutionException e) {
                 handleException(e);
             }
@@ -142,7 +142,7 @@ public CompletableFuture<WorkflowData> execute() {
     private void handleException(Exception e) {
         // TODO: better handling of getCause
         this.future.completeExceptionally(e);
-        logger.debug("<<< Completed Exceptionally {}", this.id);
+        logger.debug("<<< Completed Exceptionally {}", this.id, e.getCause());
     }
 
     @Override

diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java
@@ -42,7 +42,6 @@ private void populateMap() {
         // Replace with actual implementations such as
         // https://github.com/opensearch-project/opensearch-ai-flow-framework/pull/38
         // https://github.com/opensearch-project/opensearch-ai-flow-framework/pull/44
-        stepMap.put("create_index", new CreateIndexWorkflowStep());
         stepMap.put("fetch_model", new DemoWorkflowStep(3000));
         stepMap.put("create_ingest_pipeline", new DemoWorkflowStep(3000));
         stepMap.put("create_search_pipeline", new DemoWorkflowStep(5000));
@@ -70,8 +69,13 @@ public String getName() {
      * @return an instance of the specified type
      */
     public WorkflowStep createStep(String type) {
-        if (stepMap.containsKey(type)) {
-            return stepMap.get(type);
+        switch (type) {
+            case "create_index":
+                return new CreateIndexWorkflowStep();
+            default:
+                if (stepMap.containsKey(type)) {
+                    return stepMap.get(type);
+                }
         }
         // TODO: replace this with a FlowFrameworkException
         // https://github.com/opensearch-project/opensearch-ai-flow-framework/pull/43