Add Response Convenience (#173)

* Added Response Convenience

* Applied review comments

* Use Lombok's Value

* Renaming and adding tests.

* Add and test getTokenUsage

* Add and test getAllMessages

* Apply requested changes

* Apply requested changes

* Apply requested changes

---------

Co-authored-by: Jonas Israel <[email protected]>
Co-authored-by: Matthias Kuhr <[email protected]>

Author: Jonas-Isr

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

@@ -2,8 +2,14 @@
 
 import static lombok.AccessLevel.PACKAGE;
 
+import com.sap.ai.sdk.orchestration.client.model.ChatMessage;
 import com.sap.ai.sdk.orchestration.client.model.CompletionPostResponse;
+import com.sap.ai.sdk.orchestration.client.model.LLMChoice;
 import com.sap.ai.sdk.orchestration.client.model.LLMModuleResultSynchronous;
+import com.sap.ai.sdk.orchestration.client.model.TokenUsage;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
 import javax.annotation.Nonnull;
 import lombok.RequiredArgsConstructor;
 import lombok.Value;
@@ -24,18 +30,47 @@ public class OrchestrationChatResponse {
    */
   @Nonnull
   public String getContent() throws OrchestrationClientException {
-    final var choices =
-        ((LLMModuleResultSynchronous) originalResponse.getOrchestrationResult()).getChoices();
-
-    if (choices.isEmpty()) {
-      return "";
-    }
-
-    final var choice = choices.get(0);
+    final var choice = getCurrentChoice();
 
     if ("content_filter".equals(choice.getFinishReason())) {
       throw new OrchestrationClientException("Content filter filtered the output.");
     }
     return choice.getMessage().getContent();
   }
+
+  /**
+   * Get the token usage.
+   *
+   * @return The token usage.
+   */
+  @Nonnull
+  public TokenUsage getTokenUsage() {
+    return ((LLMModuleResultSynchronous) originalResponse.getOrchestrationResult()).getUsage();
+  }
+
+  /**
+   * Get all messages. This can be used for subsequent prompts as a message history.
+   *
+   * @return A list of all messages.
+   */
+  @Nonnull
+  public List<ChatMessage> getAllMessages() {
+    final var items = Objects.requireNonNull(originalResponse.getModuleResults().getTemplating());
+    final var messages = new ArrayList<>(items);
+    messages.add(getCurrentChoice().getMessage());
+    return messages;
+  }
+
+  /**
+   * Get current choice.
+   *
+   * @return The current choice.
+   */
+  @Nonnull
+  private LLMChoice getCurrentChoice() {
+    //    We expect choices to be defined and never empty.
+    return ((LLMModuleResultSynchronous) originalResponse.getOrchestrationResult())
+        .getChoices()
+        .get(0);
+  }
 }

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

@@ -129,9 +129,9 @@ void testTemplating() throws IOException {
 
     final var response = result.getOriginalResponse();
     assertThat(response.getRequestId()).isEqualTo("26ea36b5-c196-4806-a9a6-a686f0c6ad91");
-    assertThat(response.getModuleResults().getTemplating().get(0).getContent())
+    assertThat(result.getAllMessages().get(0).getContent())
         .isEqualTo("Reply with 'Orchestration Service is working!' in German");
-    assertThat(response.getModuleResults().getTemplating().get(0).getRole()).isEqualTo("user");
+    assertThat(result.getAllMessages().get(0).getRole()).isEqualTo("user");
     var llm = (LLMModuleResultSynchronous) response.getModuleResults().getLlm();
     assertThat(llm).isNotNull();
     assertThat(llm.getId()).isEqualTo("chatcmpl-9lzPV4kLrXjFckOp2yY454wksWBoj");
@@ -144,7 +144,7 @@ void testTemplating() throws IOException {
         .isEqualTo("Orchestration Service funktioniert!");
     assertThat(choices.get(0).getMessage().getRole()).isEqualTo("assistant");
     assertThat(choices.get(0).getFinishReason()).isEqualTo("stop");
-    var usage = llm.getUsage();
+    var usage = result.getTokenUsage();
     assertThat(usage.getCompletionTokens()).isEqualTo(7);
     assertThat(usage.getPromptTokens()).isEqualTo(19);
     assertThat(usage.getTotalTokens()).isEqualTo(26);
@@ -159,7 +159,7 @@ void testTemplating() throws IOException {
         .isEqualTo("Orchestration Service funktioniert!");
     assertThat(choices.get(0).getMessage().getRole()).isEqualTo("assistant");
     assertThat(choices.get(0).getFinishReason()).isEqualTo("stop");
-    usage = orchestrationResult.getUsage();
+    usage = result.getTokenUsage();
     assertThat(usage.getCompletionTokens()).isEqualTo(7);
     assertThat(usage.getPromptTokens()).isEqualTo(19);
     assertThat(usage.getTotalTokens()).isEqualTo(26);
@@ -380,17 +380,4 @@ void testErrorHandling() {
 
     softly.assertAll();
   }
-
-  @Test
-  void testEmptyChoicesResponse() {
-    stubFor(
-        post(urlPathEqualTo("/v2/inference/deployments/abcdef0123456789/completion"))
-            .willReturn(
-                aResponse()
-                    .withBodyFile("emptyChoicesResponse.json")
-                    .withHeader("Content-Type", "application/json")));
-    final var result = client.chatCompletion(prompt, config);
-
-    assertThat(result.getContent()).isEmpty();
-  }
 }

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

@@ -37,14 +37,14 @@ void testTemplate() {
     assertThat(controller.config.getLlmConfig()).isNotNull();
     final var modelName = controller.config.getLlmConfig().getModelName();
 
-    final var response = controller.template();
-    final var result = response.getOriginalResponse();
+    final var result = controller.template();
+    final var response = result.getOriginalResponse();
 
-    assertThat(result.getRequestId()).isNotEmpty();
-    assertThat(result.getModuleResults().getTemplating().get(0).getContent())
+    assertThat(response.getRequestId()).isNotEmpty();
+    assertThat(result.getAllMessages().get(0).getContent())
         .isEqualTo("Reply with 'Orchestration Service is working!' in German");
-    assertThat(result.getModuleResults().getTemplating().get(0).getRole()).isEqualTo("user");
-    var llm = (LLMModuleResultSynchronous) result.getModuleResults().getLlm();
+    assertThat(result.getAllMessages().get(0).getRole()).isEqualTo("user");
+    var llm = (LLMModuleResultSynchronous) response.getModuleResults().getLlm();
     assertThat(llm.getId()).isNotEmpty();
     assertThat(llm.getObject()).isEqualTo("chat.completion");
     assertThat(llm.getCreated()).isGreaterThan(1);
@@ -54,12 +54,12 @@ void testTemplate() {
     assertThat(choices.get(0).getMessage().getContent()).isNotEmpty();
     assertThat(choices.get(0).getMessage().getRole()).isEqualTo("assistant");
     assertThat(choices.get(0).getFinishReason()).isEqualTo("stop");
-    var usage = llm.getUsage();
+    var usage = result.getTokenUsage();
     assertThat(usage.getCompletionTokens()).isGreaterThan(1);
     assertThat(usage.getPromptTokens()).isGreaterThan(1);
     assertThat(usage.getTotalTokens()).isGreaterThan(1);
 
-    var orchestrationResult = ((LLMModuleResultSynchronous) result.getOrchestrationResult());
+    var orchestrationResult = ((LLMModuleResultSynchronous) response.getOrchestrationResult());
     assertThat(orchestrationResult.getObject()).isEqualTo("chat.completion");
     assertThat(orchestrationResult.getCreated()).isGreaterThan(1);
     assertThat(orchestrationResult.getModel()).isEqualTo(modelName);
@@ -68,7 +68,7 @@ void testTemplate() {
     assertThat(choices.get(0).getMessage().getContent()).isNotEmpty();
     assertThat(choices.get(0).getMessage().getRole()).isEqualTo("assistant");
     assertThat(choices.get(0).getFinishReason()).isEqualTo("stop");
-    usage = orchestrationResult.getUsage();
+    usage = result.getTokenUsage();
     assertThat(usage.getCompletionTokens()).isGreaterThan(1);
     assertThat(usage.getPromptTokens()).isGreaterThan(1);
     assertThat(usage.getTotalTokens()).isGreaterThan(1);