Skip to content

Commit 4242a37

Browse files
jonathan-buttnerjonathan-buttner
committed
Refactoring the error parsing logic
1 parent a0984c7 commit 4242a37

File tree

11 files changed

+113
-94
lines changed

11 files changed

+113
-94
lines changed

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/HttpClient.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ public void cancelled() {
137137
private void respondUsingUtilityThread(HttpResponse response, HttpRequest request, ActionListener<HttpResult> listener) {
138138
threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(() -> {
139139
try {
140-
listener.onResponse(HttpResult.create(settings.getMaxResponseSize(), response, request));
140+
listener.onResponse(HttpResult.create(settings.getMaxResponseSize(), response));
141141
} catch (Exception e) {
142142
throttlerManager.warn(
143143
logger,

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/HttpResult.java‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,16 @@
1212
import org.elasticsearch.core.Streams;
1313
import org.elasticsearch.rest.RestStatus;
1414
import org.elasticsearch.xpack.inference.common.SizeLimitInputStream;
15-
import org.elasticsearch.xpack.inference.external.request.HttpRequest;
1615

1716
import java.io.ByteArrayOutputStream;
1817
import java.io.IOException;
1918
import java.io.InputStream;
2019
import java.util.Objects;
2120

22-
public record HttpResult(HttpResponse response, byte[] body, HttpRequest request) {
21+
public record HttpResult(HttpResponse response, byte[] body) {
2322

24-
public static HttpResult create(ByteSizeValue maxResponseSize, HttpResponse response, HttpRequest request) throws IOException {
25-
return new HttpResult(response, limitBody(maxResponseSize, response), request);
23+
public static HttpResult create(ByteSizeValue maxResponseSize, HttpResponse response) throws IOException {
24+
return new HttpResult(response, limitBody(maxResponseSize, response));
2625
}
2726

2827
private static byte[] limitBody(ByteSizeValue maxResponseSize, HttpResponse response) throws IOException {
@@ -44,7 +43,6 @@ private static byte[] limitBody(ByteSizeValue maxResponseSize, HttpResponse resp
4443
public HttpResult {
4544
Objects.requireNonNull(response);
4645
Objects.requireNonNull(body);
47-
Objects.requireNonNull(request);
4846
}
4947

5048
public boolean isBodyEmpty() {

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ public static CustomModel of(CustomModel model, Map<String, Object> taskSettings
3333
}
3434

3535
public CustomModel(
36-
String modelId,
36+
String inferenceId,
3737
TaskType taskType,
3838
String service,
3939
Map<String, Object> serviceSettings,
@@ -42,26 +42,26 @@ public CustomModel(
4242
ConfigurationParseContext context
4343
) {
4444
this(
45-
modelId,
45+
inferenceId,
4646
taskType,
4747
service,
48-
CustomServiceSettings.fromMap(serviceSettings, context, taskType),
48+
CustomServiceSettings.fromMap(serviceSettings, context, taskType, inferenceId),
4949
CustomTaskSettings.fromMap(taskSettings),
5050
CustomSecretSettings.fromMap(secrets)
5151
);
5252
}
5353

5454
// should only be used for testing
5555
CustomModel(
56-
String modelId,
56+
String inferenceId,
5757
TaskType taskType,
5858
String service,
5959
CustomServiceSettings serviceSettings,
6060
CustomTaskSettings taskSettings,
6161
@Nullable CustomSecretSettings secretSettings
6262
) {
6363
this(
64-
new ModelConfigurations(modelId, taskType, service, serviceSettings, taskSettings),
64+
new ModelConfigurations(inferenceId, taskType, service, serviceSettings, taskSettings),
6565
new ModelSecrets(secretSettings),
6666
serviceSettings
6767
);

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,12 @@ public class CustomServiceSettings extends FilteredXContentObject implements Ser
6565
private static final RateLimitSettings DEFAULT_RATE_LIMIT_SETTINGS = new RateLimitSettings(10_000);
6666
private static final String RESPONSE_SCOPE = String.join(".", ModelConfigurations.SERVICE_SETTINGS, RESPONSE);
6767

68-
public static CustomServiceSettings fromMap(Map<String, Object> map, ConfigurationParseContext context, TaskType taskType) {
68+
public static CustomServiceSettings fromMap(
69+
Map<String, Object> map,
70+
ConfigurationParseContext context,
71+
TaskType taskType,
72+
String inferenceId
73+
) {
6974
ValidationException validationException = new ValidationException();
7075

7176
var textEmbeddingSettings = TextEmbeddingSettings.fromMap(map, taskType, validationException);
@@ -110,7 +115,7 @@ public static CustomServiceSettings fromMap(Map<String, Object> map, Configurati
110115
validationException
111116
);
112117

113-
var errorParser = ErrorResponseParser.fromMap(errorParserMap, RESPONSE_SCOPE, validationException);
118+
var errorParser = ErrorResponseParser.fromMap(errorParserMap, RESPONSE_SCOPE, inferenceId, validationException);
114119

115120
RateLimitSettings rateLimitSettings = RateLimitSettings.of(
116121
map,
@@ -151,6 +156,7 @@ public record TextEmbeddingSettings(
151156
@Nullable Integer maxInputTokens,
152157
@Nullable DenseVectorFieldMapper.ElementType elementType
153158
) implements ToXContentFragment, Writeable {
159+
154160
// This specifies float for the element type but null for all other settings
155161
public static final TextEmbeddingSettings DEFAULT_FLOAT = new TextEmbeddingSettings(
156162
null,

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/response/ErrorResponseParser.java‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
2525

2626
import java.io.IOException;
27+
import java.nio.charset.StandardCharsets;
2728
import java.util.Map;
2829
import java.util.Objects;
2930
import java.util.function.Function;
@@ -38,10 +39,12 @@ public class ErrorResponseParser implements ToXContentFragment, Function<HttpRes
3839
public static final String MESSAGE_PATH = "path";
3940

4041
private final String messagePath;
42+
private final String inferenceId;
4143

4244
public static ErrorResponseParser fromMap(
4345
Map<String, Object> responseParserMap,
4446
String scope,
47+
String inferenceId,
4548
ValidationException validationException
4649
) {
4750
var path = extractRequiredString(responseParserMap, MESSAGE_PATH, String.join(".", scope, ERROR_PARSER), validationException);
@@ -50,19 +53,22 @@ public static ErrorResponseParser fromMap(
5053
throw validationException;
5154
}
5255

53-
return new ErrorResponseParser(path);
56+
return new ErrorResponseParser(path, inferenceId);
5457
}
5558

56-
public ErrorResponseParser(String messagePath) {
59+
public ErrorResponseParser(String messagePath, String inferenceId) {
5760
this.messagePath = Objects.requireNonNull(messagePath);
61+
this.inferenceId = Objects.requireNonNull(inferenceId);
5862
}
5963

6064
public ErrorResponseParser(StreamInput in) throws IOException {
6165
this.messagePath = in.readString();
66+
this.inferenceId = in.readString();
6267
}
6368

6469
public void writeTo(StreamOutput out) throws IOException {
6570
out.writeString(messagePath);
71+
out.writeString(inferenceId);
6672
}
6773

6874
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
@@ -104,16 +110,19 @@ public ErrorResponse apply(HttpResult httpResult) {
104110
var errorText = toType(MapPathExtractor.extract(map, messagePath).extractedObject(), String.class, messagePath);
105111
return new ErrorResponse(errorText);
106112
} catch (Exception e) {
113+
var resultAsString = new String(httpResult.body(), StandardCharsets.UTF_8);
114+
107115
logger.info(
108116
Strings.format(
109-
"Failed to parse error object for custom service inference id [%s], message path: [%s]",
110-
httpResult.request().inferenceEntityId(),
111-
messagePath
117+
"Failed to parse error object for custom service inference id [%s], message path: [%s], result as string: [%s]",
118+
inferenceId,
119+
messagePath,
120+
resultAsString
112121
),
113122
e
114123
);
115-
}
116124

117-
return ErrorResponse.UNDEFINED_ERROR;
125+
return new ErrorResponse(Strings.format("Unable to parse the error, response body: [%s]", resultAsString));
126+
}
118127
}
119128
}

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomModelTests.java‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -74,23 +74,23 @@ public void testOverride() {
7474
}
7575

7676
public static CustomModel createModel(
77-
String modelId,
77+
String inferenceId,
7878
TaskType taskType,
7979
Map<String, Object> serviceSettings,
8080
Map<String, Object> taskSettings,
8181
@Nullable Map<String, Object> secrets
8282
) {
83-
return new CustomModel(modelId, taskType, CustomService.NAME, serviceSettings, taskSettings, secrets, null);
83+
return new CustomModel(inferenceId, taskType, CustomService.NAME, serviceSettings, taskSettings, secrets, null);
8484
}
8585

8686
public static CustomModel createModel(
87-
String modelId,
87+
String inferenceId,
8888
TaskType taskType,
8989
CustomServiceSettings serviceSettings,
9090
CustomTaskSettings taskSettings,
9191
@Nullable CustomSecretSettings secretSettings
9292
) {
93-
return new CustomModel(modelId, taskType, CustomService.NAME, serviceSettings, taskSettings, secretSettings);
93+
return new CustomModel(inferenceId, taskType, CustomService.NAME, serviceSettings, taskSettings, secretSettings);
9494
}
9595

9696
public static CustomModel getTestModel() {
@@ -102,6 +102,7 @@ public static CustomModel getTestModel(TaskType taskType, CustomResponseParser r
102102
}
103103

104104
public static CustomModel getTestModel(TaskType taskType, CustomResponseParser responseParser, String url) {
105+
var inferenceId = "inference_id";
105106
Integer dims = 1536;
106107
Integer maxInputTokens = 512;
107108
Map<String, String> headers = Map.of(HttpHeaders.AUTHORIZATION, "${" + secretSettingsKey + "}");
@@ -120,12 +121,12 @@ public static CustomModel getTestModel(TaskType taskType, CustomResponseParser r
120121
requestContentString,
121122
responseParser,
122123
new RateLimitSettings(10_000),
123-
new ErrorResponseParser("$.error.message")
124+
new ErrorResponseParser("$.error.message", inferenceId)
124125
);
125126

126127
CustomTaskSettings taskSettings = new CustomTaskSettings(Map.of(taskSettingsKey, taskSettingsValue));
127128
CustomSecretSettings secretSettings = new CustomSecretSettings(Map.of(secretSettingsKey, secretSettingsValue));
128129

129-
return CustomModelTests.createModel("service", taskType, serviceSettings, taskSettings, secretSettings);
130+
return CustomModelTests.createModel(inferenceId, taskType, serviceSettings, taskSettings, secretSettings);
130131
}
131132
}

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomRequestManagerTests.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ public void tearDown() throws Exception {
4949
}
5050

5151
public void testCreateRequest_ThrowsException_ForInvalidUrl() {
52+
var inferenceId = "inference_id";
53+
5254
var requestContentString = """
5355
{
5456
"input": ${input}
@@ -63,11 +65,11 @@ public void testCreateRequest_ThrowsException_ForInvalidUrl() {
6365
requestContentString,
6466
new RerankResponseParser("$.result.score"),
6567
new RateLimitSettings(10_000),
66-
new ErrorResponseParser("$.error.message")
68+
new ErrorResponseParser("$.error.message", inferenceId)
6769
);
6870

6971
var model = CustomModelTests.createModel(
70-
"service",
72+
inferenceId,
7173
TaskType.RERANK,
7274
serviceSettings,
7375
new CustomTaskSettings(Map.of("url", "^")),

0 commit comments

Comments
 (0)