feat: [Orchestration] Mask grounding (#371)

* fix: [Orchestration] Added location to server errors

* wip

* location only

* reduce coverage for error 500

* message + location

* Added test

* finito

* docs

* fixed bug

* Added unit test

* Updated test

* More tests

* Used lombok @with

* Update docs/guides/ORCHESTRATION_CHAT_COMPLETION.md

---------

Co-authored-by: Jonas-Isr <[email protected]>

Author: CharlesDuboisSAP

docs/guides/ORCHESTRATION_CHAT_COMPLETION.md

@@ -271,6 +271,23 @@ var response = client.chatCompletion(prompt, configWithGrounding);
 
 Please find [an example in our Spring Boot application](../../sample-code/spring-app/src/main/java/com/sap/ai/sdk/app/services/OrchestrationService.java).
 
+### Mask Grounding
+
+You can also mask both the grounding information and the prompt message:
+
+```java
+var maskingConfig =
+    DpiMasking.anonymization()
+        .withEntities(DPIEntities.SENSITIVE_DATA)
+        .withMaskGroundingEnabled()
+        .withAllowList(List.of("SAP", "Joule"));
+var maskedGroundingConfig = groundingConfig.withMaskingConfig(maskingConfig);
+
+var result = client.chatCompletion(prompt, maskedGroundingConfig);
+```
+
+Please find [an example in our Spring Boot application](../../sample-code/spring-app/src/main/java/com/sap/ai/sdk/app/services/OrchestrationService.java).
+
 ## Stream chat completion
 
 It's possible to pass a stream of chat completion delta elements, e.g. from the application backend to the frontend in real-time.

docs/release-notes/release_notes.md

@@ -14,6 +14,7 @@
 
 - [Orchestration] [Add Spring AI Chat Memory support](https://github.com/SAP/ai-sdk-java/tree/main/docs/guides/SPRING_AI_INTEGRATION.md#chat-memory)
 - [Orchestration] [Prompt templates can be consumed from registry.](https://github.com/SAP/ai-sdk-java/tree/main/docs/guides/ORCHESTRATION_CHAT_COMPLETION.md#Chat-completion-with-Templates)
+- [Orchestration] [Masking is now available on grounding.](https://github.com/SAP/ai-sdk-java/tree/main/docs/guides/ORCHESTRATION_CHAT_COMPLETION.md#mask-grounding)
 - [Orchestration] [Grounding via *help.sap.com* is enabled.](https://github.com/SAP/ai-sdk-java/tree/main/docs/guides/ORCHESTRATION_CHAT_COMPLETION.md#grounding)
 
 ### 📈 Improvements

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

@@ -5,6 +5,7 @@
 import static com.sap.ai.sdk.orchestration.model.DPIConfig.TypeEnum.SAP_DATA_PRIVACY_INTEGRATION;
 
 import com.sap.ai.sdk.orchestration.model.DPIConfig;
+import com.sap.ai.sdk.orchestration.model.DPIConfigMaskGroundingInput;
 import com.sap.ai.sdk.orchestration.model.DPIEntities;
 import com.sap.ai.sdk.orchestration.model.DPIEntityConfig;
 import com.sap.ai.sdk.orchestration.model.MaskingProviderConfig;
@@ -16,6 +17,7 @@
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.Value;
+import lombok.With;
 import lombok.val;
 
 /**
@@ -31,6 +33,8 @@
 public class DpiMasking implements MaskingProvider {
   @Nonnull DPIConfig.MethodEnum maskingMethod;
   @Nonnull List<DPIEntities> entities;
+  @With boolean maskGroundingInput;
+  @Nonnull List<String> allowList;
 
   /**
    * Build a configuration applying anonymization.
@@ -65,7 +69,7 @@ public static class Builder {
      *
      * @param entity An entity type to mask (required)
      * @param entities Additional entity types to mask (optional)
-     * @return A configured {@link DpiMasking} instance
+     * @return A new {@link DpiMasking} instance
      * @see DPIEntities
      */
     @Nonnull
@@ -74,17 +78,30 @@ public DpiMasking withEntities(
       val entitiesList = new ArrayList<DPIEntities>();
       entitiesList.add(entity);
       entitiesList.addAll(Arrays.asList(entities));
-      return new DpiMasking(maskingMethod, entitiesList);
+      return new DpiMasking(maskingMethod, entitiesList, false, List.of());
     }
   }
 
+  /**
+   * Set words that should not be masked.
+   *
+   * @param allowList List of strings that should not be masked
+   * @return A new {@link DpiMasking} instance
+   */
+  @Nonnull
+  public DpiMasking withAllowList(@Nonnull final List<String> allowList) {
+    return new DpiMasking(maskingMethod, entities, maskGroundingInput, allowList);
+  }
+
   @Nonnull
   @Override
   public MaskingProviderConfig createConfig() {
     val entitiesDTO = entities.stream().map(it -> DPIEntityConfig.create().type(it)).toList();
     return DPIConfig.create()
         .type(SAP_DATA_PRIVACY_INTEGRATION)
         .method(maskingMethod)
-        .entities(entitiesDTO);
+        .entities(entitiesDTO)
+        .maskGroundingInput(DPIConfigMaskGroundingInput.create().enabled(maskGroundingInput))
+        .allowlist(allowList);
   }
 }

orchestration/src/test/java/com/sap/ai/sdk/orchestration/OrchestrationModuleConfigTest.java

@@ -77,7 +77,11 @@ void testThrowOnEmptyFilterConfig() {
 
   @Test
   void testDpiMaskingConfig() {
-    var maskingConfig = DpiMasking.anonymization().withEntities(DPIEntities.ADDRESS);
+    var maskingConfig =
+        DpiMasking.anonymization()
+            .withEntities(DPIEntities.ADDRESS)
+            .withMaskGroundingInput(true)
+            .withAllowList(List.of("Alice"));
     var config =
         new OrchestrationModuleConfig().withLlmConfig(GPT_4O).withMaskingConfig(maskingConfig);
 
@@ -87,6 +91,8 @@ void testDpiMaskingConfig() {
     assertThat(dpiConfig.getMethod()).isEqualTo(DPIConfig.MethodEnum.ANONYMIZATION);
     assertThat(dpiConfig.getEntities()).hasSize(1);
     assertThat(dpiConfig.getEntities().get(0).getType()).isEqualTo(DPIEntities.ADDRESS);
+    assertThat(dpiConfig.getMaskGroundingInput().isEnabled()).isEqualTo(true);
+    assertThat(dpiConfig.getAllowlist()).containsExactly("Alice");
 
     var configModified = config.withMaskingConfig(maskingConfig);
     assertThat(configModified.getMaskingConfig()).isNotNull();

orchestration/src/test/java/com/sap/ai/sdk/orchestration/OrchestrationUnitTest.java

@@ -170,7 +170,13 @@ void testGrounding() throws IOException {
         GroundingModuleConfig.create()
             .type(GroundingModuleConfig.TypeEnum.DOCUMENT_GROUNDING_SERVICE)
             .config(groundingConfigConfig);
-    final var configWithGrounding = config.withGroundingConfig(groundingConfig);
+    val maskingConfig = // optional masking configuration
+        DpiMasking.anonymization()
+            .withEntities(DPIEntities.SENSITIVE_DATA)
+            .withMaskGroundingInput(true)
+            .withAllowList(List.of("SAP", "Joule"));
+    final var configWithGrounding =
+        config.withGroundingConfig(groundingConfig).withMaskingConfig(maskingConfig);
 
     final Map<String, String> inputParams =
         Map.of("query", "String used for similarity search in database");
@@ -200,6 +206,18 @@ void testGrounding() throws IOException {
             "First chunk```Second chunk```Last found chunk");
     assertThat(groundingModule.getData()).isEqualTo(groundingData);
 
+    var inputMasking = response.getOriginalResponse().getModuleResults().getInputMasking();
+    assertThat(inputMasking.getMessage())
+        .isEqualTo("Input to LLM and Grounding is masked successfully.");
+    Object data = inputMasking.getData();
+    assertThat(data)
+        .isEqualTo(
+            Map.of(
+                "masked_template",
+                "[{\"role\": \"user\", \"content\": \"Context message with embedded grounding results. First chunk```MASKED_SENSITIVE_DATA```Last found chunk\"}]",
+                "masked_grounding_input", // maskGroundingInput: true will make this field present
+                "[\"What does Joule do?\"]"));
+
     try (var requestInputStream = fileLoader.apply("groundingRequest.json")) {
       final String request = new String(requestInputStream.readAllBytes());
       verify(postRequestedFor(urlPathEqualTo("/completion")).withRequestBody(equalToJson(request)));

orchestration/src/test/resources/__files/groundingResponse.json

@@ -14,6 +14,13 @@
         "content": "Context message with embedded grounding results. First chunk```Second chunk```Last found chunk"
       }
     ],
+    "input_masking": {
+      "message": "Input to LLM and Grounding is masked successfully.",
+      "data": {
+        "masked_template": "[{\"role\": \"user\", \"content\": \"Context message with embedded grounding results. First chunk```MASKED_SENSITIVE_DATA```Last found chunk\"}]",
+        "masked_grounding_input": "[\"What does Joule do?\"]"
+      }
+    },
     "llm": {
       "id": "chatcmpl-Apz5s3CMf99jkOnxvPshH1rGLwvvU",
       "object": "chat.completion",
@@ -59,4 +66,4 @@
       "total_tokens": 380
     }
   }
-}
+}

orchestration/src/test/resources/groundingRequest.json

@@ -21,6 +21,19 @@
         "defaults" : { },
         "tools" : [ ]
       },
+      "masking_module_config" : {
+        "masking_providers" : [ {
+          "type" : "sap_data_privacy_integration",
+          "method" : "anonymization",
+          "entities" : [ {
+            "type" : "profile-sensitive-data"
+          } ],
+          "allowlist" : [ "SAP", "Joule" ],
+          "mask_grounding_input" : {
+            "enabled" : true
+          }
+        } ]
+      },
       "grounding_module_config" : {
         "type" : "document_grounding_service",
         "config" : {
@@ -53,4 +66,4 @@
     "query" : "String used for similarity search in database"
   },
   "messages_history" : [ ]
-}
+}

orchestration/src/test/resources/maskingRequest.json

@@ -33,7 +33,10 @@
                 "type": "profile-phone"
               }
             ],
