Skip to content
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

feat: Support AI extract and AI extract structured #1266

Merged
merged 8 commits into from
Oct 11, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 25 additions & 10 deletions doc/ai.md
Original file line number Diff line number Diff line change
Expand Up @@ -95,39 +95,54 @@ 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, BoxAIAgentExtract agent)`][extract-metadata-freeform] method.
[`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)),
agent
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-com.box.sdk.ai.BoxAIAgentExtract-
[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, call static
[`extractMetadataStructured extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items, BoxAIExtractMetadataTemplate template, List<BoxAIExtractField> fields, BoxAIAgentExtractStructured agent)`][extract-metadata-structured] method.

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,
null,
agent
template
);
JsonObject sourceJson = result.getSourceJson();
```

To send an AI request to extract metadata from files with a custom fields, call static
congminh1254 marked this conversation as resolved.
Show resolved Hide resolved
[`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]: 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-
[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-
12 changes: 2 additions & 10 deletions src/intTest/java/com/box/sdk/BoxAIIT.java
Original file line number Diff line number Diff line change
Expand Up @@ -235,16 +235,9 @@ public void aiExtractStructuredWithFields() throws InterruptedException {
retry(() -> {
BoxAIExtractStructuredResponse 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("firstName"));
add(new BoxAIExtractField("lastName"));
add(new BoxAIExtractField("date",
"Person date of birth",
"Birth date",
Expand Down Expand Up @@ -320,7 +313,6 @@ public void aiExtractStructuredWithMetadataTemplate() throws InterruptedExceptio
BoxAIExtractStructuredResponse response = BoxAI.extractMetadataStructured(api,
Collections.singletonList(new BoxAIItem(uploadedFile.getID(), BoxAIItem.Type.FILE)),
new BoxAIExtractMetadataTemplate(templateKey, "enterprise"),
null,
agentExtractStructured);
JsonObject sourceJson = response.getSourceJson();
assertThat(sourceJson.get("firstName").asString(), is(equalTo("John")));
Expand Down
91 changes: 91 additions & 0 deletions src/main/java/com/box/sdk/BoxAI.java
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,21 @@ public String toString() {
}
}

/**
* 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.
* @return The response from the AI.
*/
public static BoxAIResponse extractMetadataFreeform(BoxAPIConnection api,
String prompt,
List<BoxAIItem> items) {
return extractMetadataFreeform(api, prompt, items, null);
}

/**
* 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.
Expand Down Expand Up @@ -284,7 +299,83 @@ public static BoxAIResponse extractMetadataFreeform(BoxAPIConnection api,
}
}

/**
* 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.
*
* @param api The API connection to be used by the created user.
* @param items The items to be processed by the LLM, currently only files are supported.
* @param template The metadata template to be used for the request.
* @return The response from the AI.
*/
public static BoxAIExtractStructuredResponse extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items,
congminh1254 marked this conversation as resolved.
Show resolved Hide resolved
BoxAIExtractMetadataTemplate template) {
return extractMetadataStructured(api, items, template, null, null);
}

/**
* 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.
*
* @param api The API connection to be used by the created user.
* @param items The items to be processed by the LLM, currently only files are supported.
* @param template The metadata template to be used for the request.
* @param agent The AI agent configuration to be used for the request.
* @return The response from the AI.
*/
public static BoxAIExtractStructuredResponse extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items,
BoxAIExtractMetadataTemplate template,
BoxAIAgentExtractStructured agent) {
return extractMetadataStructured(api, items, template, null, agent);
}

/**
* 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.
*
* @param api The API connection to be used by the created user.
* @param items The items to be processed by the LLM, currently only files are supported.
* @param fields The fields to be extracted from the items.
* @return The response from the AI.
*/
public static BoxAIExtractStructuredResponse extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items,
List<BoxAIExtractField> fields) {
return extractMetadataStructured(api, items, null, fields, null);
}

/**
* 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.
*
* @param api The API connection to be used by the created user.
* @param items The items to be processed by the LLM, currently only files are supported.
* @param fields The fields to be extracted from the items.
* @param agent The AI agent configuration to be used for the request.
* @return The response from the AI.
*/
public static BoxAIExtractStructuredResponse extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items,
List<BoxAIExtractField> fields,
BoxAIAgentExtractStructured agent) {
return extractMetadataStructured(api, items, null, fields, agent);
}

