Fix Assistant POJO issues

Easier to make separate objects for getting, creating, and modifying Assistants in case the parameters diverge later.
TheoKanning · Nov 14, 2023 · f223175 · f223175
1 parent bc81a29
commit f223175
Show file tree

Hide file tree

Showing 8 changed files with 122 additions and 52 deletions.
diff --git a/api/src/main/java/com/theokanning/openai/assistants/Assistant.java b/api/src/main/java/com/theokanning/openai/assistants/Assistant.java
@@ -1,10 +1,16 @@
 package com.theokanning.openai.assistants;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
-import lombok.Data;
+import lombok.*;
 
+import java.util.List;
+import java.util.Map;
+
+@Builder
+@NoArgsConstructor
+@AllArgsConstructor
 @Data
-public class Assistant extends AssistantBase {
+public class Assistant {
 
     /**
      * The identifier, which can be referenced in API endpoints.
@@ -21,4 +27,41 @@ public class Assistant extends AssistantBase {
      */
     @JsonProperty("created_at")
     Integer createdAt;
+
+    /**
+     * The name of the assistant. The maximum length is 256
+     */
+    String name;
+
+    /**
+     * The description of the assistant.
+     */
+    String description;
+
+    /**
+     * ID of the model to use
+     */
+    @NonNull
+    String model;
+
+    /**
+     * The system instructions that the assistant uses.
+     */
+    String instructions;
+
+    /**
+     * A list of tools enabled on the assistant.
+     */
+    List<Tool> tools;
+
+    /**
+     * A list of file IDs attached to this assistant.
+     */
+    @JsonProperty("file_ids")
+    List<String> fileIds;
+
+    /**
+     * Set of 16 key-value pairs that can be attached to an object.
+     */
+    Map<String, String> metadata;
 }
diff --git a/api/src/main/java/com/theokanning/openai/assistants/AssistantRequest.java b/api/src/main/java/com/theokanning/openai/assistants/AssistantRequest.java
@@ -1,6 +1,52 @@
 package com.theokanning.openai.assistants;
 
 
-public class AssistantRequest extends AssistantBase {
+import com.fasterxml.jackson.annotation.JsonProperty;
+import lombok.*;
 
+import java.util.List;
+import java.util.Map;
+
+@Builder
+@NoArgsConstructor
+@AllArgsConstructor
+@Data
+public class AssistantRequest {
+
+    /**
+     * ID of the model to use
+     */
+    @NonNull
+    String model;
+
+    /**
+     * The name of the assistant. The maximum length is 256
+     */
+    String name;
+
+    /**
+     * The description of the assistant.
+     */
+    String description;
+
+    /**
+     * The system instructions that the assistant uses.
+     */
+    String instructions;
+
+    /**
+     * A list of tools enabled on the assistant.
+     */
+    List<Tool> tools;
+
+    /**
+     * A list of file IDs attached to this assistant.
+     */
+    @JsonProperty("file_ids")
+    List<String> fileIds;
+
+    /**
+     * Set of 16 key-value pairs that can be attached to an object.
+     */
+    Map<String, String> metadata;
 }
diff --git a/api/src/main/java/com/theokanning/openai/assistants/ListAssistant.java b/api/src/main/java/com/theokanning/openai/assistants/ListAssistant.java
@@ -3,7 +3,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.theokanning.openai.OpenAiResponse;
 
-public class ListAssistant<T extends AssistantBase> extends OpenAiResponse<T> {
+public class ListAssistant<T> extends OpenAiResponse<T> {
 
     @JsonProperty("first_id")
     String firstId;

diff --git a/...ning/openai/assistants/AssistantBase.java → ...ai/assistants/ModifyAssistantRequest.java b/...ning/openai/assistants/AssistantBase.java → ...ai/assistants/ModifyAssistantRequest.java
@@ -1,11 +1,8 @@
 package com.theokanning.openai.assistants;
 
+
 import com.fasterxml.jackson.annotation.JsonProperty;
-import lombok.AllArgsConstructor;
-import lombok.Builder;
-import lombok.Data;
-import lombok.NoArgsConstructor;
-import lombok.NonNull;
+import lombok.*;
 
 import java.util.List;
 import java.util.Map;
@@ -14,12 +11,11 @@
 @NoArgsConstructor
 @AllArgsConstructor
 @Data
-public class AssistantBase {
+public class ModifyAssistantRequest {
 
     /**
      * ID of the model to use
      */
-    @NonNull
     String model;
 
     /**
@@ -46,7 +42,7 @@ public class AssistantBase {
      * A list of file IDs attached to this assistant.
      */
     @JsonProperty("file_ids")
-    List<String> fields;
+    List<String> fileIds;
 
     /**
      * Set of 16 key-value pairs that can be attached to an object.

diff --git a/api/src/main/java/com/theokanning/openai/assistants/Tool.java b/api/src/main/java/com/theokanning/openai/assistants/Tool.java
@@ -1,5 +1,6 @@
 package com.theokanning.openai.assistants;
 
+import com.theokanning.openai.completion.chat.ChatFunction;
 import lombok.AllArgsConstructor;
 import lombok.Data;
 import lombok.NoArgsConstructor;
@@ -8,5 +9,13 @@
 @AllArgsConstructor
 @Data
 public class Tool {
+     /**
+      * The type of tool being defined
+      */
      AssistantToolsEnum type;
+
+     /**
+      * Function definition, only used if type is "function"
+      */
+     ChatFunction function;
 }
diff --git a/client/src/main/java/com/theokanning/openai/client/OpenAiApi.java b/client/src/main/java/com/theokanning/openai/client/OpenAiApi.java
@@ -2,12 +2,7 @@
 
 import com.theokanning.openai.DeleteResult;
 import com.theokanning.openai.OpenAiResponse;
-import com.theokanning.openai.assistants.AssistantBase;
-import com.theokanning.openai.assistants.Assistant;
-import com.theokanning.openai.assistants.AssistantFile;
-import com.theokanning.openai.assistants.AssistantFileRequest;
-import com.theokanning.openai.assistants.ListAssistant;
-import com.theokanning.openai.assistants.ListAssistantQueryRequest;
+import com.theokanning.openai.assistants.*;
 import com.theokanning.openai.audio.CreateSpeechRequest;
 import com.theokanning.openai.audio.TranscriptionResult;
 import com.theokanning.openai.audio.TranslationResult;
@@ -195,15 +190,15 @@ public interface OpenAiApi {
 
     @Headers({"OpenAI-Beta: assistants=v1"})
     @POST("/v1/assistants")
-    Single<Assistant> createAssistant(@Body AssistantBase request);
+    Single<Assistant> createAssistant(@Body AssistantRequest request);
 
     @Headers({"OpenAI-Beta: assistants=v1"})
     @GET("/v1/assistants/{assistant_id}")
     Single<Assistant> retrieveAssistant(@Path("assistant_id") String assistantId);
 
     @Headers({"OpenAI-Beta: assistants=v1"})
     @POST("/v1/assistants/{assistant_id}")
-    Single<Assistant> modifyAssistant(@Path("assistant_id") String assistantId, @Body AssistantBase request);
+    Single<Assistant> modifyAssistant(@Path("assistant_id") String assistantId, @Body ModifyAssistantRequest request);
 
     @Headers({"OpenAI-Beta: assistants=v1"})
     @DELETE("/v1/assistants/{assistant_id}")

diff --git a/service/src/main/java/com/theokanning/openai/service/OpenAiService.java b/service/src/main/java/com/theokanning/openai/service/OpenAiService.java
@@ -6,15 +6,11 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.PropertyNamingStrategy;
 import com.fasterxml.jackson.databind.node.TextNode;
+import com.sun.org.apache.xpath.internal.operations.Mod;
 import com.theokanning.openai.DeleteResult;
 import com.theokanning.openai.OpenAiError;
 import com.theokanning.openai.OpenAiHttpException;
-import com.theokanning.openai.assistants.Assistant;
-import com.theokanning.openai.assistants.AssistantBase;
-import com.theokanning.openai.assistants.AssistantFile;
-import com.theokanning.openai.assistants.AssistantFileRequest;
-import com.theokanning.openai.assistants.ListAssistant;
-import com.theokanning.openai.assistants.ListAssistantQueryRequest;
+import com.theokanning.openai.assistants.*;
 import com.theokanning.openai.audio.CreateSpeechRequest;
 import com.theokanning.openai.audio.CreateTranscriptionRequest;
 import com.theokanning.openai.audio.CreateTranslationRequest;
@@ -362,15 +358,15 @@ public ModerationResult createModeration(ModerationRequest request) {
         return execute(api.createModeration(request));
     }
 
-    public Assistant createAssistant(AssistantBase request) {
+    public Assistant createAssistant(AssistantRequest request) {
         return execute(api.createAssistant(request));
     }
 
     public Assistant retrieveAssistant(String assistantId) {
         return execute(api.retrieveAssistant(assistantId));
     }
 
-    public Assistant modifyAssistant(String assistantId, AssistantBase request) {
+    public Assistant modifyAssistant(String assistantId, ModifyAssistantRequest request) {
         return execute(api.modifyAssistant(assistantId, request));
     }
 

diff --git a/service/src/test/java/com/theokanning/openai/service/AssistantTest.java b/service/src/test/java/com/theokanning/openai/service/AssistantTest.java
@@ -1,16 +1,7 @@
 package com.theokanning.openai.service;
 
 import com.theokanning.openai.DeleteResult;
-import com.theokanning.openai.assistants.Assistant;
-import com.theokanning.openai.assistants.AssistantBase;
-import com.theokanning.openai.assistants.AssistantFile;
-import com.theokanning.openai.assistants.AssistantFileRequest;
-import com.theokanning.openai.assistants.AssistantRequest;
-import com.theokanning.openai.assistants.AssistantSortOrder;
-import com.theokanning.openai.assistants.AssistantToolsEnum;
-import com.theokanning.openai.assistants.ListAssistant;
-import com.theokanning.openai.assistants.ListAssistantQueryRequest;
-import com.theokanning.openai.assistants.Tool;
+import com.theokanning.openai.assistants.*;
 import com.theokanning.openai.file.File;
 import com.theokanning.openai.utils.TikTokensUtil;
 import org.junit.jupiter.api.AfterAll;
@@ -44,10 +35,12 @@ void retrieveAssistant() {
     void modifyAssistant() {
         Assistant createAssistantResponse = createAndValidateAssistant();
 
-        String modifiedName = MATH_TUTOR + " Modified";
-        createAssistantResponse.setName(modifiedName);//modify a field
+        String modifiedName = MATH_TUTOR + "Modified";
+        ModifyAssistantRequest modifyRequest = ModifyAssistantRequest.builder()
+                .name(modifiedName)
+                .build();
 
-        Assistant modifiedAssistantResponse = service.modifyAssistant(createAssistantResponse.getId(), createAssistantResponse);
+        Assistant modifiedAssistantResponse = service.modifyAssistant(createAssistantResponse.getId(), modifyRequest);
         assertNotNull(modifiedAssistantResponse);
         assertEquals(modifiedName, modifiedAssistantResponse.getName());
     }
@@ -123,20 +116,19 @@ private static File uploadAssistantFile() {
     }
 
     private static Assistant createAndValidateAssistant() {
-        AssistantBase assistantRequest = assistantStub();
+        AssistantRequest assistantRequest = assistantStub();
         Assistant createAssistantResponse = service.createAssistant(assistantRequest);
         validateAssistantResponse(createAssistantResponse);
 
         return createAssistantResponse;
     }
 
-
-    private static AssistantBase assistantStub() {
+    private static AssistantRequest assistantStub() {
         return AssistantRequest.builder()
                 .model(TikTokensUtil.ModelEnum.GPT_4_1106_preview.getName())
                 .name(MATH_TUTOR)
                 .instructions(ASSISTANT_INSTRUCTION)
-                .tools(Collections.singletonList(new Tool(AssistantToolsEnum.CODE_INTERPRETER)))
+                .tools(Collections.singletonList(new Tool(AssistantToolsEnum.CODE_INTERPRETER, null)))
                 .build();
     }
 
@@ -148,11 +140,4 @@ private static void validateAssistantResponse(Assistant assistantResponse) {
         assertEquals(assistantResponse.getTools().get(0).getType(),  AssistantToolsEnum.CODE_INTERPRETER);
         assertEquals(MATH_TUTOR, assistantResponse.getName());
     }
-
-    private static List<Assistant> validateListAssistants(ListAssistant<Assistant> assistants) {
-        assertNotNull(assistants);
-        List<Assistant> data = assistants.getData();
-        assertNotNull(data);
-        return data;
-    }
 }