Fixes #4123: Add anthropic claude support

neo4j-contrib · Jul 30, 2024 · 5f55ce1 · 5f55ce1
1 parent 30fe1c9
commit 5f55ce1
Show file tree

Hide file tree

Showing 8 changed files with 376 additions and 25 deletions.
diff --git a/docs/asciidoc/modules/ROOT/pages/ml/openai.adoc b/docs/asciidoc/modules/ROOT/pages/ml/openai.adoc
@@ -14,9 +14,10 @@ All the following procedures can have the following APOC config, i.e. in `apoc.c
 .Apoc configuration
 |===
 |key | description | default
-| apoc.ml.openai.type | "AZURE", "HUGGINGFACE", "OPENAI", indicates whether the API is Azure, HuggingFace or another one | "OPENAI" 
-| apoc.ml.openai.url | the OpenAI endpoint base url | https://api.openai.com/v1 
-    (or empty string if `apoc.ml.openai.type=<AZURE OR HUGGINGFACE>`)
+| apoc.ml.openai.type | "AZURE", "HUGGINGFACE", "OPENAI", indicates whether the API is Azure, HuggingFace, Anthropic or another one | "OPENAI" 
+| apoc.ml.openai.url | the OpenAI endpoint base url | `https://api.openai.com/v1` by default, 
+    `https://api.anthropic.com/v1` if `apoc.ml.openai.type=<ANTHROPIC>`,
+    or empty string if `apoc.ml.openai.type=<AZURE OR HUGGINGFACE>`
 | apoc.ml.azure.api.version | in case of `apoc.ml.openai.type=AZURE`, indicates the `api-version` to be passed after the `?api-version=` url
 |===
 
