Skip to content

[ML] Refactor inference API service tests base classes #135461

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.

Sign up for GitHub

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

Merged
merged 16 commits into from
Sep 30, 2025
Merged

[ML] Refactor inference API service tests base classes #135461

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
028d445
Refactoring openai
jonathan-buttner Sep 24, 2025
c3243fd
Splitting up parameterized tests
jonathan-buttner Sep 24, 2025
399a50b
Working tests
jonathan-buttner Sep 25, 2025
060eed9
[CI] Auto commit changes from spotless
Sep 25, 2025
ac44cc6
[CI] Update transport version definitions
Sep 25, 2025
0af841d
Merge branch 'main' of github.com:elastic/elasticsearch into ml-refac…
jonathan-buttner Sep 25, 2025
bbd0707
Merge branch 'ml-refactor-tests-abstract' of github.com:jonathan-butt…
jonathan-buttner Sep 25, 2025
9392ec0
Merge branch 'main' of github.com:elastic/elasticsearch into ml-refac…
jonathan-buttner Sep 29, 2025
cc99292
Removing deprecated function
jonathan-buttner Sep 29, 2025
0896e5c
Moving string creation and refactoring customservice chunking
jonathan-buttner Sep 30, 2025
49c1301
Removing usages of persistent function
jonathan-buttner Sep 30, 2025
3110967
Merge branch 'main' of github.com:elastic/elasticsearch into ml-refac…
jonathan-buttner Sep 30, 2025
6ec99f9
[CI] Auto commit changes from spotless
Sep 30, 2025
d52d4dc
Finishing comment
jonathan-buttner Sep 30, 2025
bbd1f3e
Merge branch 'ml-refactor-tests-abstract' of github.com:jonathan-butt…
jonathan-buttner Sep 30, 2025
dd9dde0
Merge branch 'main' of github.com:elastic/elasticsearch into ml-refac…
jonathan-buttner Sep 30, 2025
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
22 changes: 19 additions & 3 deletions ...ugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ServiceUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1079,11 +1079,27 @@ public interface EnumConstructor<E extends Enum<E>> {
E apply(String name) throws IllegalArgumentException;
}

