[ML] Adding custom headers support openai text embeddings

docs/changelog/134960.yaml

@@ -0,0 +1,5 @@
+pr: 134960
+summary: Adding custom headers support openai text embeddings
+area: Machine Learning
+type: enhancement
+issues: []

server/src/main/resources/transport/definitions/referable/inference_api_openai_embeddings_headers.csv

@@ -0,0 +1 @@
+9169000

server/src/main/resources/transport/upper_bounds/9.2.csv

@@ -1 +1 @@
-security_stats_endpoint,9168000
+inference_api_openai_embeddings_headers,9169000

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

@@ -485,9 +485,8 @@ public static InferenceServiceConfiguration get() {
 
                 configurationMap.put(
                     HEADERS,
-                    new SettingsConfiguration.Builder(EnumSet.of(TaskType.COMPLETION, TaskType.CHAT_COMPLETION)).setDescription(
-                        "Custom headers to include in the requests to OpenAI."
-                    )
+                    new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING, TaskType.COMPLETION, TaskType.CHAT_COMPLETION))
+                        .setDescription("Custom headers to include in the requests to OpenAI.")
                         .setLabel("Custom Headers")
                         .setRequired(false)
                         .setSensitive(false)

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

@@ -0,0 +1,126 @@
+/*
+ * 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.services.openai;
+
+import org.elasticsearch.common.ValidationException;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.inference.ModelConfigurations;
+import org.elasticsearch.inference.TaskSettings;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMapRemoveNulls;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalString;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.validateMapStringValues;
+import static org.elasticsearch.xpack.inference.services.openai.OpenAiServiceFields.HEADERS;
+import static org.elasticsearch.xpack.inference.services.openai.OpenAiServiceFields.USER;
+
+public abstract class OpenAiTaskSettings<T extends OpenAiTaskSettings<T>> implements TaskSettings {
+    private static final Settings EMPTY_SETTINGS = new Settings(null, null);
+
+    private final Settings taskSettings;
+
+    public OpenAiTaskSettings(Map<String, Object> map) {
+        this(fromMap(map));
+    }
+
+    public record Settings(@Nullable String user, @Nullable Map<String, String> headers) {}
+
+    public static Settings createSettings(String user, Map<String, String> stringHeaders) {
+        if (user == null && stringHeaders == null) {
+            return EMPTY_SETTINGS;
+        } else {
+            return new Settings(user, stringHeaders);
+        }
+    }
+
+    private static Settings fromMap(Map<String, Object> map) {
+        if (map.isEmpty()) {
+            return EMPTY_SETTINGS;
+        }
+
+        ValidationException validationException = new ValidationException();
+
+        String user = extractOptionalString(map, USER, ModelConfigurations.TASK_SETTINGS, validationException);
+        Map<String, Object> headers = extractOptionalMapRemoveNulls(map, HEADERS, validationException);
+        var stringHeaders = validateMapStringValues(headers, HEADERS, validationException, false, null);
+
+        if (validationException.validationErrors().isEmpty() == false) {
+            throw validationException;
+        }
+
+        return createSettings(user, stringHeaders);
+    }
+
+    public OpenAiTaskSettings(@Nullable String user, @Nullable Map<String, String> headers) {
+        this(new Settings(user, headers));
+    }
+
+    protected OpenAiTaskSettings(Settings taskSettings) {
+        this.taskSettings = Objects.requireNonNull(taskSettings);
+    }
+
+    public String user() {
+        return taskSettings.user();
+    }
+
+    public Map<String, String> headers() {
+        return taskSettings.headers();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return taskSettings.user() == null && (taskSettings.headers() == null || taskSettings.headers().isEmpty());
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+
+        if (taskSettings.user() != null) {
+            builder.field(USER, taskSettings.user());
+        }
+
+        if (taskSettings.headers() != null && taskSettings.headers().isEmpty() == false) {
+            builder.field(HEADERS, taskSettings.headers());
+        }
+
+        builder.endObject();
+
+        return builder;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        OpenAiTaskSettings<?> that = (OpenAiTaskSettings<?>) o;
+        return Objects.equals(taskSettings, that.taskSettings);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(taskSettings);
+    }
+
+    @Override
+    public T updatedTaskSettings(Map<String, Object> newSettings) {
+        Settings updatedSettings = fromMap(new HashMap<>(newSettings));
+
+        var userToUse = updatedSettings.user() == null ? taskSettings.user() : updatedSettings.user();
+        var headersToUse = updatedSettings.headers() == null ? taskSettings.headers() : updatedSettings.headers();
+        return create(userToUse, headersToUse);
+    }
+
+    protected abstract T create(@Nullable String user, @Nullable Map<String, String> headers);
+
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/completion/OpenAiChatCompletionModel.java

@@ -35,8 +35,7 @@ public static OpenAiChatCompletionModel of(OpenAiChatCompletionModel model, Map<
             return model;
         }
 
-        var requestTaskSettings = OpenAiChatCompletionRequestTaskSettings.fromMap(taskSettings);
-        return new OpenAiChatCompletionModel(model, OpenAiChatCompletionTaskSettings.of(model.getTaskSettings(), requestTaskSettings));
+        return new OpenAiChatCompletionModel(model, model.getTaskSettings().updatedTaskSettings(taskSettings));
     }
 
     public static OpenAiChatCompletionModel of(OpenAiChatCompletionModel model, UnifiedCompletionRequest request) {
@@ -73,7 +72,7 @@ public OpenAiChatCompletionModel(
             taskType,
             service,
             OpenAiChatCompletionServiceSettings.fromMap(serviceSettings, context),
-            OpenAiChatCompletionTaskSettings.fromMap(taskSettings),
+            new OpenAiChatCompletionTaskSettings(taskSettings),
             DefaultSecretSettings.fromMap(secrets)
         );
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/completion/OpenAiChatCompletionTaskSettings.java

@@ -9,100 +9,44 @@
 
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
-import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.Nullable;
-import org.elasticsearch.inference.ModelConfigurations;
-import org.elasticsearch.inference.TaskSettings;
-import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.inference.services.openai.OpenAiTaskSettings;
 
 import java.io.IOException;
-import java.util.HashMap;
 import java.util.Map;
-import java.util.Objects;
 
 import static org.elasticsearch.TransportVersions.INFERENCE_API_OPENAI_HEADERS;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMapRemoveNulls;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalString;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.validateMapStringValues;
-import static org.elasticsearch.xpack.inference.services.openai.OpenAiServiceFields.HEADERS;
-import static org.elasticsearch.xpack.inference.services.openai.OpenAiServiceFields.USER;
 
-public class OpenAiChatCompletionTaskSettings implements TaskSettings {
+public class OpenAiChatCompletionTaskSettings extends OpenAiTaskSettings<OpenAiChatCompletionTaskSettings> {
 
     public static final String NAME = "openai_completion_task_settings";
 
-    public static OpenAiChatCompletionTaskSettings fromMap(Map<String, Object> map) {
-        ValidationException validationException = new ValidationException();
-
-        String user = extractOptionalString(map, USER, ModelConfigurations.TASK_SETTINGS, validationException);
-        var headers = extractOptionalMapRemoveNulls(map, HEADERS, validationException);
-        var stringHeaders = validateMapStringValues(headers, HEADERS, validationException, false, null);
-
-        if (validationException.validationErrors().isEmpty() == false) {
-            throw validationException;
-        }
-
-        return new OpenAiChatCompletionTaskSettings(user, stringHeaders);
+    public OpenAiChatCompletionTaskSettings(Map<String, Object> map) {
+        super(map);
     }
 
-    private final String user;
-    @Nullable
-    private final Map<String, String> headers;
-
     public OpenAiChatCompletionTaskSettings(@Nullable String user, @Nullable Map<String, String> headers) {
-        this.user = user;
-        this.headers = headers;
+        super(user, headers);
     }
 
     public OpenAiChatCompletionTaskSettings(StreamInput in) throws IOException {
-        this.user = in.readOptionalString();
+        super(readTaskSettingsFromStream(in));
+    }
+
+    private static Settings readTaskSettingsFromStream(StreamInput in) throws IOException {
+        var user = in.readOptionalString();
+
+        Map<String, String> headers;
 
         if (in.getTransportVersion().onOrAfter(INFERENCE_API_OPENAI_HEADERS)) {
             headers = in.readOptionalImmutableMap(StreamInput::readString, StreamInput::readString);
         } else {
             headers = null;
         }
-    }
-
-    @Override
-    public boolean isEmpty() {
-        return user == null && (headers == null || headers.isEmpty());
-    }
-
-    public static OpenAiChatCompletionTaskSettings of(
-        OpenAiChatCompletionTaskSettings originalSettings,
-        OpenAiChatCompletionRequestTaskSettings requestSettings
-    ) {
-        var userToUse = requestSettings.user() == null ? originalSettings.user : requestSettings.user();
-        var headersToUse = requestSettings.headers() == null ? originalSettings.headers : requestSettings.headers();
-        return new OpenAiChatCompletionTaskSettings(userToUse, headersToUse);
-    }
-
-    public String user() {
-        return user;
-    }
 
-    public Map<String, String> headers() {
-        return headers;
-    }
-
-    @Override
-    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.startObject();
-
-        if (user != null) {
-            builder.field(USER, user);
-        }
-
-        if (headers != null && headers.isEmpty() == false) {
-            builder.field(HEADERS, headers);
-        }
-
-        builder.endObject();
-
-        return builder;
+        return createSettings(user, headers);
     }
 
     @Override
@@ -117,30 +61,14 @@ public TransportVersion getMinimalSupportedVersion() {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeOptionalString(user);
+        out.writeOptionalString(user());
         if (out.getTransportVersion().onOrAfter(INFERENCE_API_OPENAI_HEADERS)) {
-            out.writeOptionalMap(headers, StreamOutput::writeString, StreamOutput::writeString);
+            out.writeOptionalMap(headers(), StreamOutput::writeString, StreamOutput::writeString);
         }
     }
 
     @Override
-    public boolean equals(Object object) {
-        if (this == object) return true;
-        if (object == null || getClass() != object.getClass()) return false;
-        OpenAiChatCompletionTaskSettings that = (OpenAiChatCompletionTaskSettings) object;
-        return Objects.equals(user, that.user) && Objects.equals(headers, that.headers);
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(user, headers);
-    }
-
-    @Override
-    public TaskSettings updatedTaskSettings(Map<String, Object> newSettings) {
-        OpenAiChatCompletionRequestTaskSettings updatedSettings = OpenAiChatCompletionRequestTaskSettings.fromMap(
-            new HashMap<>(newSettings)
-        );
-        return of(this, updatedSettings);
+    protected OpenAiChatCompletionTaskSettings create(@Nullable String user, @Nullable Map<String, String> headers) {
+        return new OpenAiChatCompletionTaskSettings(user, headers);
     }
 }