-            "allowlist" : [ ]
+            "allowlist" : [ ],
+            "mask_grounding_input" : {
+              "enabled" : false
+            }
           }
         ]
       }

sample-code/spring-app/src/main/java/com/sap/ai/sdk/app/controllers/OrchestrationController.java

@@ -180,11 +180,12 @@ Object maskingPseudonymization(
     return response.getContent();
   }
 
-  @GetMapping("/grounding")
+  @GetMapping("/grounding/{maskGroundingInput}")
   @Nonnull
   Object grounding(
-      @Nullable @RequestParam(value = "format", required = false) final String format) {
-    final var response = service.grounding("What does Joule do?");
+      @Nullable @RequestParam(value = "format", required = false) final String format,
+      @PathVariable("maskGroundingInput") final boolean maskGroundingInput) {
+    final var response = service.grounding("What does Joule do?", maskGroundingInput);
     if ("json".equals(format)) {
       return response;
     }

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

@@ -316,10 +316,12 @@ public OrchestrationChatResponse maskingPseudonymization(@Nonnull final DPIEntit
    * @link <a href="https://help.sap.com/docs/sap-ai-core/sap-ai-core-service-guide/grounding">SAP
    *     AI Core: Orchestration - Grounding</a>
    * @param userMessage the user message to provide grounding for
+   * @param maskGroundingInput whether to mask the request sent to the Grounding Service
    * @return the assistant response object
    */
   @Nonnull
-  public OrchestrationChatResponse grounding(@Nonnull final String userMessage) {
+  public OrchestrationChatResponse grounding(
+      @Nonnull final String userMessage, final boolean maskGroundingInput) {
     // optional filter for collections
     val documentMetadata =
         SearchDocumentKeyValueListPair.create()
@@ -335,7 +337,13 @@ public OrchestrationChatResponse grounding(@Nonnull final String userMessage) {
 
     val groundingConfig = Grounding.create().filters(databaseFilter);
     val prompt = groundingConfig.createGroundingPrompt(userMessage);
-    val configWithGrounding = config.withGrounding(groundingConfig);
+    val maskingConfig = // optional masking configuration
+        DpiMasking.anonymization()
+            .withEntities(DPIEntities.SENSITIVE_DATA)
+            .withMaskGroundingInput(maskGroundingInput)
+            .withAllowList(List.of("SAP", "Joule"));
+    val configWithGrounding =
+        config.withGrounding(groundingConfig).withMaskingConfig(maskingConfig);
 
     return client.chatCompletion(prompt, configWithGrounding);
   }