Springboot集成WebSocket实现智能聊天【Demo】

背景

openai 目前越来越流行，其他 ai 产业也随之而来，偶然翻到 openai接口文档，就想着可以调用接口实现智能聊天，接下来就写写我怎么接入 websocket 的过程，文笔不佳，谅解。

接入WebSocket

1. webSocket依赖

<dependency>
  <groupId>org.springframework.boot</groupId>
  <artifactId>spring-boot-starter-websocket</artifactId>
  <version>${spring.boot.version}</version>  <!--填写springboot 版本 -->
</dependency>

2. 配置文件

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

@Configuration
public class WebSocketConfig {

    /**
     * 注入ServerEndpointExporter，该Bean会自动注册使用@ServerEndpoint注解申明的websocket endpoint
     */
     @Bean
     public ServerEndpointExporter serverEndpointExporter() {
         return new ServerEndpointExporter();
     }

}

3. 编写Server


import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.handsometaoa.utils.HttpClientUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
@ServerEndpoint(value = "/aiChatServer/{username}/{password}")
public class WebSocketServer {

    @Value("${ai.chat.url}")
    private String aiChatUrl;
    @Value("${ai.chat.token}")
    private String aiChatToken;


    /**
     * 记录当前在线连接信息
     */
    public static final Map<String, Session> sessionMap = new ConcurrentHashMap<String, Session>();

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("username") String username,@PathParam("password") String password) {
        // 1. todo: 校验密码（由于AI接口是计费的，所以需要校验密码）
        // 2. 将用户session放入sessionMap中
        sessionMap.put(username, session);
    }
    
    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public void onClose(Session session, @PathParam("username") String username) {
        sessionMap.remove(username);
    }

    /**
     * 收到客户端消息后调用的方法
     * 用来处理接收到消息的逻辑
     */
    @OnMessage
    public void onMessage(String message, @PathParam("username") String username) {
        log.info("服务端收到用户username={}的消息:{}", username, message);
        // 调用AI接口获取回复
        HashMap<String, String> headers = new HashMap<>();
        headers.put("Authorization", "Bearer " + aiChatToken);
        String requestBody = buildAiRequestBody(message);
        String post = HttpClientUtil.post(aiChatUrl, headers, requestBody);
        JSONObject response = JSONObject.parseObject(post);
        String aiResponseText = response.getJSONArray("choices").getJSONObject(0).getJSONObject("message").getString("content");
        // 由于是机器人对话聊天，消息发给自己
        sessionMap.get(username).getAsyncRemote().sendText(aiResponseText);
    }

    /**
     * 发生错误时调用
     */
    @OnError
    public void onError(Session session, Throwable error) {
        log.error("发生错误:{}",error.getMessage(),error);
    }

    /**
     * 服务端发送消息给客户端
     */
    private void sendMessage(String message, Session toSession) {
        try {
            log.info("服务端给客户端[{}]发送消息{}", toSession.getId(), message);
            toSession.getBasicRemote().sendText(message);
        } catch (Exception e) {
            log.error("服务端发送消息给客户端失败", e);
        }
    }

    /**
     * 服务端发送消息给所有客户端
     */
    private void sendAllMessage(String message) {
        try {
            for (Session session : sessionMap.values()) {
                log.info("服务端给客户端[{}]发送消息{}", session.getId(), message);
                session.getBasicRemote().sendText(message);
            }
        } catch (Exception e) {
            log.error("服务端发送消息给客户端失败", e);
        }
    }


    /**
     * 封装构建OpenAi请求体
     */
    private static String buildAiRequestBody(String message) {
        JSONObject body = new JSONObject();
        body.put("model","gpt-3.5-turbo");
        JSONArray jsonArray = new JSONArray();
        JSONObject messages = new JSONObject();
        messages.put("role","user");
        messages.put("content", message);
        jsonArray.add(messages);
        body.put("messages",jsonArray);
        return body.toString();
    }
}