Add outbound request writing (WIP)

Author: Max Hniebergall

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;
 
@@ -40,6 +42,10 @@ public record UnifiedCompletionRequest(
     @Nullable String user
 ) implements Writeable {
 
+    public sealed interface Content extends NamedWriteable permits ContentObjects, ContentString {
+        void toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException;
+    }
+
     @SuppressWarnings("unchecked")
     static final ConstructingObjectParser<UnifiedCompletionRequest, Void> PARSER = new ConstructingObjectParser<>(
         InferenceAction.NAME,
@@ -158,8 +164,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";
@@ -173,6 +177,17 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeCollection(contentObjects);
         }
 
+        @Override
+        public void toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
+            builder.startArray();
+            for (ContentObject contentObject : contentObjects) {
+                builder.startObject();
+                contentObject.toXContentObject(builder, params);
+                builder.endObject();
+            }
+            builder.endArray();
+        }
+
         @Override
         public String getWriteableName() {
             return NAME;
@@ -199,6 +214,14 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeString(text);
             out.writeString(type);
         }
+
+        public XContentBuilder toXContentObject(XContentBuilder builder, ToXContent.Params params) throws IOException {
+            builder.startObject();
+            builder.field("text", text);
+            builder.field("type", type);
+            builder.endObject();
+            return builder;
+        }
     }
 
     public record ContentString(String content) implements Content, NamedWriteable {
@@ -222,6 +245,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 {
@@ -437,7 +464,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/OpenAiCompletionRequestManager.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/OpenAiUnifiedCompletionRequestManager.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/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 {