Repository: LangStream/langstream (language: python)

Issue:
Switch to the OpenAI Async client
Now that the pipelines are fully async it is better to switch to the official async client for OpenAI
https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/openai/azure-ai-openai/src/main/java/com/azure/ai/openai/OpenAIAsyncClient.java


Discussion / hints:

Fix patch:
diff --git a/docker/build.sh b/docker/build.sh
index 13969472b..95cae924e 100755
--- a/docker/build.sh
+++ b/docker/build.sh
@@ -33,7 +33,7 @@ common_flags="-DskipTests -PskipPython -Dlicense.skip -Dspotless.skip"
 build_docker_image() {
   module=$1
   ./mvnw install -am -pl $module -T 1C $common_flags
-  ./mvnw package -Pdocker -pl $module -Ddocker.platforms="$(docker_platforms)" $common_flags
+  ./mvnw clean package -Pdocker -pl $module -Ddocker.platforms="$(docker_platforms)" $common_flags
   docker images | head -n 2
 }
 
diff --git a/examples/applications/openai-completions/pipeline.yaml b/examples/applications/openai-completions/pipeline.yaml
index 2af385648..db8e09877 100644
--- a/examples/applications/openai-completions/pipeline.yaml
+++ b/examples/applications/openai-completions/pipeline.yaml
@@ -47,8 +47,6 @@ pipeline:
       # with 1 chunk, then with 2 chunks....up to the min-chunks-per-message value
       # eventually we want to send bigger messages to reduce the overhead of each message on the topic
       min-chunks-per-message: 20
-      # you can turn off the streaming behavior by setting this to false
-      stream: true
       messages:
         - role: user
           content: "You are an helpful assistant. Below you can fine a question from the user. Please try to help them the best way you can.\n\n{{% value.question}}"
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/services/impl/OpenAICompletionService.java b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/services/impl/OpenAICompletionService.java
index 5d3b73eb9..d309ae380 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/services/impl/OpenAICompletionService.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/services/impl/OpenAICompletionService.java
@@ -19,11 +19,10 @@
 import static ai.langstream.api.util.ConfigurationUtils.getDouble;
 import static ai.langstream.api.util.ConfigurationUtils.getInteger;
 
-import com.azure.ai.openai.OpenAIClient;
+import com.azure.ai.openai.OpenAIAsyncClient;
 import com.azure.ai.openai.models.ChatCompletionsOptions;
 import com.azure.ai.openai.models.ChatRole;
 import com.azure.ai.openai.models.CompletionsFinishReason;
-import com.azure.core.util.IterableStream;
 import com.datastax.oss.streaming.ai.completions.ChatChoice;
 import com.datastax.oss.streaming.ai.completions.ChatCompletions;
 import com.datastax.oss.streaming.ai.completions.ChatMessage;
