Fixes 3611: ML API procedures handle null values being passed in better

build.gradle

@@ -131,7 +131,7 @@ subprojects {
 
 ext {
     // NB: due to version.json generation by parsing this file, the next line must not have any if/then/else logic
-    neo4jVersion = "5.19.0"
+    neo4jVersion = "5.18.0"
     // instead we apply the override logic here
     neo4jVersionEffective = project.hasProperty("neo4jVersionOverride") ? project.getProperty("neo4jVersionOverride") : neo4jVersion
     testContainersVersion = '1.18.3'

extended/src/main/java/apoc/ml/OpenAI.java

@@ -16,6 +16,8 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import static apoc.ExtendedApocConfig.APOC_ML_OPENAI_TYPE;
@@ -111,13 +113,26 @@ public Stream<EmbeddingResult> getEmbedding(@Name("texts") List<String> texts, @
       "model": "text-embedding-ada-002",
       "usage": { "prompt_tokens": 8, "total_tokens": 8 } }
     */
-        Stream<Object> resultStream = executeRequest(apiKey, configuration, "embeddings", "text-embedding-ada-002", "input", texts, "$.data", apocConfig, urlAccessChecker);
-        return resultStream
+        Map<Boolean, List<String>> collect = texts.stream()
+                .collect(Collectors.groupingBy(Objects::nonNull));
+
+        List<String> nonNullTexts = collect.get(true);
+
+        Stream<Object> resultStream = executeRequest(apiKey, configuration, "embeddings", "text-embedding-ada-002", "input", nonNullTexts, "$.data", apocConfig, urlAccessChecker);
+        Stream<EmbeddingResult> embeddingResultStream = resultStream
                 .flatMap(v -> ((List<Map<String, Object>>) v).stream())
                 .map(m -> {
                     Long index = (Long) m.get("index");
-                    return new EmbeddingResult(index, texts.get(index.intValue()), (List<Double>) m.get("embedding"));
+                    return new EmbeddingResult(index, nonNullTexts.get(index.intValue()), (List<Double>) m.get("embedding"));
+                });
+
+        List<String> nullTexts = collect.getOrDefault(false, List.of());
+        Stream<EmbeddingResult> nullResultStream = nullTexts.stream()
+                .map(i -> {
+                    // null text return index -1 to indicate that are not coming from `/embeddings` RestAPI
+                    return new EmbeddingResult(-1, i, List.of());
                 });
+        return Stream.concat(embeddingResultStream, nullResultStream);
     }
 
 
@@ -132,13 +147,19 @@ public Stream<MapResult> completion(@Name("prompt") String prompt, @Name("api_ke
       "usage": { "prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12 }
     }
     */
+        if (prompt == null) {
+            return Stream.of(new MapResult(null));
+        }
         return executeRequest(apiKey, configuration, "completions", "gpt-3.5-turbo-instruct", "prompt", prompt, "$", apocConfig, urlAccessChecker)
                 .map(v -> (Map<String,Object>)v).map(MapResult::new);
     }
 
     @Procedure("apoc.ml.openai.chat")
     @Description("apoc.ml.openai.chat(messages, api_key, configuration]) - prompts the completion API")
     public Stream<MapResult> chatCompletion(@Name("messages") List<Map<String, Object>> messages, @Name("api_key") String apiKey, @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) throws Exception {
+        if (messages == null) {
+            return Stream.of(new MapResult(null));
+        }
         return executeRequest(apiKey, configuration, "chat/completions", "gpt-3.5-turbo", "messages", messages, "$", apocConfig, urlAccessChecker)
                 .map(v -> (Map<String,Object>)v).map(MapResult::new);
         // https://platform.openai.com/docs/api-reference/chat/create

extended/src/main/java/apoc/ml/VertexAI.java

@@ -20,7 +20,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 
@@ -94,16 +96,30 @@ public Stream<EmbeddingResult> getEmbedding(@Name("texts") List<String> texts, @
   ]
 }
     */
+
+        Map<Boolean, List<String>> collect = texts.stream()
+                .collect(Collectors.groupingBy(Objects::nonNull));
+
+        List<String> nonNullTexts = collect.get(true);
+
         Object inputs = texts.stream().map(text -> Map.of("content", text)).toList();
         Stream<Object> resultStream = executeRequest(accessToken, project, configuration, "textembedding-gecko", inputs, List.of(), urlAccessChecker);
         AtomicInteger ai = new AtomicInteger();
-        return resultStream
+        Stream<EmbeddingResult> embeddingResultStream = resultStream
                 .flatMap(v -> ((List<Map<String, Object>>) v).stream())
                 .map(m -> {
-                    Map<String,Object> embeddings = (Map<String, Object>) ((Map)m).get("embeddings");
+                    Map<String, Object> embeddings = (Map<String, Object>) ((Map) m).get("embeddings");
                     int index = ai.getAndIncrement();
-                    return new EmbeddingResult(index, texts.get(index), (List<Double>) embeddings.get("values"));
+                    return new EmbeddingResult(index, nonNullTexts.get(index), (List<Double>) embeddings.get("values"));
                 });
+
+        List<String> nullTexts = collect.getOrDefault(false, List.of());
+        Stream<EmbeddingResult> nullResultStream = nullTexts.stream()
+                .map(text -> {
+                    // null text return index -1 to indicate that are not coming from `/embeddings` RestAPI
+                    return new EmbeddingResult(-1, text, List.of());
+                });
+        return Stream.concat(embeddingResultStream, nullResultStream);
     }
 
 
@@ -153,6 +169,10 @@ public Stream<MapResult> completion(@Name("prompt") String prompt, @Name("access
   ]
 }
  */
+        if (prompt == null) {
+            return Stream.of(new MapResult(null));
+        }
+
         Object input = List.of(Map.of("prompt",prompt));
         var parameterKeys = List.of("temperature", "topK", "topP", "maxOutputTokens");
         var resultStream = executeRequest(accessToken, project, configuration, "text-bison", input, parameterKeys, urlAccessChecker);
@@ -192,6 +212,9 @@ public Stream<MapResult> chatCompletion(@Name("messages") List<Map<String, Strin
                                             @Name(value = "context",defaultValue = "") String context,
                                             @Name(value = "examples", defaultValue = "[]") List<Map<String, Map<String,String>>> examples
                                             ) throws Exception {
+        if (messages == null) {
+            return Stream.of(new MapResult(null));
+        }
         Object inputs = List.of(Map.of("context",context, "examples",examples, "messages", messages));
         var parameterKeys = List.of("temperature", "topK", "topP", "maxOutputTokens");
         return executeRequest(accessToken, project, configuration, "chat-bison", inputs, parameterKeys, urlAccessChecker)

extended/src/main/java/apoc/ml/Watson.java

@@ -39,6 +39,9 @@ public class Watson {
     @Procedure("apoc.ml.watson.chat")
     @Description("apoc.ml.watson.chat(messages, accessToken, $configuration) - prompts the completion API")
     public Stream<MapResult> chatCompletion(@Name("messages") List<Map<String, Object>> messages, @Name("accessToken") String accessToken, @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) throws Exception {
+        if (messages == null) {
+            return Stream.of(new MapResult(null));
+        }
         String prompt = messages.stream()
                 .map(message -> {
                     Object role = message.get("role");
@@ -56,6 +59,9 @@ public Stream<MapResult> chatCompletion(@Name("messages") List<Map<String, Objec
     @Procedure("apoc.ml.watson.completion")
     @Description("apoc.ml.watson.completion(prompt, accessToken, $configuration) - prompts the completion API")
     public Stream<MapResult> completion(@Name("prompt") String prompt, @Name("accessToken") String accessToken, @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) throws Exception {
+        if (prompt == null) {
+            return Stream.of(new MapResult(null));
+        }
         return executeRequest(prompt, accessToken, configuration);
     }
 

extended/src/main/java/apoc/ml/aws/Bedrock.java

@@ -60,7 +60,9 @@ public Stream<MapResult> custom(@Name(value = "body") Map<String, Object> body,
     public Stream<MapResult> chatCompletion(
             @Name("messages") List<Map<String, Object>> messages,
             @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) {
-
+        if (messages == null) {
+            return Stream.of(new MapResult(null));
+        }
         var config = new HashMap<>(configuration);
         config.putIfAbsent(MODEL, ANTHROPIC_CLAUDE_V2);
 
@@ -94,7 +96,9 @@ private void transformOpenAiToBedrockRequestBody(Map<String, Object> message) {
     @Description("apoc.ml.bedrock.completion(prompt, $conf) - prompts the completion API")
     public Stream<MapResult> completion(@Name("prompt") String prompt,
                                        @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) {
-
+        if (prompt == null) {
+            return Stream.of(new MapResult(null));
+        }
         var config = new HashMap<>(configuration);
         config.putIfAbsent(MODEL, JURASSIC_2_ULTRA);
 
@@ -109,6 +113,7 @@ public Stream<MapResult> completion(@Name("prompt") String prompt,
     @Description("apoc.ml.bedrock.embedding([texts], $configuration) - returns the embeddings for a given text")
     public Stream<Embedding> embedding(@Name(value = "texts") List<String> texts,
                                        @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) {
+
         var config = new HashMap<>(configuration);
         config.putIfAbsent(MODEL, TITAN_EMBED_TEXT);
 

extended/src/test/java/apoc/ml/OpenAIIT.java

@@ -1,19 +1,25 @@
 package apoc.ml;
 
 import apoc.util.TestUtil;
+import apoc.util.collection.Iterators;
 import org.junit.Assume;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.neo4j.test.rule.DbmsRule;
 import org.neo4j.test.rule.ImpermanentDbmsRule;
 
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 
 import static apoc.ml.OpenAI.MODEL_CONF_KEY;
 import static apoc.ml.OpenAITestResultUtils.*;
 import static apoc.util.TestUtil.testCall;
+import static apoc.util.TestUtil.testResult;
 import static java.util.Collections.emptyMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
 public class OpenAIIT {
 
@@ -68,6 +74,19 @@ public void getEmbedding3LargeWithDimensionsRequestParameter() {
                 r -> assertEmbeddings(r, 256));
     }
 
+    @Test
+    public void getEmbeddingNull() {
+        testResult(db, "CALL apoc.ml.openai.embedding([null, 'Some Text', null, 'Other Text'], $apiKey, $conf)", Map.of("apiKey",openaiKey, "conf", emptyMap()),
+                r -> {
+                    Set<String> actual = Iterators.asSet(r.columnAs("text"));
+
+                    Set<String> expected = new HashSet<>() {{
+                        add(null); add(null); add("Some Text"); add("Other Text");
+                    }};
+                    assertEquals(expected, actual);
+                });
+    }
+
     @Test
     public void completion() {
         testCall(db, COMPLETION_QUERY,
@@ -104,4 +123,21 @@ public void chatCompletion() {
 }
          */
     }
+
+    @Test
+    public void completionNull() {
+        testCall(db, "CALL apoc.ml.openai.completion(null, $apiKey, $conf)",
+                Map.of("apiKey", openaiKey, "conf", emptyMap()),
+                (row) -> assertNull(row.get("value"))
+        );
+    }
+
+    @Test
+    public void chatCompletionNull() {
+        testCall(db, 
+                "CALL apoc.ml.openai.chat(null, $apiKey, $conf)",
+                Map.of("apiKey", openaiKey, "conf", emptyMap()),
+                (row) -> assertNull(row.get("value"))
+        );
+    }
 }

extended/src/test/java/apoc/ml/VertexAIIT.java

@@ -1,6 +1,7 @@
 package apoc.ml;
 
 import apoc.util.TestUtil;
+import apoc.util.collection.Iterators;
 import org.apache.commons.io.FileUtils;
 import org.junit.Assume;
 import org.junit.Before;
@@ -13,15 +14,20 @@
 import java.io.IOException;
 import java.util.Base64;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static apoc.ml.VertexAIHandler.ENDPOINT_CONF_KEY;
 import static apoc.ml.VertexAIHandler.MODEL_CONF_KEY;
 import static apoc.ml.VertexAIHandler.PREDICT_RESOURCE;
 import static apoc.ml.VertexAIHandler.RESOURCE_CONF_KEY;
 import static apoc.ml.VertexAIHandler.STREAM_RESOURCE;
 import static apoc.util.TestUtil.testCall;
+import static apoc.util.TestUtil.testResult;
+import static java.util.Collections.emptyMap;
+import static org.junit.Assert.assertNull;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -66,6 +72,20 @@ public void getEmbedding() {
         });
     }
 
+    @Test
+    public void getEmbeddingNull() {
+        testResult(db, "CALL apoc.ml.vertexai.embedding([null, 'Some Text', null, 'Other Text'], $apiKey, $project)",
+                parameters,
+                r -> {
+                    Set<String> actual = Iterators.asSet(r.columnAs("text"));
+
+                    Set<String> expected = new HashSet<>() {{
+                        add(null); add(null); add("Some Text"); add("Other Text");
+                    }};
+                    assertEquals(expected, actual);
+                });
+    }
+
     @Test
     public void completion() {
         testCall(db, "CALL apoc.ml.vertexai.completion('What color is the sky? Answer in one word: ', $apiKey, $project)", parameters,(row) -> {
@@ -209,4 +229,21 @@ private void assertCorrectResponse(Map<String, Object> row, String expected) {
         assertTrue(stringRow.toLowerCase().contains(expected),
                 "Actual result is: " + stringRow);
     }
+
+    @Test
+    public void completionNull() {
+        testCall(db, "CALL apoc.ml.vertexai.completion(null, $apiKey, $project)",
+                parameters,
+                (row) -> assertNull(row.get("value"))
+        );
+    }
+
+    @Test
+    public void chatCompletionNull() {
+        testCall(db,
+                "CALL apoc.ml.vertexai.chat(null, $apiKey, $project)",
+                parameters,
+                (row) -> assertNull(row.get("value"))
+        );
+    }
 }

extended/src/test/java/apoc/ml/WatsonIT.java

@@ -14,6 +14,8 @@
 import static apoc.ExtendedApocConfig.APOC_ML_WATSON_URL;
 import static apoc.ExtendedApocConfig.APOC_ML_WATSON_PROJECT_ID;
 import static apoc.util.TestUtil.testCall;
+import static java.util.Collections.emptyMap;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assume.assumeNotNull;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -130,4 +132,21 @@ private static void commonAssertions(Map<String, Object> row, String text, long
         assertEquals(inputTokenCount, result.get("input_token_count"));
         assertEquals(stopReason, result.get("stop_reason"));
     }
+
+    @Test
+    public void completionNull() {
+        testCall(db, "CALL apoc.ml.watson.completion(null, $apiKey, $conf)",
+                Map.of("apiKey", accessToken, "conf", emptyMap()),
+                (row) -> assertNull(row.get("value"))
+        );
+    }
+
+    @Test
+    public void chatCompletionNull() {
+        testCall(db,
+                "CALL apoc.ml.watson.chat(null, $apiKey, $conf)",
+                Map.of("apiKey", accessToken, "conf", emptyMap()),
+                (row) -> assertNull(row.get("value"))
+        );
+    }
 }

extended/src/test/java/apoc/ml/aws/BedrockIT.java

@@ -27,8 +27,10 @@
 import static apoc.ml.aws.BedrockUtil.*;
 import static apoc.util.TestUtil.testCall;
 import static apoc.util.TestUtil.testResult;
+import static java.util.Collections.emptyMap;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeNotNull;
@@ -291,4 +293,21 @@ private static void assertionsTitanEmbed(Map value) {
         assertNotNull(value.get("inputTextTokenCount"));
         assertNotNull(value.get("embedding"));
     }
+
+    @Test
+    public void completionNull() {
+        testCall(db, "CALL apoc.ml.bedrock.completion(null, $conf)",
+                Map.of("conf", emptyMap()),
+                (row) -> assertNull(row.get("value"))
+        );
+    }
+
+    @Test
+    public void chatCompletionNull() {
+        testCall(db,
+                "CALL apoc.ml.bedrock.chat(null, $conf)",
+                Map.of("conf", emptyMap()),
+                (row) -> assertNull(row.get("value"))
+        );
+    }
 }