fix: Fix AI dialogue history (#1289)

box · Feb 19, 2025 · 29b6519 · 29b6519
1 parent 9119adc
commit 29b6519
Show file tree

Hide file tree

Showing 13 changed files with 186 additions and 56 deletions.
diff --git a/src/intTest/java/com/box/sdk/BoxAIIT.java b/src/intTest/java/com/box/sdk/BoxAIIT.java
@@ -10,9 +10,9 @@
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.notNullValue;
 
 import com.eclipsesource.json.Json;
-import com.eclipsesource.json.JsonObject;
 import java.text.ParseException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -111,8 +111,8 @@ public void askAIMultipleItems() throws InterruptedException {
     public void askAITextGenItemWithDialogueHistory() throws ParseException, InterruptedException {
         BoxAPIConnection api = jwtApiForServiceAccount();
         String fileName = "[askAITextGenItemWithDialogueHistory] Test File.txt";
-        Date date1 = BoxDateFormat.parse("2013-05-16T15:27:57-07:00");
-        Date date2 = BoxDateFormat.parse("2013-05-16T15:26:57-07:00");
+        Date date1 = BoxDateFormat.parse("2021-01-01T00:00:00Z");
+        Date date2 = BoxDateFormat.parse("2021-02-01T00:00:00Z");
 
         BoxFile uploadedFile = uploadFileToUniqueFolder(api, fileName, "Test file");
         try {
@@ -148,28 +148,25 @@ public void askAITextGenItemWithDialogueHistory() throws ParseException, Interru
     @Test
     public void getAIAgentDefaultConfiguration() {
         BoxAPIConnection api = jwtApiForServiceAccount();
-        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.ASK,
-            "en", "openai__gpt_3_5_turbo");
+        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.ASK);
         BoxAIAgentAsk askAgent = (BoxAIAgentAsk) agent;
 
         assertThat(askAgent.getType(), is(equalTo(BoxAIAgentAsk.TYPE)));
-        assertThat(askAgent.getBasicText().getModel(), is(equalTo("openai__gpt_3_5_turbo")));
+        assertThat(askAgent.getBasicText().getModel(), is(notNullValue()));
 
-        BoxAIAgent agent2 = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.TEXT_GEN,
-            "en", "openai__gpt_3_5_turbo");
+        BoxAIAgent agent2 = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.TEXT_GEN);
         BoxAIAgentTextGen textGenAgent = (BoxAIAgentTextGen) agent2;
 
         assertThat(textGenAgent.getType(), is(equalTo(BoxAIAgentTextGen.TYPE)));
-        assertThat(textGenAgent.getBasicGen().getModel(), is(equalTo("openai__gpt_3_5_turbo")));
+        assertThat(textGenAgent.getBasicGen().getModel(), is(notNullValue()));
     }
 
     @Test
     public void askAISingleItemWithAgent() throws InterruptedException {
         BoxAPIConnection api = jwtApiForServiceAccount();
         String fileName = "[askAISingleItem] Test File.txt";
         BoxFile uploadedFile = uploadFileToUniqueFolder(api, fileName, "Test file");
-        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.ASK,
-            "en", "openai__gpt_3_5_turbo_16k");
+        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.ASK);
         BoxAIAgentAsk askAgent = (BoxAIAgentAsk) agent;
 
         try {
@@ -199,8 +196,10 @@ public void askAISingleItemWithAgent() throws InterruptedException {
     @Test
     public void aiExtract() throws InterruptedException {
         BoxAPIConnection api = jwtApiForServiceAccount();
-        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT, "en-US", null);
+        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT);
         BoxAIAgentExtract agentExtract = (BoxAIAgentExtract) agent;
+        // AI team is going to move away from supporting overriding embeddings model
+        agentExtract.getLongText().setEmbeddings(null);
 
         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.");
@@ -224,8 +223,10 @@ public void aiExtract() throws InterruptedException {
     @Test
     public void aiExtractStructuredWithFields() throws InterruptedException {
         BoxAPIConnection api = jwtApiForServiceAccount();
-        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT_STRUCTURED, "en-US", null);
+        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT_STRUCTURED);
         BoxAIAgentExtractStructured agentExtractStructured = (BoxAIAgentExtractStructured) agent;
+        // AI team is going to move away from supporting overriding embeddings model
+        agentExtractStructured.getLongText().setEmbeddings(null);
 
         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.");
@@ -259,12 +260,16 @@ public void aiExtractStructuredWithFields() throws InterruptedException {
                                 "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(response.getSourceJson().get("answer"), is(equalTo(response.getAnswer())));
+
+                assertThat(response.getAnswer().get("firstName").asString(), is(equalTo("John")));
+                assertThat(response.getAnswer().get("lastName").asString(), is(equalTo("Doe")));
+                assertThat(response.getAnswer().get("dateOfBirth").asString(), is(equalTo("1990-07-04")));
+                assertThat(response.getAnswer().get("age").asInt(), is(equalTo(34)));
+                assertThat(response.getAnswer().get("hobby").asArray().get(0).asString(), is(equalTo("guitar")));
+
+                assertThat(response.getCompletionReason(), equalTo("done"));
+                assertThat(response.getCreatedAt(), is(notNullValue()));
             }, 2, 2000);
         } finally {
             deleteFile(uploadedFile);
@@ -274,8 +279,10 @@ public void aiExtractStructuredWithFields() throws InterruptedException {
     @Test
     public void aiExtractStructuredWithMetadataTemplate() throws InterruptedException {
         BoxAPIConnection api = jwtApiForServiceAccount();
-        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT_STRUCTURED, "en-US", null);
+        BoxAIAgent agent = BoxAI.getAiAgentDefaultConfig(api, BoxAIAgent.Mode.EXTRACT_STRUCTURED);
         BoxAIAgentExtractStructured agentExtractStructured = (BoxAIAgentExtractStructured) agent;
+        // AI team is going to move away from supporting overriding embeddings model
+        agentExtractStructured.getLongText().setEmbeddings(null);
 
         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.");
@@ -312,12 +319,15 @@ public void aiExtractStructuredWithMetadataTemplate() throws InterruptedExceptio
                     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-04T00:00:00Z")));
-                assertThat(sourceJson.get("age").asInt(), is(equalTo(34)));
-                assertThat(sourceJson.get("hobby").asArray().get(0).asString(), is(equalTo("guitar")));
+                assertThat(response.getSourceJson().get("answer"), is(equalTo(response.getAnswer())));
+
+                assertThat(response.getAnswer().get("firstName").asString(), is(equalTo("John")));
+                assertThat(response.getAnswer().get("lastName").asString(), is(equalTo("Doe")));
+                assertThat(response.getAnswer().get("dateOfBirth").asString(), is(equalTo("1990-07-04T00:00:00Z")));
+                assertThat(response.getAnswer().get("age").asInt(), is(equalTo(34)));
+                assertThat(response.getAnswer().get("hobby").asArray().get(0).asString(), is(equalTo("guitar")));
+                assertThat(response.getCompletionReason(), equalTo("done"));
+                assertThat(response.getCreatedAt(), is(notNullValue()));
             }, 2, 2000);
         } finally {
             deleteFile(uploadedFile);

diff --git a/src/main/java/com/box/sdk/BoxAIAgentAsk.java b/src/main/java/com/box/sdk/BoxAIAgentAsk.java
@@ -151,10 +151,18 @@ void parseJSONMember(JsonObject.Member member) {
     public JsonObject getJSONObject() {
         JsonObject jsonObject = new JsonObject();
         JsonUtils.addIfNotNull(jsonObject, "type", this.getType());
-        JsonUtils.addIfNotNull(jsonObject, "basic_text", this.basicText.getJSONObject());
-        JsonUtils.addIfNotNull(jsonObject, "basic_text_multi", this.basicTextMulti.getJSONObject());
-        JsonUtils.addIfNotNull(jsonObject, "long_text", this.longText.getJSONObject());
-        JsonUtils.addIfNotNull(jsonObject, "long_text_multi", this.longTextMulti.getJSONObject());
+        if (this.basicText != null) {
+            jsonObject.add("basic_text", this.basicText.getJSONObject());
+        }
+        if (this.basicTextMulti != null) {
+            jsonObject.add("basic_text_multi", this.basicTextMulti.getJSONObject());
+        }
+        if (this.longText != null) {
+            jsonObject.add("long_text", this.longText.getJSONObject());
+        }
+        if (this.longTextMulti != null) {
+            jsonObject.add("long_text_multi", this.longTextMulti.getJSONObject());
+        }
         return jsonObject;
     }
 }

diff --git a/src/main/java/com/box/sdk/BoxAIAgentAskBasicText.java b/src/main/java/com/box/sdk/BoxAIAgentAskBasicText.java
@@ -175,7 +175,9 @@ void parseJSONMember(JsonObject.Member member) {
 
     public JsonObject getJSONObject() {
         JsonObject jsonObject = new JsonObject();
-        JsonUtils.addIfNotNull(jsonObject, "llm_endpoint_params", this.llmEndpointParams.getJSONObject());
+        if (this.llmEndpointParams != null) {
+            jsonObject.add("llm_endpoint_params", this.llmEndpointParams.getJSONObject());
+        }
         JsonUtils.addIfNotNull(jsonObject, "model", this.model);
         JsonUtils.addIfNotNull(jsonObject, "num_tokens_for_completion", this.numTokensForCompletion);
         JsonUtils.addIfNotNull(jsonObject, "prompt_template", this.promptTemplate);

diff --git a/src/main/java/com/box/sdk/BoxAIAgentAskLongText.java b/src/main/java/com/box/sdk/BoxAIAgentAskLongText.java
@@ -201,8 +201,12 @@ void parseJSONMember(JsonObject.Member member) {
 
     public JsonObject getJSONObject() {
         JsonObject jsonObject = new JsonObject();
-        JsonUtils.addIfNotNull(jsonObject, "embeddings", this.embeddings.getJSONObject());
-        JsonUtils.addIfNotNull(jsonObject, "llm_endpoint_params", this.llmEndpointParams.getJSONObject());
+        if (this.embeddings != null) {
+            jsonObject.add("embeddings", this.embeddings.getJSONObject());
+        }
+        if (this.llmEndpointParams != null) {
+            jsonObject.add("llm_endpoint_params", this.llmEndpointParams.getJSONObject());
+        }
         JsonUtils.addIfNotNull(jsonObject, "model", this.model);
         JsonUtils.addIfNotNull(jsonObject, "num_tokens_for_completion", this.numTokensForCompletion);
         JsonUtils.addIfNotNull(jsonObject, "prompt_template", this.promptTemplate);

diff --git a/src/main/java/com/box/sdk/BoxAIAgentEmbeddings.java b/src/main/java/com/box/sdk/BoxAIAgentEmbeddings.java
@@ -87,7 +87,9 @@ void parseJSONMember(JsonObject.Member member) {
     public JsonObject getJSONObject() {
         JsonObject jsonObject = new JsonObject();
         JsonUtils.addIfNotNull(jsonObject, "model", this.model);
-        JsonUtils.addIfNotNull(jsonObject, "strategy", this.strategy.getJSONObject());
+        if (this.strategy != null) {
+            jsonObject.add("strategy", this.strategy.getJSONObject());
+        }
         return jsonObject;
     }
 }
diff --git a/src/main/java/com/box/sdk/BoxAIAgentExtract.java b/src/main/java/com/box/sdk/BoxAIAgentExtract.java
@@ -101,8 +101,12 @@ void parseJSONMember(JsonObject.Member member) {
     public JsonObject getJSONObject() {
         JsonObject jsonObject = new JsonObject();
         JsonUtils.addIfNotNull(jsonObject, "type", this.getType());
-        JsonUtils.addIfNotNull(jsonObject, "basic_text", this.basicText.getJSONObject());
-        JsonUtils.addIfNotNull(jsonObject, "long_text", this.longText.getJSONObject());
+        if (this.basicText != null) {
+            jsonObject.add("basic_text", this.basicText.getJSONObject());
+        }
+        if (this.longText != null) {
+            jsonObject.add("long_text", this.longText.getJSONObject());
+        }
         return jsonObject;
     }
 }

diff --git a/src/main/java/com/box/sdk/BoxAIAgentExtractStructured.java b/src/main/java/com/box/sdk/BoxAIAgentExtractStructured.java
@@ -101,8 +101,12 @@ void parseJSONMember(JsonObject.Member member) {
     public JsonObject getJSONObject() {
         JsonObject jsonObject = new JsonObject();
         JsonUtils.addIfNotNull(jsonObject, "type", this.getType());
-        JsonUtils.addIfNotNull(jsonObject, "basic_text", this.basicText.getJSONObject());
-        JsonUtils.addIfNotNull(jsonObject, "long_text", this.longText.getJSONObject());
+        if (this.basicText != null) {
+            jsonObject.add("basic_text", this.basicText.getJSONObject());
+        }
+        if (this.longText != null) {
+            jsonObject.add("long_text", this.longText.getJSONObject());
+        }
         return jsonObject;
     }
 }

diff --git a/src/main/java/com/box/sdk/BoxAIAgentTextGen.java b/src/main/java/com/box/sdk/BoxAIAgentTextGen.java
@@ -64,7 +64,9 @@ void parseJSONMember(JsonObject.Member member) {
     public JsonObject getJSONObject() {
         JsonObject jsonObject = new JsonObject();
         JsonUtils.addIfNotNull(jsonObject, "type", this.getType());
-        JsonUtils.addIfNotNull(jsonObject, "basic_gen", this.basicGen.getJSONObject());
+        if (this.basicGen != null) {
+            jsonObject.add("basic_gen", this.basicGen.getJSONObject());
+        }
         return jsonObject;
     }
 }
diff --git a/src/main/java/com/box/sdk/BoxAIAgentTextGenBasicGen.java b/src/main/java/com/box/sdk/BoxAIAgentTextGenBasicGen.java
@@ -224,8 +224,12 @@ void parseJSONMember(JsonObject.Member member) {
     public JsonObject getJSONObject() {
         JsonObject jsonObject = new JsonObject();
         JsonUtils.addIfNotNull(jsonObject, "content_template", this.contentTemplate);
-        JsonUtils.addIfNotNull(jsonObject, "embeddings", this.embeddings.getJSONObject());
-        JsonUtils.addIfNotNull(jsonObject, "llm_endpoint_params", this.llmEndpointParams.getJSONObject());
+        if (this.embeddings != null) {
+            jsonObject.add("embeddings", this.embeddings.getJSONObject());
+        }
+        if (this.llmEndpointParams != null) {
+            jsonObject.add("llm_endpoint_params", this.llmEndpointParams.getJSONObject());
+        }
         JsonUtils.addIfNotNull(jsonObject, "model", this.model);
         JsonUtils.addIfNotNull(jsonObject, "num_tokens_for_completion", this.numTokensForCompletion);
         JsonUtils.addIfNotNull(jsonObject, "prompt_template", this.promptTemplate);

diff --git a/src/main/java/com/box/sdk/BoxAIDialogueEntry.java b/src/main/java/com/box/sdk/BoxAIDialogueEntry.java
@@ -94,11 +94,11 @@ public void setCreatedAt(Date createdAt) {
      */
     public JsonObject getJSONObject() {
         JsonObject itemJSON = new JsonObject()
-                .add("id", this.prompt)
-                .add("type", this.answer);
+                .add("prompt", this.prompt)
+                .add("answer", this.answer);
 
         if (this.createdAt != null) {
-            itemJSON.add("content", this.createdAt.toString());
+            itemJSON.add("created_at", BoxDateFormat.format(this.createdAt));
         }
 
         return itemJSON;

diff --git a/src/main/java/com/box/sdk/BoxAIExtractStructuredResponse.java b/src/main/java/com/box/sdk/BoxAIExtractStructuredResponse.java
@@ -2,12 +2,18 @@
 
 import com.eclipsesource.json.Json;
 import com.eclipsesource.json.JsonObject;
+import com.eclipsesource.json.JsonValue;
+import java.text.ParseException;
+import java.util.Date;
 
 /**
  * AI response to a user request.
  */
 public class BoxAIExtractStructuredResponse extends BoxJSONObject {
     private final JsonObject sourceJson;
+    private JsonObject answer;
+    private String completionReason;
+    private Date createdAt;
 
     /**
      * Constructs a BoxAIResponse object.
@@ -35,4 +41,52 @@ public BoxAIExtractStructuredResponse(String json) {
     public JsonObject getSourceJson() {
         return sourceJson;
     }
+
+    /**
+     * Gets the answer of the AI.
+     *
+     * @return the answer of the AI.
+     */
+    public JsonObject getAnswer() {
+        return answer;
+    }
+
+    /**
+     * Gets reason the response finishes.
+     *
+     * @return the reason the response finishes.
+     */
+    public String getCompletionReason() {
+        return completionReason;
+    }
+
+    /**
+     * Gets the ISO date formatted timestamp of when the answer to the prompt was created.
+     *
+     * @return The ISO date formatted timestamp of when the answer to the prompt was created.
+     */
+    public Date getCreatedAt() {
+        return createdAt;
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    void parseJSONMember(JsonObject.Member member) {
+        JsonValue value = member.getValue();
+        String memberName = member.getName();
+        try {
+            if (memberName.equals("answer")) {
+                this.answer = value.asObject();
+            } else if (memberName.equals("completion_reason")) {
+                this.completionReason = value.asString();
+            } else if (memberName.equals("created_at")) {
+                this.createdAt = BoxDateFormat.parse(value.asString());
+            }
+        } catch (ParseException e) {
+            assert false : "A ParseException indicates a bug in the SDK.";
+        }
+    }
+
 }
diff --git a/src/test/Fixtures/BoxAI/ExtractMetadataStructured200.json b/src/test/Fixtures/BoxAI/ExtractMetadataStructured200.json
@@ -1,9 +1,13 @@
 {
-  "firstName": "John",
-  "lastName": "Doe",
-  "age": 25,
-  "hobbies": [
-    "reading",
-    "travelling"
-  ]
+  "answer": {
+    "firstName": "John",
+    "lastName": "Doe",
+    "age": 25,
+    "hobbies": [
+      "reading",
+      "travelling"
+    ]
+  },
+  "completion_reason": "done",
+  "created_at": "2012-12-12T10:53:43.123-08:00"
 }