-
Notifications
You must be signed in to change notification settings - Fork 494
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add OpenAI (LLM) procedures #3575
Merged
Merged
Changes from all commits
Commits
Show all changes
9 commits
Select commit
Hold shift + click to select a range
b42ba7c
WIP
tomasonjo 01dbb37
Add completion API
tomasonjo 462077c
add chatCompletion
tomasonjo e536d76
prettificiation
tomasonjo 7f365c6
WIP
tomasonjo 99ae402
prettify
tomasonjo c9812a3
Refactoring, todo docs & tests
jexp 7d726fc
Added Tests & Docs for OpenAI procs (WIP)
jexp 469bbd1
Update openai.adoc
jexp File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,10 @@ | ||
[[ml]] | ||
= Machine Learning (ML | ||
:description: This chapter describes procedures that can be used for adding Machine Learning (ML) functionality to graph applications. | ||
|
||
The procedures described in this chapter act as wrappers around cloud based Machine Learning APIs. | ||
These procedures generate embeddings, analyze text, complete text, complete chat conversations and more. | ||
|
||
This section includes: | ||
|
||
* xref::ml/openai.adoc[] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,122 @@ | ||
[[openai-api]] | ||
= OpenAI API Access | ||
:description: This section describes procedures that can be used to access the OpenAI API. | ||
|
||
NOTE: You need to acquire an https://platform.openai.com/account/api-keys[OpenAI API key^] to use these procedures. Using them will incurr costs on your OpenAI account. | ||
|
||
== Generate Embeddings API | ||
|
||
This procedure `apoc.ml.openai.embedding` can take a list of text strings, and will return one row per string, with the embedding data as a 1536 element vector. | ||
It uses the `/embeddings/create` API which is https://platform.openai.com/docs/api-reference/embeddings/create[documented here^]. | ||
|
||
Additional configuration is passed to the API, the default model used is `text-embedding-ada-002`. | ||
|
||
.Generate Embeddings Call | ||
[source,cypher] | ||
---- | ||
CALL apoc.ml.openai.embedding(['Some Text'], $apiKey, {}) yield index, text, embedding; | ||
---- | ||
|
||
.Generate Embeddings Response | ||
[%autowidth, opts=header] | ||
|=== | ||
|index | text | embedding | ||
|0 | "Some Text" | [-0.0065358975, -7.9563365E-4, .... -0.010693862, -0.005087272] | ||
|=== | ||
|
||
.Parameters | ||
[%autowidth, opts=header] | ||
|=== | ||
|name | description | ||
| texts | List of text strings | ||
| apiKey | OpenAI API key | ||
| configuration | optional map for entries like model and other request parameters | ||
|=== | ||
|
||
|
||
.Results | ||
[%autowidth, opts=header] | ||
|=== | ||
|name | description | ||
| index | index entry in original list | ||
| text | line of text from original list | ||
| embedding | 1536 element floating point embedding vector for ada-002 model | ||
|=== | ||
|
||
== Text Completion API | ||
|
||
This procedure `apoc.ml.openai.completion` can continue/complete a given text. | ||
|
||
It uses the `/completions/create` API which is https://platform.openai.com/docs/api-reference/completions/create[documented here^]. | ||
|
||
Additional configuration is passed to the API, the default model used is `text-davinci-003`. | ||
|
||
.Text Completion Call | ||
[source,cypher] | ||
---- | ||
CALL apoc.ml.openai.completion('What color is the sky? Answer in one word: ', $apiKey, {config}) yield value; | ||
---- | ||
|
||
.Text Completion Response | ||
---- | ||
{ created=1684248202, model="text-davinci-003", id="cmpl-7GqBWwX49yMJljdmnLkWxYettZoOy", | ||
usage={completion_tokens=2, prompt_tokens=12, total_tokens=14}, | ||
choices=[{finish_reason="stop", index=0, text="Blue", logprobs=null}], object="text_completion"} | ||
---- | ||
|
||
.Parameters | ||
[%autowidth, opts=header] | ||
|=== | ||
|name | description | ||
| prompt | Text to complete | ||
| apiKey | OpenAI API key | ||
| configuration | optional map for entries like model, temperature, and other request parameters | ||
|=== | ||
|
||
.Results | ||
[%autowidth, opts=header] | ||
|=== | ||
|name | description | ||
| value | result entry from OpenAI (containing) | ||
|=== | ||
|
||
== Chat Completion API | ||
|
||
This procedure `apoc.ml.openai.chat` takes a list of maps of chat exchanges between assistant and user (with optional system message), and will return the next message in the flow. | ||
|
||
It uses the `/chat/create` API which is https://platform.openai.com/docs/api-reference/chat/create[documented here^]. | ||
|
||
Additional configuration is passed to the API, the default model used is `gpt-3.5-turbo`. | ||
|
||
.Chat Completion Call | ||
[source,cypher] | ||
---- | ||
CALL apoc.ml.openai.chat([ | ||
{role:"system", content:"Only answer with a single word"}, | ||
{role:"user", content:"What planet do humans live on?"} | ||
], $apiKey) yield value | ||
---- | ||
|
||
.Chat Completion Response | ||
---- | ||
{created=1684248203, id="chatcmpl-7GqBXZr94avd4fluYDi2fWEz7DIHL", | ||
object="chat.completion", model="gpt-3.5-turbo-0301", | ||
usage={completion_tokens=2, prompt_tokens=26, total_tokens=28}, | ||
choices=[{finish_reason="stop", index=0, message={role="assistant", content="Earth."}}]} | ||
---- | ||
|
||
.Parameters | ||
[%autowidth, opts=header] | ||
|=== | ||
|name | description | ||
| messages | List of maps of instructions with `{role:"assistant|user|system", content:"text}` | ||
| apiKey | OpenAI API key | ||
| configuration | optional map for entries like model, temperature, and other request parameters | ||
|=== | ||
|
||
.Results | ||
[%autowidth, opts=header] | ||
|=== | ||
|name | description | ||
| value | result entry from OpenAI (containing created, id, model, object, usage(tokens), choices(message, index, finish_reason)) | ||
|=== |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,114 @@ | ||
package apoc.ml; | ||
|
||
import apoc.Extended; | ||
import apoc.util.JsonUtil; | ||
import com.fasterxml.jackson.core.JsonProcessingException; | ||
import org.neo4j.procedure.Description; | ||
import org.neo4j.procedure.Name; | ||
import org.neo4j.procedure.Procedure; | ||
|
||
import java.net.MalformedURLException; | ||
import java.net.URL; | ||
import java.util.HashMap; | ||
import java.util.Map; | ||
import java.util.List; | ||
import java.util.stream.Stream; | ||
|
||
import apoc.result.MapResult; | ||
|
||
import com.fasterxml.jackson.databind.ObjectMapper; | ||
|
||
|
||
@Extended | ||
public class OpenAI { | ||
|
||
public static final String APOC_ML_OPENAI_URL = "apoc.ml.openai.url"; | ||
|
||
public static class EmbeddingResult { | ||
public final long index; | ||
public final String text; | ||
public final List<Double> embedding; | ||
|
||
public EmbeddingResult(long index, String text, List<Double> embedding) { | ||
this.index = index; | ||
this.text = text; | ||
this.embedding = embedding; | ||
} | ||
} | ||
|
||
private static Stream<Object> executeRequest(String apiKey, Map<String, Object> configuration, String path, String model, String key, Object inputs, String jsonPath) throws JsonProcessingException, MalformedURLException { | ||
if (apiKey == null || apiKey.isBlank()) | ||
throw new IllegalArgumentException("API Key must not be empty"); | ||
String endpoint = System.getProperty(APOC_ML_OPENAI_URL,"https://api.openai.com/v1/"); | ||
Map<String, Object> headers = Map.of( | ||
"Content-Type", "application/json", | ||
"Authorization", "Bearer " + apiKey | ||
); | ||
|
||
var config = new HashMap<>(configuration); | ||
config.putIfAbsent("model", model); | ||
config.put(key, inputs); | ||
|
||
String payload = new ObjectMapper().writeValueAsString(config); | ||
|
||
var url = new URL(new URL(endpoint), path).toString(); | ||
return JsonUtil.loadJson(url, headers, payload, jsonPath, true, List.of()); | ||
} | ||
|
||
@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 { | ||
// https://platform.openai.com/docs/api-reference/embeddings/create | ||
/* | ||
{ "object": "list", | ||
"data": [ | ||
{ | ||
"object": "embedding", | ||
"embedding": [ 0.0023064255, -0.009327292, .... (1536 floats total for ada-002) -0.0028842222 ], | ||
"index": 0 | ||
} | ||
], | ||
"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"); | ||
return 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")); | ||
}); | ||
} | ||
|
||
|
||
@Procedure("apoc.ml.openai.completion") | ||
@Description("apoc.ml.openai.completion(prompt, api_key, configuration) - prompts the completion API") | ||
public Stream<MapResult> completion(@Name("prompt") String prompt, @Name("api_key") String apiKey, @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) throws Exception { | ||
// https://platform.openai.com/docs/api-reference/completions/create | ||
/* | ||
{ "id": "cmpl-uqkvlQyYK7bGYrRHQ0eXlWi7", | ||
"object": "text_completion", "created": 1589478378, "model": "text-davinci-003", | ||
"choices": [ { "text": "\n\nThis is indeed a test", "index": 0, "logprobs": null, "finish_reason": "length" } ], | ||
"usage": { "prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12 } | ||
} | ||
*/ | ||
return executeRequest(apiKey, configuration, "completions", "text-davinci-003", "prompt", prompt, "$") | ||
.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 { | ||
return executeRequest(apiKey, configuration, "chat/completions", "gpt-3.5-turbo", "messages", messages, "$") | ||
.map(v -> (Map<String,Object>)v).map(MapResult::new); | ||
// https://platform.openai.com/docs/api-reference/chat/create | ||
/* | ||
{ 'id': 'chatcmpl-6p9XYPYSTTRi0xEviKjjilqrWU2Ve', 'object': 'chat.completion', 'created': 1677649420, 'model': 'gpt-3.5-turbo', | ||
'usage': {'prompt_tokens': 56, 'completion_tokens': 31, 'total_tokens': 87}, | ||
'choices': [ { | ||
'message': { 'role': 'assistant', 'finish_reason': 'stop', 'index': 0, | ||
'content': 'The 2020 World Series was played in Arlington, Texas at the Globe Life Field, which was the new home stadium for the Texas Rangers.'} | ||
} ] } | ||
*/ | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,120 @@ | ||
package apoc.ml; | ||
|
||
import apoc.util.TestUtil; | ||
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.nio.file.Paths; | ||
import java.util.List; | ||
import java.util.Map; | ||
|
||
import static apoc.ApocConfig.APOC_IMPORT_FILE_ENABLED; | ||
import static apoc.ApocConfig.apocConfig; | ||
import static apoc.util.TestUtil.getUrlFileName; | ||
import static apoc.util.TestUtil.testCall; | ||
import static org.junit.jupiter.api.Assertions.assertEquals; | ||
|
||
public class OpenAIIT { | ||
|
||
private String openaiKey; | ||
|
||
@Rule | ||
public DbmsRule db = new ImpermanentDbmsRule(); | ||
|
||
public OpenAIIT() { | ||
} | ||
|
||
@Before | ||
public void setUp() throws Exception { | ||
openaiKey = System.getenv("OPENAI_KEY"); | ||
Assume.assumeNotNull("No OPENAI_KEY environment configured", openaiKey); | ||
TestUtil.registerProcedure(db, OpenAI.class); | ||
} | ||
|
||
@Test | ||
public void getEmbedding() { | ||
testCall(db, "CALL apoc.ml.openai.embedding(['Some Text'], $apiKey)", Map.of("apiKey",openaiKey),(row) -> { | ||
System.out.println("row = " + row); | ||
assertEquals(0L, row.get("index")); | ||
assertEquals("Some Text", row.get("text")); | ||
var embedding = (List<Double>) row.get("embedding"); | ||
assertEquals(1536, embedding.size()); | ||
assertEquals(true, embedding.stream().allMatch(d -> d instanceof Double)); | ||
}); | ||
} | ||
|
||
@Test | ||
public void completion() { | ||
testCall(db, "CALL apoc.ml.openai.completion('What color is the sky? Answer in one word: ', $apiKey)", | ||
Map.of("apiKey",openaiKey),(row) -> { | ||
System.out.println("row = " + row); | ||
var result = (Map<String,Object>)row.get("value"); | ||
assertEquals(true, result.get("created") instanceof Number); | ||
assertEquals(true, result.containsKey("choices")); | ||
var finishReason = (String)((List<Map>) result.get("choices")).get(0).get("finish_reason"); | ||
assertEquals(true, finishReason.matches("stop|length")); | ||
String text = (String) ((List<Map>) result.get("choices")).get(0).get("text"); | ||
assertEquals(true, text != null && !text.isBlank()); | ||
assertEquals(true, text.toLowerCase().contains("blue")); | ||
assertEquals(true, result.containsKey("usage")); | ||
assertEquals(true, ((Map)result.get("usage")).get("prompt_tokens") instanceof Number); | ||
assertEquals("text-davinci-003", result.get("model")); | ||
assertEquals("text_completion", result.get("object")); | ||
}); | ||
} | ||
|
||
@Test | ||
public void chatCompletion() { | ||
testCall(db, """ | ||
CALL apoc.ml.openai.chat([ | ||
{role:"system", content:"Only answer with a single word"}, | ||
{role:"user", content:"What planet do humans live on?"} | ||
], $apiKey) | ||
""", Map.of("apiKey",openaiKey), (row) -> { | ||
System.out.println("row = " + row); | ||
var result = (Map<String,Object>)row.get("value"); | ||
assertEquals(true, result.get("created") instanceof Number); | ||
assertEquals(true, result.containsKey("choices")); | ||
|
||
Map message = ((List<Map<String,Map>>) result.get("choices")).get(0).get("message"); | ||
assertEquals("assistant", message.get("role")); | ||
// assertEquals("stop", message.get("finish_reason")); | ||
String text = (String) message.get("content"); | ||
assertEquals(true, text != null && !text.isBlank()); | ||
|
||
|
||
assertEquals(true, result.containsKey("usage")); | ||
assertEquals(true, ((Map)result.get("usage")).get("prompt_tokens") instanceof Number); | ||
assertEquals("gpt-3.5-turbo-0301", result.get("model")); | ||
assertEquals("chat.completion", result.get("object")); | ||
}); | ||
|
||
/* | ||
{ | ||
"id": "chatcmpl-6p9XYPYSTTRi0xEviKjjilqrWU2Ve", | ||
"object": "chat.completion", | ||
"created": 1677649420, | ||
"model": "gpt-3.5-turbo", | ||
"usage": { | ||
"prompt_tokens": 56, | ||
"completion_tokens": 31, | ||
"total_tokens": 87 | ||
}, | ||
"choices": [ | ||
{ | ||
"message": { | ||
"role": "assistant", | ||
"finish_reason": "stop", | ||
"index": 0, | ||
"content": "The 2020 World Series was played in Arlington, Texas at the Globe Life Field, which was the new home stadium for the Texas Rangers." | ||
} | ||
} | ||
] | ||
} | ||
*/ | ||
} | ||
} |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@jexp @tomasonjo does make sense to make this and the other models configurable, via
configuration
?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
not sure exactly if I understand what you mean, but we could add model parameter under the configuration map. I would leave the prompts as a separate parameter.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The model is something important, so it cannot be missed. For the user it's already in the configuration. But internally I think we should make it explicit.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah I mean something like this
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not sure we need to do this on every call, as it's already done inside of executeRequest.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ok