Merge branch 'ml-inference-unified-api-elastic' of github.com:elastic/elasticsearch into ml-inference-unified-api-elastic

Author: jonathan-buttner

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/UnifiedCompletionRequest.java

@@ -16,6 +16,8 @@
 import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ObjectParser;
 import org.elasticsearch.xcontent.ParseField;
+import org.elasticsearch.xcontent.ToXContent;
+import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParseException;
 import org.elasticsearch.xcontent.XContentParser;
 
@@ -39,6 +41,8 @@ public record UnifiedCompletionRequest(
     @Nullable String user
 ) implements Writeable {
 
+    public sealed interface Content extends NamedWriteable permits ContentObjects, ContentString {}
+
     @SuppressWarnings("unchecked")
     static final ConstructingObjectParser<UnifiedCompletionRequest, Void> PARSER = new ConstructingObjectParser<>(
         InferenceAction.NAME,
@@ -153,8 +157,6 @@ public void writeTo(StreamOutput out) throws IOException {
         }
     }
 
-    public sealed interface Content extends NamedWriteable permits ContentObjects, ContentString {}
-
     public record ContentObjects(List<ContentObject> contentObjects) implements Content, Writeable {
 
         public static final String NAME = "content_objects";
@@ -194,6 +196,7 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeString(text);
             out.writeString(type);
         }
+
     }
 
     public record ContentString(String content) implements Content, NamedWriteable {
@@ -217,6 +220,10 @@ public void writeTo(StreamOutput out) throws IOException {
         public String getWriteableName() {
             return NAME;
         }
+
+        public void toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
+            builder.value(content);
+        }
     }
 
     public record ToolCall(String id, FunctionField function, String type) implements Writeable {
@@ -432,7 +439,7 @@ public void writeTo(StreamOutput out) throws IOException {
         public record FunctionField(
             @Nullable String description,
             String name,
-            @Nullable Map<String, Object> parameters,
+            @Nullable Map<String, Object> parameters, // TODO can we parse this as a string?
             @Nullable Boolean strict
         ) implements Writeable {
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/CompletionInputs.java

@@ -15,7 +15,7 @@
 import org.elasticsearch.xpack.inference.external.http.retry.RequestSender;
 import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
 import org.elasticsearch.xpack.inference.external.openai.OpenAiChatCompletionResponseHandler;
-import org.elasticsearch.xpack.inference.external.request.openai.OpenAiChatCompletionRequest;
+import org.elasticsearch.xpack.inference.external.request.openai.OpenAiUnifiedChatCompletionRequest;
 import org.elasticsearch.xpack.inference.external.response.openai.OpenAiChatCompletionResponseEntity;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModel;
 
@@ -35,7 +35,7 @@ public static OpenAiCompletionRequestManager of(OpenAiChatCompletionModel model,
     private final OpenAiChatCompletionModel model;
 
     private OpenAiCompletionRequestManager(OpenAiChatCompletionModel model, ThreadPool threadPool) {
-        super(threadPool, model, OpenAiChatCompletionRequest::buildDefaultUri);
+        super(threadPool, model, OpenAiUnifiedChatCompletionRequest::buildDefaultUri);
         this.model = Objects.requireNonNull(model);
     }
 
@@ -46,10 +46,11 @@ public void execute(
         Supplier<Boolean> hasRequestCompletedFunction,
         ActionListener<InferenceServiceResults> listener
     ) {
-        var docsOnly = DocumentsOnlyInput.of(inferenceInputs);
-        var docsInput = docsOnly.getInputs();
-        var stream = docsOnly.stream();
-        OpenAiChatCompletionRequest request = new OpenAiChatCompletionRequest(docsInput, model, stream);
+
+        OpenAiUnifiedChatCompletionRequest request = new OpenAiUnifiedChatCompletionRequest(
+            UnifiedChatInput.of(inferenceInputs).getRequestEntity(),
+            model
+        );
 
         execute(new ExecutableInferenceRequest(requestSender, logger, request, HANDLER, hasRequestCompletedFunction, listener));
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiCompletionRequestManager.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.http.sender;
+
+import org.elasticsearch.xpack.inference.external.request.openai.OpenAiUnifiedChatCompletionRequestEntity;
+
+import java.util.Objects;
+
+public class UnifiedChatInput extends InferenceInputs {
+
+    public static UnifiedChatInput of(InferenceInputs inferenceInputs) {
+
+        if (inferenceInputs instanceof DocumentsOnlyInput docsOnly) {
+            return new UnifiedChatInput(new OpenAiUnifiedChatCompletionRequestEntity(docsOnly));
+        } else if (inferenceInputs instanceof UnifiedChatInput == false) {
+            throw createUnsupportedTypeException(inferenceInputs);
+        }
+
+        return (UnifiedChatInput) inferenceInputs;
+    }
+
+    public OpenAiUnifiedChatCompletionRequestEntity getRequestEntity() {
+        return requestEntity;
+    }
+
+    private final OpenAiUnifiedChatCompletionRequestEntity requestEntity;
+
+    public UnifiedChatInput(OpenAiUnifiedChatCompletionRequestEntity requestEntity) {
+        this.requestEntity = Objects.requireNonNull(requestEntity);
+    }
+
+    public boolean stream() {
+        return requestEntity.isStream();
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiUnifiedCompletionRequestManager.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.request;
+
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xpack.core.inference.action.UnifiedCompletionRequest;
+
+import java.util.List;
+
+public record UnifiedRequest(
+    List<UnifiedCompletionRequest.Message> messages,
+    @Nullable String model,
+    @Nullable Long maxCompletionTokens,
+    @Nullable Integer n,
+    @Nullable UnifiedCompletionRequest.Stop stop,
+    @Nullable Float temperature,
+    @Nullable UnifiedCompletionRequest.ToolChoice toolChoice,
+    @Nullable List<UnifiedCompletionRequest.Tool> tool,
+    @Nullable Float topP,
+    @Nullable String user,
+    boolean stream
+) {}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/UnifiedChatInput.java

@@ -21,35 +21,28 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
-import java.util.List;
 import java.util.Objects;
 
 import static org.elasticsearch.xpack.inference.external.request.RequestUtils.createAuthBearerHeader;
 import static org.elasticsearch.xpack.inference.external.request.openai.OpenAiUtils.createOrgHeader;
 
-public class OpenAiChatCompletionRequest implements OpenAiRequest {
+public class OpenAiUnifiedChatCompletionRequest implements OpenAiRequest {
 
     private final OpenAiAccount account;
-    private final List<String> input;
+    private final OpenAiUnifiedChatCompletionRequestEntity requestEntity;
     private final OpenAiChatCompletionModel model;
-    private final boolean stream;
 
-    public OpenAiChatCompletionRequest(List<String> input, OpenAiChatCompletionModel model, boolean stream) {
-        this.account = OpenAiAccount.of(model, OpenAiChatCompletionRequest::buildDefaultUri);
-        this.input = Objects.requireNonNull(input);
+    public OpenAiUnifiedChatCompletionRequest(OpenAiUnifiedChatCompletionRequestEntity requestEntity, OpenAiChatCompletionModel model) {
+        this.account = OpenAiAccount.of(model, OpenAiUnifiedChatCompletionRequest::buildDefaultUri);
+        this.requestEntity = Objects.requireNonNull(requestEntity);
         this.model = Objects.requireNonNull(model);
-        this.stream = stream;
     }
 
     @Override
     public HttpRequest createHttpRequest() {
         HttpPost httpPost = new HttpPost(account.uri());
 
-        ByteArrayEntity byteEntity = new ByteArrayEntity(
-            Strings.toString(
-                new OpenAiChatCompletionRequestEntity(input, model.getServiceSettings().modelId(), model.getTaskSettings().user(), stream)
-            ).getBytes(StandardCharsets.UTF_8)
-        );
+        ByteArrayEntity byteEntity = new ByteArrayEntity(Strings.toString(requestEntity).getBytes(StandardCharsets.UTF_8));
         httpPost.setEntity(byteEntity);
 
         httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType());
@@ -87,7 +80,7 @@ public String getInferenceEntityId() {
 
     @Override
     public boolean isStreaming() {
-        return stream;
+        return requestEntity.isStream();
     }
 
     public static URI buildDefaultUri() throws URISyntaxException {