Skip to content

Commit 6aa0ea5

Browse files
jonathan-buttnerjonathan-buttner
committed
Addressing feedback
1 parent 578546f commit 6aa0ea5

File tree

17 files changed

+304
-215
lines changed

17 files changed

+304
-215
lines changed

‎x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/BaseMockEISAuthServerTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525

2626
import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.getModel;
2727
import static org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings.CCM_SUPPORTED_ENVIRONMENT;
28-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.ELSER_V2_ENDPOINT_ID;
28+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.ELSER_V2_ENDPOINT_ID;
2929

3030
public class BaseMockEISAuthServerTest extends ESRestTestCase {
3131

‎x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetModelsWithElasticInferenceServiceIT.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,12 @@
1818

1919
import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.getAllModels;
2020
import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.getModels;
21-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.ELSER_V2_ENDPOINT_ID;
22-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.GP_LLM_V2_CHAT_COMPLETION_ENDPOINT_ID;
23-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.GP_LLM_V2_COMPLETION_ENDPOINT_ID;
24-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.JINA_EMBED_V3_ENDPOINT_ID;
25-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.RAINBOW_SPRINKLES_ENDPOINT_ID;
26-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.RERANK_V1_ENDPOINT_ID;
21+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.ELSER_V2_ENDPOINT_ID;
22+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.GP_LLM_V2_CHAT_COMPLETION_ENDPOINT_ID;
23+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.GP_LLM_V2_COMPLETION_ENDPOINT_ID;
24+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.JINA_EMBED_V3_ENDPOINT_ID;
25+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.RAINBOW_SPRINKLES_ENDPOINT_ID;
26+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.RERANK_V1_ENDPOINT_ID;
2727
import static org.hamcrest.Matchers.hasSize;
2828
import static org.hamcrest.Matchers.is;
2929

‎x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,8 @@ public void testGetServicesWithCompletionTaskType() throws IOException {
176176
"hugging_face",
177177
"amazon_sagemaker",
178178
"mistral",
179-
"watsonxai"
179+
"watsonxai",
180+
"elastic"
180181
).toArray()
181182
)
182183
);

‎x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MockElasticInferenceServiceAuthorizationServer.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import org.junit.runners.model.Statement;
1818

1919
import static org.elasticsearch.core.Strings.format;
20-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.getEisAuthorizationResponseWithMultipleEndpoints;
20+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.getEisAuthorizationResponseWithMultipleEndpoints;
2121

