Skip to content

[8.19] [ML] CustomService adding template validation prior to request flow (#129591) #129672

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: 8.19
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
);
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,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;
Expand All @@ -55,6 +56,7 @@
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";

Expand Down Expand Up @@ -101,12 +103,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 (IllegalStateException e) {
var validationException = new ValidationException();
validationException.addValidationError(Strings.format("Failed to validate model configuration: %s", e.getMessage()));
throw validationException;
}
}

private static ChunkingSettings extractChunkingSettings(Map<String, Object> config, TaskType taskType) {
if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
return ChunkingSettingsBuilder.fromMap(removeFromMap(config, ModelConfigurations.CHUNKING_SETTINGS));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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."
)
);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.action.support.PlainActionFuture;
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.ValidationException;
import org.elasticsearch.common.settings.SecureString;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
Expand Down Expand Up @@ -52,6 +53,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.entityAsMap;
import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl;
import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings;
Expand Down Expand Up @@ -611,6 +613,42 @@ 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)))
)
);

var config = getRequestConfigMap(settingsMap, createTaskSettingsMap(), createSecretSettingsMap());

var listener = new PlainActionFuture<Model>();
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.;"
)
);
}
}

public void testChunkedInfer_ChunkingSettingsSet() throws IOException {
var model = createInternalEmbeddingModel(
SimilarityMeasure.DOT_PRODUCT,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -264,7 +264,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() {
Expand Down