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();
    }
}


測試

webSocket 測試網站:http://www.websocket-test.com/


補充

OpenAi接口文檔 : https://platform.openai.com/docs/api-reference/chat

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章