利用 OkHttp 或者 WebClient 对接大模型的流式接口 - ToB企服应用市场:ToB评测及商务社交产业平台

<dependency>
<groupId>com.squareup.okhttp3</groupId>
<artifactId>okhttp</artifactId>
<version>4.12.0</version>在这里插入代码片
</dependency>

复制代码

package com.dolphin.bootstrap.agent.streamchat;
import com.alibaba.fastjson.JSONObject;
import com.dolphin.bootstrap.agent.BaseAgent;
import io.micrometer.common.util.StringUtils;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okio.BufferedSource;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.IOException;
import java.util.Collections;
import java.util.concurrent.TimeUnit;
/**
* 利用 OKHttp 库，实现 Stream 流式调用
*/
@Component
@Slf4j
public class OKHttpStream implements BaseAgent {
private static final Integer TIME_OUT = 5 * 60;
@Value("ai.agent.streamChatUrl")
private String streamChatUrl;
@Override
public void chatStream(SseEmitter sseEmitter, String question) {
// 1. 创建 OkHttpClient 对象，并设置超时时间
OkHttpClient client = new OkHttpClient.Builder()
.connectTimeout(TIME_OUT, TimeUnit.SECONDS)
.readTimeout(TIME_OUT, TimeUnit.SECONDS)
.writeTimeout(TIME_OUT, TimeUnit.SECONDS)
.build();
// 封装请求体参数
JSONObject params = new JSONObject();
params.put("question", question);
RequestBody requestBody = RequestBody.create(params.toJSONString(), MediaType.parse("application/json; charset=utf-8"));
// 封装请求头
Headers headers = new Headers.Builder()
.set("Content-Type","application/json")
.set("Accept","text/event-stream")
.build();
// 2. 构建 Request 对象
Request request = new Request.Builder()
.url(streamChatUrl)
.headers(headers)
.post(requestBody)
.build();
// 3. 创建 Call 对象
Call call = client.newCall(request);
// 4. 监听回调
call.enqueue(new Callback() {
@Override
public void onFailure(@NotNull Call call, @NotNull IOException e) {
log.error("进入 onFailure方法 {}", e.getMessage(), e);
sseEmitter.completeWithError(e);
}
@Override
public void onResponse(@NotNull Call call, @NotNull Response response) throws IOException {
if (response.isSuccessful()) {
String chunkMessage = "";
try (ResponseBody responseBody = response.body();) {
BufferedSource source = responseBody.source();
while (!source.exhausted()) {
chunkMessage = source.readUtf8Line();
if (StringUtils.isBlank(chunkMessage)) {
continue;
}
JSONObject jsonObject = JSONObject.parseObject(chunkMessage);
if (null != jsonObject && null != jsonObject.getJSONObject("data")) {
String answer = jsonObject.getJSONObject("data").getString("answer");
sseEmitter.send(answer);
}
}
} catch (Exception e) {
log.error("解析失败： {}", e.getMessage(), e);
sseEmitter.completeWithError(e);
}
} else {
log.error("onResponse 方法请求失败 {}", response.message());
// TODO 重新发起请求
}
}
});
BaseAgent.super.chatStream(sseEmitter, question);
}
}

复制代码

<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
<version>6.1.4</version>
</dependency>

复制代码

package com.dolphin.bootstrap.agent.streamchat;
import com.alibaba.fastjson.JSONObject;
import com.dolphin.bootstrap.agent.BaseAgent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;
import java.io.IOException;
import java.util.Collections;
/**
* 利用 WebClient 类，实现 Stream 流式调用
*/
public class WebClientStream implements BaseAgent {
private static final Logger log = LoggerFactory.getLogger(WebClientStream.class);
@Value("ai.agent.streamChatUrl")
private String streamChatUrl;
@Override
public void chatStream(SseEmitter sseEmitter, String question) {
// 创建 WebClient 客户端
WebClient webClient = WebClient.builder().baseUrl(streamChatUrl).build();
// 封装参数
JSONObject params = new JSONObject();
params.put("question", question);
// 封装请求头
HttpHeaders headers = new HttpHeaders();
headers.setContentType(MediaType.APPLICATION_JSON_UTF8);
// accept 一定要设置为 TEXT_EVENT_STREAM
headers.setAccept(Collections.singletonList(MediaType.TEXT_EVENT_STREAM));
Flux<String> eventStream = webClient
.post()
.uri("/stream")
.accept(MediaType.valueOf("text/event-stream;charset=UTF-8")) // 一定要设置
.headers(httpHeaders -> httpHeaders.addAll(headers))
.bodyValue(params.toJSONString())
.retrieve()
.bodyToFlux(String.class);
eventStream.subscribe(
data -> { // data 有什么key，value。具体看你对接agent的文档
try {
JSONObject bodyJson = JSONObject.parseObject(data);
if (bodyJson.getIntValue("code") == HttpStatus.OK.value()) {
// 事件类型
String event = bodyJson.getString("event");
// agent回答
String answer = bodyJson.getJSONObject("data").getString("answer");
if ("message".equals(event)) {
// 中间一段一段的消息
sseEmitter.send(answer);
} else if ("end".equals(event)) {
// 最后会全部返回
sseEmitter.send(answer);
}
}
} catch (IOException e) {
throw new RuntimeException(e);
}
},
error -> {
sseEmitter.completeWithError(error);
log.error("报错了:{}", error.getMessage(), error);
},
sseEmitter::complete
);
}
}

复制代码

webClient.get()
.uri("/")
.retrieve()
.onStatus(HttpStatus::is4xxClientError, response -> Mono.error(new RuntimeException("客户端错误")))
.onStatus(HttpStatus::is5xxServerError, response -> Mono.error(new RuntimeException("服务器错误")))
.bodyToFlux(String.class)
.doOnNext(data -> System.out.println("接收到数据块：" + data))
.subscribe();

复制代码