/**
* 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.
*
* @param api The API connection to be used by the created user.
* @param items The items to be processed by the LLM, currently only files are supported.
* @param template The metadata template to be used for the request.
* @param fields The fields to be extracted from the items.
* @param agent The AI agent configuration to be used for the request.
* @return The response from the AI.
*/
private static BoxAIExtractStructuredResponse extractMetadataStructured(BoxAPIConnection api, List<BoxAIItem> items,
BoxAIExtractMetadataTemplate template,
List<BoxAIExtractField> fields,
BoxAIAgentExtractStructured agent) {
Expand Down
7 changes: 7 additions & 0 deletions src/main/java/com/box/sdk/BoxAIExtractField.java
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,13 @@ public class BoxAIExtractField extends BoxJSONObject {
*/
private String prompt;

/**
* Constructs a BoxAIExtractField object with a given key.
*/
public BoxAIExtractField(String key) {
this.key = key;
}

/**
* Constructs a BoxAIExtractField object with a given type, description, display name, key, options, and prompt.
*
Expand Down
56 changes: 44 additions & 12 deletions src/test/java/com/box/sdk/BoxAITest.java
Original file line number Diff line number Diff line change
Expand Up @@ -315,23 +315,13 @@ public void testExtractMetadataFreeformSuccess() {
}

@Test
public void testExtractMetadataStructuredSuccess() {
public void testExtractMetadataStructuredWithTemplateSuccess() {
String agentString = TestUtils.getFixture("BoxAI/GetAIAgentDefaultConfigExtractStructured200");
final String fileId = "12345";
final BoxAIExtractMetadataTemplate template = new BoxAIExtractMetadataTemplate("templateId", "enterprise");
final List<BoxAIItem> items = Collections.singletonList(new BoxAIItem(fileId, BoxAIItem.Type.FILE));
final String result = TestUtils.getFixture("BoxAI/ExtractMetadataStructured200");
final BoxAIAgentExtractStructured agent = new BoxAIAgentExtractStructured(Json.parse(agentString).asObject());
final BoxAIExtractField field = new BoxAIExtractField(
"text",
"The name of the file",
"Name",
"name",
new ArrayList<BoxAIExtractFieldOption>() {{
add(new BoxAIExtractFieldOption("option 1"));
add(new BoxAIExtractFieldOption("option 2"));
}},
"What is the name of the file?");

final JsonObject expectedRequestBody = new JsonObject()
.add("items", new JsonArray().add(new JsonObject()
Expand All @@ -341,6 +331,49 @@ public void testExtractMetadataStructuredSuccess() {
.add("type", "metadata_template")
.add("template_key", template.getTemplateKey())
.add("scope", template.getScope()))
.add("ai_agent", agent.getJSONObject());

wireMockRule.stubFor(WireMock.post(WireMock.urlPathEqualTo("/2.0/ai/extract_structured"))
.withRequestBody(WireMock.equalToJson(
expectedRequestBody.toString()
))
.willReturn(WireMock.aResponse()
.withHeader("Content-Type", APPLICATION_JSON)
.withBody(result)));

BoxAIExtractStructuredResponse response = BoxAI.extractMetadataStructured(
api,
items,
template,
agent
);

assertThat(response.getSourceJson(), equalTo(Json.parse(result).asObject()));
assertThat(response.getSourceJson().get("firstName").asString(), equalTo("John"));
}

@Test
public void testExtractMetadataStructuredWithFieldSuccess() {
String agentString = TestUtils.getFixture("BoxAI/GetAIAgentDefaultConfigExtractStructured200");
final String fileId = "12345";
final List<BoxAIItem> items = Collections.singletonList(new BoxAIItem(fileId, BoxAIItem.Type.FILE));
final String result = TestUtils.getFixture("BoxAI/ExtractMetadataStructured200");
final BoxAIAgentExtractStructured agent = new BoxAIAgentExtractStructured(Json.parse(agentString).asObject());
final BoxAIExtractField field = new BoxAIExtractField(
"text",
"The name of the file",
"Name",
"name",
new ArrayList<BoxAIExtractFieldOption>() {{
add(new BoxAIExtractFieldOption("option 1"));
add(new BoxAIExtractFieldOption("option 2"));
}},
"What is the name of the file?");

final JsonObject expectedRequestBody = new JsonObject()
.add("items", new JsonArray().add(new JsonObject()
.add("type", "file")
.add("id", fileId)))
.add("ai_agent", agent.getJSONObject())
.add("fields", new JsonArray().add(field.getJSONObject()));

Expand All @@ -355,7 +388,6 @@ public void testExtractMetadataStructuredSuccess() {
BoxAIExtractStructuredResponse response = BoxAI.extractMetadataStructured(
api,
items,
template,
Collections.singletonList(field),
agent
);
Expand Down
Loading