使用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; }

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

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

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