Fix type of Orchestration Response Messages

Author: MatKuhr

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

@@ -82,7 +82,7 @@ public String chatCompletion(@Nonnull final String userPrompt)
           response.originalResponseDto());
       throw new OrchestrationClientException("Output content filter triggered");
     }
-    return response.assistantMessage().getContent();
+    return response.assistantMessage().content();
   }
 
   /**

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

@@ -7,6 +7,7 @@
 import java.util.Arrays;
 import java.util.List;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import lombok.RequiredArgsConstructor;
 import lombok.val;
 
@@ -21,8 +22,8 @@
  *     result details.
  */
 public record OrchestrationResponse(
-    @Nonnull ChatMessage assistantMessage,
-    @Nonnull List<ChatMessage> allMessages,
+    @Nonnull AssistantMessage assistantMessage,
+    @Nonnull List<Message> allMessages,
     @Nonnull FinishReason finishReason,
     @Nonnull TokenUsage tokenUsage,
     @Nonnull CompletionPostResponse originalResponseDto) {
@@ -55,11 +56,37 @@ static FinishReason fromValue(@Nonnull final String value) {
   static OrchestrationResponse fromCompletionPostResponseDTO(
       @Nonnull final CompletionPostResponse response) {
     val choice = response.getOrchestrationResult().getChoices().get(0);
-    val message = choice.getMessage();
+    val message = new AssistantMessage(choice.getMessage().getContent());
     val finishReason = FinishReason.fromValue(choice.getFinishReason());
     val tokenUsage = response.getOrchestrationResult().getUsage();
-    val allMessages = new ArrayList<>(response.getModuleResults().getTemplating());
+    val allMessages = new ArrayList<Message>();
+    response.getModuleResults().getTemplating().stream()
+        .map(OrchestrationResponse::fromChatMessage)
+        .forEach(allMessages::add);
     allMessages.add(message);
     return new OrchestrationResponse(message, allMessages, finishReason, tokenUsage, response);
   }
+
+  @Nonnull
+  static Message fromChatMessage(@Nonnull final ChatMessage chatMessage) {
+    return switch (chatMessage.getRole()) {
+      case "system" -> new SystemMessage(chatMessage.getContent());
+      case "user" -> new UserMessage(chatMessage.getContent());
+      case "assistant" -> new AssistantMessage(chatMessage.getContent());
+      default ->
+          new Message() {
+            @Nonnull
+            @Override
+            public String type() {
+              return chatMessage.getRole();
+            }
+
+            @Nullable
+            @Override
+            public String content() {
+              return chatMessage.getContent();
+            }
+          };
+    };
+  }
 }

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

@@ -15,9 +15,9 @@
 class OrchestrationResponseTest {
   @Test
   void testFromCompletionPostResponseDTO() {
-    var message1 = mock(ChatMessage.class);
-    var message2 = mock(ChatMessage.class);
-    var message3 = mock(ChatMessage.class);
+    var message1 = ChatMessage.create().role("system").content("foo");
+    var message2 = ChatMessage.create().role("user").content("bar");
+    var message3 = ChatMessage.create().role("assistant").content("baz");
     var moduleResults = ModuleResults.create().templating(List.of(message1, message2));
 
     var orchestrationResult =
@@ -37,8 +37,10 @@ void testFromCompletionPostResponseDTO() {
 
     var result = OrchestrationResponse.fromCompletionPostResponseDTO(postResponse);
 
-    assertThat(result.assistantMessage()).isSameAs(message3);
-    assertThat(result.allMessages()).containsExactly(message1, message2, message3);
+    assertThat(result.assistantMessage()).isEqualTo(new AssistantMessage("baz"));
+    assertThat(result.allMessages())
+        .containsExactly(
+            new SystemMessage("foo"), new UserMessage("bar"), new AssistantMessage("baz"));
     assertThat(result.finishReason()).isEqualTo(OrchestrationResponse.FinishReason.STOP);
     assertThat(result.originalResponseDto()).isSameAs(postResponse);
   }

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

@@ -5,9 +5,11 @@
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import com.sap.ai.sdk.foundationmodels.openai.OpenAiModel;
+import com.sap.ai.sdk.orchestration.AssistantMessage;
 import com.sap.ai.sdk.orchestration.AzureContentFilter;
 import com.sap.ai.sdk.orchestration.OrchestrationClientException;
 import com.sap.ai.sdk.orchestration.OrchestrationResponse;
+import com.sap.ai.sdk.orchestration.UserMessage;
 import java.util.List;
 import java.util.Map;
 import org.assertj.core.api.InstanceOfAssertFactories;
@@ -27,8 +29,11 @@ void testCompletion() {
     final var response = controller.completion();
 
     assertThat(response.finishReason()).isEqualTo(OrchestrationResponse.FinishReason.STOP);
-    assertThat(response.assistantMessage().getContent()).isNotEmpty();
-    assertThat(response.assistantMessage().getRole()).isEqualTo("assistant");
+    assertThat(response.assistantMessage())
+        .isInstanceOf(AssistantMessage.class)
+        .extracting(AssistantMessage::content)
+        .asString()
+        .isNotEmpty();
     assertThat(response.tokenUsage().getPromptTokens()).isPositive();
     assertThat(response.tokenUsage().getCompletionTokens()).isPositive();
     assertThat(response.tokenUsage().getTotalTokens()).isPositive();
@@ -70,16 +75,16 @@ void testTemplate() {
     var result = controller.template();
 
     var templateResult = result.allMessages().get(0);
-    assertThat(templateResult.getContent())
+    assertThat(templateResult.content())
         .isEqualTo("Reply with 'The Orchestration Service is working!' in german");
-    assertThat(templateResult.getRole()).isEqualTo("user");
+    assertThat(templateResult).isInstanceOf(UserMessage.class);
   }
 
   @Test
   void testLenientContentFilter() {
     var result = controller.filter(AzureContentFilter.Sensitivity.LOW);
     assertThat(result.finishReason()).isEqualTo(OrchestrationResponse.FinishReason.STOP);
-    assertThat(result.assistantMessage().getContent()).isNotEmpty();
+    assertThat(result.assistantMessage().content()).isNotEmpty();
 
     var filterResult = result.originalResponseDto().getModuleResults().getInputFiltering();
     assertThat(filterResult.getMessage()).contains("passed");
@@ -98,7 +103,7 @@ void testStrictContentFilter() {
   void testMasking() {
     var result = controller.masking();
 
-    assertThat(result.assistantMessage().getContent()).contains("[email protected]");
+    assertThat(result.assistantMessage().content()).contains("[email protected]");
     var maskingResult = result.originalResponseDto().getModuleResults().getInputMasking();
     var data = (Map<String, Object>) maskingResult.getData();
     var maskedMessage = ((List<Map<String, Object>>) data.get("masked_template")).get(0);