From ac4da6a7ddd3b6491a96816b942e607783632662 Mon Sep 17 00:00:00 2001 From: Jonathan Buttner Date: Tue, 17 Jun 2025 16:41:12 -0400 Subject: [PATCH 1/3] Adding template validation prior to request flow --- .../common/ValidatingSubstitutor.java | 7 ++- .../services/custom/CustomService.java | 21 +++++++++ .../services/custom/CustomServiceTests.java | 45 +++++++++++++++++++ 3 files changed, 72 insertions(+), 1 deletion(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/ValidatingSubstitutor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/ValidatingSubstitutor.java index ec0619ff06389..ab98408a5156a 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/ValidatingSubstitutor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/ValidatingSubstitutor.java @@ -51,7 +51,12 @@ private static void ensureNoMorePlaceholdersExist(String substitutedString, Stri Matcher matcher = VARIABLE_PLACEHOLDER_PATTERN.matcher(substitutedString); if (matcher.find()) { throw new IllegalStateException( - Strings.format("Found placeholder [%s] in field [%s] after replacement call", matcher.group(), field) + Strings.format( + "Found placeholder [%s] in field [%s] after replacement call, " + + "please check that all templates have a corresponding field definition.", + matcher.group(), + field + ) ); } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java index 5e9aef099f622..f7737baa35749 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java @@ -36,6 +36,7 @@ import org.elasticsearch.xpack.inference.services.SenderService; import org.elasticsearch.xpack.inference.services.ServiceComponents; import org.elasticsearch.xpack.inference.services.ServiceUtils; +import org.elasticsearch.xpack.inference.services.custom.request.CustomRequest; import java.util.EnumSet; import java.util.HashMap; @@ -94,12 +95,32 @@ public void parseRequestConfig( throwIfNotEmptyMap(serviceSettingsMap, NAME); throwIfNotEmptyMap(taskSettingsMap, NAME); + validateConfiguration(model); + parsedModelListener.onResponse(model); } catch (Exception e) { parsedModelListener.onFailure(e); } } + /** + * This does some initial validation with mock inputs to determine if any templates are missing a field to fill them. + */ + private static void validateConfiguration(CustomModel model) { + String query = null; + if (model.getTaskType() == TaskType.RERANK) { + query = "test query"; + } + + try { + new CustomRequest(query, List.of("test input"), model).createHttpRequest(); + } catch (Exception e) { + var validationException = new ValidationException(); + validationException.addValidationError(Strings.format("Failed to validate model configuration: %s", e.getMessage())); + throw validationException; + } + } + @Override public InferenceServiceConfiguration getConfiguration() { return Configuration.get(); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceTests.java index f8d650144693d..fbbab6a82564b 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceTests.java @@ -8,6 +8,7 @@ package org.elasticsearch.xpack.inference.services.custom; import org.elasticsearch.action.support.PlainActionFuture; +import org.elasticsearch.common.ValidationException; import org.elasticsearch.common.settings.SecureString; import org.elasticsearch.core.Nullable; import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; @@ -44,6 +45,7 @@ import java.util.Map; import static org.elasticsearch.xpack.inference.Utils.TIMEOUT; +import static org.elasticsearch.xpack.inference.Utils.getRequestConfigMap; import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl; import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings; import static org.elasticsearch.xpack.inference.services.custom.response.RerankResponseParser.RERANK_PARSER_DOCUMENT_TEXT; @@ -546,4 +548,47 @@ public void testInfer_HandlesSparseEmbeddingRequest_Alibaba_Format() throws IOEx ); } } + + public void testParseRequestConfig_ThrowsAValidationError_WhenReplacementDoesNotFillTemplate() throws Exception { + try (var service = createService(threadPool, clientManager)) { + + var settingsMap = new HashMap<>( + Map.of( + CustomServiceSettings.URL, + "http://www.abc.com", + CustomServiceSettings.HEADERS, + Map.of("key", "value"), + QueryParameters.QUERY_PARAMETERS, + List.of(List.of("key", "value")), + CustomServiceSettings.REQUEST, + "request body ${some_template}", + CustomServiceSettings.RESPONSE, + new HashMap<>( + Map.of( + CustomServiceSettings.JSON_PARSER, + createResponseParserMap(TaskType.COMPLETION), + CustomServiceSettings.ERROR_PARSER, + new HashMap<>(Map.of(ErrorResponseParser.MESSAGE_PATH, "$.error.message")) + ) + ) + ) + ); + + var config = getRequestConfigMap(settingsMap, createTaskSettingsMap(), createSecretSettingsMap()); + + var listener = new PlainActionFuture(); + service.parseRequestConfig("id", TaskType.COMPLETION, config, listener); + + var exception = expectThrows(ValidationException.class, () -> listener.actionGet(TIMEOUT)); + + assertThat( + exception.getMessage(), + is( + "Validation Failed: 1: Failed to validate model configuration: Found placeholder " + + "[${some_template}] in field [request] after replacement call, please check that all " + + "templates have a corresponding field definition.;" + ) + ); + } + } } From 3ad6077b801f02184e78a4cdda069a1033f1827f Mon Sep 17 00:00:00 2001 From: Jonathan Buttner Date: Wed, 18 Jun 2025 09:41:22 -0400 Subject: [PATCH 2/3] Fixing tests --- .../common/ValidatingSubstitutorTests.java | 24 ++++++++++++++++--- .../custom/request/CustomRequestTests.java | 8 ++++++- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/common/ValidatingSubstitutorTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/common/ValidatingSubstitutorTests.java index eb0224ff774dc..61563fd82de31 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/common/ValidatingSubstitutorTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/common/ValidatingSubstitutorTests.java @@ -37,20 +37,38 @@ public void testReplace_ThrowsException_WhenPlaceHolderStillExists() { var sub = new ValidatingSubstitutor(Map.of("some_key", "value", "key2", "value2"), "${", "}"); var exception = expectThrows(IllegalStateException.class, () -> sub.replace("super:${key}", "setting")); - assertThat(exception.getMessage(), is("Found placeholder [${key}] in field [setting] after replacement call")); + assertThat( + exception.getMessage(), + is( + "Found placeholder [${key}] in field [setting] after replacement call, " + + "please check that all templates have a corresponding field definition." + ) + ); } // only reports the first placeholder pattern { var sub = new ValidatingSubstitutor(Map.of("some_key", "value", "some_key2", "value2"), "${", "}"); var exception = expectThrows(IllegalStateException.class, () -> sub.replace("super, ${key}, ${key2}", "setting")); - assertThat(exception.getMessage(), is("Found placeholder [${key}] in field [setting] after replacement call")); + assertThat( + exception.getMessage(), + is( + "Found placeholder [${key}] in field [setting] after replacement call, " + + "please check that all templates have a corresponding field definition." + ) + ); } { var sub = new ValidatingSubstitutor(Map.of("some_key", "value", "key2", "value2"), "${", "}"); var exception = expectThrows(IllegalStateException.class, () -> sub.replace("super:${ \\/\tkey\"}", "setting")); - assertThat(exception.getMessage(), is("Found placeholder [${ \\/\tkey\"}] in field [setting] after replacement call")); + assertThat( + exception.getMessage(), + is( + "Found placeholder [${ \\/\tkey\"}] in field [setting] after replacement call," + + " please check that all templates have a corresponding field definition." + ) + ); } } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/request/CustomRequestTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/request/CustomRequestTests.java index 0f5daaf13af43..deb0f0a7d6597 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/request/CustomRequestTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/request/CustomRequestTests.java @@ -270,7 +270,13 @@ public void testCreateRequest_IgnoresNonStringFields_ForStringParams() throws IO var request = new CustomRequest(null, List.of("abc", "123"), model); var exception = expectThrows(IllegalStateException.class, request::createHttpRequest); - assertThat(exception.getMessage(), is("Found placeholder [${task.key}] in field [header.Accept] after replacement call")); + assertThat( + exception.getMessage(), + is( + "Found placeholder [${task.key}] in field [header.Accept] after replacement call, " + + "please check that all templates have a corresponding field definition." + ) + ); } public void testCreateRequest_ThrowsException_ForInvalidUrl() { From 53936fb131bba68ea9261f3ad29744307c9a97a3 Mon Sep 17 00:00:00 2001 From: Jonathan Buttner Date: Wed, 18 Jun 2025 12:25:53 -0400 Subject: [PATCH 3/3] Narrowing exception --- .../xpack/inference/services/custom/CustomService.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java index 220e323558606..a2878a33087df 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java @@ -7,6 +7,8 @@ package org.elasticsearch.xpack.inference.services.custom; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.elasticsearch.ElasticsearchStatusException; import org.elasticsearch.TransportVersion; import org.elasticsearch.TransportVersions; @@ -52,8 +54,10 @@ import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwUnsupportedUnifiedCompletionOperation; public class CustomService extends SenderService { + public static final String NAME = "custom"; private static final String SERVICE_NAME = "Custom"; + private static final Logger logger = LogManager.getLogger(CustomService.class); private static final EnumSet supportedTaskTypes = EnumSet.of( TaskType.TEXT_EMBEDDING, @@ -114,7 +118,7 @@ private static void validateConfiguration(CustomModel model) { try { new CustomRequest(query, List.of("test input"), model).createHttpRequest(); - } catch (Exception e) { + } catch (IllegalStateException e) { var validationException = new ValidationException(); validationException.addValidationError(Strings.format("Failed to validate model configuration: %s", e.getMessage())); throw validationException;