Need to address quoted strings

Author: jonathan-buttner

x-pack/plugin/inference/build.gradle

@@ -448,4 +448,3 @@ tasks.named("thirdPartyAudit").configure {
 tasks.named('yamlRestTest') {
   usesDefaultDistribution("to be triaged")
 }
-

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/JsonUtils.java

@@ -0,0 +1,30 @@
+/*
+ * 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.common;
+
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.json.JsonXContent;
+
+public class JsonUtils {
+
+    public static <T> String toJson(T value, String field) {
+        try {
+            XContentBuilder builder = JsonXContent.contentBuilder();
+            builder.value(value);
+            return Strings.toString(builder);
+        } catch (Exception e) {
+            throw new IllegalStateException(
+                Strings.format("Failed to serialize custom request value as JSON, field: %s, error: %s", field, e.getMessage()),
+                e
+            );
+        }
+    }
+
+    private JsonUtils() {}
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/ValidatingSubstitutor.java

@@ -0,0 +1,57 @@
+/*
+ * 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.common;
+
+import org.apache.commons.text.StringSubstitutor;
+
+import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * Substitutes placeholder values in a string that match the keys in a provided map with the map's corresponding values.
+ */
+public class ValidatingSubstitutor {
+    /**
+     * This regex pattern matches on the string {@code ${<any characters>}} excluding newlines.
+     */
+    private static final Pattern VARIABLE_PLACEHOLDER_PATTERN = Pattern.compile("\\$\\{.*?\\}");
+
+    private final StringSubstitutor substitutor;
+
+    /**
+     * @param params a map containing the placeholders as the keys, the values will be used to replace the placeholders
+     * @param prefix a string indicating the start of a placeholder
+     * @param suffix a string indicating the end of a placeholder
+     */
+    public ValidatingSubstitutor(Map<String, String> params, String prefix, String suffix) {
+        substitutor = new StringSubstitutor(params, prefix, suffix);
+    }
+
+    /**
+     * Substitutes placeholder values in a string that match the keys in a provided map with the map's corresponding values.
+     * After replacement, if the source still contains a placeholder an {@link IllegalStateException} is thrown.
+     * @param source the string that will be searched for placeholders to be replaced
+     * @param settingName a description of the source string
+     * @return a string with the placeholders replaced by string values
+     */
+    public String replace(String source, String settingName) {
+        var replacedString = substitutor.replace(source);
+        ensureNoMorePlaceholdersExist(replacedString, settingName);
+        return replacedString;
+    }
+
+    private static void ensureNoMorePlaceholdersExist(String substitutedString, String settingName) {
+        Matcher matcher = VARIABLE_PLACEHOLDER_PATTERN.matcher(substitutedString);
+        if (matcher.find()) {
+            throw new IllegalStateException(
+                String.format("Found placeholder [%s] in setting [%s] after replacement call", matcher.group(), settingName)
+            );
+        }
+    }
+}

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

@@ -54,6 +54,7 @@ public class CustomServiceSettings extends FilteredXContentObject implements Ser
     public static final String NAME = "custom_service_settings";
     public static final String URL = "url";
     public static final String HEADERS = "headers";
+    public static final String QUERY_PARAMETERS = "query_parameters";
     public static final String REQUEST = "request";
     public static final String REQUEST_CONTENT = "content";
     public static final String RESPONSE = "response";
@@ -71,6 +72,15 @@ public static CustomServiceSettings fromMap(Map<String, Object> map, Configurati
 
         String url = extractRequiredString(map, URL, ModelConfigurations.SERVICE_SETTINGS, validationException);
 
+        Map<String, Object> queryParameters = extractOptionalMap(
+            map,
+            QUERY_PARAMETERS,
+            ModelConfigurations.SERVICE_SETTINGS,
+            validationException
+        );
+        removeNullValues(queryParameters);
+        var stringQueryParameters = validateMapStringValues(queryParameters, QUERY_PARAMETERS, validationException, false);
+
         Map<String, Object> headers = extractOptionalMap(map, HEADERS, ModelConfigurations.SERVICE_SETTINGS, validationException);
         removeNullValues(headers);
         var stringHeaders = validateMapStringValues(headers, HEADERS, validationException, false);
@@ -136,6 +146,7 @@ public static CustomServiceSettings fromMap(Map<String, Object> map, Configurati
             maxInputTokens,
             url,
             stringHeaders,
+            stringQueryParameters,
             requestContentString,
             responseJsonParser,
             rateLimitSettings,
@@ -148,6 +159,7 @@ public static CustomServiceSettings fromMap(Map<String, Object> map, Configurati
     private final Integer maxInputTokens;
     private final String url;
     private final Map<String, String> headers;
+    private final Map<String, String> queryParameters;
     private final String requestContentString;
     private final CustomResponseParser responseJsonParser;
     private final RateLimitSettings rateLimitSettings;
@@ -159,6 +171,7 @@ public CustomServiceSettings(
         @Nullable Integer maxInputTokens,
         String url,
         @Nullable Map<String, String> headers,
+        @Nullable Map<String, String> queryParameters,
         String requestContentString,
         CustomResponseParser responseJsonParser,
         @Nullable RateLimitSettings rateLimitSettings,
@@ -169,6 +182,7 @@ public CustomServiceSettings(
         this.maxInputTokens = maxInputTokens;
         this.url = Objects.requireNonNull(url);
         this.headers = Collections.unmodifiableMap(Objects.requireNonNullElse(headers, Map.of()));
+        this.queryParameters = Collections.unmodifiableMap(Objects.requireNonNullElse(queryParameters, Map.of()));
         this.requestContentString = Objects.requireNonNull(requestContentString);
         this.responseJsonParser = Objects.requireNonNull(responseJsonParser);
         this.rateLimitSettings = Objects.requireNonNullElse(rateLimitSettings, DEFAULT_RATE_LIMIT_SETTINGS);
@@ -181,6 +195,7 @@ public CustomServiceSettings(StreamInput in) throws IOException {
         maxInputTokens = in.readOptionalVInt();
         url = in.readString();
         headers = in.readImmutableMap(StreamInput::readString);
+        queryParameters = in.readImmutableMap(StreamInput::readString);
         requestContentString = in.readString();
         responseJsonParser = in.readNamedWriteable(CustomResponseParser.class);
         rateLimitSettings = new RateLimitSettings(in);
@@ -214,6 +229,10 @@ public Map<String, String> getHeaders() {
         return headers;
     }
 
+    public Map<String, String> getQueryParameters() {
+        return queryParameters;
+    }
+
     public String getRequestContentString() {
         return requestContentString;
     }
@@ -267,6 +286,10 @@ public XContentBuilder toXContentFragmentOfExposedFields(XContentBuilder builder
             builder.field(HEADERS, headers);
         }
 
+        if (queryParameters.isEmpty() == false) {
+            builder.field(QUERY_PARAMETERS, queryParameters);
+        }
+
         builder.startObject(REQUEST);
         {
             builder.field(REQUEST_CONTENT, requestContentString);
@@ -302,6 +325,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalVInt(maxInputTokens);
         out.writeString(url);
         out.writeMap(headers, StreamOutput::writeString, StreamOutput::writeString);
+        out.writeMap(queryParameters, StreamOutput::writeString, StreamOutput::writeString);
         out.writeString(requestContentString);
         out.writeNamedWriteable(responseJsonParser);
         rateLimitSettings.writeTo(out);
@@ -318,6 +342,7 @@ public boolean equals(Object o) {
             && Objects.equals(maxInputTokens, that.maxInputTokens)
             && Objects.equals(url, that.url)
             && Objects.equals(headers, that.headers)
+            && Objects.equals(queryParameters, that.queryParameters)
             && Objects.equals(requestContentString, that.requestContentString)
             && Objects.equals(responseJsonParser, that.responseJsonParser)
             && Objects.equals(rateLimitSettings, that.rateLimitSettings)
@@ -332,6 +357,7 @@ public int hashCode() {
             maxInputTokens,
             url,
             headers,
+            queryParameters,
             requestContentString,
             responseJsonParser,
             rateLimitSettings,

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

@@ -29,6 +29,7 @@
 public class CustomTaskSettings implements TaskSettings {
     public static final String NAME = "custom_task_settings";
 
+    public static final String PARAMETERS = "json_parameters";
     public static final String PARAMETERS = "parameters";
 
     static final CustomTaskSettings EMPTY_SETTINGS = new CustomTaskSettings(new HashMap<>());
@@ -41,7 +42,13 @@ public static CustomTaskSettings fromMap(Map<String, Object> map) {
 
         Map<String, Object> parameters = extractOptionalMap(map, PARAMETERS, ModelConfigurations.TASK_SETTINGS, validationException);
         removeNullValues(parameters);
-        validateMapValues(parameters, List.of(String.class, Integer.class, Double.class, Float.class, Boolean.class), PARAMETERS, validationException, false);
+        validateMapValues(
+            parameters,
+            List.of(String.class, Integer.class, Double.class, Float.class, Boolean.class),
+            PARAMETERS,
+            validationException,
+            false
+        );
 
         if (validationException.validationErrors().isEmpty() == false) {
             throw validationException;

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

@@ -7,44 +7,35 @@
 
 package org.elasticsearch.xpack.inference.services.custom.request;
 
-import org.apache.commons.text.StringSubstitutor;
 import org.apache.http.HttpHeaders;
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.client.methods.HttpRequestBase;
 import org.apache.http.entity.StringEntity;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentType;
-import org.elasticsearch.xcontent.json.JsonXContent;
+import org.elasticsearch.xpack.inference.common.ValidatingSubstitutor;
 import org.elasticsearch.xpack.inference.external.request.HttpRequest;
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.services.custom.CustomModel;
 import org.elasticsearch.xpack.inference.services.custom.CustomServiceSettings;
 
-import java.io.IOException;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 
+import static org.elasticsearch.xpack.inference.common.JsonUtils.toJson;
 import static org.elasticsearch.xpack.inference.services.custom.CustomServiceSettings.REQUEST_CONTENT;
 import static org.elasticsearch.xpack.inference.services.custom.CustomServiceSettings.URL;
 
 public class CustomRequest implements Request {
-    /**
-     * This regex pattern matches on the string {@code "${<any characters>}"}
-     */
-    private static final Pattern VARIABLE_PLACEHOLDER_PATTERN = Pattern.compile("\\$\\{.*?\\}");
-
     private static final String QUERY = "query";
     private static final String INPUT = "input";
 
     private final URI uri;
-    private final StringSubstitutor substitutor;
+    private final ValidatingSubstitutor substitutor;
     private final CustomModel model;
 
     public CustomRequest(String query, List<String> input, CustomModel model) {
@@ -60,28 +51,21 @@ public CustomRequest(String query, List<String> input, CustomModel model) {
 
         jsonParams.put(INPUT, toJson(input, INPUT));
 
-        substitutor = new StringSubstitutor(jsonParams, "${", "}");
+        substitutor = new ValidatingSubstitutor(jsonParams, "${", "}");
         uri = buildUri();
     }
 
-    // default for testing
-    static <T> String toJson(T value, String field) {
-        try {
-            XContentBuilder builder = JsonXContent.contentBuilder();
-            // TODO test this, I think it'll write the quotes for us so we don't need to include them in the content string
-            builder.value(value);
-            return Strings.toString(builder);
-        } catch (IOException e) {
-            throw new IllegalStateException(Strings.format("Failed to serialize custom request value as json, field: %s", field), e);
-        }
-    }
-
     private static void addJsonStringParams(Map<String, String> jsonStringParams, Map<String, ?> params) {
         for (var entry : params.entrySet()) {
             jsonStringParams.put(entry.getKey(), toJson(entry.getValue(), entry.getKey()));
         }
     }
 
+    private URI buildUri() {
+        String replacedUrl = substitutor.replace(model.getServiceSettings().getUrl(), URL);
+        return URI.create(replacedUrl);
+    }
+
     @Override
     public HttpRequest createHttpRequest() {
         HttpPost httpRequest = new HttpPost(uri);
@@ -97,15 +81,13 @@ private void setHeaders(HttpRequestBase httpRequest) {
         httpRequest.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType());
 
         for (var entry : model.getServiceSettings().getHeaders().entrySet()) {
-            String replacedHeadersValue = substitutor.replace(entry.getValue());
-            placeholderValidation(replacedHeadersValue, Strings.format("header.%s", entry.getKey()));
+            String replacedHeadersValue = substitutor.replace(entry.getValue(), Strings.format("header.%s", entry.getKey()));
             httpRequest.setHeader(entry.getKey(), replacedHeadersValue);
         }
     }
 
     private void setRequestContent(HttpPost httpRequest) {
-        String replacedRequestContentString = substitutor.replace(model.getServiceSettings().getRequestContentString());
-        placeholderValidation(replacedRequestContentString, REQUEST_CONTENT);
+        String replacedRequestContentString = substitutor.replace(model.getServiceSettings().getRequestContentString(), REQUEST_CONTENT);
         StringEntity stringEntity = new StringEntity(replacedRequestContentString, StandardCharsets.UTF_8);
         httpRequest.setEntity(stringEntity);
     }
@@ -133,19 +115,4 @@ public Request truncate() {
     public boolean[] getTruncationInfo() {
         return null;
     }
-
-    // default for testing
-    URI buildUri() {
-        String replacedUrl = substitutor.replace(model.getServiceSettings().getUrl());
-        placeholderValidation(replacedUrl, URL);
-        return URI.create(replacedUrl);
-    }
-
-    // default for testing
-    static void placeholderValidation(String substitutedString, String settingName) {
-        Matcher matcher = VARIABLE_PLACEHOLDER_PATTERN.matcher(substitutedString);
-        if (matcher.find()) {
-            throw new IllegalStateException(String.format("Found placeholder in [%s] after replacement call", settingName));
-        }
-    }
 }