Fixing tests

Author: jonathan-buttner

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

@@ -30,17 +30,19 @@ public class CompletionResponseParser extends BaseCustomResponseParser<ChatCompl
 
     private final String completionResultPath;
 
-<<<<<<< HEAD
-    public static CompletionResponseParser fromMap(Map<String, Object> responseParserMap, String scope, ValidationException validationException) {
-        var path = extractRequiredString(responseParserMap, COMPLETION_PARSER_RESULT, String.join(".", scope, JSON_PARSER), validationException);
-
-        if (path == null) {
-=======
-    public static CompletionResponseParser fromMap(Map<String, Object> responseParserMap, ValidationException validationException) {
-        var path = extractRequiredString(responseParserMap, COMPLETION_PARSER_RESULT, JSON_PARSER, validationException);
+    public static CompletionResponseParser fromMap(
+        Map<String, Object> responseParserMap,
+        String scope,
+        ValidationException validationException
+    ) {
+        var path = extractRequiredString(
+            responseParserMap,
+            COMPLETION_PARSER_RESULT,
+            String.join(".", scope, JSON_PARSER),
+            validationException
+        );
 
         if (validationException.validationErrors().isEmpty() == false) {
->>>>>>> 23b7a31406ea7fe4502a8b582915cb066bf8b86e
             throw validationException;
         }
 
@@ -88,35 +90,22 @@ public String getWriteableName() {
 
     @Override
     public ChatCompletionResults transform(Map<String, Object> map) {
-<<<<<<< HEAD
-        var extractedField = MapPathExtractor.extract(map, completionResultPath);
-
-        validateNonNull(extractedField);
-
-        if (extractedField instanceof List<?> extractedList) {
-            var completionList = validateAndCastList(extractedList, (obj) -> toType(obj, String.class));
-=======
         var result = MapPathExtractor.extract(map, completionResultPath);
         var extractedField = result.extractedObject();
 
         validateNonNull(extractedField, completionResultPath);
 
         if (extractedField instanceof List<?> extractedList) {
             var completionList = castList(extractedList, (obj, fieldName) -> toType(obj, String.class, fieldName), completionResultPath);
->>>>>>> 23b7a31406ea7fe4502a8b582915cb066bf8b86e
             return new ChatCompletionResults(completionList.stream().map(ChatCompletionResults.Result::new).toList());
         } else if (extractedField instanceof String extractedString) {
             return new ChatCompletionResults(List.of(new ChatCompletionResults.Result(extractedString)));
         } else {
             throw new IllegalArgumentException(
                 Strings.format(
-<<<<<<< HEAD
-                    "Extracted field is an invalid type, expected a list or a string but received [%s]",
-=======
                     "Extracted field [%s] from path [%s] is an invalid type, expected a list or a string but received [%s]",
                     result.getArrayFieldName(0),
                     completionResultPath,
->>>>>>> 23b7a31406ea7fe4502a8b582915cb066bf8b86e
                     extractedField.getClass().getSimpleName()
                 )
             );

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/common/MapPathExtractorTests.java

@@ -111,9 +111,12 @@ public void testExtract_IteratesSparseEmbeddingStyleMap_ExtractsMaps() {
         assertThat(
             MapPathExtractor.extract(input, "$.result.sparse_embeddings[*].embedding[*]"),
             is(
-                List.of(
-                    List.of(Map.of("tokenId", 6, "weight", 0.123d), Map.of("tokenId", 100, "weight", -123d)),
-                    List.of(Map.of("tokenId", 7, "weight", 0.456d), Map.of("tokenId", 200, "weight", -456d))
+                new MapPathExtractor.Result(
+                    List.of(
+                        List.of(Map.of("tokenId", 6, "weight", 0.123d), Map.of("tokenId", 100, "weight", -123d)),
+                        List.of(Map.of("tokenId", 7, "weight", 0.456d), Map.of("tokenId", 200, "weight", -456d))
+                    ),
+                    List.of("result.sparse_embeddings", "result.sparse_embeddings.embedding")
                 )
             )
         );
@@ -143,7 +146,12 @@ public void testExtract_IteratesSparseEmbeddingStyleMap_ExtractsFieldFromMap() {
 
         assertThat(
             MapPathExtractor.extract(input, "$.result.sparse_embeddings[*].embedding[*].tokenId"),
-            is(List.of(List.of(6, 100), List.of(7, 200)))
+            is(
+                new MapPathExtractor.Result(
+                    List.of(List.of(6, 100), List.of(7, 200)),
+                    List.of("result.sparse_embeddings", "result.sparse_embeddings.embedding", "result.sparse_embeddings.embedding.tokenId")
+                )
+            )
         );
     }
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/response/CompletionResponseParserTests.java

@@ -39,7 +39,11 @@ public static CompletionResponseParser createRandom() {
     public void testFromMap() {
         var validation = new ValidationException();
 
-        var parser = CompletionResponseParser.fromMap(new HashMap<>(Map.of(COMPLETION_PARSER_RESULT, "$.result[*].text")), "scope", validation);
+        var parser = CompletionResponseParser.fromMap(
+            new HashMap<>(Map.of(COMPLETION_PARSER_RESULT, "$.result[*].text")),
+            "scope",
+            validation
+        );
 
         assertThat(parser, is(new CompletionResponseParser("$.result[*].text")));
     }