feat: Spring AI Orchestration streaming (#292)

* Added stream to Spring AI Orchestration

* Fixed

* Formatting

* Added e2e test + docs

* Removed unused field

* Update orchestration/src/main/java/com/sap/ai/sdk/orchestration/spring/OrchestrationSpringChatDelta.java

Co-authored-by: Alexander Dümont <[email protected]>

* Remove comments

* Fixed mistake

* Fixed mistake

---------

Co-authored-by: SAP Cloud SDK Bot <[email protected]>
Co-authored-by: Alexander Dümont <[email protected]>

Author: 3 people

docs/guides/SPRING_AI_INTEGRATION.md

@@ -50,8 +50,7 @@ Prompt prompt = new Prompt("What is the capital of France?", opts);
 ChatResponse response = client.call(prompt);
 ```
 
-Please
-find [an example in our Spring Boot application](../../sample-code/spring-app/src/main/java/com/sap/ai/sdk/app/controllers/SpringAiOrchestrationController.java).
+Please find [an example in our Spring Boot application](../../sample-code/spring-app/src/main/java/com/sap/ai/sdk/app/controllers/SpringAiOrchestrationController.java).
 
 ## Orchestration Masking
 
@@ -76,3 +75,27 @@ ChatResponse response = client.call(prompt);
 
 Please
 find [an example in our Spring Boot application](../../sample-code/spring-app/src/main/java/com/sap/ai/sdk/app/controllers/SpringAiOrchestrationController.java).
+
+## Stream chat completion
+
+It's possible to pass a stream of chat completion delta elements, e.g. from the application backend
+to the frontend in real-time.
+
+```java
+ChatModel client = new OrchestrationChatModel();
+OrchestrationModuleConfig config = new OrchestrationModuleConfig().withLlmConfig(GPT_35_TURBO);
+OrchestrationChatOptions opts = new OrchestrationChatOptions(config);
+
+Prompt prompt =
+    new Prompt(
+        "Can you give me the first 100 numbers of the Fibonacci sequence?", opts);
+Flux<ChatResponse> flux = client.stream(prompt);
+
+// also possible to keep only the chat completion text
+Flux<String> responseFlux = 
+    flux.map(chatResponse -> chatResponse.getResult().getOutput().getContent());
+```
+
+_Note: A Spring endpoint can return `Flux` instead of `ResponseEntity`._
+
+Please find [an example in our Spring Boot application](../../sample-code/spring-app/src/main/java/com/sap/ai/sdk/app/controllers/SpringAiOrchestrationController.java).

orchestration/pom.xml

@@ -59,6 +59,11 @@
       <artifactId>spring-ai-core</artifactId>
       <optional>true</optional>
     </dependency>
+    <dependency>
+      <groupId>io.projectreactor</groupId>
+      <artifactId>reactor-core</artifactId>
+      <optional>true</optional>
+    </dependency>
     <dependency>
       <groupId>org.apache.httpcomponents.core5</groupId>
       <artifactId>httpcore5</artifactId>

orchestration/src/main/java/com/sap/ai/sdk/orchestration/spring/OrchestrationChatModel.java

@@ -1,14 +1,19 @@
 package com.sap.ai.sdk.orchestration.spring;
 
+import static com.sap.ai.sdk.orchestration.OrchestrationClient.toCompletionPostRequest;
+
 import com.google.common.annotations.Beta;
 import com.sap.ai.sdk.orchestration.AssistantMessage;
+import com.sap.ai.sdk.orchestration.OrchestrationChatCompletionDelta;
 import com.sap.ai.sdk.orchestration.OrchestrationClient;
+import com.sap.ai.sdk.orchestration.OrchestrationClientException;
 import com.sap.ai.sdk.orchestration.OrchestrationPrompt;
 import com.sap.ai.sdk.orchestration.SystemMessage;
 import com.sap.ai.sdk.orchestration.UserMessage;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
+import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -17,6 +22,7 @@
 import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.chat.model.ChatResponse;
 import org.springframework.ai.chat.prompt.Prompt;
+import reactor.core.publisher.Flux;
 
 /**
  * Spring AI integration for the orchestration service.
@@ -52,6 +58,47 @@ public ChatResponse call(@Nonnull final Prompt prompt) {
         "Please add OrchestrationChatOptions to the Prompt: new Prompt(\"message\", new OrchestrationChatOptions(config))");
   }
 
+  @Override
+  @Nonnull
+  public Flux<ChatResponse> stream(@Nonnull final Prompt prompt) {
+
+    if (prompt.getOptions() instanceof OrchestrationChatOptions options) {
+
+      val orchestrationPrompt = toOrchestrationPrompt(prompt);
+      val request = toCompletionPostRequest(orchestrationPrompt, options.getConfig());
+      val stream = client.streamChatCompletionDeltas(request);
+
+      final Flux<OrchestrationChatCompletionDelta> flux =
+          Flux.generate(
+              stream::iterator,
+              (iterator, sink) -> {
+                if (iterator.hasNext()) {
+                  sink.next(iterator.next());
+                } else {
+                  sink.complete();
+                }
+                return iterator;
+              });
+      return flux.map(
+          delta -> {
+            throwOnContentFilter(stream, delta);
+            return new OrchestrationSpringChatDelta(delta);
+          });
+    }
+    throw new IllegalArgumentException(
+        "Please add OrchestrationChatOptions to the Prompt: new Prompt(\"message\", new OrchestrationChatOptions(config))");
+  }
+
+  private static void throwOnContentFilter(
+      @Nonnull final Stream<OrchestrationChatCompletionDelta> stream,
+      @Nonnull final OrchestrationChatCompletionDelta delta) {
+    final String finishReason = delta.getFinishReason();
+    if (finishReason != null && finishReason.equals("content_filter")) {
+      stream.close();
+      throw new OrchestrationClientException("Content filter filtered the output.");
+    }
+  }
+
   @Nonnull
   private OrchestrationPrompt toOrchestrationPrompt(@Nonnull final Prompt prompt) {
     val messages = toOrchestrationMessages(prompt.getInstructions());

orchestration/src/main/java/com/sap/ai/sdk/orchestration/spring/OrchestrationSpringChatDelta.java

@@ -0,0 +1,83 @@
+package com.sap.ai.sdk.orchestration.spring;
+
+import com.google.common.annotations.Beta;
+import com.sap.ai.sdk.orchestration.OrchestrationChatCompletionDelta;
+import com.sap.ai.sdk.orchestration.model.LLMChoice;
+import com.sap.ai.sdk.orchestration.model.LLMModuleResultSynchronous;
+import com.sap.ai.sdk.orchestration.model.TokenUsage;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import lombok.EqualsAndHashCode;
+import lombok.Value;
+import lombok.val;
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.metadata.DefaultUsage;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.model.Generation;
+
+/**
+ * Response from the orchestration service in a Spring AI {@link ChatResponse}.
+ *
+ * @since 1.2.0
+ */
+@Beta
+@Value
+@EqualsAndHashCode(callSuper = true)
+public class OrchestrationSpringChatDelta extends ChatResponse {
+
+  OrchestrationSpringChatDelta(@Nonnull final OrchestrationChatCompletionDelta delta) {
+    super(
+        toGenerations((LLMModuleResultSynchronous) delta.getOrchestrationResult()),
+        toChatResponseMetadata((LLMModuleResultSynchronous) delta.getOrchestrationResult()));
+  }
+
+  @Nonnull
+  static List<Generation> toGenerations(@Nonnull final LLMModuleResultSynchronous result) {
+    return result.getChoices().stream().map(OrchestrationSpringChatDelta::toGeneration).toList();
+  }
+
+  @Nonnull
+  static Generation toGeneration(@Nonnull final LLMChoice choice) {
+    val metadata = ChatGenerationMetadata.builder().finishReason(choice.getFinishReason());
+    metadata.metadata("index", choice.getIndex());
+    if (!choice.getLogprobs().isEmpty()) {
+      metadata.metadata("logprobs", choice.getLogprobs());
+    }
+    return new Generation(new AssistantMessage(getContent(choice)), metadata.build());
+  }
+
+  @Nonnull
+  private static String getContent(@Nonnull final LLMChoice choice) {
+    return choice.getCustomField("delta") instanceof Map<?, ?> delta
+            && delta.get("content") instanceof String content
+        ? content
+        : "";
+  }
+
+  @Nonnull
+  static ChatResponseMetadata toChatResponseMetadata(
+      @Nonnull final LLMModuleResultSynchronous orchestrationResult) {
+    val metadataBuilder = ChatResponseMetadata.builder();
+
+    metadataBuilder
+        .id(orchestrationResult.getId())
+        .model(orchestrationResult.getModel())
+        .keyValue("object", orchestrationResult.getObject())
+        .keyValue("created", orchestrationResult.getCreated());
+    if (orchestrationResult.getUsage() != null) {
+      metadataBuilder.usage(toDefaultUsage(orchestrationResult.getUsage()));
+    }
+    return metadataBuilder.build();
+  }
+
+  @Nonnull
+  private static DefaultUsage toDefaultUsage(@Nonnull final TokenUsage usage) {
+    return new DefaultUsage(
+        usage.getPromptTokens().longValue(),
+        usage.getCompletionTokens().longValue(),
+        usage.getTotalTokens().longValue());
+  }
+}

orchestration/src/main/java/com/sap/ai/sdk/orchestration/spring/OrchestrationSpringChatResponse.java

@@ -5,14 +5,13 @@
 import com.sap.ai.sdk.orchestration.model.LLMChoice;
 import com.sap.ai.sdk.orchestration.model.LLMModuleResultSynchronous;
 import com.sap.ai.sdk.orchestration.model.TokenUsage;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import javax.annotation.Nonnull;
 import lombok.EqualsAndHashCode;
 import lombok.Value;
 import lombok.val;
 import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
 import org.springframework.ai.chat.metadata.ChatResponseMetadata;
 import org.springframework.ai.chat.metadata.DefaultUsage;
 import org.springframework.ai.chat.model.ChatResponse;
@@ -43,21 +42,18 @@ public class OrchestrationSpringChatResponse extends ChatResponse {
 
   @Nonnull
   static List<Generation> toGenerations(@Nonnull final LLMModuleResultSynchronous result) {
-    return result.getChoices().stream()
-        .map(OrchestrationSpringChatResponse::toAssistantMessage)
-        .map(Generation::new)
-        .toList();
+    return result.getChoices().stream().map(OrchestrationSpringChatResponse::toGeneration).toList();
   }
 
   @Nonnull
-  static AssistantMessage toAssistantMessage(@Nonnull final LLMChoice choice) {
-    final Map<String, Object> metadata = new HashMap<>();
-    metadata.put("finish_reason", choice.getFinishReason());
-    metadata.put("index", choice.getIndex());
+  static Generation toGeneration(@Nonnull final LLMChoice choice) {
+    val metadata = ChatGenerationMetadata.builder().finishReason(choice.getFinishReason());
+    metadata.metadata("index", choice.getIndex());
     if (!choice.getLogprobs().isEmpty()) {
-      metadata.put("logprobs", choice.getLogprobs());
+      metadata.metadata("logprobs", choice.getLogprobs());
     }
-    return new AssistantMessage(choice.getMessage().getContent(), metadata);
+    val message = new AssistantMessage(choice.getMessage().getContent());
+    return new Generation(message, metadata.build());
   }
 
   @Nonnull

orchestration/src/test/java/com/sap/ai/sdk/orchestration/spring/OrchestrationChatDeltaTest.java

@@ -0,0 +1,63 @@
+package com.sap.ai.sdk.orchestration.spring;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.sap.ai.sdk.orchestration.model.LLMChoice;
+import com.sap.ai.sdk.orchestration.model.LLMModuleResultSynchronous;
+import com.sap.ai.sdk.orchestration.model.ResponseChatMessage;
+import com.sap.ai.sdk.orchestration.model.TokenUsage;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.metadata.EmptyUsage;
+import org.springframework.ai.chat.model.Generation;
+
+class OrchestrationChatDeltaTest {
+
+  @Test
+  void testToGeneration() {
+    var choice =
+        LLMChoice.create()
+            .index(0)
+            .message(ResponseChatMessage.create().role("wrong").content("wrong"))
+            .finishReason("stop");
+    // this will be fixed once the spec is fixed
+    choice.setCustomField("delta", Map.of("content", "Hello, world!"));
+
+    Generation generation = OrchestrationSpringChatDelta.toGeneration(choice);
+
+    assertThat(generation.getOutput().getContent()).isEqualTo("Hello, world!");
+    assertThat(generation.getMetadata().getFinishReason()).isEqualTo("stop");
+    assertThat(generation.getMetadata().<Integer>get("index")).isEqualTo(0);
+  }
+
+  @Test
+  void testToChatResponseMetadata() {
+    var moduleResult =
+        LLMModuleResultSynchronous.create()
+            .id("test-id")
+            ._object("test-object")
+            .created(123456789)
+            .model("test-model")
+            .choices(List.of())
+            .usage(TokenUsage.create().completionTokens(20).promptTokens(10).totalTokens(30));
+
+    var metadata = OrchestrationSpringChatDelta.toChatResponseMetadata(moduleResult);
+
+    assertThat(metadata.getId()).isEqualTo("test-id");
+    assertThat(metadata.getModel()).isEqualTo("test-model");
+    assertThat(metadata.<String>get("object")).isEqualTo("test-object");
+    assertThat(metadata.<Integer>get("created")).isEqualTo(123456789);
+
+    var usage = metadata.getUsage();
+
+    assertThat(usage.getPromptTokens()).isEqualTo(10L);
+    assertThat(usage.getGenerationTokens()).isEqualTo(20L);
+    assertThat(usage.getTotalTokens()).isEqualTo(30L);
+
+    // delta without token usage
+    moduleResult.usage(null);
+    metadata = OrchestrationSpringChatDelta.toChatResponseMetadata(moduleResult);
+    assertThat(metadata.getUsage()).isInstanceOf(EmptyUsage.class);
+  }
+}