Introduce enrich and enrichOutput FuncDSL methods

Signed-off-by: Matheus Cruz <matheuscruz.dev@gmail.com>

Author: mcruzdev

experimental/fluent/func/src/main/java/io/serverlessworkflow/fluent/func/dsl/EnrichWithModelBiFunction.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2020-Present The Serverless Workflow Specification Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.serverlessworkflow.fluent.func.dsl;
+
+import io.serverlessworkflow.impl.WorkflowModel;
+
+/**
+ * A function that enriches a typed value by combining it with the root workflow input.
+ *
+ * <p>This is useful when you need to merge a task output or state value with the original workflow
+ * input, where the value is typed and the root input is a {@link WorkflowModel}.
+ *
+ * @param <T> The type of the input value to be enriched
+ * @param <R> The type of the enriched result
+ */
+@FunctionalInterface
+public interface EnrichWithModelBiFunction<T, R> {
+  /**
+   * Applies this function to enrich a typed value with the root workflow input.
+   *
+   * @param value the typed value to be enriched (for example a task output or state value)
+   * @param rootInput the original workflow input as {@link WorkflowModel}
+   * @return the enriched result combining both inputs
+   */
+  R apply(T value, WorkflowModel rootInput);
+}

experimental/fluent/func/src/main/java/io/serverlessworkflow/fluent/func/dsl/FuncDSL.java

@@ -32,6 +32,7 @@
 import io.serverlessworkflow.fluent.spec.configurers.AuthenticationConfigurer;
 import io.serverlessworkflow.impl.TaskContextData;
 import io.serverlessworkflow.impl.WorkflowContextData;
+import io.serverlessworkflow.impl.WorkflowModel;
 import java.net.URI;
 import java.util.Collection;
 import java.util.List;
@@ -279,6 +280,176 @@ public static FuncPredicateEventConfigurer event(String type) {
     return OPS.event(type);
   }
 
