001/*
002 *  Copyright (c) 2023-2026, Agents-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.agentsflex.core.model.client;
017
018import com.agentsflex.core.message.AiMessage;
019import com.agentsflex.core.model.chat.ChatContext;
020import com.agentsflex.core.model.chat.ChatModel;
021import com.agentsflex.core.model.chat.StreamResponseListener;
022import com.agentsflex.core.model.chat.response.AiMessageResponse;
023import com.agentsflex.core.parser.AiMessageParser;
024import com.agentsflex.core.util.StringUtil;
025import com.alibaba.fastjson2.JSON;
026import com.alibaba.fastjson2.JSONObject;
027
028import java.util.concurrent.atomic.AtomicBoolean;
029
030public class BaseStreamClientListener implements StreamClientListener {
031
032    private final StreamResponseListener streamResponseListener;
033    private final ChatContext chatContext;
034    private final AiMessageParser<JSONObject> messageParser;
035    private final StreamContext context;
036    private final AiMessage fullMessage = new AiMessage();
037    private final AtomicBoolean finishedFlag = new AtomicBoolean(false);
038    private final AtomicBoolean stoppedFlag = new AtomicBoolean(false);
039    private final AtomicBoolean isFailure = new AtomicBoolean(false);
040
041    public BaseStreamClientListener(
042        ChatModel chatModel,
043        ChatContext chatContext,
044        StreamClient client,
045        StreamResponseListener streamResponseListener,
046        AiMessageParser<JSONObject> messageParser) {
047        this.streamResponseListener = streamResponseListener;
048        this.chatContext = chatContext;
049        this.messageParser = messageParser;
050        this.context = new StreamContext(chatModel, chatContext, client);
051    }
052
053    @Override
054    public void onStart(StreamClient client) {
055        streamResponseListener.onStart(context);
056    }
057
058    @Override
059    public void onMessage(StreamClient client, String response) {
060        if (StringUtil.noText(response) || "[DONE]".equalsIgnoreCase(response.trim()) || finishedFlag.get()) {
061            notifyLastMessageAndStop(response);
062            return;
063        }
064
065        try {
066            JSONObject jsonObject = JSON.parseObject(response);
067            AiMessage delta = messageParser.parse(jsonObject, chatContext);
068
069            //合并 增量 delta 到 fullMessage
070            fullMessage.merge(delta);
071
072            // 设置 delta 全内容
073            delta.setFullContent(fullMessage.getContent());
074            delta.setFullReasoningContent(fullMessage.getReasoningContent());
075
076            //输出内容
077            AiMessageResponse resp = new AiMessageResponse(chatContext, response, delta);
078            streamResponseListener.onMessage(context, resp);
079        } catch (Exception err) {
080            onFailure(this.context.getClient(), err);
081            onStop(this.context.getClient());
082        }
083    }
084
085    private void notifyLastMessage(String response) {
086        if (finishedFlag.compareAndSet(false, true)) {
087            finalizeFullMessage();
088            AiMessageResponse resp = new AiMessageResponse(chatContext, response, fullMessage);
089            streamResponseListener.onMessage(context, resp);
090        }
091    }
092
093    private void notifyLastMessageAndStop(String response) {
094        try {
095            notifyLastMessage(response);
096        } finally {
097            if (stoppedFlag.compareAndSet(false, true)) {
098                streamResponseListener.onStop(context);
099            }
100        }
101    }
102
103
104    @Override
105    public void onStop(StreamClient client) {
106        try {
107            if (!isFailure.get()) {
108                // onStop 在 sse 的 onClosed 中会被调用，可以用于在 onMessage 出现异常时进行兜底
109                notifyLastMessage(null);
110            }
111        } finally {
112            if (stoppedFlag.compareAndSet(false, true)) {
113                streamResponseListener.onStop(context);
114            }
115        }
116    }
117
118    private void finalizeFullMessage() {
119        String currentContent = fullMessage.getContent();
120        String currentReasoningContent = fullMessage.getReasoningContent();
121
122        fullMessage.setFullContent(currentContent);
123        fullMessage.setContent(null);
124
125        fullMessage.setFullReasoningContent(currentReasoningContent);
126        fullMessage.setReasoningContent(null);
127
128        fullMessage.setFinished(true);
129
130        context.setFullMessage(fullMessage);
131    }
132
133
134    @Override
135    public void onFailure(StreamClient client, Throwable throwable) {
136        if (isFailure.compareAndSet(false, true)) {
137            context.setThrowable(throwable);
138            streamResponseListener.onFailure(context, throwable);
139        }
140    }
141
142}