Skip to content

Commit

Permalink
feat: Support AI extract and AI extract structured (#1266)
Browse files Browse the repository at this point in the history
  • Loading branch information
@congminh1254
congminh1254 authored Oct 11, 2024
1 parent 3cb2c7c commit 7ba90b9
Show file tree
Hide file tree
Showing 17 changed files with 1,096 additions and 45 deletions.
60 changes: 59 additions & 1 deletion doc/ai.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ an answer based on the provided prompt and items.
- [Send AI request](#send-ai-request)
- [Send AI text generation request](#send-ai-text-generation-request)
- [Get AI Agent default configuration](#get-ai-agent-default-configuration)
- [Extract metadata freeform](#extract-metadata-freeform)
- [Extract metadata structured](#extract-metadata-structured)

<!-- END doctoc generated TOC please keep comment here to allow auto update -->

Expand Down Expand Up @@ -87,4 +89,60 @@ BoxAIAgentConfig config = BoxAI.getAiAgentDefaultConfig(
);
```

[get-ai-agent-default-config]: http://opensource.box.com/box-java-sdk/javadoc/com/box/sdk/BoxAI.html#getAiAgentDefaultConfig-com.box.sdk.BoxAPIConnection-com.box.sdk.ai.BoxAIAgent.Mode-java.lang.String-java.lang.String-
[get-ai-agent-default-config]: http://opensource.box.com/box-java-sdk/javadoc/com/box/sdk/BoxAI.html#getAiAgentDefaultConfig-com.box.sdk.BoxAPIConnection-com.box.sdk.ai.BoxAIAgent.Mode-java.lang.String-java.lang.String-

Extract metadata freeform
--------------------------

To send an AI request to supported Large Language Models (LLMs) and extract metadata in form of key-value pairs, call static
[`extractMetadataFreeform(BoxAPIConnection api, String prompt, List<BoxAIItem> items)`][extract-metadata-freeform] method.
In the request you have to provide a prompt, a list of items that your prompt refers to and an optional agent configuration.

<!-- sample post_ai_extract -->
```java
BoxAIResponse response = BoxAI.extractMetadataFreeform(
api,
"firstName, lastName, location, yearOfBirth, company",
Collections.singletonList(new BoxAIItem("123456", BoxAIItem.Type.FILE))
);
```

[extract-metadata-freeform]: https://opensource.box.com/box-java-sdk/javadoc/com/box/sdk/BoxAI.html#extractMetadataFreeform-com.box.sdk.BoxAPIConnection-java.lang.String-java.util.List-

Extract metadata structured
--------------------------

Sends an AI request to supported Large Language Models (LLMs) and returns extracted metadata as a set of key-value pairs. For this request, you need to use an already defined metadata template or define a schema yourself.

To send an AI request to extract metadata from files with a predefined metadata template, call static
[`extractMetadataStructured extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items, BoxAIExtractMetadataTemplate template)`][extract-metadata-structured-metadata-template] method.

<!-- sample post_ai_extract_structured -->
```java
BoxAIExtractMetadataTemplate template = new BoxAIExtractMetadataTemplate("templateKey", "enterprise");
BoxAIExtractStructuredResponse result = BoxAI.extractMetadataStructured(
api,
Collections.singletonList(new BoxAIItem("123456", BoxAIItem.Type.FILE)),
template
);
JsonObject sourceJson = result.getSourceJson();
```

To send an AI request to extract metadata from files with custom fields, call static
[`extractMetadataStructured extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items, List<BoxAIExtractField> fields)`][extract-metadata-structured-fields] method.

<!-- sample post_ai_extract_structured_fields -->
```java
List<BoxAIExtractField> fields = new ArrayList<>();
fields.add(new BoxAIExtractField("firstName"));

BoxAIExtractStructuredResponse result = BoxAI.extractMetadataStructured(
api,
Collections.singletonList(new BoxAIItem("123456", BoxAIItem.Type.FILE)),
fields
);
JsonObject sourceJson = result.getSourceJson();
```

[extract-metadata-structured-metadata-template]: https://opensource.box.com/box-java-sdk/javadoc/com/box/sdk/BoxAI.html#extractMetadataStructured-com.box.sdk.BoxAPIConnection-java.util.List-com.box.sdk.ai.metadata.BoxAIExtractMetadataTemplate-
[extract-metadata-structured-fields]: https://opensource.box.com/box-java-sdk/javadoc/com/box/sdk/BoxAI.html#extractMetadataStructured-com.box.sdk.BoxAPIConnection-java.util.List-java.util.List-
165 changes: 150 additions & 15 deletions src/intTest/java/com/box/sdk/BoxAIIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;

import com.eclipsesource.json.Json;
import com.eclipsesource.json.JsonObject;
import java.text.ParseException;
import java.util.ArrayList;
import java.util.Collections;
Expand Down Expand Up @@ -49,10 +51,10 @@ public void askAISingleItem() throws InterruptedException {
// and 412 is returned
retry(() -> {
BoxAIResponse response = BoxAI.sendAIRequest(
api,
"What is the name of the file?",
Collections.singletonList(new BoxAIItem(uploadedFileInfo.getID(), BoxAIItem.Type.FILE)),
BoxAI.Mode.SINGLE_ITEM_QA
api,
"What is the name of the file?",
Collections.singletonList(new BoxAIItem(uploadedFileInfo.getID(), BoxAIItem.Type.FILE)),
BoxAI.Mode.SINGLE_ITEM_QA
);
assertThat(response.getAnswer(), containsString("Test file"));
assert response.getCreatedAt().before(new Date(System.currentTimeMillis()));
Expand Down Expand Up @@ -86,10 +88,10 @@ public void askAIMultipleItems() throws InterruptedException {
// and 412 is returned
retry(() -> {
BoxAIResponse response = BoxAI.sendAIRequest(
api,
"What is the content of these files?",
items,
BoxAI.Mode.MULTIPLE_ITEM_QA
api,
"What is the content of these files?",
items,
BoxAI.Mode.MULTIPLE_ITEM_QA
);
assertThat(response.getAnswer(), containsString("Test file"));
assert response.getCreatedAt().before(new Date(System.currentTimeMillis()));
Expand All @@ -111,7 +113,7 @@ public void askAITextGenItemWithDialogueHistory() throws ParseException, Interru
Date date1 = BoxDateFormat.parse("2013-05-16T15:27:57-07:00");
Date date2 = BoxDateFormat.parse("2013-05-16T15:26:57-07:00");

BoxFile uploadedFile = uploadFileToUniqueFolder(api, fileName, "Test file");
BoxFile uploadedFile = uploadFileToUniqueFolder(api, fileName, "Test file");
try {
// When a file has been just uploaded, AI service may not be ready to return text response
// and 412 is returned
Expand All @@ -121,16 +123,16 @@ public void askAITextGenItemWithDialogueHistory() throws ParseException, Interru

List<BoxAIDialogueEntry> dialogueHistory = new ArrayList<>();
dialogueHistory.add(
new BoxAIDialogueEntry("What is the name of the file?", "Test file", date1)
new BoxAIDialogueEntry("What is the name of the file?", "Test file", date1)
);
dialogueHistory.add(
new BoxAIDialogueEntry("What is the size of the file?", "10kb", date2)
new BoxAIDialogueEntry("What is the size of the file?", "10kb", date2)
);
BoxAIResponse response = BoxAI.sendAITextGenRequest(
api,
"What is the name of the file?",
Collections.singletonList(new BoxAIItem(uploadedFileInfo.getID(), BoxAIItem.Type.FILE)),
dialogueHistory
api,
"What is the name of the file?",
Collections.singletonList(new BoxAIItem(uploadedFileInfo.getID(), BoxAIItem.Type.FILE)),
dialogueHistory
);
assertThat(response.getAnswer(), containsString("name"));
assert response.getCreatedAt().before(new Date(System.currentTimeMillis()));
Expand Down Expand Up @@ -192,4 +194,137 @@ public void askAISingleItemWithAgent() throws InterruptedException {
deleteFile(uploadedFile);
}
}

@Test
public void aiExtract() throws InterruptedException {
BoxAPIConnection api = jwtApiForServiceAccount();
BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT, "en-US", null);
BoxAIAgentExtract agentExtract = (BoxAIAgentExtract) agent;

BoxFile uploadedFile = uploadFileToUniqueFolder(api, "[aiExtract] Test File.txt",
"My name is John Doe. I live in San Francisco. I was born in 1990. I work at Box.");

try {
// When a file has been just uploaded, AI service may not be ready to return text response
// and 412 is returned
retry(() -> {
BoxAIResponse response = BoxAI.extractMetadataFreeform(api,
"firstName, lastName, location, yearOfBirth, company",
Collections.singletonList(new BoxAIItem(uploadedFile.getID(), BoxAIItem.Type.FILE)),
agentExtract);
assertThat(response.getAnswer(), containsString("John"));
assertThat(response.getCompletionReason(), equalTo("done"));
}, 2, 2000);
} finally {
deleteFile(uploadedFile);
}
}

@Test
public void aiExtractStructuredWithFields() throws InterruptedException {
BoxAPIConnection api = jwtApiForServiceAccount();
BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT_STRUCTURED, "en-US", null);
BoxAIAgentExtractStructured agentExtractStructured = (BoxAIAgentExtractStructured) agent;

BoxFile uploadedFile = uploadFileToUniqueFolder(api, "[aiExtractStructuredWithFields] Test File.txt",
"My name is John Doe. I was born in 4th July 1990. I am 34 years old. My hobby is guitar and books.");

try {
// When a file has been just uploaded, AI service may not be ready to return text response
// and 412 is returned
retry(() -> {
BoxAIExtractStructuredResponse response = BoxAI.extractMetadataStructured(api,
Collections.singletonList(new BoxAIItem(uploadedFile.getID(), BoxAIItem.Type.FILE)),
new ArrayList<BoxAIExtractField>() {{
add(new BoxAIExtractField("firstName"));
add(new BoxAIExtractField("lastName"));
add(new BoxAIExtractField("date",
"Person date of birth",
"Birth date",
"dateOfBirth",
null,
"What is the date of your birth?"));
add(new BoxAIExtractField("float",
"Person age",
"Age",
"age",
null,
"How old are you?"));
add(new BoxAIExtractField("multiSelect",
"Person hobby",
"Hobby",
"hobby",
new ArrayList<BoxAIExtractFieldOption>() {{
add(new BoxAIExtractFieldOption("guitar"));
add(new BoxAIExtractFieldOption("books"));
}},
"What is your hobby?"));
}},
agentExtractStructured);
JsonObject sourceJson = response.getSourceJson();
assertThat(sourceJson.get("firstName").asString(), is(equalTo("John")));
assertThat(sourceJson.get("lastName").asString(), is(equalTo("Doe")));
assertThat(sourceJson.get("dateOfBirth").asString(), is(equalTo("1990-07-04")));
assertThat(sourceJson.get("age").asInt(), is(equalTo(34)));
assertThat(sourceJson.get("hobby").asArray().get(0).asString(), is(equalTo("guitar")));
assertThat(sourceJson.get("hobby").asArray().get(1).asString(), is(equalTo("books")));
}, 2, 2000);
} finally {
deleteFile(uploadedFile);
}
}

@Test
public void aiExtractStructuredWithMetadataTemplate() throws InterruptedException {
BoxAPIConnection api = jwtApiForServiceAccount();
BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT_STRUCTURED, "en-US", null);
BoxAIAgentExtractStructured agentExtractStructured = (BoxAIAgentExtractStructured) agent;

BoxFile uploadedFile = uploadFileToUniqueFolder(api, "[aiExtractStructuredWithMetadataTemplate] Test File.txt",
"My name is John Doe. I was born in 4th July 1990. I am 34 years old. My hobby is guitar and books.");
String templateKey = "key" + java.util.UUID.randomUUID().toString();
MetadataTemplate template = MetadataTemplate.createMetadataTemplate(api,
"enterprise",
templateKey,
templateKey,
false,
new ArrayList<MetadataTemplate.Field>() {{
add(new MetadataTemplate.Field(Json.parse(
"{\"key\":\"firstName\",\"displayName\":\"First name\","
+ "\"description\":\"Person first name\",\"type\":\"string\"}").asObject()));
add(new MetadataTemplate.Field(Json.parse(
"{\"key\":\"lastName\",\"displayName\":\"Last name\","
+ "\"description\":\"Person last name\",\"type\":\"string\"}").asObject()));
add(new MetadataTemplate.Field(Json.parse(
"{\"key\":\"dateOfBirth\",\"displayName\":\"Birth date\","
+ "\"description\":\"Person date of birth\",\"type\":\"date\"}").asObject()));
add(new MetadataTemplate.Field(Json.parse(
"{\"key\":\"age\",\"displayName\":\"Age\","
+ "\"description\":\"Person age\",\"type\":\"float\"}").asObject()));
add(new MetadataTemplate.Field(Json.parse(
"{\"key\":\"hobby\",\"displayName\":\"Hobby\","
+ "\"description\":\"Person hobby\",\"type\":\"multiSelect\"}").asObject()));
}});

try {
// When a file has been just uploaded, AI service may not be ready to return text response
// and 412 is returned
retry(() -> {
BoxAIExtractStructuredResponse response = BoxAI.extractMetadataStructured(api,
Collections.singletonList(new BoxAIItem(uploadedFile.getID(), BoxAIItem.Type.FILE)),
new BoxAIExtractMetadataTemplate(templateKey, "enterprise"),
agentExtractStructured);
JsonObject sourceJson = response.getSourceJson();
assertThat(sourceJson.get("firstName").asString(), is(equalTo("John")));
assertThat(sourceJson.get("lastName").asString(), is(equalTo("Doe")));
assertThat(sourceJson.get("dateOfBirth").asString(), is(equalTo("1990-07-04")));
assertThat(sourceJson.get("age").asInt(), is(equalTo(34)));
assertThat(sourceJson.get("hobby").asArray().get(0).asString(), is(equalTo("guitar")));
assertThat(sourceJson.get("hobby").asArray().get(1).asString(), is(equalTo("books")));
}, 2, 2000);
} finally {
deleteFile(uploadedFile);
MetadataTemplate.deleteMetadataTemplate(api, template.getScope(), template.getTemplateKey());
}
}
}
Loading

0 comments on commit 7ba90b9

Please sign in to comment.