Skip to content

Commit

Permalink
feat: Support AI extract and AI extract structured
Browse files Browse the repository at this point in the history
  • Loading branch information
@congminh1254
congminh1254 committed Oct 7, 2024
1 parent 3cb2c7c commit 468c98a
Show file tree
Hide file tree
Showing 15 changed files with 895 additions and 45 deletions.
48 changes: 47 additions & 1 deletion doc/ai.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,4 +87,50 @@ 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 extracts metadata in form of key-value pairs, call static
[` extractMetadataFreeform(BoxAPIConnection api, String prompt, List<BoxAIItem> items, BoxAIAgentExtract agent)`][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
BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT, "en-US", null);
BoxAIAgentExtract agentExtract = (BoxAIAgentExtract) agent;

BoxAIResponse response = BoxAI.extractMetadataFreeform(
api,
"What is the content of the file?",
Collections.singletonList(new BoxAIItem("123456", BoxAIItem.Type.FILE)),
agentExtract
);
```

[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-com.box.sdk.ai.BoxAIAgentExtract-

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 a define a schema yourself.
To send an AI request to extract metadata from files, call static
[`extractMetadataStructured extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items, BoxAIExtractMetadataTemplate template, List<BoxAIExtractField> fields, BoxAIAgentExtractStructured agent)`][extract-metadata-structured] method.

<!-- sample post_ai_extract_structured -->
```java
BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT_STRUCTURED, "en-US", null);
BoxAIAgentExtractStructured agentExtractStructured = (BoxAIAgentExtractStructured) agent;
BoxAIExtractMetadataTemplate template = new BoxAIExtractMetadataTemplate("templateKey", "enterprise");

JsonObject result = BoxAI.extractMetadataStructured(
api,
Collections.singletonList(new BoxAIItem("123456", BoxAIItem.Type.FILE)),
template,
null,
agentExtractStructured
);
```

[extract-metadata-structured]: 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.BoxAIExtractMetadataTemplate-java.util.List-com.box.sdk.ai.BoxAIAgentExtractStructured-
171 changes: 156 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,143 @@ 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(() -> {
JsonObject response = BoxAI.extractMetadataStructured(api,
Collections.singletonList(new BoxAIItem(uploadedFile.getID(), BoxAIItem.Type.FILE)),
null,
new ArrayList<BoxAIExtractField>() {{
add(new BoxAIExtractField("string",
"Person first name",
"First name",
"firstName",
null,
"What is the your first name?"));
add(new BoxAIExtractField("string",
"Person last name", "Last name", "lastName", null, "What is the your last name?"));
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<String>() {{
add("guitar");
add("books");
}},
"What is your hobby?"));
}},
agentExtractStructured);
assertThat(response.get("firstName").asString(), is(equalTo("John")));
assertThat(response.get("lastName").asString(), is(equalTo("Doe")));
assertThat(response.get("dateOfBirth").asString(), is(equalTo("1990-07-04")));
assertThat(response.get("age").asInt(), is(equalTo(34)));
assertThat(response.get("hobby").asArray().get(0).asString(), is(equalTo("guitar")));
assertThat(response.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(() -> {
JsonObject response = BoxAI.extractMetadataStructured(api,
Collections.singletonList(new BoxAIItem(uploadedFile.getID(), BoxAIItem.Type.FILE)),
new BoxAIExtractMetadataTemplate(templateKey, "enterprise"),
null,
agentExtractStructured);
assertThat(response.get("firstName").asString(), is(equalTo("John")));
assertThat(response.get("lastName").asString(), is(equalTo("Doe")));
assertThat(response.get("dateOfBirth").asString(), is(equalTo("1990-07-04")));
assertThat(response.get("age").asInt(), is(equalTo(34)));
assertThat(response.get("hobby").asArray().get(0).asString(), is(equalTo("guitar")));
assertThat(response.get("hobby").asArray().get(1).asString(), is(equalTo("books")));
}, 2, 2000);
} finally {
deleteFile(uploadedFile);
MetadataTemplate.deleteMetadataTemplate(api, template.getScope(), template.getTemplateKey());
}
}
}
81 changes: 81 additions & 0 deletions src/main/java/com/box/sdk/BoxAI.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,14 @@ public final class BoxAI {
* AI agent default config url.
*/
public static final URLTemplate AI_AGENT_DEFAULT_CONFIG_URL = new URLTemplate("ai_agent_default");
/**
* AI extract metadata freeform url.
*/
public static final URLTemplate EXTRACT_METADATA_FREEFORM_URL = new URLTemplate("ai/extract");
/**
* AI extract metadata structured url.
*/
public static final URLTemplate EXTRACT_METADATA_STRUCTURED_URL = new URLTemplate("ai/extract_structured");

private BoxAI() {
}
Expand Down Expand Up @@ -239,4 +247,77 @@ public String toString() {
return this.mode;
}
}

/**
* Sends an AI request to supported Large Language Models (LLMs) and extracts metadata in form of key-value pairs.
* Freeform metadata extraction does not require any metadata template setup before sending the request.
*
* @param api the API connection to be used by the created user.
* @param prompt The prompt provided by the client to be answered by the LLM.
* @param items The items to be processed by the LLM, currently only files are supported.
* @param agent The AI agent configuration to be used for the request.
* @return The response from the AI.
*/
public static BoxAIResponse extractMetadataFreeform(BoxAPIConnection api,
String prompt,
List<BoxAIItem> items,
BoxAIAgentExtract agent) {
URL url = EXTRACT_METADATA_FREEFORM_URL.build(api.getBaseURL());

JsonObject requestJSON = new JsonObject();
JsonArray itemsJSON = new JsonArray();
for (BoxAIItem item : items) {
itemsJSON.add(item.getJSONObject());
}
requestJSON.add("items", itemsJSON);
requestJSON.add("prompt", prompt);
if (agent != null) {
requestJSON.add("ai_agent", agent.getJSONObject());
}

BoxJSONRequest req = new BoxJSONRequest(api, url, HttpMethod.POST);
req.setBody(requestJSON.toString());

try (BoxJSONResponse response = req.send()) {
JsonObject responseJSON = Json.parse(response.getJSON()).asObject();
return new BoxAIResponse(responseJSON);
}
}

public static JsonObject extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items,
BoxAIExtractMetadataTemplate template,
List<BoxAIExtractField> fields,
BoxAIAgentExtractStructured agent) {
URL url = EXTRACT_METADATA_STRUCTURED_URL.build(api.getBaseURL());

JsonObject requestJSON = new JsonObject();
JsonArray itemsJSON = new JsonArray();
for (BoxAIItem item : items) {
itemsJSON.add(item.getJSONObject());
}
requestJSON.add("items", itemsJSON);

if (template != null) {
requestJSON.add("metadata_template", template.getJSONObject());
}

if (fields != null) {
JsonArray fieldsJSON = new JsonArray();
for (BoxAIExtractField field : fields) {
fieldsJSON.add(field.getJSONObject());
}
requestJSON.add("fields", fieldsJSON);
}

if (agent != null) {
requestJSON.add("ai_agent", agent.getJSONObject());
}

BoxJSONRequest req = new BoxJSONRequest(api, url, HttpMethod.POST);
req.setBody(requestJSON.toString());

try (BoxJSONResponse response = req.send()) {
return Json.parse(response.getJSON()).asObject();
}
}
}
Loading

0 comments on commit 468c98a

Please sign in to comment.