2222
public class MockElasticInferenceServiceAuthorizationServer implements TestRule {
2323

‎x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationPoller;
2929
import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationTaskExecutor;
3030
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings;
31-
import org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests;
31+
import org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests;
3232
import org.junit.After;
3333
import org.junit.AfterClass;
3434
import org.junit.Before;
@@ -43,12 +43,12 @@
4343
import java.util.stream.Collectors;
4444

4545
import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl;
46-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.EIS_EMPTY_RESPONSE;
47-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.ELSER_V2_ENDPOINT_ID;
48-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.JINA_EMBED_V3_ENDPOINT_ID;
49-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.RAINBOW_SPRINKLES_ENDPOINT_ID;
50-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.RERANK_V1_ENDPOINT_ID;
51-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.getEisRainbowSprinklesAuthorizationResponse;
46+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.EIS_EMPTY_RESPONSE;
47+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.ELSER_V2_ENDPOINT_ID;
48+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.JINA_EMBED_V3_ENDPOINT_ID;
49+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.RAINBOW_SPRINKLES_ENDPOINT_ID;
50+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.RERANK_V1_ENDPOINT_ID;
51+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.getEisRainbowSprinklesAuthorizationResponse;
5252
import static org.hamcrest.Matchers.empty;
5353
import static org.hamcrest.Matchers.is;
5454
import static org.hamcrest.Matchers.not;
@@ -66,7 +66,7 @@ public class AuthorizationTaskExecutorIT extends ESSingleNodeTestCase {
6666

6767
private static final MockWebServer webServer = new MockWebServer();
6868
private static String gatewayUrl;
69-
private static AuthorizationResponseEntityTests.EisAuthorizationResponse chatCompletionResponse;
69+
private static ElasticInferenceServiceAuthorizationResponseEntityTests.EisAuthorizationResponse chatCompletionResponse;
7070

7171
private ModelRegistry modelRegistry;
7272
private AuthorizationTaskExecutor authorizationTaskExecutor;
@@ -273,7 +273,7 @@ public void testCreatesChatCompletion_AndThenCreatesTextEmbedding() throws Excep
273273
assertChatCompletionEndpointExists();
274274

275275
// Simulate that a text embedding model is now authorized
276-
var jinaEmbedResponse = AuthorizationResponseEntityTests.getEisJinaEmbedAuthorizationResponse(gatewayUrl);
276+
var jinaEmbedResponse = ElasticInferenceServiceAuthorizationResponseEntityTests.getEisJinaEmbedAuthorizationResponse(gatewayUrl);
277277
webServer.enqueue(new MockResponse().setResponseCode(200).setBody(jinaEmbedResponse.responseJson()));
278278

279279
restartPollingTaskAndWaitForAuthResponse();

‎x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorMultipleNodesIT.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettings;
2121
import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationPoller;
2222
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings;
23-
import org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests;
23+
import org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests;
2424
import org.junit.AfterClass;
2525
import org.junit.Before;
2626
import org.junit.BeforeClass;
@@ -35,9 +35,9 @@
3535
import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl;
3636
import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.cancelAuthorizationTask;
3737
import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.waitForTask;
38-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.EIS_EMPTY_RESPONSE;
39-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.RAINBOW_SPRINKLES_ENDPOINT_ID;
40-
import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.getEisRainbowSprinklesAuthorizationResponse;
38+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.EIS_EMPTY_RESPONSE;
39+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.RAINBOW_SPRINKLES_ENDPOINT_ID;
40+
import static org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests.getEisRainbowSprinklesAuthorizationResponse;
4141
import static org.hamcrest.CoreMatchers.is;
4242
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
4343
import static org.hamcrest.Matchers.not;
@@ -55,7 +55,7 @@ public class AuthorizationTaskExecutorMultipleNodesIT extends ESIntegTestCase {
5555
private static final String AUTH_TASK_ACTION = AuthorizationPoller.TASK_NAME + "[c]";
5656
private static final MockWebServer webServer = new MockWebServer();
5757
private static String gatewayUrl;
58-
private static AuthorizationResponseEntityTests.EisAuthorizationResponse chatCompletionResponse;
58+
private static ElasticInferenceServiceAuthorizationResponseEntityTests.EisAuthorizationResponse chatCompletionResponse;
5959

6060
@BeforeClass
6161
public static void initClass() throws IOException {

‎x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/CCMServiceIT.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMModel;
2424
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMService;
2525
import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings;
26-
import org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests;
26+
import org.elasticsearch.xpack.inference.services.elastic.response.ElasticInferenceServiceAuthorizationResponseEntityTests;
2727
import org.junit.After;
2828
import org.junit.AfterClass;
2929
import org.junit.Before;
@@ -49,7 +49,7 @@ public class CCMServiceIT extends CCMSingleNodeIT {
4949

5050
private static final MockWebServer webServer = new MockWebServer();
5151
private static String gatewayUrl;
52-
private static AuthorizationResponseEntityTests.EisAuthorizationResponse chatCompletionResponse;
52+
private static ElasticInferenceServiceAuthorizationResponseEntityTests.EisAuthorizationResponse chatCompletionResponse;
5353

5454
private AuthorizationTaskExecutor authorizationTaskExecutor;
5555
private ModelRegistry modelRegistry;
@@ -79,7 +79,9 @@ public static void initClass() throws IOException {
7979

8080
webServer.start();
8181
gatewayUrl = getUrl(webServer);
82-
chatCompletionResponse = AuthorizationResponseEntityTests.getEisRainbowSprinklesAuthorizationResponse(gatewayUrl);
82+
chatCompletionResponse = ElasticInferenceServiceAuthorizationResponseEntityTests.getEisRainbowSprinklesAuthorizationResponse(
83+
gatewayUrl
84+
);
8385
}
8486

8587
@Before

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportGetInferenceServicesAction.java‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
import org.elasticsearch.xpack.core.inference.action.GetInferenceServicesAction;
2727
import org.elasticsearch.xpack.inference.external.http.sender.Sender;
2828
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService;
29-
import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationModel;
29+
import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationModel;
3030
import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationRequestHandler;
3131

3232
import java.util.ArrayList;
@@ -122,7 +122,7 @@ private void getServiceConfigurationsForServicesAndEis(
122122
ArrayList<Map.Entry<String, InferenceService>> availableServices,
123123
@Nullable TaskType requestedTaskType
124124
) {
125-
SubscribableListener.<AuthorizationModel>newForked(authModelListener -> {
125+
SubscribableListener.<ElasticInferenceServiceAuthorizationModel>newForked(authModelListener -> {
126126
// Executing on a separate thread because there's a chance the authorization call needs to do some initialization for the Sender
127127
threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(() -> getEisAuthorization(authModelListener, eisSender));
128128
}).<List<InferenceServiceConfiguration>>andThen((configurationListener, authorizationModel) -> {
@@ -133,12 +133,14 @@ private void getServiceConfigurationsForServicesAndEis(
133133
return;
134134
}
135135

136-
var config = ElasticInferenceService.createConfiguration(authorizationModel.getTaskTypes());
136+
// If there was a requested task type and the authorization response from EIS doesn't support it, we'll exclude EIS as a valid
137+
// service
137138
if (requestedTaskType != null && authorizationModel.getTaskTypes().contains(requestedTaskType) == false) {
138139
configurationListener.onResponse(serviceConfigs);
139140
return;
140141
}
141142

143+
var config = ElasticInferenceService.createConfiguration(authorizationModel.getTaskTypes());
142144
serviceConfigs.add(config);
143145
serviceConfigs.sort(Comparator.comparing(InferenceServiceConfiguration::getService));
144146
configurationListener.onResponse(serviceConfigs);
@@ -150,14 +152,14 @@ private void getServiceConfigurationsForServicesAndEis(
150152
);
151153
}
152154

153-
private void getEisAuthorization(ActionListener<AuthorizationModel> listener, Sender sender) {
155+
private void getEisAuthorization(ActionListener<ElasticInferenceServiceAuthorizationModel> listener, Sender sender) {
154156
var disabledServiceListener = listener.delegateResponse((delegate, e) -> {
155157
logger.warn(
156158
"Failed to retrieve authorization information from the "
157159
+ "Elastic Inference Service while determining service configurations. Marking service as disabled.",
158160
e
159161
);
160-
delegate.onResponse(AuthorizationModel.empty());
162+
delegate.onResponse(ElasticInferenceServiceAuthorizationModel.unauthorized());
161163
});
162164

163165
eisAuthorizationRequestHandler.getAuthorization(disabledServiceListener, sender);

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/AuthorizationPoller.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -306,15 +306,15 @@ private void shouldSendAuthRequest(ActionListener<Consumer<ActionListener<Void>>
306306
}
307307

308308
private void sendRequest(ActionListener<Void> listener) {
309-
SubscribableListener.<AuthorizationModel>newForked(
309+
SubscribableListener.<ElasticInferenceServiceAuthorizationModel>newForked(
310310
authModelListener -> authorizationHandler.getAuthorization(authModelListener, sender)
311311
)
312312
.andThenApply(this::getNewInferenceEndpointsToStore)
313313
.<Void>andThen((storeListener, newInferenceIds) -> storePreconfiguredModels(newInferenceIds, storeListener))
314314
.addListener(listener);
315315
}
316316

317-
private List<? extends Model> getNewInferenceEndpointsToStore(AuthorizationModel authModel) {
317+
private List<? extends Model> getNewInferenceEndpointsToStore(ElasticInferenceServiceAuthorizationModel authModel) {
318318
logger.debug("Received authorization response, {}", authModel);
319319

320320
var scopedAuthModel = authModel.newLimitedToTaskTypes(EnumSet.copyOf(IMPLEMENTED_TASK_TYPES));

0 commit comments

Comments
 (0)