public static String parsePersistedConfigErrorMsg(String inferenceEntityId, String serviceName) {
/**
* Create an exception for when the task type is not valid for the service.
*/
public static ElasticsearchStatusException createInvalidTaskTypeException(
String inferenceEntityId,
String serviceName,
TaskType taskType,
ConfigurationParseContext parseContext
) {
var message = parseContext == ConfigurationParseContext.PERSISTENT
? parsePersistedConfigErrorMsg(inferenceEntityId, serviceName, taskType)
: TaskType.unsupportedTaskTypeErrorMsg(taskType, serviceName);
return new ElasticsearchStatusException(message, RestStatus.BAD_REQUEST);
}

private static String parsePersistedConfigErrorMsg(String inferenceEntityId, String serviceName, TaskType taskType) {
return format(
"Failed to parse stored model [%s] for [%s] service, please delete and add the service again",
"Failed to parse stored model [%s] for [%s] service, error: [%s]. Please delete and add the service again",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would deleting and adding the service again actually help if the task type was unsupported?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I think deleting is probably the only solution here. I think this would only occur if the the inference endpoint got corrupted somehow. Basically this is saying that the persisted inference endpoint is stating it is leverage a particular task type that is not supported. The request context parsing should prevent getting into that scenario. But if we had a regression or the endpoint was corrupted somehow we could.

inferenceEntityId,
serviceName
serviceName,
TaskType.unsupportedTaskTypeErrorMsg(taskType, serviceName)
);
}

Expand Down
38 changes: 7 additions & 31 deletions .../inference/src/main/java/org/elasticsearch/xpack/inference/services/ai21/Ai21Service.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@

package org.elasticsearch.xpack.inference.services.ai21;

import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.TransportVersion;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.cluster.service.ClusterService;
Expand All @@ -27,7 +26,6 @@
import org.elasticsearch.inference.SettingsConfiguration;
import org.elasticsearch.inference.TaskType;
import org.elasticsearch.inference.configuration.SettingsConfigurationFieldType;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
import org.elasticsearch.xpack.inference.external.http.sender.GenericRequestManager;
Expand Down Expand Up @@ -55,7 +53,7 @@
import java.util.Set;

import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidTaskTypeException;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrDefaultEmpty;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrThrowIfNull;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwIfNotEmptyMap;
Expand Down Expand Up @@ -178,14 +176,7 @@ public void parseRequestConfig(
Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);

Ai21Model model = createModel(
modelId,
taskType,
serviceSettingsMap,
serviceSettingsMap,
TaskType.unsupportedTaskTypeErrorMsg(taskType, NAME),
ConfigurationParseContext.REQUEST
);
Ai21Model model = createModel(modelId, taskType, serviceSettingsMap, serviceSettingsMap, ConfigurationParseContext.REQUEST);

throwIfNotEmptyMap(config, NAME);
throwIfNotEmptyMap(serviceSettingsMap, NAME);
Expand All @@ -208,21 +199,15 @@ public Ai21Model parsePersistedConfigWithSecrets(
removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
Map<String, Object> secretSettingsMap = removeFromMapOrDefaultEmpty(secrets, ModelSecrets.SECRET_SETTINGS);

return createModelFromPersistent(
modelId,
taskType,
serviceSettingsMap,
secretSettingsMap,
parsePersistedConfigErrorMsg(modelId, NAME)
);
return createModelFromPersistent(modelId, taskType, serviceSettingsMap, secretSettingsMap);
}

@Override
public Ai21Model parsePersistedConfig(String modelId, TaskType taskType, Map<String, Object> config) {
Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);

return createModelFromPersistent(modelId, taskType, serviceSettingsMap, null, parsePersistedConfigErrorMsg(modelId, NAME));
return createModelFromPersistent(modelId, taskType, serviceSettingsMap, null);
}

@Override
Expand All @@ -240,32 +225,23 @@ private static Ai21Model createModel(
TaskType taskType,
Map<String, Object> serviceSettings,
@Nullable Map<String, Object> secretSettings,
String failureMessage,
ConfigurationParseContext context
) {
switch (taskType) {
case CHAT_COMPLETION, COMPLETION:
return new Ai21ChatCompletionModel(modelId, taskType, NAME, serviceSettings, secretSettings, context);
default:
throw new ElasticsearchStatusException(failureMessage, RestStatus.BAD_REQUEST);
throw createInvalidTaskTypeException(modelId, NAME, taskType, context);
}
}

private Ai21Model createModelFromPersistent(
String inferenceEntityId,
TaskType taskType,
Map<String, Object> serviceSettings,
Map<String, Object> secretSettings,
String failureMessage
Map<String, Object> secretSettings
) {
return createModel(
inferenceEntityId,
taskType,
serviceSettings,
secretSettings,
failureMessage,
ConfigurationParseContext.PERSISTENT
);
return createModel(inferenceEntityId, taskType, serviceSettings, secretSettings, ConfigurationParseContext.PERSISTENT);
}