+  /**
+   * Create an input transformation that enriches the typed last state with the root workflow input
+   * as WorkflowModel.
+   *
+   * <p>This is useful when you want to combine the last task output (with its actual type) with the
+   * original workflow input as WorkflowModel.
+   *
+   * <p>Example usage:
+   *
+   * <pre>{@code
+   * function("processData", (Long input) -> input + 5, Long.class),
+   * function("combineData", (Long enrichedValue) -> enrichedValue, Long.class)
+   *     .inputFrom(enrich((lastState, rootInputModel) -> {
+   *         Long rootInput = rootInputModel.as(Long.class).orElse(0L);
+   *         return lastState + rootInput;
+   *     }, Long.class))
+   * }</pre>
+   *
+   * @param fn the enrichment function that receives typed lastState and WorkflowModel rootInput
+   * @param lastStateClass the class of the last state type
+   * @param <T> the type of the last state
+   * @param <R> the type of the enriched result
+   * @return a JavaContextFunction that can be used with inputFrom
+   */
+  public static <T, R> JavaContextFunction<T, R> enrich(
+      EnrichWithModelBiFunction<T, R> fn, Class<T> lastStateClass) {
+    return (lastState, workflowContext) -> {
+      Objects.requireNonNull(lastStateClass, "lastStateClass must not be null");
+      final WorkflowModel rootInput = workflowContext.instanceData().input();
+      final WorkflowModel lastStateModel =
+          workflowContext.definition().application().modelFactory().fromAny(lastState);
+      final T typedLastState =
+          lastStateModel
+              .as(lastStateClass)
+              .orElseThrow(
+                  () ->
+                      new IllegalArgumentException(
+                          "Cannot convert lastState of type "
+                              + (lastState != null ? lastState.getClass().getName() : "null")
+                              + " to "
+                              + lastStateClass.getName()));
+      return fn.apply(typedLastState, rootInput);
+    };
+  }
+
+  /**
+   * Create an input transformation that uses only the root workflow input as WorkflowModel.
+   *
+   * <p>This is useful when you want to transform the task input based solely on the original
+   * workflow input, ignoring the last state.
+   *
+   * <p>Example usage:
+   *
+   * <pre>{@code
+   * function("processData", (Long input) -> input * 2, Long.class)
+   *     .inputFrom(enrich(rootInput -> rootInput.asNumber().orElseThrow()))
+   * }</pre>
+   *
+   * @param fn the function that receives the root workflow input and returns the enriched result
+   * @param <R> the type of the enriched result
+   * @return a JavaContextFunction that can be used with inputFrom
+   */
+  public static <R> JavaContextFunction<Object, R> enrich(Function<WorkflowModel, R> fn) {
+    return (lastState, workflowContext) -> {
+      final WorkflowModel rootInput = workflowContext.instanceData().input();
+      return fn.apply(rootInput);
+    };
+  }
+
+  /**
+   * Create an output transformation that enriches the typed task output with the root workflow
+   * input as WorkflowModel.
+   *
+   * <p>This is useful when you want to combine the task output (with its actual type) with the
+   * original workflow input as WorkflowModel.
+   *
+   * <p>Example usage:
+   *
+   * <pre>{@code
+   * function("processData", (Long input) -> input + 5, Long.class)
+   *     .outputAs(FuncDSL.enrichOutput((taskOutput, rootInputModel) -> {
+   *         Long rootInput = rootInputModel.as(Long.class).orElse(0L);
+   *         return taskOutput + rootInput;
+   *     }, Long.class))
+   * }</pre>
+   *
+   * @param fn the enrichment function that receives typed taskOutput and WorkflowModel rootInput
+   * @param taskOutputClass the class of the task output type
+   * @param <T> the type of the task output
+   * @param <R> the type of the enriched result
+   * @return a JavaContextFunction that can be used with outputAs
+   */
+  public static <T, R> JavaContextFunction<T, R> enrichOutput(
+      EnrichWithModelBiFunction<T, R> fn, Class<T> taskOutputClass) {
+    return (taskOutput, workflowContext) -> {
+      Objects.requireNonNull(taskOutputClass, "taskOutputClass must not be null");
+      final WorkflowModel rootInput = workflowContext.instanceData().input();
+      final T typedTaskOutput = convertTaskOutput(taskOutput, taskOutputClass);
+      return fn.apply(typedTaskOutput, rootInput);
+    };
+  }
+
+  private static <T> T convertTaskOutput(Object taskOutput, Class<T> taskOutputClass) {
+    if (taskOutput == null) {
+      return null;
+    }
+    if (taskOutputClass.isInstance(taskOutput)) {
+      return taskOutputClass.cast(taskOutput);
+    }
+    if (Number.class.isAssignableFrom(taskOutputClass) && taskOutput instanceof Number) {
+      Number number = (Number) taskOutput;
+      if (taskOutputClass == Long.class) {
+        @SuppressWarnings("unchecked")
+        T converted = (T) Long.valueOf(number.longValue());
+        return converted;
+      }
+      if (taskOutputClass == Integer.class) {
+        @SuppressWarnings("unchecked")
+        T converted = (T) Integer.valueOf(number.intValue());
+        return converted;
+      }
+      if (taskOutputClass == Short.class) {
+        @SuppressWarnings("unchecked")
+        T converted = (T) Short.valueOf(number.shortValue());
+        return converted;
+      }
+      if (taskOutputClass == Byte.class) {
+        @SuppressWarnings("unchecked")
+        T converted = (T) Byte.valueOf(number.byteValue());
+        return converted;
+      }
+      if (taskOutputClass == Double.class) {
+        @SuppressWarnings("unchecked")
+        T converted = (T) Double.valueOf(number.doubleValue());
+        return converted;
+      }
+      if (taskOutputClass == Float.class) {
+        @SuppressWarnings("unchecked")
+        T converted = (T) Float.valueOf(number.floatValue());
+        return converted;
+      }
+    }
+    // Fallback to the original behavior for incompatible types.
+    return taskOutputClass.cast(taskOutput);
+  }
+
+  /**
+   * Create an output transformation that uses only the root workflow input as WorkflowModel.
+   *
+   * <p>This is useful when you want to transform the task output based solely on the original
+   * workflow input, ignoring the actual task output.
+   *
+   * <p>Example usage:
+   *
+   * <pre>{@code
+   * function("processData", (Long input) -> input * 2, Long.class)
+   *     .outputAs(FuncDSL.enrichOutput(rootInput -> rootInput.asNumber().orElseThrow()))
+   * }</pre>
+   *
+   * @param fn the function that receives the root workflow input and returns the enriched result
+   * @param <R> the type of the enriched result
+   * @return a JavaContextFunction that can be used with outputAs
+   */
+  public static <R> JavaContextFunction<Object, R> enrichOutput(Function<WorkflowModel, R> fn) {
+    return (taskOutput, workflowContext) -> {
+      final WorkflowModel rootInput = workflowContext.instanceData().input();
+      return fn.apply(rootInput);
+    };
+  }
+
   /**
    * Create a {@link FuncCallStep} that calls a simple Java {@link Function} with explicit input
    * type.

impl/test/pom.xml

@@ -94,6 +94,18 @@
             <artifactId>grpc-netty</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.serverlessworkflow</groupId>
+            <artifactId>serverlessworkflow-experimental-fluent-func</artifactId>
+            <scope>test</scope>
+            <version>${project.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.serverlessworkflow</groupId>
+            <artifactId>serverlessworkflow-experimental-lambda</artifactId>
+            <scope>test</scope>
+            <version>${project.version}</version>
+        </dependency>
     </dependencies>
     <build>
         <resources>