[ML] Adding headers support for OpenAI chat completion (#134504)

* Adding headers support for openai chat completion and completion

* Update docs/changelog/134504.yaml

* [CI] Auto commit changes from spotless

* Adjusting test name

* Updating the changelog

* Adding headers to the configuration

* [CI] Auto commit changes from spotless

* [CI] Update transport version definitions

* Fixing transport versions

---------

Co-authored-by: elasticsearchmachine <infra-root+elasticsearchmachine@elastic.co>

Author: jonathan-buttner

docs/changelog/134504.yaml

@@ -0,0 +1,5 @@
+pr: 134504
+summary: Adding headers support for OpenAI chat completion
+area: Machine Learning
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -324,6 +324,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INDEX_SOURCE = def(9_158_0_00);
     public static final TransportVersion MAX_HEAP_SIZE_PER_NODE_IN_CLUSTER_INFO = def(9_159_0_00);
     public static final TransportVersion TIMESERIES_DEFAULT_LIMIT = def(9_160_0_00);
+    public static final TransportVersion INFERENCE_API_OPENAI_HEADERS = def(9_161_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java

@@ -841,7 +841,19 @@ public <V> Map<String, V> readImmutableMap(Writeable.Reader<V> valueReader) thro
     }
 
     /**
-     * Read a {@link Map} using the given key and value readers. The return Map is immutable.
+     * Read an optional {@link Map} using the given key and value readers. The returned Map is immutable.
+     *
+     * @param keyReader Method to read a key. Must not return null.
+     * @param valueReader Method to read a value. Must not return null.
+     * @return The immutable map or null if not present
+     */
+    public <K, V> Map<K, V> readOptionalImmutableMap(Writeable.Reader<K> keyReader, Writeable.Reader<V> valueReader) throws IOException {
+        final boolean present = readBoolean();
+        return present ? readImmutableMap(keyReader, valueReader) : null;
+    }
+
+    /**
+     * Read a {@link Map} using the given key and value readers. The returned Map is immutable.
      *
      * @param keyReader Method to read a key. Must not return null.
      * @param valueReader Method to read a value. Must not return null.

server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java

@@ -643,6 +643,26 @@ public final <K extends Writeable, V extends Writeable> void writeMap(final Map<
         writeMap(map, StreamOutput::writeWriteable, StreamOutput::writeWriteable);
     }
 
+    /**
+     * Write an optional {@link Map} of {@code K}-type keys to {@code V}-type.
+     * <pre><code>
+     * Map&lt;String, String&gt; map = ...;
+     * out.writeMap(map, StreamOutput::writeString, StreamOutput::writeString);
+     * </code></pre>
+     *
+     * @param keyWriter The key writer
+     * @param valueWriter The value writer
+     */
+    public final <K, V> void writeOptionalMap(final Map<K, V> map, final Writer<K> keyWriter, final Writer<V> valueWriter)
+        throws IOException {
+        if (map == null) {
+            writeBoolean(false);
+        } else {
+            writeBoolean(true);
+            writeMap(map, keyWriter, valueWriter);
+        }
+    }
+
     /**
      * Write a {@link Map} of {@code K}-type keys to {@code V}-type.
      * <pre><code>

server/src/main/java/org/elasticsearch/inference/configuration/SettingsConfigurationFieldType.java

@@ -13,7 +13,8 @@ public enum SettingsConfigurationFieldType {
     STRING("str"),
     INTEGER("int"),
     LIST("list"),
-    BOOLEAN("bool");
+    BOOLEAN("bool"),
+    MAP("map");
 
     private final String value;
 

server/src/test/java/org/elasticsearch/common/io/stream/AbstractStreamTests.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.core.CheckedConsumer;
 import org.elasticsearch.core.CheckedFunction;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.test.ESTestCase;
@@ -674,22 +675,63 @@ public void testObjectArrayIsWriteable() throws IOException {
         assertNotWriteable(new Object[] { new Unwriteable() }, Unwriteable.class);
     }
 
-    public void assertImmutableMapSerialization(Map<String, Integer> expected) throws IOException {
-        final BytesStreamOutput output = new BytesStreamOutput();
-        output.writeMap(expected, StreamOutput::writeString, StreamOutput::writeVInt);
-        final BytesReference bytesReference = output.bytes();
+    public void testImmutableMapSerialization() throws IOException {
+        CheckedBiConsumer<BytesStreamOutput, Map<String, Integer>, IOException> writer = (out, map) -> out.writeMap(
+            map,
+            StreamOutput::writeString,
+            StreamOutput::writeVInt
+        );
+        CheckedFunction<StreamInput, Map<String, Integer>, IOException> reader = in -> in.readImmutableMap(
+            StreamInput::readString,
+            StreamInput::readVInt
+        );
 
-        final StreamInput input = getStreamInput(bytesReference);
-        Map<String, Integer> got = input.readImmutableMap(StreamInput::readString, StreamInput::readVInt);
+        assertOptionalImmutableMapSerialization(Map.of(), writer, reader);
+        assertOptionalImmutableMapSerialization(Map.of("a", 1), writer, reader);
+        assertOptionalImmutableMapSerialization(Map.of("a", 1, "b", 2), writer, reader);
+    }
+
+    public void testOptionalImmutableMapSerialization() throws IOException {
+        CheckedBiConsumer<BytesStreamOutput, Map<String, Integer>, IOException> writer = (out, map) -> out.writeOptionalMap(
+            map,
+            StreamOutput::writeString,
+            StreamOutput::writeVInt
+        );
+        CheckedFunction<StreamInput, Map<String, Integer>, IOException> reader = in -> in.readOptionalImmutableMap(
+            StreamInput::readString,
+            StreamInput::readVInt
+        );
+
+        assertOptionalImmutableMapSerialization(null, writer, reader);
+        assertOptionalImmutableMapSerialization(Map.of(), writer, reader);
+        assertOptionalImmutableMapSerialization(Map.of("a", 1), writer, reader);
+        assertOptionalImmutableMapSerialization(Map.of("a", 1, "b", 2), writer, reader);
+    }
+
+    public void assertOptionalImmutableMapSerialization(
+        @Nullable Map<String, Integer> expected,
+        CheckedBiConsumer<BytesStreamOutput, Map<String, Integer>, IOException> writer,
+        CheckedFunction<StreamInput, Map<String, Integer>, IOException> reader
+    ) throws IOException {
+        var got = writeThenReadImmutableMap(expected, writer, reader);
         assertThat(got, equalTo(expected));
 
-        expectThrows(UnsupportedOperationException.class, () -> got.put("blah", 1));
+        if (got != null) {
+            expectThrows(UnsupportedOperationException.class, () -> got.put("blah", 1));
+        }
     }
 
-    public void testImmutableMapSerialization() throws IOException {
-        assertImmutableMapSerialization(Map.of());
-        assertImmutableMapSerialization(Map.of("a", 1));
-        assertImmutableMapSerialization(Map.of("a", 1, "b", 2));
+    private <K, V, E extends IOException> Map<K, V> writeThenReadImmutableMap(
+        @Nullable Map<K, V> expected,
+        CheckedBiConsumer<BytesStreamOutput, Map<K, V>, E> writer,
+        CheckedFunction<StreamInput, Map<K, V>, E> reader
+    ) throws IOException {
+        final BytesStreamOutput output = new BytesStreamOutput();
+        writer.accept(output, expected);
+        final BytesReference bytesReference = output.bytes();
+
+        final StreamInput input = getStreamInput(bytesReference);
+        return reader.apply(input);
     }
 
     public <T> void assertImmutableListSerialization(List<T> expected, Writeable.Reader<T> reader, Writeable.Writer<T> writer)

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

@@ -532,7 +532,6 @@ public static Map<String, Object> extractRequiredMap(
     public static Map<String, Object> extractOptionalMap(
         Map<String, Object> map,
         String settingName,
-        String scope,
         ValidationException validationException
     ) {
         int initialValidationErrorCount = validationException.validationErrors().size();
@@ -545,6 +544,14 @@ public static Map<String, Object> extractOptionalMap(
         return optionalField;
     }
 
+    public static Map<String, Object> extractOptionalMapRemoveNulls(
+        Map<String, Object> map,
+        String settingName,
+        ValidationException validationException
+    ) {
+        return removeNullValues(extractOptionalMap(map, settingName, validationException));
+    }
+
     public static List<Tuple<String, String>> extractOptionalListOfStringTuples(
         Map<String, Object> map,
         String settingName,
@@ -626,6 +633,20 @@ private static void validateString(
         }
     }
 
+    public static Map<String, String> validateMapStringValues(
+        Map<String, ?> map,
+        String settingName,
+        ValidationException validationException,
+        boolean censorValue,
+        @Nullable Map<String, String> defaultValue
+    ) {
+        if (map == null) {
+            return defaultValue;
+        }
+
+        return validateMapStringValues(map, settingName, validationException, censorValue);
+    }
+
     /**
      * Validates that each value in the map is a {@link String} and returns a new map of {@code Map<String, String>}.
      */

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

@@ -23,8 +23,7 @@
 import java.util.Objects;
 
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.convertMapStringsToSecureString;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMap;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeNullValues;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMapRemoveNulls;
 
 public class CustomSecretSettings implements SecretSettings {
     public static final String NAME = "custom_secret_settings";
@@ -37,8 +36,7 @@ public static CustomSecretSettings fromMap(@Nullable Map<String, Object> map) {
 
         ValidationException validationException = new ValidationException();
 
-        Map<String, Object> requestSecretParamsMap = extractOptionalMap(map, SECRET_PARAMETERS, NAME, validationException);
-        removeNullValues(requestSecretParamsMap);
+        Map<String, Object> requestSecretParamsMap = extractOptionalMapRemoveNulls(map, SECRET_PARAMETERS, validationException);
         var secureStringMap = convertMapStringsToSecureString(requestSecretParamsMap, SECRET_PARAMETERS, validationException);
 
         if (validationException.validationErrors().isEmpty() == false) {

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

@@ -42,13 +42,12 @@
 import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS;
 import static org.elasticsearch.xpack.inference.services.ServiceFields.MAX_INPUT_TOKENS;
 import static org.elasticsearch.xpack.inference.services.ServiceFields.SIMILARITY;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMap;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMapRemoveNulls;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalPositiveInteger;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractRequiredMap;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractRequiredString;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractSimilarity;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeAsType;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeNullValues;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwIfNotEmptyMap;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.validateMapStringValues;
 
@@ -75,8 +74,7 @@ public static CustomServiceSettings fromMap(Map<String, Object> map, Configurati
 
         var queryParams = QueryParameters.fromMap(map, validationException);
 
-        Map<String, Object> headers = extractOptionalMap(map, HEADERS, ModelConfigurations.SERVICE_SETTINGS, validationException);
-        removeNullValues(headers);
+        Map<String, Object> headers = extractOptionalMapRemoveNulls(map, HEADERS, validationException);
         var stringHeaders = validateMapStringValues(headers, HEADERS, validationException, false);
 
         String requestContentString = extractRequiredString(map, REQUEST, ModelConfigurations.SERVICE_SETTINGS, validationException);

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

@@ -12,7 +12,6 @@
 import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.TaskSettings;
 import org.elasticsearch.xcontent.XContentBuilder;
 
@@ -22,8 +21,7 @@
 import java.util.Map;
 import java.util.Objects;
 
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMap;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeNullValues;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMapRemoveNulls;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.validateMapValues;
 
 public class CustomTaskSettings implements TaskSettings {
@@ -39,8 +37,7 @@ public static CustomTaskSettings fromMap(Map<String, Object> map) {
             return EMPTY_SETTINGS;
         }
 
-        Map<String, Object> parameters = extractOptionalMap(map, PARAMETERS, ModelConfigurations.TASK_SETTINGS, validationException);
-        removeNullValues(parameters);
+        Map<String, Object> parameters = extractOptionalMapRemoveNulls(map, PARAMETERS, validationException);
         validateMapValues(
             parameters,
             List.of(String.class, Integer.class, Double.class, Float.class, Boolean.class),