使用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
{@Beanpublic 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];/*** 连接建立成功调用的方法*/@OnOpenpublic void onOpen(Session session) throws Exception{SemaphoreUtils.tryAcquire(socketSemaphore);LOGGER.info("\n 成功连接 - {}", session);}/*** 连接关闭时处理*/@OnClosepublic void onClose(Session session){LOGGER.info("\n 关闭连接 - {}", session);// 获取到信号量则需释放SemaphoreUtils.release(socketSemaphore);}/*** 抛出异常时处理*/@OnErrorpublic 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);}/*** 服务器接收到客户端消息时调用的方法*/@OnMessagepublic 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. 发出请求,调用大模型APIRequestBody 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;}

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

 

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

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