智慧工作流：Spring AI高效批次化提示訪問方案

linkcxt發表於2024-05-11
SpringAI

基於SpringAI搭建系統，依靠執行緒池\負載均衡等技術進行請求最佳化，用於解決科研&開發過程中對GPT介面進行批次化介面請求中出現的問題。

github地址：https://github.com/linkcao/springai-wave

大語言模型介面以OpenAI的GPT 3.5為例，JDK版本為17，其他依賴版本可見倉庫pom.xml

擬解決的問題

在處理大量提示文字時，存在以下挑戰：

  1. API金鑰請求限制： 大部分AI服務提供商對API金鑰的請求次數有限制，單個金鑰每分鐘只能傳送有限數量的請求。
  2. 處理速度慢： 大量的提示文字需要逐條傳送請求，處理速度較慢，影響效率。
  3. 結果儲存和分析困難： 處理完成的結果需要儲存到本地資料庫中，並進行後續的資料分析，但這一過程相對複雜。

解決方案

為了解決上述問題，本文提出了一種基於Spring框架的批次化提示訪問方案，如下圖所示:

image-20240511160521257

其中具體包括以下步驟:

  1. 多執行緒處理提示文字： 將每個提示文字看作一個獨立的任務，採用執行緒池的方式進行多執行緒處理，提高處理效率。
  2. 動態分配API金鑰： 線上程池初始化時，透過讀取本地資料庫中儲存的API金鑰資訊，動態分配每個執行緒單元所攜帶的金鑰，實現負載均衡。
  3. 結果儲存和管理： 在請求完成後，將每個請求的問題和回答儲存到本地資料庫中，以便後續的資料分析和管理。
  4. 狀態實時更新： 將整個批次請求任務區分為進行中、失敗和完成狀態，並透過資料庫儲存狀態碼實時更新任務狀態，方便監控和管理。

關鍵程式碼示例

  1. 多執行緒非同步請求提示資訊(所在包: ChatService)
    // 執行緒池初始化
	private static final ExecutorService executor = Executors.newFixedThreadPool(10);
    /**
     * 多執行緒請求提示
     * @param prompts
     * @param user
     * @param task
     * @return
     */
    @Async
    public CompletableFuture<Void> processPrompts(List<String> prompts， Users user， Task task) {
        for (int i = 0; i < prompts.size();i++) {
            int finalI = i;
            // 提交任務
            executor.submit(() -> processPrompt(prompts.get(finalI)， user， finalI));
        }
        // 設定批次任務狀態
        task.setStatus(TaskStatus.COMPLETED);
        taskService.setTask(task);
        return CompletableFuture.completedFuture(null);
    }

  • 如上所示，利用了Spring框架的@Async註解和執行緒池的功能，實現了多執行緒非同步處理提示資訊。

  • 首先，使用了ExecutorService建立了一個固定大小的執行緒池，以便同時處理多個提示文字。

  • 然後，透過CompletableFuture來實現非同步任務的管理。

  • 在處理每個提示文字時，透過executor.submit()方法提交一個任務給執行緒池，讓執行緒池來處理。

  • 處理完成後，將批次任務的狀態設定為已完成，並更新任務狀態。

  • 一個執行緒任務需要繫結請求的使用者以及所在的批次任務，當前任務所分配的key由任務所在佇列的下標決定。

  1. 處理單條提示資訊(所在包: ChatService)
    /**
     * 處理單條提示文字
     * @param prompt 提示文字
     * @param user 使用者
     * @param index 所在佇列下標
     */
    public void processPrompt(String prompt， Users user， int index) {
        // 獲取Api Key
        OpenAiApi openAiApi = getApiByIndex(user， index);
        assert openAiApi != null;
        ChatClient client = new OpenAiChatClient(openAiApi);
        // 提示文字請求
        String response = client.call(prompt);
        // 日誌記錄
        log.info("提示資訊" + prompt );
        log.info("輸出" + response );
        // 回答儲存資料庫
        saveQuestionAndAnswer(user， prompt， response);
    }
  • 首先根據任務佇列的下標獲取對應的API金鑰
  • 然後利用該金鑰建立一個與AI服務進行通訊的客戶端。
  • 接著，使用客戶端傳送提示文字請求，並獲取AI模型的回答。
  • 最後，將問題和回答儲存到本地資料庫和日誌中，以便後續的資料分析和管理。
  1. Api Key 負載均衡(所在包: ChatService)
    /**
     * 採用任務下標分配key的方式進行負載均衡
     * @param index 任務下標
     * @return OpenAiApi
     */
    private OpenAiApi getApiByIndex(int index){
        List<KeyInfo> keyInfoList = keyRepository.findAll();
        if (keyInfoList.isEmpty()) {
            return null;
        }
        // 根據任務佇列下標分配 Key
        KeyInfo keyInfo = keyInfoList.get(index % keyInfoList.size());
        return new OpenAiApi(keyInfo.getApi()，keyInfo.getKeyValue());
    }
  • 首先從本地資料庫中獲取所有可用的API金鑰資訊
  • 然後根據任務佇列的下標來動態分配API金鑰。
  • 確保每個執行緒單元都攜帶了不同的API金鑰，避免了因為某個金鑰請求次數達到限制而導致的請求失敗問題。
  1. 依靠執行緒池批次請求GPT整體方法(所在包: ChatController)
