use ConstructingObjectParser for response parsing

Author: brendan-jugan-elastic

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/elastic/ElasticInferenceServiceDenseTextEmbeddingsResponseEntity.java

@@ -7,30 +7,23 @@
 
 package org.elasticsearch.xpack.inference.external.response.elastic;
 
-import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
 import org.elasticsearch.common.xcontent.XContentParserUtils;
+import org.elasticsearch.xcontent.ConstructingObjectParser;
+import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.XContentFactory;
-import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.request.Request;
 
 import java.io.IOException;
-import java.util.Collections;
 import java.util.List;
 
-import static org.elasticsearch.common.xcontent.XContentParserUtils.ensureExpectedToken;
-import static org.elasticsearch.common.xcontent.XContentParserUtils.parseList;
-import static org.elasticsearch.xpack.inference.external.response.XContentUtils.moveToFirstToken;
-import static org.elasticsearch.xpack.inference.external.response.XContentUtils.positionParserAtTokenAfterField;
+import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
 
 public class ElasticInferenceServiceDenseTextEmbeddingsResponseEntity {
 
-    private static final String FAILED_TO_FIND_FIELD_TEMPLATE =
-        "Failed to find required field [%s] in Elastic Inference Service dense text embeddings response";
-
     /**
      * Parses the Elastic Inference Service Dense Text Embeddings response.
      *
@@ -64,43 +57,51 @@ public class ElasticInferenceServiceDenseTextEmbeddingsResponseEntity {
      *     </code>
      * </pre>
      */
-
     public static TextEmbeddingFloatResults fromResponse(Request request, HttpResult response) throws IOException {
-        var parserConfig = XContentParserConfiguration.EMPTY.withDeprecationHandler(LoggingDeprecationHandler.INSTANCE);
-
-        try (XContentParser jsonParser = XContentFactory.xContent(XContentType.JSON).createParser(parserConfig, response.body())) {
-            moveToFirstToken(jsonParser);
-
-            XContentParser.Token token = jsonParser.currentToken();
-            ensureExpectedToken(XContentParser.Token.START_OBJECT, token, jsonParser);
+        try (var p = XContentFactory.xContent(XContentType.JSON).createParser(XContentParserConfiguration.EMPTY, response.body())) {
+            return EmbeddingFloatResult.PARSER.apply(p, null).toTextEmbeddingFloatResults();
+        }
+    }
 
-            positionParserAtTokenAfterField(jsonParser, "data", FAILED_TO_FIND_FIELD_TEMPLATE);
+    public record EmbeddingFloatResult(List<EmbeddingFloatResultEntry> embeddingResults) {
+        @SuppressWarnings("unchecked")
+        public static final ConstructingObjectParser<EmbeddingFloatResult, Void> PARSER = new ConstructingObjectParser<>(
+            EmbeddingFloatResult.class.getSimpleName(),
+            true,
+            args -> new EmbeddingFloatResult((List<EmbeddingFloatResultEntry>) args[0])
+        );
 
-            List<TextEmbeddingFloatResults.Embedding> parsedEmbeddings = parseList(
-                jsonParser,
-                (parser, index) -> ElasticInferenceServiceDenseTextEmbeddingsResponseEntity.parseTextEmbeddingObject(parser)
+        static {
+            // Custom field declaration to handle array of arrays format
+            PARSER.declareField(
+                constructorArg(),
+                (parser, context) -> {
+                    return XContentParserUtils.parseList(parser, (p, index) -> {
+                        List<Float> embedding = XContentParserUtils.parseList(p, (innerParser, innerIndex) -> innerParser.floatValue());
+                        return EmbeddingFloatResultEntry.fromFloatArray(embedding);
+                    });
+                },
+                new ParseField("data"),
+                org.elasticsearch.xcontent.ObjectParser.ValueType.OBJECT_ARRAY
             );
-
-            if (parsedEmbeddings.isEmpty()) {
-                return new TextEmbeddingFloatResults(Collections.emptyList());
-            }
-
-            return new TextEmbeddingFloatResults(parsedEmbeddings);
         }
-    }
 
-    private static TextEmbeddingFloatResults.Embedding parseTextEmbeddingObject(XContentParser parser) throws IOException {
-        List<Float> embeddingValueList = parseList(
-            parser,
-            ElasticInferenceServiceDenseTextEmbeddingsResponseEntity::parseEmbeddingFloatValueList
-        );
-        return TextEmbeddingFloatResults.Embedding.of(embeddingValueList);
+        public TextEmbeddingFloatResults toTextEmbeddingFloatResults() {
+            return new TextEmbeddingFloatResults(
+                embeddingResults.stream().map(entry -> TextEmbeddingFloatResults.Embedding.of(entry.embedding)).toList()
+            );
+        }
     }
 
-    private static float parseEmbeddingFloatValueList(XContentParser parser) throws IOException {
-        XContentParser.Token token = parser.currentToken();
-        XContentParserUtils.ensureExpectedToken(XContentParser.Token.VALUE_NUMBER, token, parser);
-        return parser.floatValue();
+    /**
+     * Represents a single embedding entry in the response.
+     * For the Elastic Inference Service, each entry is just an array of floats (no wrapper object).
+     * This is a simpler wrapper that just holds the float array.
+     */
+    public record EmbeddingFloatResultEntry(List<Float> embedding) {
+        public static EmbeddingFloatResultEntry fromFloatArray(List<Float> floats) {
+            return new EmbeddingFloatResultEntry(floats);
+        }
     }
 
     private ElasticInferenceServiceDenseTextEmbeddingsResponseEntity() {}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/action/ElasticInferenceServiceActionCreatorTests.java

@@ -48,6 +48,7 @@
 import static org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests.createSender;
 import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings;
 import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
@@ -407,10 +408,7 @@ public void testSend_FailsFromInvalidResponseFormat_ForDenseTextEmbeddingsAction
             );
 
             var thrownException = expectThrows(ElasticsearchException.class, () -> listener.actionGet(TIMEOUT));
-            assertThat(
-                thrownException.getMessage(),
-                is("Failed to parse object: expecting token of type [START_ARRAY] but found [START_OBJECT]")
-            );
+            assertThat(thrownException.getMessage(), containsString("[EmbeddingFloatResult] failed to parse field [data]"));
 
             assertThat(webServer.requests(), hasSize(1));
             assertNull(webServer.requests().get(0).getUri().getQuery());