From 1b8a4677f362d234d834120deac4880d7ae89a50 Mon Sep 17 00:00:00 2001
From: Junjie <fallin.jie@qq.com>
Date: 星期四, 12 三月 2026 17:03:59 +0800
Subject: [PATCH] #
---
src/main/java/com/zy/ai/service/LlmSpringAiClientService.java | 35 +++++++++++++++++++++++++----------
1 files changed, 25 insertions(+), 10 deletions(-)
diff --git a/src/main/java/com/zy/ai/service/LlmSpringAiClientService.java b/src/main/java/com/zy/ai/service/LlmSpringAiClientService.java
index 746aa20..0079565 100644
--- a/src/main/java/com/zy/ai/service/LlmSpringAiClientService.java
+++ b/src/main/java/com/zy/ai/service/LlmSpringAiClientService.java
@@ -22,6 +22,7 @@
import java.util.List;
import java.util.Locale;
import java.util.Map;
+import java.util.function.Consumer;
@Service
public class LlmSpringAiClientService {
@@ -43,19 +44,33 @@
legacy);
}
- public Flux<String> streamCompletion(String baseUrl, String apiKey, ChatCompletionRequest req) {
+ public Flux<String> streamCompletion(String baseUrl, String apiKey, ChatCompletionRequest req, Consumer<ChatCompletionResponse.Usage> usageConsumer) {
OpenAiApi api = buildOpenAiApi(baseUrl, apiKey);
OpenAiApi.ChatCompletionRequest springReq = buildSpringAiRequest(req, true);
return api.chatCompletionStream(springReq)
- .flatMapIterable(chunk -> chunk == null || chunk.choices() == null
- ? List.<OpenAiApi.ChatCompletionChunk.ChunkChoice>of()
- : chunk.choices())
- .map(OpenAiApi.ChatCompletionChunk.ChunkChoice::delta)
- .filter(delta -> delta != null)
- .handle((delta, sink) -> {
- String text = extractSpringAiContent(delta);
- if (text != null && !text.isEmpty()) {
- sink.next(text);
+ .handle((chunk, sink) -> {
+ if (chunk == null) {
+ return;
+ }
+ if (chunk.usage() != null && usageConsumer != null) {
+ ChatCompletionResponse.Usage usage = new ChatCompletionResponse.Usage();
+ usage.setPromptTokens(chunk.usage().promptTokens());
+ usage.setCompletionTokens(chunk.usage().completionTokens());
+ usage.setTotalTokens(chunk.usage().totalTokens());
+ usageConsumer.accept(usage);
+ }
+ List<OpenAiApi.ChatCompletionChunk.ChunkChoice> choices = chunk.choices();
+ if (choices == null || choices.isEmpty()) {
+ return;
+ }
+ for (OpenAiApi.ChatCompletionChunk.ChunkChoice choice : choices) {
+ if (choice == null || choice.delta() == null) {
+ continue;
+ }
+ String text = extractSpringAiContent(choice.delta());
+ if (text != null && !text.isEmpty()) {
+ sink.next(text);
+ }
}
});
}
--
Gitblit v1.9.1