separate out unified request and combine inputs

Author: Max Hniebergall

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/CompletionInputs.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/request/UnifiedRequest.java

@@ -7,132 +7,99 @@
 
 package org.elasticsearch.xpack.inference.external.request.openai;
 
-import org.elasticsearch.common.Strings;
-import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.core.inference.action.UnifiedCompletionRequest;
 import org.elasticsearch.xpack.inference.external.http.sender.DocumentsOnlyInput;
+import org.elasticsearch.xpack.inference.external.request.UnifiedRequest;
 
 import java.io.IOException;
 import java.util.List;
-import java.util.Objects;
 
 public class OpenAiUnifiedChatCompletionRequestEntity implements ToXContentObject {
 
-    private static final String MESSAGES_FIELD = "messages";
-    private static final String MODEL_FIELD = "model";
-
+    public static final String NAME_FIELD = "name";
+    public static final String TOOL_CALL_ID_FIELD = "tool_call_id";
+    public static final String TOOL_CALLS_FIELD = "tool_calls";
+    public static final String ID_FIELD = "id";
+    public static final String FUNCTION_FIELD = "function";
+    public static final String ARGUMENTS_FIELD = "arguments";
+    public static final String DESCRIPTION_FIELD = "description";
+    public static final String PARAMETERS_FIELD = "parameters";
+    public static final String STRICT_FIELD = "strict";
+    public static final String TOP_P_FIELD = "top_p";
+    public static final String USER_FIELD = "user";
+    public static final String STREAM_FIELD = "stream";
     private static final String NUMBER_OF_RETURNED_CHOICES_FIELD = "n";
-
+    private static final String MODEL_FIELD = "model";
+    public static final String MESSAGES_FIELD = "messages";
     private static final String ROLE_FIELD = "role";
-    private static final String USER_FIELD = "user";
     private static final String CONTENT_FIELD = "content";
-    private static final String STREAM_FIELD = "stream";
     private static final String MAX_COMPLETION_TOKENS_FIELD = "max_completion_tokens";
     private static final String STOP_FIELD = "stop";
     private static final String TEMPERATURE_FIELD = "temperature";
     private static final String TOOL_CHOICE_FIELD = "tool_choice";
     private static final String TOOL_FIELD = "tool";
-    private static final String TOP_P_FIELD = "top_p";
+    private static final String TEXT_FIELD = "text";
+    private static final String TYPE_FIELD = "type";
 
-    private final String user;
+    private final UnifiedRequest unifiedRequest;
 
-    public boolean isStream() {
-        return stream;
+    public OpenAiUnifiedChatCompletionRequestEntity(UnifiedRequest unifiedRequest) {
+        this.unifiedRequest = unifiedRequest;
     }
 
-    private final boolean stream;
-    private final Long maxCompletionTokens;
-    private final Integer n;
-    private final UnifiedCompletionRequest.Stop stop;
-    private final Float temperature;
-    private final UnifiedCompletionRequest.ToolChoice toolChoice;
-    private final List<UnifiedCompletionRequest.Tool> tool;
-    private final Float topP;
-    private final List<UnifiedCompletionRequest.Message> messages;
-    private final String model;
-
     public OpenAiUnifiedChatCompletionRequestEntity(DocumentsOnlyInput input) {
-        this(convertDocumentsOnlyInputToMessages(input), null, null, null, null, null, null, null, null, null);
+        this(new UnifiedRequest(convertDocumentsOnlyInputToMessages(input), null, null, null, null, null, null, null, null, null, true));
     }
 
     private static List<UnifiedCompletionRequest.Message> convertDocumentsOnlyInputToMessages(DocumentsOnlyInput input) {
         return input.getInputs()
             .stream()
-            .map(doc -> new UnifiedCompletionRequest.Message(new UnifiedCompletionRequest.ContentString(doc), "user", null, null, null))
+            .map(doc -> new UnifiedCompletionRequest.Message(new UnifiedCompletionRequest.ContentString(doc), USER_FIELD, null, null, null))
             .toList();
     }
 
-    public OpenAiUnifiedChatCompletionRequestEntity(
-        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
-    ) {
-        Objects.requireNonNull(messages);
-        Objects.requireNonNull(model);
-
-        this.user = user;
-        this.stream = true; // always stream in unified API
-        this.maxCompletionTokens = maxCompletionTokens;
-        this.n = n;
-        this.stop = stop;
-        this.temperature = temperature;
-        this.toolChoice = toolChoice;
-        this.tool = tool;
-        this.topP = topP;
-        this.messages = messages;
-        this.model = model;
-
-    }
-
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
         builder.startArray(MESSAGES_FIELD);
         {
-            for (UnifiedCompletionRequest.Message message : messages) {
+            for (UnifiedCompletionRequest.Message message : unifiedRequest.messages()) {
                 builder.startObject();
                 {
                     switch (message.content()) {
                         case UnifiedCompletionRequest.ContentString contentString -> builder.field(CONTENT_FIELD, contentString.content());
                         case UnifiedCompletionRequest.ContentObjects contentObjects -> {
                             for (UnifiedCompletionRequest.ContentObject contentObject : contentObjects.contentObjects()) {
                                 builder.startObject(CONTENT_FIELD);
-                                builder.field("text", contentObject.text());
-                                builder.field("type", contentObject.type());
+                                builder.field(TEXT_FIELD, contentObject.text());
+                                builder.field(TYPE_FIELD, contentObject.type());
                                 builder.endObject();
                             }
                         }
                     }
 
                     builder.field(ROLE_FIELD, message.role());
                     if (message.name() != null) {
-                        builder.field("name", message.name());
+                        builder.field(NAME_FIELD, message.name());
                     }
                     if (message.toolCallId() != null) {
-                        builder.field("tool_call_id", message.toolCallId());
+                        builder.field(TOOL_CALL_ID_FIELD, message.toolCallId());
                     }
                     if (message.toolCalls() != null) {
-                        builder.startArray("tool_calls");
+                        builder.startArray(TOOL_CALLS_FIELD);
                         for (UnifiedCompletionRequest.ToolCall toolCall : message.toolCalls()) {
                             builder.startObject();
                             {
-                                builder.field("id", toolCall.id());
-                                builder.startObject("function");
+                                builder.field(ID_FIELD, toolCall.id());
+                                builder.startObject(FUNCTION_FIELD);
                                 {
-                                    builder.field("arguments", toolCall.function().arguments());
-                                    builder.field("name", toolCall.function().name());
+                                    builder.field(ARGUMENTS_FIELD, toolCall.function().arguments());
+                                    builder.field(NAME_FIELD, toolCall.function().name());
                                 }
                                 builder.endObject();
-                                builder.field("type", toolCall.type());
+                                builder.field(TYPE_FIELD, toolCall.type());
                             }
                             builder.endObject();
                         }
@@ -144,65 +111,69 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         }
         builder.endArray();
 
-        if (model != null) {
-            builder.field(MODEL_FIELD, model);
+        if (unifiedRequest.model() != null) {
+            builder.field(MODEL_FIELD, unifiedRequest.model());
         }
-        if (maxCompletionTokens != null) {
-            builder.field(MAX_COMPLETION_TOKENS_FIELD, maxCompletionTokens);
+        if (unifiedRequest.maxCompletionTokens() != null) {
+            builder.field(MAX_COMPLETION_TOKENS_FIELD, unifiedRequest.maxCompletionTokens());
         }
-        if (n != null) {
-            builder.field(NUMBER_OF_RETURNED_CHOICES_FIELD, n);
+        if (unifiedRequest.n() != null) {
+            builder.field(NUMBER_OF_RETURNED_CHOICES_FIELD, unifiedRequest.n());
         }
-        if (stop != null) {
-            switch (stop) {
+        if (unifiedRequest.stop() != null) {
+            switch (unifiedRequest.stop()) {
                 case UnifiedCompletionRequest.StopString stopString -> builder.field(STOP_FIELD, stopString.value());
                 case UnifiedCompletionRequest.StopValues stopValues -> builder.field(STOP_FIELD, stopValues.values());
             }
         }
-        if (temperature != null) {
-            builder.field(TEMPERATURE_FIELD, temperature);
+        if (unifiedRequest.temperature() != null) {
+            builder.field(TEMPERATURE_FIELD, unifiedRequest.temperature());
         }
-        if (toolChoice != null) {
-            if (toolChoice instanceof UnifiedCompletionRequest.ToolChoiceString) {
-                builder.field(TOOL_CHOICE_FIELD, ((UnifiedCompletionRequest.ToolChoiceString) toolChoice).value());
-            } else if (toolChoice instanceof UnifiedCompletionRequest.ToolChoiceObject) {
+        if (unifiedRequest.toolChoice() != null) {
+            if (unifiedRequest.toolChoice() instanceof UnifiedCompletionRequest.ToolChoiceString) {
+                builder.field(TOOL_CHOICE_FIELD, ((UnifiedCompletionRequest.ToolChoiceString) unifiedRequest.toolChoice()).value());
+            } else if (unifiedRequest.toolChoice() instanceof UnifiedCompletionRequest.ToolChoiceObject) {
                 builder.startObject(TOOL_CHOICE_FIELD);
                 {
-                    builder.field("type", ((UnifiedCompletionRequest.ToolChoiceObject) toolChoice).type());
-                    builder.startObject("function");
+                    builder.field(TYPE_FIELD, ((UnifiedCompletionRequest.ToolChoiceObject) unifiedRequest.toolChoice()).type());
+                    builder.startObject(FUNCTION_FIELD);
                     {
-                        builder.field("name", ((UnifiedCompletionRequest.ToolChoiceObject) toolChoice).function().name());
+                        builder.field(
+                            NAME_FIELD,
+                            ((UnifiedCompletionRequest.ToolChoiceObject) unifiedRequest.toolChoice()).function().name()
+                        );
                     }
                     builder.endObject();
                 }
                 builder.endObject();
             }
         }
-        if (tool != null) {
+        if (unifiedRequest.tool() != null) {
             builder.startArray(TOOL_FIELD);
-            for (UnifiedCompletionRequest.Tool t : tool) {
+            for (UnifiedCompletionRequest.Tool t : unifiedRequest.tool()) {
                 builder.startObject();
                 {
-                    builder.field("type", t.type());
-                    builder.startObject("function");
+                    builder.field(TYPE_FIELD, t.type());
+                    builder.startObject(FUNCTION_FIELD);
                     {
-                        builder.field("description", t.function().description());
-                        builder.field("name", t.function().name());
-                        builder.field("parameters", t.function().parameters());
-                        builder.field("strict", t.function().strict());
+                        builder.field(DESCRIPTION_FIELD, t.function().description());
+                        builder.field(NAME_FIELD, t.function().name());
+                        builder.field(PARAMETERS_FIELD, t.function().parameters());
+                        builder.field(STRICT_FIELD, t.function().strict());
                     }
                     builder.endObject();
                 }
                 builder.endObject();
             }
             builder.endArray();
         }
-        if (topP != null) {
-            builder.field(TOP_P_FIELD, topP);
+        if (unifiedRequest.topP() != null) {
+            builder.field(TOP_P_FIELD, unifiedRequest.topP());
         }
-        if (Strings.isNullOrEmpty(user) == false) {
-            builder.field(USER_FIELD, user);
+        if (unifiedRequest.user() != null && unifiedRequest.user().isEmpty() == false) {
+            builder.field(USER_FIELD, unifiedRequest.user());
         }
+        builder.field(STREAM_FIELD, unifiedRequest.stream());
         builder.endObject();
         return builder;
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/openai/OpenAiUnifiedChatCompletionRequestEntity.java

@@ -21,12 +21,13 @@
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.rest.RestStatus;
-import org.elasticsearch.xpack.inference.external.http.sender.CompletionInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.DocumentsOnlyInput;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.QueryAndDocsInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
+import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
+import org.elasticsearch.xpack.inference.external.request.openai.OpenAiUnifiedChatCompletionRequestEntity;
 
 import java.io.IOException;
 import java.util.EnumSet;
@@ -73,7 +74,7 @@ public void infer(
     private static InferenceInputs createInput(Model model, List<String> input, @Nullable String query, boolean stream) {
         return switch (model.getTaskType()) {
             // TODO implement parameters
-            case COMPLETION -> new CompletionInputs(null);
+            case COMPLETION -> new UnifiedChatInput(null);
             case RERANK -> new QueryAndDocsInputs(query, input, stream);
             case TEXT_EMBEDDING -> new DocumentsOnlyInput(input, stream);
             default -> throw new ElasticsearchStatusException(
@@ -84,9 +85,14 @@ private static InferenceInputs createInput(Model model, List<String> input, @Nul
     }
 
     @Override
-    public void completionInfer(Model model, Object parameters, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
+    public void completionInfer(
+        Model model,
+        OpenAiUnifiedChatCompletionRequestEntity parameters,
+        TimeValue timeout,
+        ActionListener<InferenceServiceResults> listener
+    ) {
         init();
-        doUnifiedCompletionInfer(model, new CompletionInputs(parameters), timeout, listener);
+        doUnifiedCompletionInfer(model, new UnifiedChatInput(parameters), timeout, listener);
     }
 
     @Override
@@ -116,7 +122,7 @@ protected abstract void doInfer(
 
     protected abstract void doUnifiedCompletionInfer(
         Model model,
-        CompletionInputs inputs,
+        UnifiedChatInput inputs,
         TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     );

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java

@@ -35,11 +35,11 @@
 import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
 import org.elasticsearch.xpack.inference.external.action.SingleInputSenderExecutableAction;
 import org.elasticsearch.xpack.inference.external.action.openai.OpenAiActionCreator;
-import org.elasticsearch.xpack.inference.external.http.sender.CompletionInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.DocumentsOnlyInput;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.OpenAiCompletionRequestManager;
+import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
 import org.elasticsearch.xpack.inference.services.SenderService;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
@@ -266,7 +266,7 @@ public void doInfer(
     @Override
     public void doUnifiedCompletionInfer(
         Model model,
-        CompletionInputs inputs,
+        UnifiedChatInput inputs,
         TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {