Supporting DPICustomEntity and option field replacement_strategy

Author: n-o-u-r-h-a-n

orchestration/src/main/java/com/sap/ai/sdk/orchestration/DpiMasking.java

@@ -6,8 +6,10 @@
 
 import com.sap.ai.sdk.orchestration.model.DPIConfig;
 import com.sap.ai.sdk.orchestration.model.DPIConfigMaskGroundingInput;
+import com.sap.ai.sdk.orchestration.model.DPICustomEntity;
 import com.sap.ai.sdk.orchestration.model.DPIEntities;
 import com.sap.ai.sdk.orchestration.model.DPIEntityConfig;
+import com.sap.ai.sdk.orchestration.model.DPIMethodConstant;
 import com.sap.ai.sdk.orchestration.model.DPIStandardEntity;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -33,6 +35,7 @@
 public class DpiMasking implements MaskingProvider {
   @Nonnull DPIConfig.MethodEnum maskingMethod;
   @Nonnull List<DPIEntities> entities;
+  @Nonnull List<DPICustomEntity> customEntities;
   @With boolean maskGroundingInput;
   @Nonnull List<String> allowList;
 
@@ -78,7 +81,27 @@ public DpiMasking withEntities(
       val entitiesList = new ArrayList<DPIEntities>();
       entitiesList.add(entity);
       entitiesList.addAll(Arrays.asList(entities));
-      return new DpiMasking(maskingMethod, entitiesList, false, List.of());
+      return new DpiMasking(maskingMethod, entitiesList, List.of(), false, List.of());
+    }
+
+    /**
+     * Specifies a custom regex pattern for masking.
+     *
+     * @param regex The regex pattern to match
+     * @param replacement The replacement string
+     * @return A new {@link DpiMasking} instance
+     */
+    @Nonnull
+    public DpiMasking withRegex(@Nonnull final String regex, @Nonnull final String replacement) {
+      val customEntity =
+          DPICustomEntity.create()
+              .regex(regex)
+              .replacementStrategy(
+                  DPIMethodConstant.create()
+                      .method(DPIMethodConstant.MethodEnum.CONSTANT)
+                      .value(replacement));
+
+      return new DpiMasking(maskingMethod, List.of(), List.of(customEntity), false, List.of());
     }
   }
 
@@ -90,14 +113,18 @@ public DpiMasking withEntities(
    */
   @Nonnull
   public DpiMasking withAllowList(@Nonnull final List<String> allowList) {
-    return new DpiMasking(maskingMethod, entities, maskGroundingInput, allowList);
+    return new DpiMasking(maskingMethod, entities, customEntities, maskGroundingInput, allowList);
   }
 
   @Nonnull
   @Override
   public DPIConfig createConfig() {
     val entitiesDTO =
-        entities.stream().map(it -> (DPIEntityConfig) DPIStandardEntity.create().type(it)).toList();
+        new ArrayList<>(
+            entities.stream()
+                .map(it -> (DPIEntityConfig) DPIStandardEntity.create().type(it))
+                .toList());
+    entitiesDTO.addAll(customEntities);
     return DPIConfig.create()
         .type(SAP_DATA_PRIVACY_INTEGRATION)
         .method(maskingMethod)

sample-code/spring-app/src/main/java/com/sap/ai/sdk/app/services/OrchestrationService.java

@@ -262,9 +262,9 @@ public OrchestrationChatResponse maskingAnonymization(@Nonnull final DPIEntities
     val userMessage =
         Message.user(
             """
-                            I think the SDK is good, but could use some further enhancements.
-                            My architect Alice and manager Bob pointed out that we need the grounding capabilities, which aren't supported yet.
-                            """);
+    I think the SDK is good, but could use some further enhancements.
+    My architect Alice and manager Bob pointed out that we need the grounding capabilities, which aren't supported yet.
+    """);
 
     val prompt = new OrchestrationPrompt(systemMessage, userMessage);
     val maskingConfig = DpiMasking.anonymization().withEntities(entity);
@@ -273,6 +273,34 @@ public OrchestrationChatResponse maskingAnonymization(@Nonnull final DPIEntities
     return client.chatCompletion(prompt, configWithMasking);
   }
 
+  /**
+   * Let the orchestration service evaluate the feedback on the AI SDK provided by a hypothetical
+   * user. Anonymize any patient IDs given as they are not relevant for judging the sentiment of the
+   * feedback.
+   *
+   * @link <a
+   *     href="https://help.sap.com/docs/sap-ai-core/sap-ai-core-service-guide/data-masking">SAP AI
+   *     Core: Orchestration - Data Masking</a>
+   * @return the assistant response object
+   */
+  @Nonnull
+  public OrchestrationChatResponse maskingCustomAnonymization() {
+    val systemMessage = Message.system("Repeat this phrase");
+    val userMessage =
+        Message.user(
+            """
+       The patient id is patient_id_123.
+        """);
+
+    val prompt = new OrchestrationPrompt(systemMessage, userMessage);
+    val regex = "patient_id_[0-9]+";
+    val replacement = "REDACTED_ID";
+    val maskingConfig = DpiMasking.anonymization().withRegex(regex, replacement);
+    val configWithMasking = config.withMaskingConfig(maskingConfig);
+
+    return client.chatCompletion(prompt, configWithMasking);
+  }
+
   /**
    * Chat request to OpenAI through the Orchestration deployment under a specific resource group.
    *

sample-code/spring-app/src/test/java/com/sap/ai/sdk/app/controllers/OrchestrationTest.java

@@ -148,6 +148,25 @@ void testMaskingAnonymization() {
     assertThat(result.getIntermediateResults().getOutputUnmasking()).isEmpty();
   }
 
+  @Test
+  void testMaskingCustomAnonymization() {
+    var response = service.maskingCustomAnonymization();
+    var result = response.getOriginalResponse();
+    var llmChoice = result.getFinalResult().getChoices().get(0);
+    assertThat(llmChoice.getFinishReason()).isEqualTo("stop");
+
+    var maskingResult = result.getIntermediateResults().getInputMasking();
+    assertThat(maskingResult.getMessage()).isNotEmpty();
+    var data = (Map<String, Object>) maskingResult.getData();
+    var maskedMessage = (String) data.get("masked_template");
+    assertThat(maskedMessage)
+        .describedAs("The masked input should replace patient IDs with REDACTED_ID")
+        .doesNotContain("patient_id_123")
+        .contains("REDACTED_ID");
+
+    assertThat(result.getIntermediateResults().getOutputUnmasking()).isEmpty();
+  }
+
   @SuppressWarnings("unchecked")
   @Test
   void testMaskingPseudonymization() {