package xyz.wbsite.ai;

import dev.langchain4j.agent.tool.P;
import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.service.tool.DefaultToolExecutor;
import dev.langchain4j.service.tool.ToolExecutor;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.function.Consumer;

/**
 * Tool调用基本用法，需要自行实现AI回调消息
 */
public class Base_Tool_Example {

    public static void main(String[] args) {
        // 初始化消息列表
        List<ChatMessage> chatMessages = new ArrayList<>();
        chatMessages.add(UserMessage.from("请问，泰州市的天气怎么样？"));

        // 定义工具
        Object weatherTools = new Object() {
            @Tool("返回某一城市的天气情况")
            public String getWeather(@P("应返回天气预报的城市") String city) {
                return "天气阴转多云，1~6℃";
            }
        };
        List<ToolSpecification> toolSpecifications = ToolSpecifications.toolSpecificationsFrom(weatherTools);

        // 构建请求
        ChatRequest chatRequest = ChatRequest.builder()
                .messages(chatMessages)
                .parameters(ChatRequestParameters.builder()
                        .toolSpecifications(toolSpecifications)
                        .build())
                .build();

        // 调用LLM
        ChatResponse chatResponse = Helper.getToolChatModel().chat(chatRequest);
        AiMessage aiMessage = chatResponse.aiMessage();
        chatMessages.add(aiMessage);
        // 判断是否需要调用工具
        if (aiMessage.hasToolExecutionRequests()) {
            System.out.println("LLM决定调用工具");
            System.out.println(chatResponse.aiMessage());
            List<ToolExecutionRequest> toolExecutionRequests = chatResponse.aiMessage().toolExecutionRequests();
            toolExecutionRequests.forEach(new Consumer<ToolExecutionRequest>() {
                @Override
                public void accept(ToolExecutionRequest toolExecutionRequest) {
                    ToolExecutor toolExecutor = new DefaultToolExecutor(weatherTools, toolExecutionRequest);
                    String result = toolExecutor.execute(toolExecutionRequest, UUID.randomUUID().toString());
                    ToolExecutionResultMessage toolExecutionResultMessages = ToolExecutionResultMessage.from(toolExecutionRequest, result);
                    chatMessages.add(toolExecutionResultMessages);
                }
            });
        }

        // 再次调用LLM
        ChatRequest chatRequest2 = ChatRequest.builder()
                .messages(chatMessages)
                .parameters(ChatRequestParameters.builder()
                        .toolSpecifications(toolSpecifications)
                        .build())
                .build();

        // 返回最终回答
        ChatResponse finalChatResponse = Helper.getToolChatModel().chat(chatRequest2);
        System.out.println(finalChatResponse.aiMessage().text());
    }
}