package com.demo.ai.config;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.deepseek.api.DeepSeekApi;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.Resource;

@RequiredArgsConstructor
@Configuration
public class AiConfig {
    @Value("classpath:/prompts/system-message.st")
    private Resource systemResource;
    private final ChatMemoryRepository chatMemoryRepository;

    @Bean("chatClient")
    public ChatClient chatClient(ChatClient.Builder chatClientBuilder, ToolCallbackProvider[] toolCallbackProviders) {
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(chatMemoryRepository)
                .maxMessages(50)
                .build();
        return chatClientBuilder
                .defaultOptions(ChatOptions.builder()
                        .model(DeepSeekApi.ChatModel.DEEPSEEK_CHAT.getValue())
                        .temperature(0.7)
                        .build())
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build(), new SimpleLoggerAdvisor())
                .defaultSystem(systemResource)
                .defaultToolCallbacks(toolCallbackProviders)
                .build();
    }

    @Bean("thinkChatClient")
    public ChatClient thinkChatClient(ChatClient.Builder chatClientBuilder, ToolCallbackProvider[] toolCallbackProviders) {
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(chatMemoryRepository)
                .maxMessages(50)
                .build();
        return chatClientBuilder
                .defaultOptions(ChatOptions.builder()
                        .model(DeepSeekApi.ChatModel.DEEPSEEK_REASONER.getValue())
                        .temperature(0.6)
                        .build())
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build(), new SimpleLoggerAdvisor())
                .defaultSystem(systemResource)
                .defaultToolCallbacks(toolCallbackProviders)
                .build();
    }
}

package com.demo.ai.service;

import com.demo.base.excpetion.ServiceException;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

@Service
public class ChatService {
    private final ChatClient chatClient;
    private final ChatClient thinkChatClient;

    public ChatService(@Qualifier("chatClient") ChatClient chatClient,
                       @Qualifier("thinkChatClient") ChatClient thinkChatClient) {
        this.chatClient = chatClient;
        this.thinkChatClient = thinkChatClient;
    }

    public SseEmitter chat(String question, String conversationId, Boolean isThinking) {
        ChatClient client = isThinking ? thinkChatClient : chatClient;
        return chat(question, conversationId, client);
    }

    private SseEmitter chat(String question, String conversationId, ChatClient client) {
        SseEmitter emitter = new SseEmitter();
        client.prompt()
                .user(question)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
                .stream()
                .chatResponse()
                .doOnNext(message -> {
                    try {
                        emitter.send(message.getResult());
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                        throw new ServiceException("send error");
                    }
                })
                .doOnComplete(emitter::complete)
                .doOnError(emitter::completeWithError)
                .subscribe();
        return emitter;
    }

}