Custom interface base message classes

Author: rpanackal

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

@@ -3,9 +3,23 @@
 import com.sap.ai.sdk.orchestration.client.model.ChatMessage;
 import javax.annotation.Nonnull;
 
-public class AssistantMessage extends ChatMessage {
+/**
+ * Represents a chat message as 'assistant' to the orchestration service.
+ *
+ * @param content
+ */
+public record AssistantMessage(@Nonnull String content) implements Message {
 
-  public AssistantMessage(@Nonnull String content) {
-    this.role("assistant").content(content);
+  /** The role of the assistant. */
+  public static final String ROLE = "assistant";
+
+  /**
+   * Converts the message to a serializable ChatMessage object.
+   *
+   * @return the corresponding {@code ChatMessage} object.
+   */
+  @Nonnull
+  public ChatMessage toChatMessage() {
+    return new ChatMessage().role(ROLE).content(content);
   }
 }

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

@@ -0,0 +1,16 @@
+package com.sap.ai.sdk.orchestration;
+
+import com.sap.ai.sdk.orchestration.client.model.ChatMessage;
+import javax.annotation.Nonnull;
+
+/** Interface representing convenience wrappers of chat message to the orchestration service. */
+public sealed interface Message permits UserMessage, AssistantMessage, SystemMessage {
+
+  /**
+   * Converts the message to a serializable ChatMessage object.
+   *
+   * @return the corresponding {@code ChatMessage} object.
+   */
+  @Nonnull
+  ChatMessage toChatMessage();
+}

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

@@ -32,7 +32,7 @@ public class OrchestrationPrompt {
    * @param message A user message.
    */
   public OrchestrationPrompt(@Nonnull final String message) {
-    messages.add(new UserMessage(message));
+    messages.add(new UserMessage(message).toChatMessage());
   }
 
   /**
@@ -41,10 +41,9 @@ public OrchestrationPrompt(@Nonnull final String message) {
    * @param message The first message.
    * @param messages Optionally, more messages.
    */
-  public OrchestrationPrompt(
-      @Nonnull final ChatMessage message, @Nonnull final ChatMessage... messages) {
-    this.messages.add(message);
-    this.messages.addAll(Arrays.asList(messages));
+  public OrchestrationPrompt(@Nonnull final Message message, @Nonnull final Message... messages) {
+    this.messages.add(message.toChatMessage());
+    this.messages.addAll(Arrays.stream(messages).map(Message::toChatMessage).toList());
   }
 
   /**
@@ -53,9 +52,9 @@ public OrchestrationPrompt(
    * @param inputParams The input parameters as entries of template variables and their contents.
    */
   public OrchestrationPrompt(
-      @Nonnull final Map<String, String> inputParams, @Nonnull final ChatMessage... messages) {
+      @Nonnull final Map<String, String> inputParams, @Nonnull final Message... messages) {
     this.templateParameters.putAll(inputParams);
-    this.messages.addAll(Arrays.asList(messages));
+    this.messages.addAll(Arrays.stream(messages).map(Message::toChatMessage).toList());
   }
 
   /**
@@ -64,9 +63,9 @@ public OrchestrationPrompt(
    * @param messagesHistory The chat history to add.
    */
   @Nonnull
-  public OrchestrationPrompt messageHistory(@Nonnull final List<ChatMessage> messagesHistory) {
+  public OrchestrationPrompt messageHistory(@Nonnull final List<Message> messagesHistory) {
     this.messagesHistory.clear();
-    this.messagesHistory.addAll(messagesHistory);
+    this.messagesHistory.addAll(messagesHistory.stream().map(Message::toChatMessage).toList());
     return this;
   }
 }

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

@@ -3,9 +3,23 @@
 import com.sap.ai.sdk.orchestration.client.model.ChatMessage;
 import javax.annotation.Nonnull;
 
-public class SystemMessage extends ChatMessage {
+/**
+ * Represents a chat message as 'system' to the orchestration service.
+ *
+ * @param content
+ */
+public record SystemMessage(@Nonnull String content) implements Message {
 
-  public SystemMessage(@Nonnull String content) {
-    this.role("system").content(content);
+  /** The role of the assistant. */
+  public static final String ROLE = "system";
+
+  /**
+   * Converts the message to a serializable ChatMessage object.
+   *
+   * @return the corresponding {@code ChatMessage} object.
+   */
+  @Nonnull
+  public ChatMessage toChatMessage() {
+    return new ChatMessage().role(ROLE).content(content);
   }
 }

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

@@ -3,9 +3,23 @@
 import com.sap.ai.sdk.orchestration.client.model.ChatMessage;
 import javax.annotation.Nonnull;
 
-public class UserMessage extends ChatMessage {
+/**
+ * Represents a chat message as 'user' to the orchestration service.
+ *
+ * @param content
+ */
+public record UserMessage(@Nonnull String content) implements Message {
 
-  public UserMessage(@Nonnull String content) {
-    this.role("user").content(content);
+  /** The role of the assistant. */
+  public static final String ROLE = "user";
+
+  /**
+   * Converts the message to a serializable ChatMessage object.
+   *
+   * @return the corresponding {@code ChatMessage} object.
+   */
+  @Nonnull
+  public ChatMessage toChatMessage() {
+    return new ChatMessage().role(ROLE).content(content);
   }
 }

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

@@ -33,7 +33,9 @@ void testEmptyTemplateConfig() {
     var systemMessage = new SystemMessage("foo");
     var userMessage = new UserMessage("Hello");
 
-    var expected = new Template().template(List.of(systemMessage, userMessage));
+    var expected =
+        new Template()
+            .template(List.of(systemMessage.toChatMessage(), userMessage.toChatMessage()));
 
     var prompt = new OrchestrationPrompt(systemMessage, userMessage);
     var actual =
@@ -54,10 +56,18 @@ void testMergingTemplateConfig() {
     var userMessage = new UserMessage("Hello ");
     var userMessage2 = new UserMessage("World");
 
-    var expected = new Template().template(List.of(systemMessage, userMessage, userMessage2));
+    var expected =
+        new Template()
+            .template(
+                List.of(
+                    systemMessage.toChatMessage(),
+                    userMessage.toChatMessage(),
+                    userMessage2.toChatMessage()));
 
     var prompt = new OrchestrationPrompt(userMessage2);
-    var templateConfig = new Template().template(List.of(systemMessage, userMessage));
+    var templateConfig =
+        new Template()
+            .template(List.of(systemMessage.toChatMessage(), userMessage.toChatMessage()));
     var actual = ConfigToRequestTransformer.toTemplateModuleConfig(prompt, templateConfig);
 
     assertThat(actual).isEqualTo(expected);
@@ -72,6 +82,6 @@ void testMessagesHistory() {
         ConfigToRequestTransformer.toCompletionPostRequest(
             prompt, new OrchestrationModuleConfig().withLlmConfig(CUSTOM_GPT_35));
 
-    assertThat(actual.getMessagesHistory()).containsExactly(systemMessage);
+    assertThat(actual.getMessagesHistory()).containsExactly(systemMessage.toChatMessage());
   }
 }

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

@@ -29,7 +29,6 @@
 import com.github.tomakehurst.wiremock.junit5.WireMockTest;
 import com.github.tomakehurst.wiremock.stubbing.Scenario;
 import com.sap.ai.sdk.core.AiCoreService;
-import com.sap.ai.sdk.orchestration.client.model.ChatMessage;
 import com.sap.ai.sdk.orchestration.client.model.CompletionPostRequest;
 import com.sap.ai.sdk.orchestration.client.model.DPIEntities;
 import com.sap.ai.sdk.orchestration.client.model.GenericModuleResult;
@@ -257,7 +256,7 @@ void messagesHistory() throws IOException {
                     .withBodyFile("templatingResponse.json")
                     .withHeader("Content-Type", "application/json")));
 
-    final List<ChatMessage> messagesHistory =
+    final List<Message> messagesHistory =
         List.of(
             new UserMessage("What is the capital of France?"),
             new AssistantMessage("The capital of France is Paris."));

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

@@ -6,6 +6,7 @@
 import com.sap.ai.sdk.orchestration.AzureContentFilter;
 import com.sap.ai.sdk.orchestration.AzureFilterThreshold;
 import com.sap.ai.sdk.orchestration.DpiMasking;
+import com.sap.ai.sdk.orchestration.Message;
 import com.sap.ai.sdk.orchestration.OrchestrationChatResponse;
 import com.sap.ai.sdk.orchestration.OrchestrationClient;
 import com.sap.ai.sdk.orchestration.OrchestrationModuleConfig;
@@ -72,7 +73,7 @@ public OrchestrationChatResponse template() {
   @GetMapping("/messagesHistory")
   @Nonnull
   public OrchestrationChatResponse messagesHistory() {
-    final List<ChatMessage> messagesHistory =
+    final List<Message> messagesHistory =
         List.of(
             new UserMessage("What is the capital of France?"),
             new AssistantMessage("The capital of France is Paris."));