feat: 更新文件搜索工具类及测试用例 #61

Lambdua · Sep 19, 2024 · b682ba7 · b682ba7
1 parent 8c4affc
commit b682ba7
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 10 deletions.
diff --git a/api/src/main/java/com/theokanning/openai/assistants/assistant/FileSearch.java b/api/src/main/java/com/theokanning/openai/assistants/assistant/FileSearch.java
@@ -0,0 +1,30 @@
+package com.theokanning.openai.assistants.assistant;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
+/**
+ * @author LiangTao
+ * @date 2024年09月19 10:53
+ **/
+@Data
+@NoArgsConstructor
+@AllArgsConstructor
+@Builder
+public class FileSearch {
+    /**
+     * The maximum number of results the file search tool should output. The default is 20 for gpt-4* models and 5 for gpt-3.5-turbo. This number should be between 1 and 50 inclusive.
+     */
+    @JsonProperty("max_num_results")
+    Integer maxNumResults;
+
+    /**
+     * The ranking options for the file search. If not specified, the file search tool will use the auto ranker and a score_threshold of 0.
+     */
+    @JsonProperty("ranking_options")
+    FileSearchRankingOptions rankingOptions;
+
+}
diff --git a/api/src/main/java/com/theokanning/openai/assistants/assistant/FileSearchRankingOptions.java b/api/src/main/java/com/theokanning/openai/assistants/assistant/FileSearchRankingOptions.java
@@ -0,0 +1,30 @@
+package com.theokanning.openai.assistants.assistant;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import lombok.AllArgsConstructor;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
+import javax.validation.constraints.NotNull;
+
+/**
+ * @author LiangTao
+ * @date 2024年09月19 10:36
+ **/
+@Data
+@NoArgsConstructor
+@AllArgsConstructor
+public class FileSearchRankingOptions {
+    /**
+     * The ranker to use for the file search. If not specified will use the auto ranker.
+     */
+    String ranker;
+
+    /**
+     * The score threshold for the file search. All values must be a floating point number between 0 and 1.
+     */
+    @NotNull
+    @JsonProperty("score_threshold")
+    Double scoreThreshold;
+
+}
diff --git a/api/src/main/java/com/theokanning/openai/assistants/assistant/FileSearchTool.java b/api/src/main/java/com/theokanning/openai/assistants/assistant/FileSearchTool.java
@@ -16,9 +16,9 @@
 public class FileSearchTool implements Tool {
     final String type = "file_search";
 
-    /**
-     * The maximum number of results the file search tool should output. The default is 20 for gpt-4* models and 5 for gpt-3.5-turbo. This number should be between 1 and 50 inclusive.
-     */
-    @JsonProperty("max_num_results")
-    Integer maxNumResults;
+
+    @JsonProperty("file_search")
+    FileSearch fileSearch;
+
+
 }
diff --git a/example/src/main/java/example/AssistantExample.java b/example/src/main/java/example/AssistantExample.java
@@ -192,7 +192,7 @@ static void fileSearchExample() throws UnsupportedEncodingException {
                 .name("file search assistant")
                 .instructions("你是一个中国传统音乐教授,负责根据用户的需求解答问题")
                 //add file search tool to assistant
-                .tools(Collections.singletonList(new FileSearchTool(1)))
+                .tools(Collections.singletonList(new FileSearchTool()))
                 .temperature(0D)
                 .build();
         Assistant assistant = service.createAssistant(assistantRequest);

diff --git a/service/src/test/java/com/theokanning/openai/service/assistants/AssistantFileSearchTest.java b/service/src/test/java/com/theokanning/openai/service/assistants/AssistantFileSearchTest.java
@@ -1,9 +1,7 @@
 package com.theokanning.openai.service.assistants;
 
 import com.theokanning.openai.ListSearchParameters;
-import com.theokanning.openai.assistants.assistant.Assistant;
-import com.theokanning.openai.assistants.assistant.AssistantRequest;
-import com.theokanning.openai.assistants.assistant.FileSearchTool;
+import com.theokanning.openai.assistants.assistant.*;
 import com.theokanning.openai.assistants.message.MessageListSearchParameters;
 import com.theokanning.openai.assistants.message.MessageRequest;
 import com.theokanning.openai.assistants.run.Run;
@@ -49,7 +47,7 @@ void fileSearchExample(){
                 .name("file search assistant")
                 .instructions("你是一个中国传统音乐教授,负责根据用户的需求解答问题")
                 //add file search tool to assistant
-                .tools(Collections.singletonList(new FileSearchTool()))
+                .tools(Collections.singletonList(new FileSearchTool(new FileSearch(1,new FileSearchRankingOptions("auto", 0.5D)))))
                 .temperature(0.3D)
                 .build();
         Assistant assistant = service.createAssistant(assistantRequest);