package com.vincent.rsf.server.ai.service.impl.chat;
|
|
import com.fasterxml.jackson.databind.ObjectMapper;
|
import com.vincent.rsf.framework.exception.CoolException;
|
import com.vincent.rsf.server.ai.dto.AiChatDoneDto;
|
import com.vincent.rsf.server.ai.dto.AiChatStatusDto;
|
import com.vincent.rsf.server.ai.enums.AiErrorCategory;
|
import com.vincent.rsf.server.ai.exception.AiChatException;
|
import lombok.RequiredArgsConstructor;
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
|
import org.springframework.ai.chat.metadata.Usage;
|
import org.springframework.http.MediaType;
|
import org.springframework.stereotype.Component;
|
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
|
import java.io.IOException;
|
import java.time.Instant;
|
import java.util.LinkedHashMap;
|
import java.util.Map;
|
import java.util.concurrent.atomic.AtomicReference;
|
|
@Slf4j
|
@Component
|
@RequiredArgsConstructor
|
public class AiSseEventPublisher {
|
|
private final ObjectMapper objectMapper;
|
|
public void markFirstToken(AtomicReference<Long> firstTokenAtRef, SseEmitter emitter, String requestId,
|
Long sessionId, String model, long startedAt, AiThinkingTraceEmitter thinkingTraceEmitter) {
|
if (!firstTokenAtRef.compareAndSet(null, System.currentTimeMillis())) {
|
return;
|
}
|
if (thinkingTraceEmitter != null) {
|
thinkingTraceEmitter.startAnswer();
|
}
|
emitSafely(emitter, "status", AiChatStatusDto.builder()
|
.requestId(requestId)
|
.sessionId(sessionId)
|
.status("FIRST_TOKEN")
|
.model(model)
|
.timestamp(Instant.now().toEpochMilli())
|
.elapsedMs(System.currentTimeMillis() - startedAt)
|
.firstTokenLatencyMs(resolveFirstTokenLatency(startedAt, firstTokenAtRef.get()))
|
.build());
|
}
|
|
public void emitDone(SseEmitter emitter, String requestId, ChatResponseMetadata metadata, String fallbackModel,
|
Long sessionId, long startedAt, Long firstTokenAt) {
|
Usage usage = metadata == null ? null : metadata.getUsage();
|
emitStrict(emitter, "done", AiChatDoneDto.builder()
|
.requestId(requestId)
|
.sessionId(sessionId)
|
.model(metadata != null && metadata.getModel() != null && !metadata.getModel().isBlank() ? metadata.getModel() : fallbackModel)
|
.elapsedMs(System.currentTimeMillis() - startedAt)
|
.firstTokenLatencyMs(resolveFirstTokenLatency(startedAt, firstTokenAt))
|
.promptTokens(usage == null ? null : usage.getPromptTokens())
|
.completionTokens(usage == null ? null : usage.getCompletionTokens())
|
.totalTokens(usage == null ? null : usage.getTotalTokens())
|
.build());
|
}
|
|
public AiChatStatusDto buildTerminalStatus(String requestId, Long sessionId, String status, String model, long startedAt, Long firstTokenAt) {
|
return AiChatStatusDto.builder()
|
.requestId(requestId)
|
.sessionId(sessionId)
|
.status(status)
|
.model(model)
|
.timestamp(Instant.now().toEpochMilli())
|
.elapsedMs(System.currentTimeMillis() - startedAt)
|
.firstTokenLatencyMs(resolveFirstTokenLatency(startedAt, firstTokenAt))
|
.build();
|
}
|
|
public Long resolveFirstTokenLatency(long startedAt, Long firstTokenAt) {
|
return firstTokenAt == null ? null : Math.max(0L, firstTokenAt - startedAt);
|
}
|
|
public Map<String, String> buildMessagePayload(String... keyValues) {
|
Map<String, String> payload = new LinkedHashMap<>();
|
if (keyValues == null || keyValues.length == 0) {
|
return payload;
|
}
|
if (keyValues.length % 2 != 0) {
|
throw new CoolException("消息载荷参数必须成对出现");
|
}
|
for (int i = 0; i < keyValues.length; i += 2) {
|
payload.put(keyValues[i], keyValues[i + 1] == null ? "" : keyValues[i + 1]);
|
}
|
return payload;
|
}
|
|
public void emitStrict(SseEmitter emitter, String eventName, Object payload) {
|
try {
|
String data = objectMapper.writeValueAsString(payload);
|
emitter.send(SseEmitter.event()
|
.name(eventName)
|
.data(data, MediaType.APPLICATION_JSON));
|
} catch (IOException e) {
|
throw new AiChatException("AI_SSE_EMIT_ERROR", AiErrorCategory.STREAM, "SSE_EMIT", "SSE 输出失败: " + e.getMessage(), e);
|
}
|
}
|
|
public void emitSafely(SseEmitter emitter, String eventName, Object payload) {
|
try {
|
emitStrict(emitter, eventName, payload);
|
} catch (Exception e) {
|
log.warn("AI SSE event emit skipped, eventName={}, message={}", eventName, e.getMessage());
|
}
|
}
|
}
|
