使用websocket后端接入文心一言

最近再写项目练手，想着最近大模型那么火，也想接入项目来玩一玩，于是去了解了一下相关的api和通信协议，最后选择了文心一言进行集成，国内的相对稳定。ERNIE-Bot-turbo - 千帆大模型平台 | 百度智能云文档 (baidu.com)

使用websocket进行双向通信，因为http是单向协议显然不适合这样的场景。简单介绍一下websocket它是一个双向的通信协议，一旦通信双方建立联系，就可以互相发送消息。

http和websocket通信过程图

现在就开始进行配置吧，首先我们需要配置好websocket的的相关依赖已经配置serverEndpoint的扫描。

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
        </dependency>

@Configuration
public class WebSocketConfig
{
    @Bean
    public ServerEndpointExporter serverEndpointExporter()
    {
        return new ServerEndpointExporter();
    }
}

如此就将准备工作做好了，现在配置一下访问处理的server

@Component
@ServerEndpoint("/websocket/message")
public class WebSocketServer
{
    /**
     * 实现日志打印
     */
    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketServer.class);
 
    public static int socketMaxOnlineCount = 100;
 
    private static Semaphore socketSemaphore = new Semaphore(socketMaxOnlineCount);
 
    private List<BaiduChatMessage> userMessage[]=new List[2000];
 
    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session) throws Exception
    {
        SemaphoreUtils.tryAcquire(socketSemaphore);
        LOGGER.info("\n 成功连接 - {}", session);
    }
 
    /**
     * 连接关闭时处理
     */
    @OnClose
    public void onClose(Session session)
    {
        LOGGER.info("\n 关闭连接 - {}", session);
        // 获取到信号量则需释放
        SemaphoreUtils.release(socketSemaphore);
    }
 
    /**
     * 抛出异常时处理
     */
    @OnError
    public void onError(Session session, Throwable exception) throws Exception
    {
        if (session.isOpen())
        {
            // 关闭连接
            session.close();
        }
        String sessionId = session.getId();
        LOGGER.info("\n 连接异常 - {}", sessionId);
        LOGGER.info("\n 异常信息 - {}", exception);
        // 获取到信号量则需释放
        SemaphoreUtils.release(socketSemaphore);
    }
 
    /**
     * 服务器接收到客户端消息时调用的方法
     */
    @OnMessage
    public void onMessage(String message, Session session) throws IOException {
        // 首先，接收到一条消息
        LOGGER.info("\n 收到消息 - {}", message);
        // 1. 调用大模型API，把上下文和这次问题传入，得到回复
        BigModelService bigModelService = new BigModelService();
        boolean f = bigModelService.callModelAPI(session,message,userMessage);
        if (!f) {
            session.getBasicRemote().sendText("抱歉，似乎出了点问题，请联系管理员");
            return;
        }
    }
}

都使用注解进行控制，其中这里最重要的就是 OnMessage，这个是得到websocket消息进行的处理，是最核心的方法，本次主要的业务实现焦距于本方法。

接下来就是将问题发送给文心一言的接口然后得到回应了

下面粘贴一下主代码

public boolean callModelAPI(Session session, String message, Map<String,List<BaiduChatMessage>> userMassage) {
        String records = "{}";
        ErnieBotTurboParam param = JSONObject.parseObject(records, ErnieBotTurboParam.class);
        //未将信息存入数据库，使用会话做一个简单的持久化
        if(userMassage.get(session.getId()) == null)
            userMassage.put(session.getId(),new ArrayList<>());
        //存入最新消息
        userMassage.get(session.getId()).add(BaiduChatMessage.builder().role("user").content(message).build());
        // 1.2.2 把messages重新设置到param中
        param.setMessages(userMassage.get(session.getId()));
        try {
            // 2. 发出请求，调用大模型API
            RequestBody body = RequestBody.create(MediaType.parse("application/json"), JSONObject.toJSONString(param));
            Request request = new Request.Builder()
                    .url("https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" + getAccessToken())
                    .method("POST", body)
                    .addHeader("Content-Type", "application/json")
                    .build();
            Response response = HTTP_CLIENT.newCall(request).execute();
 
            if (response.isSuccessful()) {
                //对返回信息进行处理
                try (ResponseBody responseBody = response.body()) {
                    if (responseBody != null) {
                        InputStream inputStream = responseBody.byteStream();
                        // 以流的方式处理响应内容
                        byte[] buffer = new byte[2048];
                        int bytesRead;
                        String ss="";
                        while ((bytesRead = inputStream.read(buffer)) != -1) {
                            // 在控制台输出每个数据块
                            String s[] = new String(buffer, 0, bytesRead).split("data: ");
                            int n=1;
                            while (n<s.length) {
                                System.out.println(s[n]);
                                TurboResponse turboResponse = JSONObject.parseObject(s[n++], TurboResponse.class);
                                ss+=turboResponse.getResult();
                                session.getBasicRemote().sendText(turboResponse.getResult());
                            }
                        }
                        userMassage.get(session.getId()).add(BaiduChatMessage.builder().role("assistant").content(ss).build());
                    }
                }
                return true;
            } else {
                LOGGER.error("调用大模型API失败: {}", response.message());
            }
        } catch (IOException e) {
            LOGGER.error("调用大模型API发生异常:", e);
        }
        return false;
    }

如此便实现了基本的通信服务，下面来看看效果

效果只能说够用 ，也还不错吧。

期待这篇文章可以帮助到大家