/**
     * 依靠執行緒池批次請求GPT
     * @param promptFile 傳入的批次提示檔案，每一行為一個提示語句
     * @param username 呼叫的使用者
     * @return 處理狀態
     */
    @PostMapping("/batch")
    public String batchPrompt(MultipartFile promptFile， String username){
        if (promptFile.isEmpty()) {
            return "上傳的檔案為空";
        }
        // 批次請求任務
        Task task = new Task();
        try {
            BufferedReader reader = new BufferedReader(new InputStreamReader(promptFile.getInputStream()));
            List<String> prompts = new ArrayList<>();
            String line;
            while ((line = reader.readLine()) != null) {
                prompts.add(line);
            }
            // 使用者資訊請求
            Users user = userService.findByUsername(username);
            // 任務狀態設定
            task.setFileName(promptFile.getName());
            task.setStartTime(LocalDateTime.now());
            task.setUserId(user.getUserId());
            task.setStatus(TaskStatus.PROCESSING);
            // 執行緒池處理
            chatService.processPrompts(prompts， user， task);
            return "檔案上傳成功，已開始批次處理提示";
        } catch ( IOException e) {
            // 處理失敗
            e.printStackTrace();
            task.setStatus(TaskStatus.FAILED);
            return "上傳檔案時出錯：" + e.getMessage();
        } finally {
            // 任務狀態儲存
            taskService.setTask(task);
        }
    }
  • 首先，接收使用者上傳的批次提示檔案和使用者名稱資訊。
  • 然後，讀取檔案中的每一行提示文字，並將它們儲存在一個列表中。
  • 接著，根據使用者名稱資訊找到對應的使用者，並建立一個任務物件來跟蹤批次處理的狀態。
  • 最後，呼叫ChatService中的processPrompts()方法來處理提示文字，並返回處理狀態給使用者。

資料庫ER圖

所有資訊都與使用者ID強繫結，便於管理和查詢，ER圖如下所示:

image-20240511165330676

演示示例

  1. 透過postman攜帶批次請求檔案和username資訊進行Post請求訪問localhost:8080/batch介面：

image-20240511165636797

  1. 在實際應用中，可以根據具體需求對提示文字進行定製和擴充套件，以滿足不同場景下的需求，演示所攜帶的請求檔案內容如下:
請回答1+2=？
請回答8*12=？
請回答12*9=？
請回答321-12=？
請回答12/4=？
請回答32%2=？
  1. 最終返回的資料庫結果，左為問題庫，右為回答庫：

image-20240511165910247

  • 問題庫和答案庫透過question_id和user_id進行繫結，由於一個問題可以讓GPT回答多次，因此兩者的關係為多對一，將問題和答案分在兩個獨立的表中也便於後續的垂域定製和擴充套件。

相關文章