From c7c9657e05614cb088914e64af65762fb59a0e6d Mon Sep 17 00:00:00 2001 From: Nathalie Jonathan Date: Tue, 28 Jul 2026 16:29:27 -0700 Subject: [PATCH] Add ML gRPC support Signed-off-by: Nathalie Jonathan --- CHANGELOG.md | 1 + java-client-grpc/build.gradle.kts | 11 +- .../client/transport/grpc/GrpcTransport.java | 91 ++++ .../grpc/ml/MlExecuteAgentStreamRequest.java | 73 +++ .../transport/grpc/ml/MlInferenceResult.java | 50 ++ .../client/transport/grpc/ml/MlMessage.java | 74 +++ .../client/transport/grpc/ml/MlOutput.java | 98 ++++ .../transport/grpc/ml/MlParameters.java | 114 +++++ .../grpc/ml/MlPredictStreamRequest.java | 73 +++ .../transport/grpc/ml/MlStreamChunk.java | 105 +++++ .../transport/grpc/ml/MlStreamStatus.java | 58 +++ .../translation/MlStreamRequestConverter.java | 104 +++++ .../MlStreamResponseConverter.java | 75 +++ .../transport/grpc/MlStreamingTest.java | 438 ++++++++++++++++++ .../grpc/translation/TranslationTest.java | 184 +++++++- .../integTest/grpc/GrpcMlStreamIT.java | 141 ++++++ 16 files changed, 1687 insertions(+), 3 deletions(-) create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlExecuteAgentStreamRequest.java create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlInferenceResult.java create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlMessage.java create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlOutput.java create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlParameters.java create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlPredictStreamRequest.java create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlStreamChunk.java create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlStreamStatus.java create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/translation/MlStreamRequestConverter.java create mode 100644 java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/translation/MlStreamResponseConverter.java create mode 100644 java-client-grpc/src/test/java/org/opensearch/client/transport/grpc/MlStreamingTest.java create mode 100644 java-client-grpc/src/test/java11/org/opensearch/client/opensearch/integTest/grpc/GrpcMlStreamIT.java diff --git a/CHANGELOG.md b/CHANGELOG.md index b9ca65d96..b5b0e1980 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) - Added `equals()` and `hashCode()` implementations to `FieldValue` ([#1998](https://github.com/opensearch-project/opensearch-java/pull/1998)) - Add document lifecycle guide and runnable sample ([#2017](https://github.com/opensearch-project/opensearch-java/pull/2017)) - Add transparent gRPC transport with HybridTransport (bulk over gRPC, REST fallback), translation layer, TLS, basic auth, AWS SigV4, and JWT support ([#2062](https://github.com/opensearch-project/opensearch-java/pull/2062)) +- Add gRPC ML streaming support for the ML Commons `PredictModelStream` and `ExecuteAgentStream` APIs, exposed as `GrpcTransport.predictModelStream()` and `GrpcTransport.executeAgentStream()` ([#2070](https://github.com/opensearch-project/opensearch-java/pull/2070)) ### Fixed diff --git a/java-client-grpc/build.gradle.kts b/java-client-grpc/build.gradle.kts index 951ae9f01..e3e353693 100644 --- a/java-client-grpc/build.gradle.kts +++ b/java-client-grpc/build.gradle.kts @@ -30,7 +30,7 @@ java { val opensearchVersion = "3.5.0-SNAPSHOT" val grpcVersion = "1.68.0" val protobufVersion = "3.25.5" -val opensearchProtobufVersion = "1.2.0" +val opensearchProtobufVersion = "1.6.0" dependencies { // Depend on java-client core @@ -53,6 +53,7 @@ dependencies { // Test dependencies testImplementation("io.grpc", "grpc-testing", grpcVersion) + testImplementation("io.grpc", "grpc-inprocess", grpcVersion) testImplementation("junit", "junit", "4.13.2") testImplementation("org.opensearch.client", "opensearch-rest-client", opensearchVersion) testImplementation("software.amazon.awssdk", "sdk-core", "[2.21,3.0)") @@ -60,11 +61,18 @@ dependencies { testImplementation("software.amazon.awssdk", "http-auth-aws", "[2.21,3.0)") } +val runtimeJavaVersion = (System.getProperty("runtime.java")?.toInt())?.let(JavaVersion::toVersion) ?: JavaVersion.current() + tasks.test { systemProperty("tests.security.manager", "false") } val unitTest = tasks.register("unitTest") { + // Without these the task is NO-SOURCE and silently runs nothing; skipped on Java 8, where java-client's Jackson 3 cannot compile. + if (runtimeJavaVersion >= JavaVersion.VERSION_11) { + testClassesDirs = sourceSets.test.get().output.classesDirs + classpath = sourceSets.test.get().runtimeClasspath + } filter { excludeTestsMatching("org.opensearch.client.opensearch.integTest.*") } @@ -86,7 +94,6 @@ val integrationTest = tasks.register("integrationTest") { } // Integration tests require Java 21+ and live in src/test/java11 -val runtimeJavaVersion = (System.getProperty("runtime.java")?.toInt())?.let(JavaVersion::toVersion) ?: JavaVersion.current() if (runtimeJavaVersion >= JavaVersion.VERSION_21) { val java21: SourceSet = sourceSets.create("java21") { java { diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/GrpcTransport.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/GrpcTransport.java index 83f09cf54..569e91c73 100644 --- a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/GrpcTransport.java +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/GrpcTransport.java @@ -11,8 +11,11 @@ import io.grpc.ManagedChannel; import io.grpc.StatusRuntimeException; import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.Iterator; import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; import javax.annotation.Nullable; import org.opensearch.client.json.JsonpMapper; import org.opensearch.client.opensearch.core.BulkRequest; @@ -21,10 +24,17 @@ import org.opensearch.client.transport.OpenSearchTransport; import org.opensearch.client.transport.TransportException; import org.opensearch.client.transport.TransportOptions; +import org.opensearch.client.transport.grpc.ml.MlExecuteAgentStreamRequest; +import org.opensearch.client.transport.grpc.ml.MlPredictStreamRequest; +import org.opensearch.client.transport.grpc.ml.MlStreamChunk; import org.opensearch.client.transport.grpc.translation.BulkRequestConverter; import org.opensearch.client.transport.grpc.translation.BulkResponseConverter; import org.opensearch.client.transport.grpc.translation.GrpcStatusConverter; +import org.opensearch.client.transport.grpc.translation.MlStreamRequestConverter; +import org.opensearch.client.transport.grpc.translation.MlStreamResponseConverter; +import org.opensearch.protobufs.PredictResponse; import org.opensearch.protobufs.services.DocumentServiceGrpc; +import org.opensearch.protobufs.services.MLServiceGrpc; /** * Pure gRPC transport for OpenSearch. Implements {@link OpenSearchTransport} and routes @@ -42,6 +52,10 @@ * OpenSearchClient client = new OpenSearchClient(transport); * client.bulk(bulkRequest); // goes over gRPC * } + *

+ * In addition to the {@link OpenSearchTransport} contract, this transport exposes the + * ML server-streaming APIs directly via {@link #predictModelStream} and + * {@link #executeAgentStream}. */ public class GrpcTransport implements OpenSearchTransport { @@ -67,6 +81,7 @@ public static boolean isEndpointSupported(Endpoint endpoint) { private final ManagedChannel channel; private final DocumentServiceGrpc.DocumentServiceBlockingStub documentStub; + private final MLServiceGrpc.MLServiceBlockingStub mlStub; private final JsonpMapper jsonpMapper; private final GrpcTransportOptions grpcOptions; private final TransportOptions transportOptions; @@ -81,6 +96,7 @@ public static boolean isEndpointSupported(Endpoint endpoint) { ) { this.channel = channel; this.documentStub = channel != null ? DocumentServiceGrpc.newBlockingStub(channel) : null; + this.mlStub = channel != null ? MLServiceGrpc.newBlockingStub(channel) : null; this.jsonpMapper = jsonpMapper; this.grpcOptions = grpcOptions; this.transportOptions = transportOptions; @@ -238,6 +254,81 @@ private BulkResponse performBulk(BulkRequest request) throws TransportException } } + // ─── ML Streaming APIs ─────────────────────────────────────────────────────── + + /** + * Invokes the ML streaming predict API ({@code MLService/PredictModelStream}) and + * returns an iterator over the response chunks as the server produces them. + * + * @param request the streaming predict request + * @return an iterator over the streamed chunks + */ + public Iterator predictModelStream(MlPredictStreamRequest request) { + org.opensearch.protobufs.MlPredictModelStreamRequest protoRequest = MlStreamRequestConverter.toProto(request); + return stream(() -> requireMlStub().predictModelStream(protoRequest)); + } + + /** + * Invokes the ML streaming agent execute API ({@code MLService/ExecuteAgentStream}) + * and returns an iterator over the response chunks as the server produces them. + * + * @param request the streaming agent execute request + * @return an iterator over the streamed chunks + */ + public Iterator executeAgentStream(MlExecuteAgentStreamRequest request) { + org.opensearch.protobufs.MlExecuteAgentStreamRequest protoRequest = MlStreamRequestConverter.toProto(request); + return stream(() -> requireMlStub().executeAgentStream(protoRequest)); + } + + private MLServiceGrpc.MLServiceBlockingStub requireMlStub() { + if (mlStub == null) { + throw new IllegalStateException("gRPC channel is not available; ML streaming requires a connected transport"); + } + return mlStub; + } + + private Iterator stream(Supplier> rpc) { + return new Iterator() { + private Iterator delegate; + + @Override + public boolean hasNext() { + try { + return delegate().hasNext(); + } catch (StatusRuntimeException e) { + throw convert(e); + } + } + + @Override + public MlStreamChunk next() { + try { + return MlStreamResponseConverter.fromProto(delegate().next()); + } catch (StatusRuntimeException e) { + throw convert(e); + } + } + + private Iterator delegate() { + if (delegate == null) { + delegate = rpc.get(); + } + return delegate; + } + }; + } + + private static RuntimeException convert(StatusRuntimeException e) { + try { + GrpcStatusConverter.throwConverted(e); + } catch (TransportException converted) { + return new UncheckedIOException(converted); + } catch (RuntimeException converted) { + return converted; + } + return new UncheckedIOException(new TransportException("gRPC error: " + e.getMessage(), e)); + } + // ─── Builder ───────────────────────────────────────────────────────────────── /** diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlExecuteAgentStreamRequest.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlExecuteAgentStreamRequest.java new file mode 100644 index 000000000..1fec86a8b --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlExecuteAgentStreamRequest.java @@ -0,0 +1,73 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.ml; + +import javax.annotation.Nullable; + +/** + * Request for the ML streaming agent execute API ({@code MLService/ExecuteAgentStream}). + */ +public final class MlExecuteAgentStreamRequest { + + private final String agentId; + private final MlParameters parameters; + + private MlExecuteAgentStreamRequest(Builder builder) { + if (builder.agentId == null || builder.agentId.isEmpty()) { + throw new IllegalArgumentException("agentId is required"); + } + this.agentId = builder.agentId; + this.parameters = builder.parameters; + } + + /** + * The ID of the agent to execute. Never null. + */ + public String agentId() { + return agentId; + } + + /** + * The request parameters, or null if unset. + */ + @Nullable + public MlParameters parameters() { + return parameters; + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private String agentId; + private MlParameters parameters; + + public Builder agentId(String agentId) { + this.agentId = agentId; + return this; + } + + public Builder parameters(MlParameters parameters) { + this.parameters = parameters; + return this; + } + + /** + * Configures the parameters using a builder function. + */ + public Builder parameters(java.util.function.Function fn) { + return parameters(fn.apply(MlParameters.builder()).build()); + } + + public MlExecuteAgentStreamRequest build() { + return new MlExecuteAgentStreamRequest(this); + } + } +} diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlInferenceResult.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlInferenceResult.java new file mode 100644 index 000000000..f85649c33 --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlInferenceResult.java @@ -0,0 +1,50 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.ml; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * One inference result of a streamed ML response, holding the outputs the model + * produced for this chunk. + */ +public final class MlInferenceResult { + + private final List output; + + private MlInferenceResult(Builder builder) { + this.output = Collections.unmodifiableList(new ArrayList<>(builder.output)); + } + + /** + * The outputs of this inference result. Never null; empty when the server sent none. + */ + public List output() { + return output; + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private final List output = new ArrayList<>(); + + public Builder addOutput(MlOutput out) { + this.output.add(out); + return this; + } + + public MlInferenceResult build() { + return new MlInferenceResult(this); + } + } +} diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlMessage.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlMessage.java new file mode 100644 index 000000000..d2a6b5bb7 --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlMessage.java @@ -0,0 +1,74 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.ml; + +import javax.annotation.Nullable; + +/** + * A single chat message in an ML streaming request. + */ +public final class MlMessage { + + private final String role; + private final String content; + + private MlMessage(Builder builder) { + this.role = builder.role; + this.content = builder.content; + } + + /** + * Creates a message with the given role and content. + * + * @param role the message role (e.g., {@code "user"}, {@code "assistant"}, {@code "system"}) + * @param content the message content + */ + public static MlMessage of(String role, String content) { + return builder().role(role).content(content).build(); + } + + /** + * The message role, or null if unset. + */ + @Nullable + public String role() { + return role; + } + + /** + * The message content, or null if unset. + */ + @Nullable + public String content() { + return content; + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private String role; + private String content; + + public Builder role(String role) { + this.role = role; + return this; + } + + public Builder content(String content) { + this.content = content; + return this; + } + + public MlMessage build() { + return new MlMessage(this); + } + } +} diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlOutput.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlOutput.java new file mode 100644 index 000000000..d11837939 --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlOutput.java @@ -0,0 +1,98 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.ml; + +import javax.annotation.Nullable; + +/** + * A single output entry of a streamed ML inference result. + */ +public final class MlOutput { + + private final String name; + private final String result; + private final String content; + private final Boolean isLast; + + private MlOutput(Builder builder) { + this.name = builder.name; + this.result = builder.result; + this.content = builder.content; + this.isLast = builder.isLast; + } + + /** + * The output name (e.g., {@code "response"}), or null if unset. + */ + @Nullable + public String name() { + return name; + } + + /** + * The output as a raw string, or null if unset. + */ + @Nullable + public String result() { + return result; + } + + /** + * The incremental content of this chunk (protobuf {@code data_as_map.content}), + * or null if unset. + */ + @Nullable + public String content() { + return content; + } + + /** + * Whether the server marked this as the final chunk + * (protobuf {@code data_as_map.is_last}), or null if unset. + */ + @Nullable + public Boolean isLast() { + return isLast; + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private String name; + private String result; + private String content; + private Boolean isLast; + + public Builder name(String name) { + this.name = name; + return this; + } + + public Builder result(String result) { + this.result = result; + return this; + } + + public Builder content(String content) { + this.content = content; + return this; + } + + public Builder isLast(Boolean isLast) { + this.isLast = isLast; + return this; + } + + public MlOutput build() { + return new MlOutput(this); + } + } +} diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlParameters.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlParameters.java new file mode 100644 index 000000000..059228e16 --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlParameters.java @@ -0,0 +1,114 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.ml; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import javax.annotation.Nullable; + +/** + * Parameters for an ML streaming request, mirroring the {@code parameters} object + * of the REST ML predict and agent execute APIs. + */ +public final class MlParameters { + + private final List messages; + private final String inputs; + private final String question; + private final String llmInterface; + + private MlParameters(Builder builder) { + this.messages = Collections.unmodifiableList(new ArrayList<>(builder.messages)); + this.inputs = builder.inputs; + this.question = builder.question; + this.llmInterface = builder.llmInterface; + } + + /** + * The chat messages. Never null; empty when unset. + */ + public List messages() { + return messages; + } + + /** + * The raw {@code inputs} value, or null if unset. + */ + @Nullable + public String inputs() { + return inputs; + } + + /** + * The {@code question} value, or null if unset. + */ + @Nullable + public String question() { + return question; + } + + /** + * The LLM interface identifier (protobuf field {@code _llm_interface}), + * or null if unset. + */ + @Nullable + public String llmInterface() { + return llmInterface; + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private final List messages = new ArrayList<>(); + private String inputs; + private String question; + private String llmInterface; + + /** + * Appends a single chat message. + */ + public Builder addMessage(MlMessage message) { + this.messages.add(message); + return this; + } + + /** + * Appends a single chat message built from a role and content. + */ + public Builder addMessage(String role, String content) { + return addMessage(MlMessage.of(role, content)); + } + + public Builder inputs(String inputs) { + this.inputs = inputs; + return this; + } + + public Builder question(String question) { + this.question = question; + return this; + } + + /** + * Sets the LLM interface identifier, e.g. {@code "openai/v1/chat/completions"}. + * Maps to the protobuf {@code _llm_interface} field. + */ + public Builder llmInterface(String llmInterface) { + this.llmInterface = llmInterface; + return this; + } + + public MlParameters build() { + return new MlParameters(this); + } + } +} diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlPredictStreamRequest.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlPredictStreamRequest.java new file mode 100644 index 000000000..22e055572 --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlPredictStreamRequest.java @@ -0,0 +1,73 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.ml; + +import javax.annotation.Nullable; + +/** + * Request for the ML streaming predict API ({@code MLService/PredictModelStream}). + */ +public final class MlPredictStreamRequest { + + private final String modelId; + private final MlParameters parameters; + + private MlPredictStreamRequest(Builder builder) { + if (builder.modelId == null || builder.modelId.isEmpty()) { + throw new IllegalArgumentException("modelId is required"); + } + this.modelId = builder.modelId; + this.parameters = builder.parameters; + } + + /** + * The ID of the model to invoke. Never null. + */ + public String modelId() { + return modelId; + } + + /** + * The request parameters, or null if unset. + */ + @Nullable + public MlParameters parameters() { + return parameters; + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private String modelId; + private MlParameters parameters; + + public Builder modelId(String modelId) { + this.modelId = modelId; + return this; + } + + public Builder parameters(MlParameters parameters) { + this.parameters = parameters; + return this; + } + + /** + * Configures the parameters using a builder function. + */ + public Builder parameters(java.util.function.Function fn) { + return parameters(fn.apply(MlParameters.builder()).build()); + } + + public MlPredictStreamRequest build() { + return new MlPredictStreamRequest(this); + } + } +} diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlStreamChunk.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlStreamChunk.java new file mode 100644 index 000000000..117317b9d --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlStreamChunk.java @@ -0,0 +1,105 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.ml; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import javax.annotation.Nullable; + +/** + * A single chunk of a streamed ML response. + */ +public final class MlStreamChunk { + + /** Name of the output that carries the incremental streamed payload. */ + private static final String RESPONSE_OUTPUT_NAME = "response"; + + private final MlStreamStatus status; + private final List inferenceResults; + + private MlStreamChunk(Builder builder) { + this.status = builder.status; + this.inferenceResults = Collections.unmodifiableList(new ArrayList<>(builder.inferenceResults)); + } + + /** + * The status of this chunk, or null if the server did not set one. + */ + @Nullable + public MlStreamStatus status() { + return status; + } + + /** + * The inference results of this chunk. Never null; empty when the server sent none. + */ + public List inferenceResults() { + return inferenceResults; + } + + /** + * The incremental text of this chunk — the {@code content} of the output named + * {@code "response"} — or null if this chunk carries none. + */ + @Nullable + public String content() { + MlOutput out = contentOutput(); + return out != null ? out.content() : null; + } + + /** + * Whether the server marked this as the final chunk, taken from the output named + * {@code "response"}. Null when the server did not set the flag. + */ + @Nullable + public Boolean isLast() { + MlOutput out = contentOutput(); + return out != null ? out.isLast() : null; + } + + @Nullable + private MlOutput contentOutput() { + MlOutput fallback = null; + for (MlInferenceResult result : inferenceResults) { + for (MlOutput out : result.output()) { + if (RESPONSE_OUTPUT_NAME.equals(out.name())) { + return out; + } + if (fallback == null && out.content() != null) { + fallback = out; + } + } + } + return fallback; + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private MlStreamStatus status; + private final List inferenceResults = new ArrayList<>(); + + public Builder status(MlStreamStatus status) { + this.status = status; + return this; + } + + public Builder addInferenceResult(MlInferenceResult result) { + this.inferenceResults.add(result); + return this; + } + + public MlStreamChunk build() { + return new MlStreamChunk(this); + } + } +} diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlStreamStatus.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlStreamStatus.java new file mode 100644 index 000000000..7fcbc2bf1 --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/ml/MlStreamStatus.java @@ -0,0 +1,58 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.ml; + +/** + * Status of a streamed ML response chunk. + */ +public enum MlStreamStatus { + /** Status was not set or is not recognized by this client version. */ + UNSPECIFIED, + /** The streaming task was cancelled. */ + CANCELLED, + /** The streaming task completed successfully. */ + COMPLETED, + /** The streaming task completed, but with an error. */ + COMPLETED_WITH_ERROR, + /** The streaming task was created. */ + CREATED, + /** The streaming task failed. */ + FAILED, + /** The streaming task is in progress; more chunks may follow. */ + RUNNING; + + /** + * Maps a protobuf {@code Status} to this enum. Unrecognized values — + * including statuses added by newer servers — map to {@link #UNSPECIFIED}. + * + * @param status the protobuf status + * @return the corresponding client status + */ + public static MlStreamStatus fromProto(org.opensearch.protobufs.Status status) { + if (status == null) { + return UNSPECIFIED; + } + switch (status) { + case STATUS_CANCELLED: + return CANCELLED; + case STATUS_COMPLETED: + return COMPLETED; + case STATUS_COMPLETED_WITH_ERROR: + return COMPLETED_WITH_ERROR; + case STATUS_CREATED: + return CREATED; + case STATUS_FAILED: + return FAILED; + case STATUS_RUNNING: + return RUNNING; + default: + return UNSPECIFIED; + } + } +} diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/translation/MlStreamRequestConverter.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/translation/MlStreamRequestConverter.java new file mode 100644 index 000000000..4c1d8bd64 --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/translation/MlStreamRequestConverter.java @@ -0,0 +1,104 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.translation; + +import org.opensearch.client.transport.grpc.ml.MlExecuteAgentStreamRequest; +import org.opensearch.client.transport.grpc.ml.MlMessage; +import org.opensearch.client.transport.grpc.ml.MlParameters; +import org.opensearch.client.transport.grpc.ml.MlPredictStreamRequest; +import org.opensearch.protobufs.MLExecuteAgentStreamRequestBody; +import org.opensearch.protobufs.MLPredictModelStreamRequestBody; +import org.opensearch.protobufs.Messages; +import org.opensearch.protobufs.Parameters; + +/** + * Converts opensearch-java ML streaming requests to their protobuf equivalents + * for transmission over gRPC. + */ +public final class MlStreamRequestConverter { + + private MlStreamRequestConverter() { + // utility class + } + + /** + * Converts a client {@link MlPredictStreamRequest} to a protobuf + * {@code MlPredictModelStreamRequest}. + * + * @param request the client request + * @return the protobuf request ready to send over gRPC + */ + public static org.opensearch.protobufs.MlPredictModelStreamRequest toProto(MlPredictStreamRequest request) { + org.opensearch.protobufs.MlPredictModelStreamRequest.Builder builder = org.opensearch.protobufs.MlPredictModelStreamRequest + .newBuilder(); + + builder.setModelId(request.modelId()); + + if (request.parameters() != null) { + builder.setMlPredictModelStreamRequestBody( + MLPredictModelStreamRequestBody.newBuilder().setParameters(toProto(request.parameters())).build() + ); + } + + return builder.build(); + } + + /** + * Converts a client {@link MlExecuteAgentStreamRequest} to a protobuf + * {@code MlExecuteAgentStreamRequest}. + * + * @param request the client request + * @return the protobuf request ready to send over gRPC + */ + public static org.opensearch.protobufs.MlExecuteAgentStreamRequest toProto(MlExecuteAgentStreamRequest request) { + org.opensearch.protobufs.MlExecuteAgentStreamRequest.Builder builder = org.opensearch.protobufs.MlExecuteAgentStreamRequest + .newBuilder(); + + builder.setAgentId(request.agentId()); + + if (request.parameters() != null) { + builder.setMlExecuteAgentStreamRequestBody( + MLExecuteAgentStreamRequestBody.newBuilder().setParameters(toProto(request.parameters())).build() + ); + } + + return builder.build(); + } + + /** + * Converts client {@link MlParameters} to a protobuf {@code Parameters}. + * Only fields that were set are mapped. + */ + private static Parameters toProto(MlParameters parameters) { + Parameters.Builder builder = Parameters.newBuilder(); + + for (MlMessage message : parameters.messages()) { + Messages.Builder messageBuilder = Messages.newBuilder(); + if (message.role() != null) { + messageBuilder.setRole(message.role()); + } + if (message.content() != null) { + messageBuilder.setContent(message.content()); + } + builder.addMessages(messageBuilder.build()); + } + + if (parameters.inputs() != null) { + builder.setInputs(parameters.inputs()); + } + if (parameters.question() != null) { + builder.setQuestion(parameters.question()); + } + if (parameters.llmInterface() != null) { + builder.setXLlmInterface(parameters.llmInterface()); + } + + return builder.build(); + } +} diff --git a/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/translation/MlStreamResponseConverter.java b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/translation/MlStreamResponseConverter.java new file mode 100644 index 000000000..bbcd7e5b4 --- /dev/null +++ b/java-client-grpc/src/main/java/org/opensearch/client/transport/grpc/translation/MlStreamResponseConverter.java @@ -0,0 +1,75 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc.translation; + +import org.opensearch.client.transport.grpc.ml.MlInferenceResult; +import org.opensearch.client.transport.grpc.ml.MlOutput; +import org.opensearch.client.transport.grpc.ml.MlStreamChunk; +import org.opensearch.client.transport.grpc.ml.MlStreamStatus; +import org.opensearch.protobufs.DataAsMap; +import org.opensearch.protobufs.InferenceResults; +import org.opensearch.protobufs.Output; +import org.opensearch.protobufs.PredictResponse; + +/** + * Converts protobuf {@code PredictResponse} chunks received from an ML server stream + * to opensearch-java {@link MlStreamChunk} objects. + */ +public final class MlStreamResponseConverter { + + private MlStreamResponseConverter() { + // utility class + } + + /** + * Converts one protobuf {@code PredictResponse} to an {@link MlStreamChunk}. + * + * @param response the streamed protobuf response + * @return the converted chunk + */ + public static MlStreamChunk fromProto(PredictResponse response) { + MlStreamChunk.Builder builder = MlStreamChunk.builder(); + + if (response.hasStatus()) { + builder.status(MlStreamStatus.fromProto(response.getStatus())); + } + + for (InferenceResults results : response.getInferenceResultsList()) { + MlInferenceResult.Builder resultBuilder = MlInferenceResult.builder(); + for (Output output : results.getOutputList()) { + resultBuilder.addOutput(convertOutput(output)); + } + builder.addInferenceResult(resultBuilder.build()); + } + + return builder.build(); + } + + private static MlOutput convertOutput(Output output) { + MlOutput.Builder builder = MlOutput.builder(); + + if (output.hasName()) { + builder.name(output.getName()); + } + if (output.hasResult()) { + builder.result(output.getResult()); + } + if (output.hasDataAsMap()) { + DataAsMap dataAsMap = output.getDataAsMap(); + if (dataAsMap.hasContent()) { + builder.content(dataAsMap.getContent()); + } + if (dataAsMap.hasIsLast()) { + builder.isLast(dataAsMap.getIsLast()); + } + } + + return builder.build(); + } +} diff --git a/java-client-grpc/src/test/java/org/opensearch/client/transport/grpc/MlStreamingTest.java b/java-client-grpc/src/test/java/org/opensearch/client/transport/grpc/MlStreamingTest.java new file mode 100644 index 000000000..288ae0c96 --- /dev/null +++ b/java-client-grpc/src/test/java/org/opensearch/client/transport/grpc/MlStreamingTest.java @@ -0,0 +1,438 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.transport.grpc; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import io.grpc.ManagedChannel; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.stub.StreamObserver; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.opensearch.client.json.JsonpMapper; +import org.opensearch.client.json.jackson3.JacksonJsonpMapper; +import org.opensearch.client.opensearch._types.OpenSearchException; +import org.opensearch.client.transport.TransportException; +import org.opensearch.client.transport.grpc.ml.MlExecuteAgentStreamRequest; +import org.opensearch.client.transport.grpc.ml.MlOutput; +import org.opensearch.client.transport.grpc.ml.MlPredictStreamRequest; +import org.opensearch.client.transport.grpc.ml.MlStreamChunk; +import org.opensearch.client.transport.grpc.ml.MlStreamStatus; +import org.opensearch.protobufs.DataAsMap; +import org.opensearch.protobufs.InferenceResults; +import org.opensearch.protobufs.MlPredictModelStreamRequest; +import org.opensearch.protobufs.Output; +import org.opensearch.protobufs.PredictResponse; +import org.opensearch.protobufs.services.MLServiceGrpc; + +/** + * Tests for the ML gRPC streaming APIs on {@link GrpcTransport}. + */ +public class MlStreamingTest { + + private final JsonpMapper mapper = new JacksonJsonpMapper(); + + private Server server; + private ManagedChannel channel; + private GrpcTransport transport; + private FakeMlService service; + + @Before + public void setUp() throws IOException { + String name = InProcessServerBuilder.generateName(); + service = new FakeMlService(); + server = InProcessServerBuilder.forName(name).directExecutor().addService(service).build().start(); + channel = InProcessChannelBuilder.forName(name).directExecutor().build(); + transport = GrpcTransport.builder("localhost", 9400).jsonpMapper(mapper).channel(channel).build(); + } + + @After + public void tearDown() throws IOException { + if (transport != null) { + transport.close(); + } + if (server != null) { + server.shutdownNow(); + } + } + + // ═══ Predict streaming ═══════════════════════════════════════════════════════ + + @Test + public void testPredictStreamYieldsAllChunks() { + service.responses.add(runningChunk("Hello", false)); + service.responses.add(runningChunk(" world", true)); + + List content = collectContent(transport.predictModelStream(predictRequest())); + + assertEquals(2, content.size()); + assertEquals("Hello", content.get(0)); + assertEquals(" world", content.get(1)); + } + + @Test + public void testPredictStreamConvertsChunkFields() { + service.responses.add(runningChunk("chunk", true)); + + Iterator chunks = transport.predictModelStream(predictRequest()); + assertTrue(chunks.hasNext()); + MlStreamChunk chunk = chunks.next(); + + assertEquals(MlStreamStatus.RUNNING, chunk.status()); + assertEquals("chunk", chunk.content()); + assertEquals(Boolean.TRUE, chunk.isLast()); + assertEquals(1, chunk.inferenceResults().size()); + assertEquals("response", chunk.inferenceResults().get(0).output().get(0).name()); + assertFalse(chunks.hasNext()); + } + + @Test + public void testPredictStreamSendsModelIdAndParameters() { + service.responses.add(runningChunk("x", true)); + + drain( + transport.predictModelStream( + MlPredictStreamRequest.builder() + .modelId("model-1") + .parameters(p -> p.addMessage("user", "Hi").llmInterface("openai/v1/chat/completions")) + .build() + ) + ); + + MlPredictModelStreamRequest sent = service.lastPredictRequest; + assertNotNull(sent); + assertEquals("model-1", sent.getModelId()); + assertTrue(sent.hasMlPredictModelStreamRequestBody()); + assertEquals("user", sent.getMlPredictModelStreamRequestBody().getParameters().getMessages(0).getRole()); + assertEquals("Hi", sent.getMlPredictModelStreamRequestBody().getParameters().getMessages(0).getContent()); + assertEquals("openai/v1/chat/completions", sent.getMlPredictModelStreamRequestBody().getParameters().getXLlmInterface()); + } + + @Test + public void testPredictStreamOmitsBodyWhenNoParameters() { + service.responses.add(runningChunk("x", true)); + + drain(transport.predictModelStream(MlPredictStreamRequest.builder().modelId("model-1").build())); + + assertFalse(service.lastPredictRequest.hasMlPredictModelStreamRequestBody()); + } + + @Test + public void testPredictStreamIsLazy() { + service.responses.add(runningChunk("x", true)); + + transport.predictModelStream(predictRequest()); + + // No terminal operation — the RPC must not have been issued yet. + assertNull(service.lastPredictRequest); + } + + @Test + public void testEmptyStreamYieldsNoChunks() { + assertFalse(transport.predictModelStream(predictRequest()).hasNext()); + } + + // ═══ Agent streaming ═════════════════════════════════════════════════════════ + + @Test + public void testExecuteAgentStreamYieldsChunks() { + service.responses.add(runningChunk("agent says", true)); + + List content = collectContent( + transport.executeAgentStream( + MlExecuteAgentStreamRequest.builder().agentId("agent-1").parameters(p -> p.question("What is OpenSearch?")).build() + ) + ); + + assertEquals(1, content.size()); + assertEquals("agent says", content.get(0)); + assertEquals("agent-1", service.lastAgentRequest.getAgentId()); + assertEquals("What is OpenSearch?", service.lastAgentRequest.getMlExecuteAgentStreamRequestBody().getParameters().getQuestion()); + } + + /** + * Agent chunks lead with {@code memory_id} and {@code parent_interaction_id} outputs + * that carry only a {@code result}, so the content shortcuts must select the + * {@code response} output by name rather than taking the first output positionally. + */ + @Test + public void testAgentChunkWithLeadingMetadataOutputsExposesContent() { + service.responses.add(agentChunk("mem-1", "parent-1", "tool call", false)); + + Iterator chunks = transport.executeAgentStream( + MlExecuteAgentStreamRequest.builder().agentId("agent-1").parameters(p -> p.question("Q")).build() + ); + + assertTrue(chunks.hasNext()); + MlStreamChunk chunk = chunks.next(); + + assertEquals("tool call", chunk.content()); + assertEquals(Boolean.FALSE, chunk.isLast()); + + // The metadata outputs remain reachable for conversation continuity. + List outputs = chunk.inferenceResults().get(0).output(); + assertEquals(3, outputs.size()); + assertEquals("memory_id", outputs.get(0).name()); + assertEquals("mem-1", outputs.get(0).result()); + assertEquals("parent_interaction_id", outputs.get(1).name()); + assertEquals("parent-1", outputs.get(1).result()); + } + + /** + * A chunk whose only outputs are metadata (no content anywhere) must report null + * rather than picking up a metadata output's absent content. + */ + @Test + public void testMetadataOnlyChunkHasNullContent() { + service.responses.add( + PredictResponse.newBuilder() + .addInferenceResults( + InferenceResults.newBuilder().addOutput(Output.newBuilder().setName("memory_id").setResult("mem-1").build()).build() + ) + .build() + ); + + Iterator chunks = transport.executeAgentStream(MlExecuteAgentStreamRequest.builder().agentId("agent-1").build()); + + MlStreamChunk chunk = chunks.next(); + assertNull(chunk.content()); + assertNull(chunk.isLast()); + } + + /** + * When the {@code response} output carries {@code is_last} but no content, both + * shortcuts must still read from it — not fall through to some other output — so they + * never report a content and an is_last flag that came from different outputs. + */ + @Test + public void testContentAndIsLastComeFromTheSameOutput() { + service.responses.add( + PredictResponse.newBuilder() + .addInferenceResults( + InferenceResults.newBuilder() + // A non-response output that does have content. + .addOutput( + Output.newBuilder() + .setName("other") + .setDataAsMap(DataAsMap.newBuilder().setContent("not the answer").build()) + .build() + ) + // The response output: terminal marker, no content of its own. + .addOutput( + Output.newBuilder().setName("response").setDataAsMap(DataAsMap.newBuilder().setIsLast(true).build()).build() + ) + .build() + ) + .build() + ); + + MlStreamChunk chunk = transport.executeAgentStream(MlExecuteAgentStreamRequest.builder().agentId("agent-1").build()).next(); + + assertEquals(Boolean.TRUE, chunk.isLast()); + assertNull(chunk.content()); + } + + @Test + public void testExecuteAgentStreamOmitsBodyWhenNoParameters() { + service.responses.add(runningChunk("x", true)); + + drain(transport.executeAgentStream(MlExecuteAgentStreamRequest.builder().agentId("agent-1").build())); + + assertFalse(service.lastAgentRequest.hasMlExecuteAgentStreamRequestBody()); + } + + // ═══ Error handling ══════════════════════════════════════════════════════════ + + @Test + public void testNotFoundBecomesOpenSearchException() { + service.error = new StatusRuntimeException(Status.NOT_FOUND.withDescription("no such model")); + + try { + drain(transport.predictModelStream(predictRequest())); + fail("Should throw"); + } catch (OpenSearchException e) { + assertEquals(404, e.status()); + } + } + + @Test + public void testUnavailableBecomesUncheckedTransportException() { + service.error = new StatusRuntimeException(Status.UNAVAILABLE.withDescription("down")); + + try { + drain(transport.predictModelStream(predictRequest())); + fail("Should throw"); + } catch (UncheckedIOException e) { + assertTrue(e.getCause() instanceof TransportException); + } + } + + @Test + public void testUnimplementedBecomesUnsupportedOperation() { + service.error = new StatusRuntimeException(Status.UNIMPLEMENTED.withDescription("not built")); + + try { + drain(transport.executeAgentStream(MlExecuteAgentStreamRequest.builder().agentId("a").build())); + fail("Should throw"); + } catch (UnsupportedOperationException e) { + assertTrue(e.getMessage().contains("not built")); + } + } + + @Test + public void testErrorAfterChunksPreservesDeliveredChunks() { + service.responses.add(runningChunk("partial", false)); + service.error = new StatusRuntimeException(Status.INTERNAL.withDescription("boom")); + + Iterator chunks = transport.predictModelStream(predictRequest()); + assertEquals("partial", chunks.next().content()); + + try { + chunks.hasNext(); + fail("Should throw"); + } catch (UncheckedIOException e) { + assertTrue(e.getCause() instanceof TransportException); + } + } + + @Test + public void testStreamingWithoutChannelThrows() throws IOException { + GrpcTransport noChannel = new GrpcTransport(null, mapper, GrpcTransportOptions.defaults(), null); + try { + noChannel.predictModelStream(predictRequest()).hasNext(); + fail("Should throw"); + } catch (IllegalStateException e) { + assertTrue(e.getMessage().contains("channel is not available")); + } finally { + noChannel.close(); + } + } + + // ═══ Request validation ══════════════════════════════════════════════════════ + + @Test(expected = IllegalArgumentException.class) + public void testPredictRequestRequiresModelId() { + MlPredictStreamRequest.builder().build(); + } + + @Test(expected = IllegalArgumentException.class) + public void testAgentRequestRequiresAgentId() { + MlExecuteAgentStreamRequest.builder().build(); + } + + // ═══ Helpers ═════════════════════════════════════════════════════════════════ + + private static MlPredictStreamRequest predictRequest() { + return MlPredictStreamRequest.builder().modelId("model-1").parameters(p -> p.addMessage("user", "Hi")).build(); + } + + /** + * An agent-shaped chunk: {@code memory_id} and {@code parent_interaction_id} outputs + * carrying only a {@code result}, followed by the {@code response} output that holds + * the actual payload. This is what {@code MLChatAgentRunner} emits. + */ + private static PredictResponse agentChunk(String memoryId, String parentInteractionId, String content, boolean isLast) { + return PredictResponse.newBuilder() + .addInferenceResults( + InferenceResults.newBuilder() + .addOutput(Output.newBuilder().setName("memory_id").setResult(memoryId).build()) + .addOutput(Output.newBuilder().setName("parent_interaction_id").setResult(parentInteractionId).build()) + .addOutput( + Output.newBuilder() + .setName("response") + .setDataAsMap(DataAsMap.newBuilder().setContent(content).setIsLast(isLast).build()) + .build() + ) + .build() + ) + .build(); + } + + private static PredictResponse runningChunk(String content, boolean isLast) { + return PredictResponse.newBuilder() + .setStatus(org.opensearch.protobufs.Status.STATUS_RUNNING) + .addInferenceResults( + InferenceResults.newBuilder() + .addOutput( + Output.newBuilder() + .setName("response") + .setDataAsMap(DataAsMap.newBuilder().setContent(content).setIsLast(isLast).build()) + .build() + ) + .build() + ) + .build(); + } + + private static List collectContent(Iterator chunks) { + List content = new ArrayList<>(); + while (chunks.hasNext()) { + content.add(chunks.next().content()); + } + return content; + } + + private static void drain(Iterator chunks) { + while (chunks.hasNext()) { + chunks.next(); + } + } + + /** + * In-process MLService that replays a canned list of chunks, optionally followed by + * an error, and records the requests it received. + */ + private static final class FakeMlService extends MLServiceGrpc.MLServiceImplBase { + final List responses = new ArrayList<>(); + StatusRuntimeException error; + MlPredictModelStreamRequest lastPredictRequest; + org.opensearch.protobufs.MlExecuteAgentStreamRequest lastAgentRequest; + + @Override + public void predictModelStream(MlPredictModelStreamRequest request, StreamObserver observer) { + lastPredictRequest = request; + replay(observer); + } + + @Override + public void executeAgentStream( + org.opensearch.protobufs.MlExecuteAgentStreamRequest request, + StreamObserver observer + ) { + lastAgentRequest = request; + replay(observer); + } + + private void replay(StreamObserver observer) { + for (PredictResponse response : responses) { + observer.onNext(response); + } + if (error != null) { + observer.onError(error); + } else { + observer.onCompleted(); + } + } + } +} diff --git a/java-client-grpc/src/test/java/org/opensearch/client/transport/grpc/translation/TranslationTest.java b/java-client-grpc/src/test/java/org/opensearch/client/transport/grpc/translation/TranslationTest.java index a89355141..ef18cbdfa 100644 --- a/java-client-grpc/src/test/java/org/opensearch/client/transport/grpc/translation/TranslationTest.java +++ b/java-client-grpc/src/test/java/org/opensearch/client/transport/grpc/translation/TranslationTest.java @@ -30,14 +30,27 @@ import org.opensearch.client.opensearch.core.bulk.BulkOperation; import org.opensearch.client.opensearch.core.bulk.OperationType; import org.opensearch.client.transport.TransportException; +import org.opensearch.client.transport.grpc.ml.MlExecuteAgentStreamRequest; +import org.opensearch.client.transport.grpc.ml.MlMessage; +import org.opensearch.client.transport.grpc.ml.MlParameters; +import org.opensearch.client.transport.grpc.ml.MlPredictStreamRequest; +import org.opensearch.client.transport.grpc.ml.MlStreamChunk; +import org.opensearch.client.transport.grpc.ml.MlStreamStatus; +import org.opensearch.protobufs.DataAsMap; import org.opensearch.protobufs.ErrorCause; +import org.opensearch.protobufs.InferenceResults; import org.opensearch.protobufs.Item; +import org.opensearch.protobufs.Messages; +import org.opensearch.protobufs.Output; +import org.opensearch.protobufs.Parameters; +import org.opensearch.protobufs.PredictResponse; import org.opensearch.protobufs.ResponseItem; import org.opensearch.protobufs.ShardInfo; /** * Combined unit tests for the gRPC translation layer: - * BulkRequestConverter, BulkResponseConverter, FieldMappingUtil, GrpcStatusConverter. + * BulkRequestConverter, BulkResponseConverter, FieldMappingUtil, GrpcStatusConverter, + * MlStreamRequestConverter, MlStreamResponseConverter. */ public class TranslationTest { @@ -394,6 +407,175 @@ public void testIngestTook() { assertEquals(Long.valueOf(25L), BulkResponseConverter.fromProto(proto).ingestTook()); } + // ═══ MlStreamRequestConverter Tests ══════════════════════════════════════════ + + @Test + public void testPredictRequestModelIdOnly() { + org.opensearch.protobufs.MlPredictModelStreamRequest proto = MlStreamRequestConverter.toProto( + MlPredictStreamRequest.builder().modelId("model-1").build() + ); + assertEquals("model-1", proto.getModelId()); + assertFalse(proto.hasMlPredictModelStreamRequestBody()); + } + + @Test + public void testPredictRequestWithMessages() { + org.opensearch.protobufs.MlPredictModelStreamRequest proto = MlStreamRequestConverter.toProto( + MlPredictStreamRequest.builder() + .modelId("model-1") + .parameters(p -> p.addMessage("system", "Be brief").addMessage("user", "Hi")) + .build() + ); + Parameters params = proto.getMlPredictModelStreamRequestBody().getParameters(); + assertEquals(2, params.getMessagesCount()); + assertEquals("system", params.getMessages(0).getRole()); + assertEquals("Be brief", params.getMessages(0).getContent()); + assertEquals("user", params.getMessages(1).getRole()); + } + + @Test + public void testPredictRequestOnlySetsProvidedParameters() { + org.opensearch.protobufs.MlPredictModelStreamRequest proto = MlStreamRequestConverter.toProto( + MlPredictStreamRequest.builder().modelId("model-1").parameters(p -> p.inputs("raw text")).build() + ); + Parameters params = proto.getMlPredictModelStreamRequestBody().getParameters(); + assertTrue(params.hasInputs()); + assertEquals("raw text", params.getInputs()); + assertFalse(params.hasQuestion()); + assertFalse(params.hasXLlmInterface()); + assertEquals(0, params.getMessagesCount()); + } + + @Test + public void testPredictRequestEmptyParametersStillSetsBody() { + // An explicitly-set but empty parameters object must still produce a body, + // distinguishing "no parameters" from "empty parameters". + org.opensearch.protobufs.MlPredictModelStreamRequest proto = MlStreamRequestConverter.toProto( + MlPredictStreamRequest.builder().modelId("model-1").parameters(MlParameters.builder().build()).build() + ); + assertTrue(proto.hasMlPredictModelStreamRequestBody()); + } + + @Test + public void testAgentRequestWithQuestion() { + org.opensearch.protobufs.MlExecuteAgentStreamRequest proto = MlStreamRequestConverter.toProto( + MlExecuteAgentStreamRequest.builder().agentId("agent-1").parameters(p -> p.question("What is OpenSearch?")).build() + ); + assertEquals("agent-1", proto.getAgentId()); + assertEquals("What is OpenSearch?", proto.getMlExecuteAgentStreamRequestBody().getParameters().getQuestion()); + } + + @Test + public void testAgentRequestAgentIdOnly() { + org.opensearch.protobufs.MlExecuteAgentStreamRequest proto = MlStreamRequestConverter.toProto( + MlExecuteAgentStreamRequest.builder().agentId("agent-1").build() + ); + assertFalse(proto.hasMlExecuteAgentStreamRequestBody()); + } + + @Test + public void testMessageWithOnlyContent() { + org.opensearch.protobufs.MlPredictModelStreamRequest proto = MlStreamRequestConverter.toProto( + MlPredictStreamRequest.builder() + .modelId("m") + .parameters(p -> p.addMessage(MlMessage.builder().content("no role").build())) + .build() + ); + Messages message = proto.getMlPredictModelStreamRequestBody().getParameters().getMessages(0); + assertFalse(message.hasRole()); + assertEquals("no role", message.getContent()); + } + + // ═══ MlStreamResponseConverter Tests ═════════════════════════════════════════ + + @Test + public void testChunkWithDataAsMap() { + MlStreamChunk chunk = MlStreamResponseConverter.fromProto( + PredictResponse.newBuilder() + .setStatus(org.opensearch.protobufs.Status.STATUS_RUNNING) + .addInferenceResults( + InferenceResults.newBuilder() + .addOutput( + Output.newBuilder() + .setName("response") + .setDataAsMap(DataAsMap.newBuilder().setContent("Hello").setIsLast(false).build()) + .build() + ) + .build() + ) + .build() + ); + assertEquals(MlStreamStatus.RUNNING, chunk.status()); + assertEquals("Hello", chunk.content()); + assertEquals(Boolean.FALSE, chunk.isLast()); + assertEquals("response", chunk.inferenceResults().get(0).output().get(0).name()); + } + + @Test + public void testChunkWithResultField() { + MlStreamChunk chunk = MlStreamResponseConverter.fromProto( + PredictResponse.newBuilder() + .addInferenceResults(InferenceResults.newBuilder().addOutput(Output.newBuilder().setResult("plain text").build()).build()) + .build() + ); + assertEquals("plain text", chunk.inferenceResults().get(0).output().get(0).result()); + assertNull(chunk.content()); + } + + @Test + public void testChunkWithoutStatus() { + assertNull(MlStreamResponseConverter.fromProto(PredictResponse.newBuilder().build()).status()); + } + + @Test + public void testEmptyChunkHasNoContent() { + MlStreamChunk chunk = MlStreamResponseConverter.fromProto(PredictResponse.newBuilder().build()); + assertTrue(chunk.inferenceResults().isEmpty()); + assertNull(chunk.content()); + assertNull(chunk.isLast()); + } + + @Test + public void testChunkWithEmptyOutputList() { + MlStreamChunk chunk = MlStreamResponseConverter.fromProto( + PredictResponse.newBuilder().addInferenceResults(InferenceResults.newBuilder().build()).build() + ); + assertEquals(1, chunk.inferenceResults().size()); + assertNull(chunk.content()); + assertNull(chunk.isLast()); + } + + @Test + public void testChunkWithMultipleOutputs() { + MlStreamChunk chunk = MlStreamResponseConverter.fromProto( + PredictResponse.newBuilder() + .addInferenceResults( + InferenceResults.newBuilder() + .addOutput(Output.newBuilder().setDataAsMap(DataAsMap.newBuilder().setContent("first").build()).build()) + .addOutput(Output.newBuilder().setDataAsMap(DataAsMap.newBuilder().setContent("second").build()).build()) + .build() + ) + .build() + ); + assertEquals(2, chunk.inferenceResults().get(0).output().size()); + assertEquals("first", chunk.content()); + } + + @Test + public void testAllStatusValuesMap() { + assertEquals(MlStreamStatus.CANCELLED, MlStreamStatus.fromProto(org.opensearch.protobufs.Status.STATUS_CANCELLED)); + assertEquals(MlStreamStatus.COMPLETED, MlStreamStatus.fromProto(org.opensearch.protobufs.Status.STATUS_COMPLETED)); + assertEquals( + MlStreamStatus.COMPLETED_WITH_ERROR, + MlStreamStatus.fromProto(org.opensearch.protobufs.Status.STATUS_COMPLETED_WITH_ERROR) + ); + assertEquals(MlStreamStatus.CREATED, MlStreamStatus.fromProto(org.opensearch.protobufs.Status.STATUS_CREATED)); + assertEquals(MlStreamStatus.FAILED, MlStreamStatus.fromProto(org.opensearch.protobufs.Status.STATUS_FAILED)); + assertEquals(MlStreamStatus.RUNNING, MlStreamStatus.fromProto(org.opensearch.protobufs.Status.STATUS_RUNNING)); + assertEquals(MlStreamStatus.UNSPECIFIED, MlStreamStatus.fromProto(org.opensearch.protobufs.Status.STATUS_UNSPECIFIED)); + assertEquals(MlStreamStatus.UNSPECIFIED, MlStreamStatus.fromProto(null)); + } + // ═══ Helper ══════════════════════════════════════════════════════════════════ public static class Doc { diff --git a/java-client-grpc/src/test/java11/org/opensearch/client/opensearch/integTest/grpc/GrpcMlStreamIT.java b/java-client-grpc/src/test/java11/org/opensearch/client/opensearch/integTest/grpc/GrpcMlStreamIT.java new file mode 100644 index 000000000..53d21a6f0 --- /dev/null +++ b/java-client-grpc/src/test/java11/org/opensearch/client/opensearch/integTest/grpc/GrpcMlStreamIT.java @@ -0,0 +1,141 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.client.opensearch.integTest.grpc; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import org.junit.After; +import org.junit.Assume; +import org.junit.Test; +import org.opensearch.client.json.jackson3.JacksonJsonpMapper; +import org.opensearch.client.transport.grpc.GrpcTransport; +import org.opensearch.client.transport.grpc.GrpcTransportOptions; +import org.opensearch.client.transport.grpc.ml.MlExecuteAgentStreamRequest; +import org.opensearch.client.transport.grpc.ml.MlPredictStreamRequest; +import org.opensearch.client.transport.grpc.ml.MlStreamChunk; + +/** + * End-to-end integration tests for the ML gRPC streaming APIs + */ +public class GrpcMlStreamIT extends AbstractGrpcIT { + + private GrpcTransport transport; + + @After + public void closeTransport() throws IOException { + if (transport != null) { + transport.close(); + transport = null; + } + } + + // ─── Tests ─────────────────────────────────────────────────────────────────── + + /** + * Streams a chat completion from a deployed model and verifies that chunks arrive + * and carry incremental content. + */ + @Test + public void testPredictModelStream() throws IOException { + assumeGrpcSupported(); + String modelId = requireId("tests.ml.model.id", "OPENSEARCH_ML_MODEL_ID"); + + Iterator chunks = grpcTransport().predictModelStream( + MlPredictStreamRequest.builder() + .modelId(modelId) + .parameters(p -> p.addMessage("user", "What is OpenSearch? Answer in one sentence.")) + .build() + ); + + List received = drain(chunks); + + assertFalse("Should receive at least one chunk", received.isEmpty()); + assertTrue("At least one chunk should carry content", received.stream().anyMatch(c -> c.content() != null)); + + // Every chunk must convert cleanly, even status-only ones. + for (MlStreamChunk chunk : received) { + assertNotNull(chunk.inferenceResults()); + } + } + + /** + * Streams an agent execution and verifies that chunks arrive. + */ + @Test + public void testExecuteAgentStream() throws IOException { + assumeGrpcSupported(); + String agentId = requireId("tests.ml.agent.id", "OPENSEARCH_ML_AGENT_ID"); + + Iterator chunks = grpcTransport().executeAgentStream( + MlExecuteAgentStreamRequest.builder().agentId(agentId).parameters(p -> p.question("What is OpenSearch?")).build() + ); + + assertFalse("Should receive at least one chunk", drain(chunks).isEmpty()); + } + + /** + * Verifies that a nonexistent model surfaces a server error during iteration rather + * than hanging or returning an empty stream. + */ + @Test + public void testPredictStreamUnknownModelFails() throws IOException { + assumeGrpcSupported(); + + Iterator chunks = grpcTransport().predictModelStream( + MlPredictStreamRequest.builder().modelId("nonexistent-model-id").parameters(p -> p.addMessage("user", "Hi")).build() + ); + + try { + drain(chunks); + // Some server configurations report the failure as a terminal chunk rather + // than a gRPC error status; either behavior is acceptable here — what matters + // is that the call completes instead of hanging. + } catch (RuntimeException e) { + assertNotNull(e.getMessage()); + } + } + + // ─── Helpers ───────────────────────────────────────────────────────────────── + + private GrpcTransport grpcTransport() { + if (transport == null) { + transport = GrpcTransport.builder(getGrpcHost(), getGrpcPort()) + .jsonpMapper(new JacksonJsonpMapper()) + .grpcOptions(GrpcTransportOptions.builder().maxRetries(2).build()) + .build(); + } + return transport; + } + + private static String requireId(String systemProperty, String environmentVariable) { + String value = System.getProperty(systemProperty); + if (value == null || value.isBlank()) { + value = System.getenv(environmentVariable); + } + Assume.assumeTrue( + "Set -D" + systemProperty + " or " + environmentVariable + " to run ML streaming integration tests", + value != null && !value.isBlank() + ); + return value; + } + + private static List drain(Iterator chunks) { + List received = new ArrayList<>(); + while (chunks.hasNext()) { + received.add(chunks.next()); + } + return received; + } +}