package com.webchat.aigc.llm;

import com.webchat.aigc.config.properties.AliBaiLianLLMPropertiesConfig;
import com.webchat.aigc.config.properties.DeepseekLLMPropertiesConfig;
import com.webchat.common.util.llm.AliBaiLianAIClient;
import com.webchat.common.util.llm.DeepSeekAIClient;
import com.webchat.domain.vo.llm.ChatCompletionMessage;
import com.webchat.domain.vo.llm.ChatCompletionRequest;
import com.webchat.domain.vo.llm.ChatCompletionResponse;
import com.webchat.domain.vo.llm.ChatCompletionStreamChoice;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.List;

/**
 * @Author 程序员七七
 * @webSite https://www.coderutil.com
 * @Date 2024/10/29 23:15
 * @description
 */
@Service
public class AliBaiLianAIService extends AbstractLLMChatService {

    @Autowired
    private AliBaiLianLLMPropertiesConfig aliBaiLianLLMPropertiesConfig;

    /**
     * 同步对话
     *
     * @param messageList
     * @return
     * @throws Exception
     */
    public ChatCompletionResponse chat(List<ChatCompletionMessage> messageList) throws Exception {
        DeepSeekAIClient client = new DeepSeekAIClient(aliBaiLianLLMPropertiesConfig.getApiKey());
        final List<ChatCompletionMessage> messages = messageList;
        return client.ChatCompletion(new ChatCompletionRequest(
                aliBaiLianLLMPropertiesConfig.getModel(),
                messages,
                2000,
                0.3f,
                1
        ));
    }

    /**
     * 流式对话
     *
     * @param emitter
     * @param messageList
     * @throws Exception
     */
    public String chat(SseEmitter emitter, List<ChatCompletionMessage> messageList) throws Exception {

        // 用于记录流式推理完整结果，返回用于，用于对话消息持久化
        StringBuilder aiMessage = new StringBuilder();

        AliBaiLianAIClient client = new AliBaiLianAIClient(aliBaiLianLLMPropertiesConfig.getApiKey());
        final List<ChatCompletionMessage> messages = messageList;
        try {
            client.ChatCompletionStream(new ChatCompletionRequest(
                    aliBaiLianLLMPropertiesConfig.getModel(),
                    messages,
                    2000,
                    0.3f,
                    1
            )).subscribe(
                    streamResponse -> {
                        if (streamResponse.getChoices().isEmpty()) {
                            return;
                        }
                        for (ChatCompletionStreamChoice choice : streamResponse.getChoices()) {
                            String finishReason = choice.getFinishReason();
                            if (finishReason != null) {
                                emitter.send("finished");
                                continue;
                            }
                            String responseContent = choice.getDelta().getContent();
                            if (responseContent.equals("```") || responseContent.equals("json") ) {
                                continue;
                            }
                            responseContent = responseContent.replaceAll("\n", "<br>");
                            System.out.println(responseContent);
                            emitter.send(responseContent);
                            aiMessage.append(responseContent);
                        }
                    },
                    error -> {
                        error.printStackTrace();
                    },
                    () -> {
//                        emitter.complete(); // 完成事件流发送
                    }
            );
        } catch (Exception e) {
            e.printStackTrace();
        }
        return aiMessage.toString();
    }
}