/**
Expand Down
18 changes: 5 additions & 13 deletions .../elasticsearch/xpack/inference/services/alibabacloudsearch/AlibabaCloudSearchService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@

package org.elasticsearch.xpack.inference.services.alibabacloudsearch;

import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.TransportVersion;
import org.elasticsearch.TransportVersions;
import org.elasticsearch.action.ActionListener;
Expand All @@ -32,7 +31,6 @@
import org.elasticsearch.inference.SimilarityMeasure;
import org.elasticsearch.inference.TaskType;
import org.elasticsearch.inference.configuration.SettingsConfigurationFieldType;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsBuilder;
import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
import org.elasticsearch.xpack.inference.external.http.sender.EmbeddingsInput;
Expand All @@ -59,7 +57,7 @@
import java.util.Map;

import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidTaskTypeException;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMap;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrDefaultEmpty;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrThrowIfNull;
Expand Down Expand Up @@ -135,7 +133,6 @@ public void parseRequestConfig(
taskSettingsMap,
chunkingSettings,
serviceSettingsMap,
TaskType.unsupportedTaskTypeErrorMsg(taskType, NAME),
ConfigurationParseContext.REQUEST
);

Expand Down Expand Up @@ -165,8 +162,7 @@ private static AlibabaCloudSearchModel createModelWithoutLoggingDeprecations(
Map<String, Object> serviceSettings,
Map<String, Object> taskSettings,
ChunkingSettings chunkingSettings,
@Nullable Map<String, Object> secretSettings,
String failureMessage
@Nullable Map<String, Object> secretSettings
) {
return createModel(
inferenceEntityId,
Expand All @@ -175,7 +171,6 @@ private static AlibabaCloudSearchModel createModelWithoutLoggingDeprecations(
taskSettings,
chunkingSettings,
secretSettings,
failureMessage,
ConfigurationParseContext.PERSISTENT
);
}
Expand All @@ -187,7 +182,6 @@ private static AlibabaCloudSearchModel createModel(
Map<String, Object> taskSettings,
ChunkingSettings chunkingSettings,
@Nullable Map<String, Object> secretSettings,
String failureMessage,
ConfigurationParseContext context
) {
return switch (taskType) {
Expand Down Expand Up @@ -229,7 +223,7 @@ private static AlibabaCloudSearchModel createModel(
secretSettings,
context
);
default -> throw new ElasticsearchStatusException(failureMessage, RestStatus.BAD_REQUEST);
default -> throw createInvalidTaskTypeException(inferenceEntityId, NAME, taskType, context);
};
}

Expand All @@ -255,8 +249,7 @@ public AlibabaCloudSearchModel parsePersistedConfigWithSecrets(
serviceSettingsMap,
taskSettingsMap,
chunkingSettings,
secretSettingsMap,
parsePersistedConfigErrorMsg(inferenceEntityId, NAME)
secretSettingsMap
);
}

Expand All @@ -276,8 +269,7 @@ public AlibabaCloudSearchModel parsePersistedConfig(String inferenceEntityId, Ta
serviceSettingsMap,
taskSettingsMap,
chunkingSettings,
null,
parsePersistedConfigErrorMsg(inferenceEntityId, NAME)
null
);
}

Expand Down
8 changes: 2 additions & 6 deletions ...n/java/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@

import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidTaskTypeException;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMap;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrDefaultEmpty;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrThrowIfNull;
Expand Down Expand Up @@ -204,7 +204,6 @@ public void parseRequestConfig(
taskSettingsMap,
chunkingSettings,
serviceSettingsMap,
TaskType.unsupportedTaskTypeErrorMsg(taskType, NAME),
ConfigurationParseContext.REQUEST
);

Expand Down Expand Up @@ -241,7 +240,6 @@ public Model parsePersistedConfigWithSecrets(
taskSettingsMap,
chunkingSettings,
secretSettingsMap,
parsePersistedConfigErrorMsg(modelId, NAME),
ConfigurationParseContext.PERSISTENT
);
}
Expand All @@ -263,7 +261,6 @@ public Model parsePersistedConfig(String modelId, TaskType taskType, Map<String,
taskSettingsMap,
chunkingSettings,
null,
parsePersistedConfigErrorMsg(modelId, NAME),
ConfigurationParseContext.PERSISTENT
);
}
Expand All @@ -285,7 +282,6 @@ private static AmazonBedrockModel createModel(
Map<String, Object> taskSettings,
ChunkingSettings chunkingSettings,
@Nullable Map<String, Object> secretSettings,
String failureMessage,
ConfigurationParseContext context
) {
switch (taskType) {
Expand Down Expand Up @@ -318,7 +314,7 @@ private static AmazonBedrockModel createModel(
checkChatCompletionProviderForTopKParameter(model);
return model;
}
default -> throw new ElasticsearchStatusException(failureMessage, RestStatus.BAD_REQUEST);
default -> throw createInvalidTaskTypeException(inferenceEntityId, NAME, taskType, context);
}
}

Expand Down
30 changes: 5 additions & 25 deletions .../src/main/java/org/elasticsearch/xpack/inference/services/anthropic/AnthropicService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@

package org.elasticsearch.xpack.inference.services.anthropic;

import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.TransportVersion;
import org.elasticsearch.TransportVersions;
import org.elasticsearch.action.ActionListener;
Expand All @@ -28,7 +27,6 @@
import org.elasticsearch.inference.SettingsConfiguration;
import org.elasticsearch.inference.TaskType;
import org.elasticsearch.inference.configuration.SettingsConfigurationFieldType;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
Expand All @@ -48,7 +46,7 @@

import static org.elasticsearch.xpack.inference.services.ServiceFields.MODEL_ID;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidTaskTypeException;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrDefaultEmpty;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrThrowIfNull;
import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwIfNotEmptyMap;
Expand Down Expand Up @@ -94,7 +92,6 @@ public void parseRequestConfig(
serviceSettingsMap,
taskSettingsMap,
serviceSettingsMap,
TaskType.unsupportedTaskTypeErrorMsg(taskType, NAME),
ConfigurationParseContext.REQUEST
);

Expand All @@ -113,16 +110,14 @@ private static AnthropicModel createModelFromPersistent(
TaskType taskType,
Map<String, Object> serviceSettings,
Map<String, Object> taskSettings,
@Nullable Map<String, Object> secretSettings,
String failureMessage
@Nullable Map<String, Object> secretSettings
) {
return createModel(
inferenceEntityId,
taskType,
serviceSettings,
taskSettings,
secretSettings,
failureMessage,
ConfigurationParseContext.PERSISTENT
);
}
Expand All @@ -133,7 +128,6 @@ private static AnthropicModel createModel(
Map<String, Object> serviceSettings,
Map<String, Object> taskSettings,
@Nullable Map<String, Object> secretSettings,
String failureMessage,
ConfigurationParseContext context
) {
return switch (taskType) {
Expand All @@ -146,7 +140,7 @@ private static AnthropicModel createModel(
secretSettings,
context
);
default -> throw new ElasticsearchStatusException(failureMessage, RestStatus.BAD_REQUEST);
default -> throw createInvalidTaskTypeException(inferenceEntityId, NAME, taskType, context);
};
}

Expand All @@ -161,29 +155,15 @@ public AnthropicModel parsePersistedConfigWithSecrets(
Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
Map<String, Object> secretSettingsMap = removeFromMapOrDefaultEmpty(secrets, ModelSecrets.SECRET_SETTINGS);

return createModelFromPersistent(
inferenceEntityId,
taskType,
serviceSettingsMap,
taskSettingsMap,
secretSettingsMap,
parsePersistedConfigErrorMsg(inferenceEntityId, NAME)
);
return createModelFromPersistent(inferenceEntityId, taskType, serviceSettingsMap, taskSettingsMap, secretSettingsMap);
}

@Override
public AnthropicModel parsePersistedConfig(String inferenceEntityId, TaskType taskType, Map<String, Object> config) {
Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);

return createModelFromPersistent(
inferenceEntityId,
taskType,
serviceSettingsMap,
taskSettingsMap,
null,
parsePersistedConfigErrorMsg(inferenceEntityId, NAME)
);
return createModelFromPersistent(inferenceEntityId, taskType, serviceSettingsMap, taskSettingsMap, null);
}

@Override
Expand Down
Loading