Message classes extending generated model

Author: rpanackal

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

@@ -3,12 +3,9 @@
 import com.sap.ai.sdk.orchestration.client.model.ChatMessage;
 import javax.annotation.Nonnull;
 
-public record AssistantMessage(@Nonnull String content) implements Message {
+public class AssistantMessage extends ChatMessage {
 
-  public static final String ROLE = "assistant";
-
-  @Nonnull
-  public ChatMessage toChatMessage() {
-    return new ChatMessage().role(ROLE).content(content);
+  public AssistantMessage(@Nonnull String content) {
+    this.role("assistant").content(content);
   }
 }

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

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

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

@@ -3,12 +3,9 @@
 import com.sap.ai.sdk.orchestration.client.model.ChatMessage;
 import javax.annotation.Nonnull;
 
-public record SystemMessage(@Nonnull String content) implements Message {
+public class SystemMessage extends ChatMessage {
 
-  public static final String ROLE = "system";
-
-  @Nonnull
-  public ChatMessage toChatMessage() {
-    return new ChatMessage().role(ROLE).content(content);
+  public SystemMessage(@Nonnull String content) {
+    this.role("system").content(content);
   }
 }

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

@@ -3,12 +3,9 @@
 import com.sap.ai.sdk.orchestration.client.model.ChatMessage;
 import javax.annotation.Nonnull;
 
-public record UserMessage(@Nonnull String content) implements Message {
+public class UserMessage extends ChatMessage {
 
-  public static final String ROLE = "user";
-
-  @Nonnull
-  public ChatMessage toChatMessage() {
-    return new ChatMessage().role(ROLE).content(content);
+  public UserMessage(@Nonnull String content) {
+    this.role("user").content(content);
   }
 }

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

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

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

@@ -29,6 +29,7 @@
 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;
@@ -256,7 +257,7 @@ void messagesHistory() throws IOException {
                     .withBodyFile("templatingResponse.json")
                     .withHeader("Content-Type", "application/json")));
 
-    final List<Message> messagesHistory =
+    final List<ChatMessage> messagesHistory =
         List.of(
             new UserMessage("What is the capital of France?"),
             new AssistantMessage("The capital of France is Paris."));

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

@@ -6,7 +6,6 @@
 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;
@@ -73,7 +72,7 @@ public OrchestrationChatResponse template() {
   @GetMapping("/messagesHistory")
   @Nonnull
   public OrchestrationChatResponse messagesHistory() {
-    final List<Message> messagesHistory =
+    final List<ChatMessage> messagesHistory =
         List.of(
             new UserMessage("What is the capital of France?"),
             new AssistantMessage("The capital of France is Paris."));