Adding inference id to error parser log message

jonathan-buttner · jonathan-buttner · commit 8d1bd22793aa · 2025-05-08T12:03:53.000-04:00
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/StreamingHttpResult.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/StreamingHttpResult.java
@@ -31,7 +31,7 @@ public void onSubscribe(Flow.Subscription subscription) {
 
             @Override
             public void onNext(byte[] item) {
-                subscriber.onNext(new HttpResult(response(), item, httpRequest));
+                subscriber.onNext(new HttpResult(response(), item));
             }
 
             @Override
@@ -70,7 +70,7 @@ public void onError(Throwable throwable) {
 
             @Override
             public void onComplete() {
-                fullResponse.onResponse(new HttpResult(response, stream.toByteArray(), httpRequest));
+                fullResponse.onResponse(new HttpResult(response, stream.toByteArray()));
             }
         });
     }
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/BaseResponseHandler.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/BaseResponseHandler.java
@@ -36,7 +36,7 @@ public abstract class BaseResponseHandler implements ResponseHandler {
     public static final String METHOD_NOT_ALLOWED = "Received a method not allowed status code";
 
     protected final String requestType;
-    private final ResponseParser parseFunction;
+    protected final ResponseParser parseFunction;
     private final BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction;
     private final boolean canHandleStreamingResponses;
 
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceSettings.java
@@ -151,12 +151,20 @@ public record TextEmbeddingSettings(
         @Nullable Integer maxInputTokens,
         @Nullable DenseVectorFieldMapper.ElementType elementType
     ) implements ToXContentFragment, Writeable {
+        // This specifies float for the element type but null for all other settings
+        public static final TextEmbeddingSettings DEFAULT_FLOAT = new TextEmbeddingSettings(
+            null,
+            null,
+            null,
+            DenseVectorFieldMapper.ElementType.FLOAT
+        );
 
-        public static final TextEmbeddingSettings EMPTY = new TextEmbeddingSettings(null, null, null, null);
+        // This refers to settings that are not related to the text embedding task type (all the settings should be null)
+        public static final TextEmbeddingSettings NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS = new TextEmbeddingSettings(null, null, null, null);
 
         public static TextEmbeddingSettings fromMap(Map<String, Object> map, TaskType taskType, ValidationException validationException) {
             if (taskType != TaskType.TEXT_EMBEDDING) {
-                return EMPTY;
+                return NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS;
             }
 
             SimilarityMeasure similarity = extractSimilarity(map, ModelConfigurations.SERVICE_SETTINGS, validationException);
@@ -207,7 +215,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     private final ErrorResponseParser errorParser;
 
     public CustomServiceSettings(
-        @Nullable TextEmbeddingSettings textEmbeddingSettings,
+        TextEmbeddingSettings textEmbeddingSettings,
         String url,
         @Nullable Map<String, String> headers,
         @Nullable QueryParameters queryParameters,
@@ -216,7 +224,7 @@ public CustomServiceSettings(
         @Nullable RateLimitSettings rateLimitSettings,
         ErrorResponseParser errorParser
     ) {
-        this.textEmbeddingSettings = textEmbeddingSettings == null ? TextEmbeddingSettings.EMPTY : textEmbeddingSettings;
+        this.textEmbeddingSettings = Objects.requireNonNull(textEmbeddingSettings);
         this.url = Objects.requireNonNull(url);
         this.headers = Collections.unmodifiableMap(Objects.requireNonNullElse(headers, Map.of()));
         this.queryParameters = Objects.requireNonNullElse(queryParameters, QueryParameters.EMPTY);
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomRequestManagerTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomRequestManagerTests.java
@@ -56,7 +56,7 @@ public void testCreateRequest_ThrowsException_ForInvalidUrl() {
             """;
 
         var serviceSettings = new CustomServiceSettings(
-            CustomServiceSettings.TextEmbeddingSettings.EMPTY,
+            CustomServiceSettings.TextEmbeddingSettings.NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS,
             "${url}",
             null,
             null,
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceSettingsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceSettingsTests.java
@@ -200,7 +200,7 @@ public void testFromMap_WithOptionalsNotSpecified() {
             settings,
             is(
                 new CustomServiceSettings(
-                    CustomServiceSettings.TextEmbeddingSettings.EMPTY,
+                    CustomServiceSettings.TextEmbeddingSettings.DEFAULT_FLOAT,
                     url,
                     Map.of(),
                     null,
@@ -664,7 +664,7 @@ public void testFromMap_ReturnsError_IfTaskTypeIsInvalid() {
 
     public void testXContent() throws IOException {
         var entity = new CustomServiceSettings(
-            CustomServiceSettings.TextEmbeddingSettings.EMPTY,
+            CustomServiceSettings.TextEmbeddingSettings.NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS,
             "http://www.abc.com",
             Map.of("key", "value"),
             null,
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
@@ -141,14 +141,8 @@ private static SenderService createService(ThreadPool threadPool, HttpClientMana
     }
 
     private static Map<String, Object> createServiceSettingsMap(TaskType taskType) {
-        return new HashMap<>(
+        var settingsMap = new HashMap<>(
             Map.of(
-                ServiceFields.SIMILARITY,
-                SimilarityMeasure.DOT_PRODUCT.toString(),
-                ServiceFields.DIMENSIONS,
-                1536,
-                ServiceFields.MAX_INPUT_TOKENS,
-                512,
                 CustomServiceSettings.URL,
                 "http://www.abc.com",
                 CustomServiceSettings.HEADERS,
@@ -168,6 +162,17 @@ private static Map<String, Object> createServiceSettingsMap(TaskType taskType) {
                 )
             )
         );
+
+        if (taskType == TaskType.TEXT_EMBEDDING) {
+            settingsMap.putAll(Map.of(ServiceFields.SIMILARITY,
+                SimilarityMeasure.DOT_PRODUCT.toString(),
+                ServiceFields.DIMENSIONS,
+                1536,
+                ServiceFields.MAX_INPUT_TOKENS,
+                512));
+        }
+
+        return settingsMap;
     }
 
     private static Map<String, Object> createResponseParserMap(TaskType taskType) {
@@ -229,7 +234,7 @@ private static CustomModel createInternalEmbeddingModel(
             CustomService.NAME,
             new CustomServiceSettings(
                 new CustomServiceSettings.TextEmbeddingSettings(
-                    SimilarityMeasure.DOT_PRODUCT,
+                    similarityMeasure,
                     123,
                     456,
                     DenseVectorFieldMapper.ElementType.FLOAT
@@ -253,7 +258,7 @@ private static CustomModel createCustomModel(TaskType taskType, CustomResponsePa
             taskType,
             CustomService.NAME,
             new CustomServiceSettings(
-                CustomServiceSettings.TextEmbeddingSettings.EMPTY,
+                getDefaultTextEmbeddingSettings(taskType),
                 url,
                 Map.of("key", "value"),
                 QueryParameters.EMPTY,
@@ -267,6 +272,12 @@ private static CustomModel createCustomModel(TaskType taskType, CustomResponsePa
         );
     }
 
+    private static CustomServiceSettings.TextEmbeddingSettings getDefaultTextEmbeddingSettings(TaskType taskType) {
+        return taskType == TaskType.TEXT_EMBEDDING
+            ? CustomServiceSettings.TextEmbeddingSettings.DEFAULT_FLOAT
+            : CustomServiceSettings.TextEmbeddingSettings.NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS;
+    }
+
     public void testInfer_HandlesTextEmbeddingRequest_OpenAI_Format() throws IOException {
         try (var service = createService(threadPool, clientManager)) {
             String responseJson = """
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
@@ -100,7 +100,7 @@ public void testCreateRequest_QueryParametersAreEscaped_AndEncoded() {
             """;
 
         var serviceSettings = new CustomServiceSettings(
-            CustomServiceSettings.TextEmbeddingSettings.EMPTY,
+            CustomServiceSettings.TextEmbeddingSettings.NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS,
             "http://www.elastic.co",
             null,
             // escaped characters retrieved from here: https://docs.microfocus.com/OMi/10.62/Content/OMi/ExtGuide/ExtApps/URL_encoding.htm
@@ -203,7 +203,7 @@ public void testCreateRequest_HandlesQuery() throws IOException {
             """;
 
         var serviceSettings = new CustomServiceSettings(
-            CustomServiceSettings.TextEmbeddingSettings.EMPTY,
+            CustomServiceSettings.TextEmbeddingSettings.NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS,
             "http://www.elastic.co",
             null,
             null,
@@ -245,7 +245,7 @@ public void testCreateRequest_IgnoresNonStringFields_ForStringParams() throws IO
             """;
 
         var serviceSettings = new CustomServiceSettings(
-            CustomServiceSettings.TextEmbeddingSettings.EMPTY,
+            CustomServiceSettings.TextEmbeddingSettings.NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS,
             "http://www.elastic.co",
             Map.of(HttpHeaders.ACCEPT, Strings.format("${task.key}")),
             null,
@@ -276,7 +276,7 @@ public void testCreateRequest_ThrowsException_ForInvalidUrl() {
             """;
 
         var serviceSettings = new CustomServiceSettings(
-            CustomServiceSettings.TextEmbeddingSettings.EMPTY,
+            CustomServiceSettings.TextEmbeddingSettings.NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS,
             "${url}",
             Map.of(HttpHeaders.ACCEPT, Strings.format("${task.key}")),
             null,
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/response/ErrorResponseParserTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/response/ErrorResponseParserTests.java
@@ -17,6 +17,7 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -26,6 +27,7 @@
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.sameInstance;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class ErrorResponseParserTests extends ESTestCase {
 
@@ -81,7 +83,7 @@ public void testErrorResponse_ExtractsError() throws IOException {
             }""");
 
         var parser = new ErrorResponseParser("$.error.message");
-        var error = parser.apply(result);
+        var error = parser.apply(getMockRequest(), result);
         assertThat(error, is(new ErrorResponse("test_error_message")));
     }
 
@@ -98,7 +100,7 @@ public void testFromResponse_WithOtherFieldsPresent() throws IOException {
             """;
 
         var parser = new ErrorResponseParser("$.error.message");
-        var error = parser.apply(getMockResult(responseJson));
+        var error = parser.apply(getMockRequest(), getMockResult(responseJson));
 
         assertThat(error, is(new ErrorResponse("You didn't provide an API key")));
     }
@@ -113,7 +115,7 @@ public void testFromResponse_noMessage() throws IOException {
             """;
 
         var parser = new ErrorResponseParser("$.error.message");
-        var error = parser.apply(getMockResult(responseJson));
+        var error = parser.apply(getMockRequest(), getMockResult(responseJson));
 
         assertThat(error, sameInstance(ErrorResponse.UNDEFINED_ERROR));
         assertThat(error.getErrorMessage(), is(""));
@@ -125,7 +127,7 @@ public void testErrorResponse_ReturnsUndefinedObjectIfNoError() throws IOExcepti
             {"noerror":true}""");
 
         var parser = new ErrorResponseParser("$.error.message");
-        var error = parser.apply(mockResult);
+        var error = parser.apply(getMockRequest(), mockResult);
 
         assertThat(error, sameInstance(ErrorResponse.UNDEFINED_ERROR));
     }
@@ -134,12 +136,23 @@ public void testErrorResponse_ReturnsUndefinedObjectIfNotJson() {
         var result = new HttpResult(mock(HttpResponse.class), Strings.toUTF8Bytes("not a json string"));
 
         var parser = new ErrorResponseParser("$.error.message");
-        var error = parser.apply(result);
+        var error = parser.apply(getMockRequest(), result);
         assertThat(error, sameInstance(ErrorResponse.UNDEFINED_ERROR));
     }
 
     private static HttpResult getMockResult(String jsonString) throws IOException {
         var response = mock(HttpResponse.class);
         return new HttpResult(response, Strings.toUTF8Bytes(XContentHelper.stripWhitespace(jsonString)));
     }
+
+    private static Request getMockRequest() {
+        return getMockRequest("id");
+    }
+
+    private static Request getMockRequest(String inferenceId) {
+        var mockRequest = mock(Request.class);
+        when(mockRequest.getInferenceEntityId()).thenReturn(inferenceId);
+
+        return mockRequest;
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@ public void onSubscribe(Flow.Subscription subscription) {`
`31`	`31`
`32`	`32`	`@Override`
`33`	`33`	`public void onNext(byte[] item) {`
`34`		`- subscriber.onNext(new HttpResult(response(), item, httpRequest));`
	`34`	`+ subscriber.onNext(new HttpResult(response(), item));`
`35`	`35`	`}`
`36`	`36`
`37`	`37`	`@Override`
`@@ -70,7 +70,7 @@ public void onError(Throwable throwable) {`
`70`	`70`
`71`	`71`	`@Override`
`72`	`72`	`public void onComplete() {`
`73`		`- fullResponse.onResponse(new HttpResult(response, stream.toByteArray(), httpRequest));`
	`73`	`+ fullResponse.onResponse(new HttpResult(response, stream.toByteArray()));`
`74`	`74`	`}`
`75`	`75`	`});`
`76`	`76`	`}`