Spring AI教程（二）Chat API之基于数据库的多Key轮询

基于数据库的多Key轮询

在之前的文章中我们所使用的Key都是一个，但事实上，官方对Key会有一定的请求限制，在实际业务场景下，我们也不可能通过一个Key来保证我们的系统稳定运行，因为一旦超过请求限制，就会出现无法请求AI的情况，这时，就需要考虑实现一种多Key轮询进行请求，从而保证我们的系统不会出现因为达到请求限制而无法运行的情况。

可惜的是，Spring AI目前还不支持多Key轮询的方式来调用大语言模型，因此需要我们自己实现。本篇将结合我对Spring AI的源码理解，来实现一个基于数据库的多Key轮询的接口。

6.1 核心类 OpenAiApi

实现多Key轮询的方式并不难，我们将API和Key的信息存储在数据库中，每次发起请求时，通过向数据库中查询API和Key的数据，手动构建一个ChatClient或StreamChatClient进行调用即可。拿OpenAI为例，就是要手动构建OpenAiChatClient对象，因为OpenAiChatClient实现了ChatClient和StreamChatClient接口。

通过观察OpenAiChatClient的构造方法，发现有一个核心类:OpenAiApi,该类的对象创建恰好就需要API和Key。

因此我们很容易想通创建OpenAiChatClient的流程：

从数据库中查询API和Key的信息；
利用API和Key构建OpenAiApi对象；
通过OpenAiApi对象再构建OpenAiChatClient；
再根据实际业务场景选择ChatClient和StreamChatClient进行调用；

6.2 环境准备

为便于演示，我新建了一个spring-ai-key-polling-demo的模块。

数据库使用的是Pgvector，它是Postgresql的扩展，可作为向量数据库，当然，本节还未涉及到向量数据库的使用，因此可以根据自己的喜好选择一个数据库；
ORM框架：Spring Data JPA，演示起来比较方便，可根据喜好选择一款ORM框架

核心依赖：

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-openai-spring-boot-starter</artifactId>
</dependency>
<dependency>
  <groupId>org.springframework.boot</groupId>
  <artifactId>spring-boot-starter-data-jpa</artifactId>
</dependency>
<dependency>
  <groupId>org.postgresql</groupId>
  <artifactId>postgresql</artifactId>
</dependency>
<dependency>
  <groupId>org.projectlombok</groupId>
  <artifactId>lombok</artifactId>
</dependency>

目录结构如下：

6.3 多Key轮询的基本实现

(1) 定义实体类

这里我定义了一个KeyInfo用来存储我们的Key和API信息,属性较为简单，实际项目中需要结合实际的业务场景添加一些额外的属性，如是否禁用、创建时间等等。

package com.ningning0111.entity;

import jakarta.annotation.Nullable;
import jakarta.persistence.*;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;

@Data
@AllArgsConstructor
@NoArgsConstructor
@Entity
@Table
public class KeyInfo {
    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;
    @Column(nullable = false)
    private String key;
    @Column(nullable = false,columnDefinition = "VARCHAR(50) DEFAULT 'https://api.openai.com'")
    private String api;

    // 描述这个Key干嘛的 可空
    private String description;

}

(2) 定义Repository

由于使用的是JPA，这里还需要创建Repository,对于Mybatis，需要创建Mapper。

package com.ningning0111.repository;

import com.ningning0111.entity.KeyInfo;
import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.stereotype.Repository;

@Repository
public interface KeyInfoRepository extends JpaRepository<KeyInfo,Long> {
}

(3) 配置文件

接着我们需要配置application.yml文件：

server:
  port: 8321

spring:
  datasource:
    driver-class-name: org.postgresql.Driver
    username: postgres
    password: postgres
    url: jdbc:postgresql://localhost/demo
  jpa:
    hibernate:
      # 自动创建表
      ddl-auto: create
    show-sql: true
    open-in-view: true
    
  autoconfigure:
    exclude: org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration

因为我们的Key和API都是从数据库中获取的，因此配置文件中就不需要配置了，为了防止相关属性因为无法自动配置而造成Spring启动不起来，就需要将OpenAi的自动装配排除。

(3) Service类

package com.ningning0111.service;

import com.ningning0111.entity.KeyInfo;
import com.ningning0111.repository.KeyInfoRepository;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.StreamingChatClient;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Random;

@RequiredArgsConstructor
@Service
public class ChatService {
    private final KeyInfoRepository repository;

    // 初始化一些key 这些key应该是可调用的
    @PostConstruct
    public void initData() {
        KeyInfo keyInfo = new KeyInfo();
        keyInfo.setKey("sk-W9kYxxxxxxxxxxxxxxxxxxxfAd460353Dc7a");
        keyInfo.setApi("https://api.mnzdna.xyz");
        keyInfo.setDescription("测试API和Key，请填写自己的Key");
        repository.save(keyInfo);
    }

    // 阻塞式
    public ChatClient getChatClient() {
        OpenAiApi openAiApi = randomGetApi();
        assert openAiApi != null;
        return new OpenAiChatClient(openAiApi);
    }

    // 流式
    public StreamingChatClient getStreamChatClient() {
        OpenAiApi openAiApi = randomGetApi();
        assert openAiApi != null;
        return new OpenAiChatClient(openAiApi);
    }

    // 随机获取一个OpenAiApi
    private OpenAiApi randomGetApi(){
        List<KeyInfo> keyInfoList = repository.findAll();
        // 如果数据库中没有KeyInfo对象，则返回null
        if (keyInfoList.isEmpty()) {
            return null;
        }
        // 随机选择一个KeyInfo对象
        Random random = new Random();
        KeyInfo randomKeyInfo = keyInfoList.get(random.nextInt(keyInfoList.size()));
        return new OpenAiApi(randomKeyInfo.getApi(),randomKeyInfo.getKey());
    }
}

(4) Controller类

package com.ningning0111.controller;

import com.ningning0111.service.ChatService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.ChatClient;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

@RestController
@RequiredArgsConstructor
public class ChatController {
    private final ChatService chatService;
    @GetMapping("/demo")
    public String chatDemo(String prompt){
        ChatClient chatClient = chatService.getChatClient();
        String response = chatClient.call(prompt);
        return response;
    }
}