#

Junjie

3 天以前 338f3b81425ab96d8c856909a775124af5365e3c

 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);
                        }
                    }
                });
    }