@@ -249,3 +250,53 @@ CALL apoc.ml.openai.chat([{"role": "user", "content": "Explain the importance of
     '<apiKey>',
     {endpoint: 'https://api.groq.com/openai/v1', model: 'mixtral-8x7b-32768'})
 ----
+
+Another alternative is to use the https://docs.anthropic.com/en/api/getting-started[Anthropic API].
+
+We can use the `apoc.ml.openai.chat` procedure to leverage the https://docs.anthropic.com/en/api/messages[Anthropic Messages API], that is:
+
+[source,cypher]
+----
+CALL apoc.ml.openai.chat([
+      { content: "What planet do humans live on?", role: "user" },
+      { content: "Only answer with a single word", role: "assistant" }
+    ],
+    $anthropicApiKey,
+    {apiType: 'ANTHROPIC'}
+)
+----
+
+.Default Anthropic body parameters
+[%autowidth, opts=header]
+|===
+|name | description
+| max_tokens | 1000
+| model | claude-3-5-sonnet-20240620
+|===
+
+Also, we can use the `apoc.ml.openai.completion` procedure to leverage the https://docs.anthropic.com/en/api/complete[Anthropic Complete API]:
+
+[source,cypher]
+----
+CALL apoc.ml.openai.completion('What color is the sky?', 
+    $anthropicApiKey,
+    {apiType: 'ANTHROPIC'}
+)
+----
+
+
+.Default Anthropic body parameters
+[%autowidth, opts=header]
+|===
+|name | description
+| max_tokens_to_sample | 1000
+| model | claude-2.1
+|===
+
+[NOTE]
+====
+At the moment Anthropic does not support embedding API.
+
+And at the time, payload with https://docs.anthropic.com/en/api/messages-streaming[`stream: true`] is not supported, since the result of apoc.ml.openai must be a JSON.
+====
+
diff --git a/extended/src/main/java/apoc/ml/MLUtil.java b/extended/src/main/java/apoc/ml/MLUtil.java
@@ -7,4 +7,7 @@ public class MLUtil {
     public static final String API_VERSION_CONF_KEY = "apiVersion";
     public static final String REGION_CONF_KEY = "region";
     public static final String MODEL_CONF_KEY = "model";
+    public static final String MAX_TOKENS = "max_tokens";
+    public static final String MAX_TOKENS_TO_SAMPLE = "max_tokens_to_sample";
+    public static final String ANTHROPIC_VERSION = "anthropic-version";
 }
diff --git a/extended/src/main/java/apoc/ml/OpenAI.java b/extended/src/main/java/apoc/ml/OpenAI.java
@@ -64,34 +64,22 @@ static Stream<Object> executeRequest(String apiKey, Map<String, Object> configur
         );
         OpenAIRequestHandler.Type type = OpenAIRequestHandler.Type.valueOf(apiTypeString.toUpperCase(Locale.ENGLISH));
 
-        var config = new HashMap<>(configuration);
-        // we remove these keys from config, since the json payload is calculated starting from the config map
-        Stream.of(ENDPOINT_CONF_KEY, API_TYPE_CONF_KEY, API_VERSION_CONF_KEY, APIKEY_CONF_KEY).forEach(config::remove);
-
-        switch (type) {
-            case MIXEDBREAD_CUSTOM -> {
-                // no payload manipulation, taken from the configuration as-is
-            }
-            case HUGGINGFACE -> {
-                config.putIfAbsent("inputs", inputs);
-                jsonPath = "$[0]";
-            }
-            default -> {
-                config.putIfAbsent(MODEL_CONF_KEY, model);
-                config.put(key, inputs);
-            }
-        }
-
+        var configForPayload = new HashMap<>(configuration);
+        // we remove these keys from configPayload, since the json payload is calculated starting from the configPayload map
+        Stream.of(ENDPOINT_CONF_KEY, API_TYPE_CONF_KEY, API_VERSION_CONF_KEY, APIKEY_CONF_KEY).forEach(configForPayload::remove);
+
+        final Map<String, Object> headers = new HashMap<>();
+
+        handleAPIProvider(type, configuration, path, model, key, inputs, configForPayload, headers);
+
+        path = (String) configuration.getOrDefault(PATH_CONF_KEY, path);
         OpenAIRequestHandler apiType = type.get();
 
         jsonPath = (String) configuration.getOrDefault(JSON_PATH_CONF_KEY, jsonPath);
-        path = (String) configuration.getOrDefault(PATH_CONF_KEY, path);
-
-        final Map<String, Object> headers = new HashMap<>();
         headers.put("Content-Type", "application/json");
         apiType.addApiKey(headers, apiKey);
 
-        String payload = JsonUtil.OBJECT_MAPPER.writeValueAsString(config);
+        String payload = JsonUtil.OBJECT_MAPPER.writeValueAsString(configForPayload);
 
         // new URL(endpoint), path) can produce a wrong path, since endpoint can have for example embedding,
         // eg: https://my-resource.openai.azure.com/openai/deployments/apoc-embeddings-model
@@ -100,6 +88,45 @@ static Stream<Object> executeRequest(String apiKey, Map<String, Object> configur
         return JsonUtil.loadJson(url, headers, payload, jsonPath, true, List.of(), urlAccessChecker);
     }
 
+    private static void handleAPIProvider(OpenAIRequestHandler.Type type,
+                                          Map<String, Object> configuration,
+                                          String path,
+                                          String model,
+                                          String key,
+                                          Object inputs,
+                                          HashMap<String, Object> configForPayload,
+                                          Map<String, Object> headers) {
+        switch (type) {
+            case MIXEDBREAD_CUSTOM -> {
+                // no payload manipulation, taken from the configuration as-is
+            }
+            case HUGGINGFACE -> {
+                configForPayload.putIfAbsent("inputs", inputs);
+                configuration.putIfAbsent(JSON_PATH_CONF_KEY, "$[0]");
+            }
+            case ANTHROPIC -> {
+                headers.putIfAbsent(ANTHROPIC_VERSION, configuration.getOrDefault(ANTHROPIC_VERSION, "2023-06-01"));
+
+                if (path.equals("completions")) {
+                    configuration.putIfAbsent(PATH_CONF_KEY, "complete");
+                    configForPayload.putIfAbsent(MAX_TOKENS_TO_SAMPLE, 1000);
+                    configForPayload.putIfAbsent(MODEL_CONF_KEY, "claude-2.1");
+                } else {
+                    configuration.putIfAbsent(PATH_CONF_KEY, "messages");
+                    configForPayload.putIfAbsent(MAX_TOKENS, 1000);
+                    configForPayload.putIfAbsent(MODEL_CONF_KEY, "claude-3-5-sonnet-20240620");
+                }
+
+                configForPayload.remove(ANTHROPIC_VERSION);
+                configForPayload.put(key, inputs);
+            }
+            default -> {
+                configForPayload.putIfAbsent(MODEL_CONF_KEY, model);
+                configForPayload.put(key, inputs);
+            }
+        }
+    }
+
     @Procedure("apoc.ml.openai.embedding")
     @Description("apoc.openai.embedding([texts], api_key, configuration) - returns the embeddings for a given text")
     public Stream<EmbeddingResult> getEmbedding(@Name("texts") List<String> texts, @Name("api_key") String apiKey, @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) throws Exception {

diff --git a/extended/src/main/java/apoc/ml/OpenAIRequestHandler.java b/extended/src/main/java/apoc/ml/OpenAIRequestHandler.java
@@ -48,6 +48,7 @@ enum Type {
         HUGGINGFACE(new OpenAi(null)),
         MIXEDBREAD_EMBEDDING(new OpenAi(MIXEDBREAD_BASE_URL)),
         MIXEDBREAD_CUSTOM(new Custom()),
+        ANTHROPIC(new Anthropic()),
         OPENAI(new OpenAi("https://api.openai.com/v1"));
 
         private final OpenAIRequestHandler handler;
@@ -94,6 +95,18 @@ public void addApiKey(Map<String, Object> headers, String apiKey) {
         }
     }
 
+    static class Anthropic extends OpenAi {
+
+        public Anthropic() {
+            super("https://api.anthropic.com/v1");
+        }
+
+        @Override
+        public void addApiKey(Map<String, Object> headers, String apiKey) {
+            headers.put("x-api-key", apiKey);
+        }
+    }
+
     static class Custom extends OpenAi {
 
         public Custom() {