fix: Unmasking Output Parsing (#165)

* fix: Unmasking Output Parsing

* Formatting

* Adjust test

* Adjust test

---------

Co-authored-by: SAP Cloud SDK Bot <[email protected]>

Author: MatKuhr

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

@@ -0,0 +1,25 @@
+package com.sap.ai.sdk.orchestration;
+
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.sap.ai.sdk.orchestration.client.model.LLMChoice;
+import com.sap.ai.sdk.orchestration.client.model.LLMModuleResultSynchronous;
+import lombok.AccessLevel;
+import lombok.NoArgsConstructor;
+
+@NoArgsConstructor(access = AccessLevel.PRIVATE)
+final class JacksonMixins {
+  /** Mixin to enforce a specific subtype to be deserialized always. */
+  @JsonTypeInfo(use = JsonTypeInfo.Id.NONE)
+  @JsonDeserialize(as = LLMModuleResultSynchronous.class)
+  interface LLMModuleResultMixIn {}
+
+  /** Mixin to enforce a specific subtype to be deserialized always. */
+  @JsonTypeInfo(use = JsonTypeInfo.Id.NONE)
+  @JsonDeserialize(as = LLMChoice.class)
+  interface ModuleResultsOutputUnmaskingInnerMixIn {}
+
+  /** Mixin to suppress @JsonTypeInfo for oneOf interfaces. */
+  @JsonTypeInfo(use = JsonTypeInfo.Id.NONE)
+  interface NoTypeInfoMixin {}
+}

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

@@ -45,11 +45,13 @@ public class OrchestrationClient {
             .visibility(PropertyAccessor.GETTER, JsonAutoDetect.Visibility.NONE)
             .visibility(PropertyAccessor.SETTER, JsonAutoDetect.Visibility.NONE)
             .serializationInclusion(JsonInclude.Include.NON_NULL)
-            .mixIn(LLMModuleResult.class, LLMModuleResultMixIn.class)
-            .mixIn(ModuleResultsOutputUnmaskingInner.class, NoTypeInfoMixin.class)
-            .mixIn(FilterConfig.class, NoTypeInfoMixin.class)
-            .mixIn(MaskingProviderConfig.class, NoTypeInfoMixin.class)
-            .mixIn(TemplatingModuleConfig.class, NoTypeInfoMixin.class)
+            .mixIn(LLMModuleResult.class, JacksonMixins.LLMModuleResultMixIn.class)
+            .mixIn(
+                ModuleResultsOutputUnmaskingInner.class,
+                JacksonMixins.ModuleResultsOutputUnmaskingInnerMixIn.class)
+            .mixIn(FilterConfig.class, JacksonMixins.NoTypeInfoMixin.class)
+            .mixIn(MaskingProviderConfig.class, JacksonMixins.NoTypeInfoMixin.class)
+            .mixIn(TemplatingModuleConfig.class, JacksonMixins.NoTypeInfoMixin.class)
             .build();
   }
 

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

@@ -298,7 +298,7 @@ void messagesHistory() throws IOException {
   }
 
   @Test
-  void maskingAnonymization() throws IOException {
+  void maskingPseudonymization() throws IOException {
     stubFor(
         post(urlPathEqualTo("/v2/inference/deployments/abcdef0123456789/completion"))
             .willReturn(
@@ -307,7 +307,7 @@ void maskingAnonymization() throws IOException {
                     .withHeader("Content-Type", "application/json")));
 
     final var maskingConfig =
-        createMaskingConfig(DPIConfig.MethodEnum.ANONYMIZATION, DPIEntities.PHONE);
+        createMaskingConfig(DPIConfig.MethodEnum.PSEUDONYMIZATION, DPIEntities.PHONE);
 
     final var result = client.chatCompletion(prompt, config.withMaskingConfig(maskingConfig));
 
@@ -316,9 +316,7 @@ void maskingAnonymization() throws IOException {
     assertThat(inputMasking.getMessage()).isEqualTo("Input to LLM is masked successfully.");
     assertThat(inputMasking.getData()).isNotNull();
     final var choices = ((LLMModuleResultSynchronous) result.getOrchestrationResult()).getChoices();
-    assertThat(choices.get(0).getMessage().getContent())
-        .isEqualTo(
-            "I'm sorry, I cannot provide information about specific individuals, including their nationality.");
+    assertThat(choices.get(0).getMessage().getContent()).contains("Hi Mallory");
 
     // verify that the request is sent correctly
     try (var requestInputStream = fileLoader.apply("maskingRequest.json")) {

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

@@ -1,74 +1,84 @@
 {
-  "request_id": "4c1b423d-f3a9-443e-9876-091888b0d585",
+  "request_id": "c252c73e-849f-4fa7-8b35-b16f3434a0da",
   "module_results": {
-    "grounding": null,
     "templating": [
+      {
+        "role": "system",
+        "content": "Please write an initial response to the below user feedback, stating that we are working on the feedback and will get back to them soon.\nPlease make sure to address the user in person and end with \"Best regards, the AI SDK team\".\n"
+      },
       {
         "role": "user",
-        "content": "What is the nationality of Patrick Morgan +49 (970) 333-3833"
+        "content": "Username: Mallory\nuserEmail: [email protected]\nDate: 2022-01-01\n\nI think the SDK is good, but could use some further enhancements.\nMy architect Alice and manager Bob pointed out that we need the grounding capabilities, which aren't supported yet.\n"
       }
     ],
     "input_masking": {
       "message": "Input to LLM is masked successfully.",
       "data": {
         "masked_template": [
+          {
+            "role": "system",
+            "content": "Please write an initial response to the below user feedback, stating that we are working on the feedback and will get back to them soon.\nPlease make sure to address the user in person and end with \"Best regards, the AI SDK team\".\n"
+          },
           {
             "role": "user",
-            "content": "What is the nationality of Patrick Morgan MASKED_PHONE_NUMBER"
+            "content": "Username: MASKED_PERSON_2\nuserEmail: MASKED_PERSON_3MASKED_EMAIL_1\nDate: 2022-01-01\n\nI think the SDK is good, but could use some further enhancements.\nMy architect MASKED_PERSON_1 and manager MASKED_PERSON_4 pointed out that we need the grounding capabilities, which aren't supported yet.\n"
           }
         ]
       }
     },
-    "input_filtering": null,
     "llm": {
+      "id": "chatcmpl-AUr7GVVoQbg52GRbADpkx4hQvDBiz",
       "object": "chat.completion",
-      "id": "chatcmpl-ADVDtCV54jp7BCKxeyWs3MWPFNx3N",
-      "created": 1727781649,
+      "created": 1731917382,
       "model": "gpt-35-turbo",
-      "system_fingerprint": "fp_e49e4201a9",
+      "system_fingerprint": "fp_808245b034",
       "choices": [
         {
           "index": 0,
           "message": {
             "role": "assistant",
-            "content": "I'm sorry, I cannot provide information about specific individuals, including their nationality."
-          },
-          "logprobs": {
+            "content": "Hi MASKED_PERSON_2,\n\nThank you for your feedback on the SDK. We appreciate your insights and are currently working on further enhancements, including the grounding capabilities that your architect and manager have highlighted. We will take your suggestions into consideration and aim to address them in our future updates.\n\nBest regards, the AI SDK team"
           },
           "finish_reason": "stop"
         }
       ],
       "usage": {
-        "completion_tokens": 16,
-        "prompt_tokens": 18,
-        "total_tokens": 34
+        "completion_tokens": 64,
+        "prompt_tokens": 134,
+        "total_tokens": 198
       }
     },
-    "output_filtering": null,
-    "output_unmasking": []
+    "output_unmasking": [
+      {
+        "index": 0,
+        "message": {
+          "role": "assistant",
+          "content": "Hi Mallory,\n\nThank you for your feedback on the SDK. We appreciate your insights and are currently working on further enhancements, including the grounding capabilities that your architect and manager have highlighted. We will take your suggestions into consideration and aim to address them in our future updates.\n\nBest regards, the AI SDK team"
+        },
+        "finish_reason": "stop"
+      }
+    ]
   },
   "orchestration_result": {
+    "id": "chatcmpl-AUr7GVVoQbg52GRbADpkx4hQvDBiz",
     "object": "chat.completion",
-    "id": "chatcmpl-ADVDtCV54jp7BCKxeyWs3MWPFNx3N",
-    "created": 1727781649,
+    "created": 1731917382,
     "model": "gpt-35-turbo",
-    "system_fingerprint": "fp_e49e4201a9",
+    "system_fingerprint": "fp_808245b034",
     "choices": [
       {
         "index": 0,
         "message": {
           "role": "assistant",
-          "content": "I'm sorry, I cannot provide information about specific individuals, including their nationality."
-        },
-        "logprobs": {
+          "content": "Hi Mallory,\n\nThank you for your feedback on the SDK. We appreciate your insights and are currently working on further enhancements, including the grounding capabilities that your architect and manager have highlighted. We will take your suggestions into consideration and aim to address them in our future updates.\n\nBest regards, the AI SDK team"
         },
         "finish_reason": "stop"
       }
     ],
     "usage": {
-      "completion_tokens": 16,
-      "prompt_tokens": 18,
-      "total_tokens": 34
+      "completion_tokens": 64,
+      "prompt_tokens": 134,
+      "total_tokens": 198
     }
   }
 }

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

@@ -23,7 +23,7 @@
         "masking_providers": [
           {
             "type": "sap_data_privacy_integration",
-            "method": "anonymization",
+            "method": "pseudonymization",
             "entities": [
               {
                 "type": "profile-phone"