@@ -36,18 +35,22 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
+import reactor.core.publisher.Flux;
 
 @Slf4j
 public class OpenAICompletionService implements CompletionsService {
 
-    private OpenAIClient client;
+    private OpenAIAsyncClient client;
 
-    public OpenAICompletionService(OpenAIClient client) {
+    public OpenAICompletionService(OpenAIAsyncClient client) {
         this.client = client;
     }
 
     @Override
+    @SneakyThrows
     public CompletableFuture<ChatCompletions> getChatCompletions(
             List<ChatMessage> messages,
             StreamingChunksConsumer streamingChunksConsumer,
@@ -78,13 +81,14 @@ public CompletableFuture<ChatCompletions> getChatCompletions(
         // it works even if the streamingChunksConsumer is null
         if (chatCompletionsOptions.isStream()) {
             CompletableFuture<?> finished = new CompletableFuture<>();
-            IterableStream<com.azure.ai.openai.models.ChatCompletions> model =
+            Flux<com.azure.ai.openai.models.ChatCompletions> flux =
                     client.getChatCompletionsStream(
                             (String) options.get("model"), chatCompletionsOptions);
+            Stream<com.azure.ai.openai.models.ChatCompletions> model = flux.toStream();
             ChatCompletionsConsumer chatCompletionsConsumer =
                     new ChatCompletionsConsumer(
                             streamingChunksConsumer, minChunksPerMessage, finished);
-            model.stream().forEach(chatCompletionsConsumer);
+            model.forEach(chatCompletionsConsumer);
             return finished.thenApply(
                     ___ -> {
                         result.setChoices(
@@ -96,8 +100,9 @@ public CompletableFuture<ChatCompletions> getChatCompletions(
                     });
         } else {
             com.azure.ai.openai.models.ChatCompletions chatCompletions =
-                    client.getChatCompletions(
-                            (String) options.get("model"), chatCompletionsOptions);
+                    client.getChatCompletions((String) options.get("model"), chatCompletionsOptions)
+                            .toFuture()
+                            .get();
             result.setChoices(
                     chatCompletions.getChoices().stream()
                             .map(c -> new ChatChoice(convertMessage(c)))
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/services/impl/OpenAIServiceProvider.java b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/services/impl/OpenAIServiceProvider.java
index 74c24037a..12c6fa006 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/services/impl/OpenAIServiceProvider.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/services/impl/OpenAIServiceProvider.java
@@ -16,7 +16,7 @@
 package ai.langstream.ai.agents.services.impl;
 
 import ai.langstream.ai.agents.services.ServiceProviderProvider;
-import com.azure.ai.openai.OpenAIClient;
+import com.azure.ai.openai.OpenAIAsyncClient;
 import com.datastax.oss.streaming.ai.model.config.OpenAIConfig;
 import com.datastax.oss.streaming.ai.services.ServiceProvider;
 import com.datastax.oss.streaming.ai.util.TransformFunctionUtil;
@@ -33,7 +33,7 @@ public ServiceProvider createImplementation(Map<String, Object> agentConfigurati
         OpenAIConfig config =
                 TransformFunctionUtil.convertFromMap(
                         (Map<String, Object>) agentConfiguration.get("openai"), OpenAIConfig.class);
-        OpenAIClient client = TransformFunctionUtil.buildOpenAIClient(config);
+        OpenAIAsyncClient client = TransformFunctionUtil.buildOpenAsyncAIClient(config);
         return new com.datastax.oss.streaming.ai.services.OpenAIServiceProvider(client);
     }
 }
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/embeddings/OpenAIEmbeddingsService.java b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/embeddings/OpenAIEmbeddingsService.java
index 7dde0e607..1e6dfdd10 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/embeddings/OpenAIEmbeddingsService.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/embeddings/OpenAIEmbeddingsService.java
@@ -15,8 +15,7 @@
  */
 package com.datastax.oss.streaming.ai.embeddings;
 
-import com.azure.ai.openai.OpenAIClient;
-import com.azure.ai.openai.models.Embeddings;
+import com.azure.ai.openai.OpenAIAsyncClient;
 import com.azure.ai.openai.models.EmbeddingsOptions;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
@@ -26,10 +25,10 @@
 @Slf4j
 public class OpenAIEmbeddingsService implements EmbeddingsService {
 
-    private final OpenAIClient openAIClient;
+    private final OpenAIAsyncClient openAIClient;
     private final String model;
 
-    public OpenAIEmbeddingsService(OpenAIClient openAIClient, String model) {
+    public OpenAIEmbeddingsService(OpenAIAsyncClient openAIClient, String model) {
         this.openAIClient = openAIClient;
         this.model = model;
     }
@@ -38,11 +37,24 @@ public OpenAIEmbeddingsService(OpenAIClient openAIClient, String model) {
     public CompletableFuture<List<List<Double>>> computeEmbeddings(List<String> texts) {
         try {
             EmbeddingsOptions embeddingsOptions = new EmbeddingsOptions(texts);
-            Embeddings embeddings = openAIClient.getEmbeddings(model, embeddingsOptions);
-            return CompletableFuture.completedFuture(
-                    embeddings.getData().stream()
-                            .map(embedding -> embedding.getEmbedding())
-                            .collect(Collectors.toList()));
+            CompletableFuture<List<List<Double>>> result =
+                    openAIClient
+                            .getEmbeddings(model, embeddingsOptions)
+                            .toFuture()
+                            .thenApply(
+                                    embeddings ->
+                                            embeddings.getData().stream()
+                                                    .map(embedding -> embedding.getEmbedding())
+                                                    .collect(Collectors.toList()));
+
+            // we need to wait, in order to guarantee ordering
+            // TODO: we should not wait, but instead use an ordered executor (per key)
+            try {
+                result.join();
+            } catch (Throwable err) {
+                log.error("Cannot compute embeddings", err);
+            }
+            return result;
         } catch (RuntimeException err) {
             log.error("Cannot compute embeddings", err);
             return CompletableFuture.failedFuture(err);
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/services/OpenAIServiceProvider.java b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/services/OpenAIServiceProvider.java
index 96e654aef..78042481d 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/services/OpenAIServiceProvider.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/services/OpenAIServiceProvider.java
@@ -16,7 +16,7 @@
 package com.datastax.oss.streaming.ai.services;
 
 import ai.langstream.ai.agents.services.impl.OpenAICompletionService;
-import com.azure.ai.openai.OpenAIClient;
+import com.azure.ai.openai.OpenAIAsyncClient;
 import com.datastax.oss.streaming.ai.completions.CompletionsService;
 import com.datastax.oss.streaming.ai.embeddings.EmbeddingsService;
 import com.datastax.oss.streaming.ai.embeddings.OpenAIEmbeddingsService;
@@ -26,13 +26,13 @@
 
 public class OpenAIServiceProvider implements ServiceProvider {
 
-    private final OpenAIClient client;
+    private final OpenAIAsyncClient client;
 
     public OpenAIServiceProvider(TransformStepConfig config) {
-        client = TransformFunctionUtil.buildOpenAIClient(config.getOpenai());
+        client = TransformFunctionUtil.buildOpenAsyncAIClient(config.getOpenai());
     }
 
-    public OpenAIServiceProvider(OpenAIClient client) {
+    public OpenAIServiceProvider(OpenAIAsyncClient client) {
         this.client = client;
     }
 
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/util/TransformFunctionUtil.java b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/util/TransformFunctionUtil.java
index 65d265286..b6c4b1cf2 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/util/TransformFunctionUtil.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/util/TransformFunctionUtil.java
@@ -15,6 +15,7 @@
  */
 package com.datastax.oss.streaming.ai.util;
 
+import com.azure.ai.openai.OpenAIAsyncClient;
 import com.azure.ai.openai.OpenAIClient;
 import com.azure.ai.openai.OpenAIClientBuilder;
 import com.azure.core.credential.AzureKeyCredential;
@@ -125,6 +126,32 @@ public static OpenAIClient buildOpenAIClient(OpenAIConfig openAIConfig) {
         return openAIClientBuilder.buildClient();
     }
 
+    public static OpenAIAsyncClient buildOpenAsyncAIClient(OpenAIConfig openAIConfig) {
+        if (openAIConfig == null) {
+            return null;
+        }
+        OpenAIClientBuilder openAIClientBuilder = new OpenAIClientBuilder();
+        if (openAIConfig.getProvider() == OpenAIProvider.AZURE) {
+            openAIClientBuilder.credential(new AzureKeyCredential(openAIConfig.getAccessKey()));
+        } else {
+            openAIClientBuilder.credential(new KeyCredential(openAIConfig.getAccessKey()));
+        }
+        if (openAIConfig.getUrl() != null && !openAIConfig.getUrl().isEmpty()) {
+            openAIClientBuilder.endpoint(openAIConfig.getUrl());
+
+            // this is for testing only
+            if (openAIConfig.getUrl().startsWith("http://localhost")) {
+                HttpPipeline httpPipeline =
+                        new HttpPipelineBuilder()
+                                .httpClient(new MockHttpClient(openAIConfig.getAccessKey()))
+                                .build();
+                openAIClientBuilder.pipeline(httpPipeline);
+            }
+        }
+
+        return openAIClientBuilder.buildAsyncClient();
+    }
+
     public static QueryStepDataSource buildDataSource(Map<String, Object> dataSourceConfig) {
         if (dataSourceConfig == null) {
             return new QueryStepDataSource() {};


Test patch:
diff --git a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/pulsar/functions/transforms/GenAITest.java b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/pulsar/functions/transforms/GenAITest.java
index 46ce4265b..c2097b2ee 100644
--- a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/pulsar/functions/transforms/GenAITest.java
+++ b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/pulsar/functions/transforms/GenAITest.java
@@ -25,10 +25,9 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import com.azure.ai.openai.OpenAIClient;
+import com.azure.ai.openai.OpenAIAsyncClient;
 import com.azure.ai.openai.models.ChatCompletions;
 import com.azure.ai.openai.models.ChatCompletionsOptions;
-import com.azure.core.util.IterableStream;
 import com.datastax.oss.streaming.ai.datasource.QueryStepDataSource;
 import com.datastax.oss.streaming.ai.services.OpenAIServiceProvider;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -49,6 +48,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.mockito.ArgumentCaptor;
+import reactor.core.publisher.Flux;
 
 @Slf4j
 public class GenAITest {
@@ -186,7 +186,7 @@ void testChatCompletions() throws Exception {
                 new Gson().fromJson(userConfig, new TypeToken<Map<String, Object>>() {}.getType());
         TransformFunction transformFunction = spy(new TransformFunction());
 
-        OpenAIClient client = mock(OpenAIClient.class);
+        OpenAIAsyncClient client = mock(OpenAIAsyncClient.class);
 
         String completion =
                 (""
@@ -205,12 +205,9 @@ void testChatCompletions() throws Exception {
         when(client.getChatCompletionsStream(eq("test-model"), any()))
                 .thenAnswer(
                         a ->
-                                IterableStream.of(
-                                        List.of(
-                                                new ObjectMapper()
-                                                        .readValue(
-                                                                completion,
-                                                                ChatCompletions.class))));
+                                Flux.just(
+                                        new ObjectMapper()
+                                                .readValue(completion, ChatCompletions.class)));
         when(transformFunction.buildServiceProvider(any()))
                 .thenReturn(new OpenAIServiceProvider(client));
 
@@ -251,7 +248,7 @@ void testChatCompletionsWithLogField() throws Exception {
                 new Gson().fromJson(userConfig, new TypeToken<Map<String, Object>>() {}.getType());
         TransformFunction transformFunction = spy(new TransformFunction());
 
-        OpenAIClient client = mock(OpenAIClient.class);
+        OpenAIAsyncClient client = mock(OpenAIAsyncClient.class);
 
         String completion =
                 (""
@@ -270,12 +267,9 @@ void testChatCompletionsWithLogField() throws Exception {
         when(client.getChatCompletionsStream(eq("test-model"), any()))
                 .thenAnswer(
                         a ->
-                                IterableStream.of(
-                                        List.of(
-                                                new ObjectMapper()
-                                                        .readValue(
-                                                                completion,
-                                                                ChatCompletions.class))));
+                                Flux.just(
+                                        new ObjectMapper()
+                                                .readValue(completion, ChatCompletions.class)));
         when(transformFunction.buildServiceProvider(any()))
                 .thenReturn(new OpenAIServiceProvider(client));
 
diff --git a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/ChatCompletionsStepTest.java b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/ChatCompletionsStepTest.java
index 57fd90de4..afd9fce91 100644
--- a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/ChatCompletionsStepTest.java
+++ b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/ChatCompletionsStepTest.java
@@ -24,10 +24,9 @@
 import static org.mockito.Mockito.when;
 
 import ai.langstream.ai.agents.services.impl.OpenAICompletionService;
-import com.azure.ai.openai.OpenAIClient;
+import com.azure.ai.openai.OpenAIAsyncClient;
 import com.azure.ai.openai.models.ChatCompletions;
 import com.azure.ai.openai.models.ChatCompletionsOptions;
-import com.azure.core.util.IterableStream;
 import com.datastax.oss.streaming.ai.completions.ChatMessage;
 import com.datastax.oss.streaming.ai.model.config.ChatCompletionsConfig;
 import com.fasterxml.jackson.databind.JsonNode;
@@ -60,6 +59,8 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.mockito.ArgumentCaptor;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 public class ChatCompletionsStepTest {
 
@@ -78,20 +79,15 @@ public class ChatCompletionsStepTest {
                     .replace("'", "\"");
     private static final ObjectMapper mapper = new ObjectMapper();
     private OpenAICompletionService completionService;
-    private OpenAIClient openAIClient;
+    private OpenAIAsyncClient openAIClient;
 
     @BeforeEach
     void setup() throws Exception {
-        openAIClient = mock(OpenAIClient.class);
+        openAIClient = mock(OpenAIAsyncClient.class);
         when(openAIClient.getChatCompletions(eq("test-model"), any()))
-                .thenReturn(mapper.readValue(COMPLETION, ChatCompletions.class));
+                .thenReturn(Mono.just(mapper.readValue(COMPLETION, ChatCompletions.class)));
         when(openAIClient.getChatCompletionsStream(eq("test-model"), any()))
-                .thenAnswer(
-                        a ->
-                                IterableStream.of(
-                                        List.of(
-                                                mapper.readValue(
-                                                        COMPLETION, ChatCompletions.class))));
+                .thenAnswer(a -> Flux.just(mapper.readValue(COMPLETION, ChatCompletions.class)));
         this.completionService = new OpenAICompletionService(openAIClient);
     }