package com.webchat.aigc.llm;

import com.webchat.aigc.config.properties.KimiLLMPropertiesConfig;
import com.webchat.common.util.llm.MoonShotAIClient;
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 00:22
 * @description
 */
@Service
public class KimiAIService extends AbstractLLMChatService {

    @Autowired
    private KimiLLMPropertiesConfig kimiLLMPropertiesConfig;

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

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

        StringBuilder aiMessage = new StringBuilder();

        MoonShotAIClient client = new MoonShotAIClient(kimiLLMPropertiesConfig.getApiKey());


        final List<ChatCompletionMessage> messages = messageList;
        try {
            client.ChatCompletionStream(new ChatCompletionRequest(
                    kimiLLMPropertiesConfig.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 != null) {
                                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();
    }
}