package com.agentsflex.llm.spark.test;

import com.agentsflex.core.document.Document;
import com.agentsflex.core.llm.Llm;
import com.agentsflex.core.llm.exception.LlmException;
import com.agentsflex.core.llm.response.AiMessageResponse;
import com.agentsflex.core.message.HumanMessage;
import com.agentsflex.core.prompt.FunctionPrompt;
import com.agentsflex.core.prompt.HistoriesPrompt;
import com.agentsflex.core.store.VectorData;
import com.agentsflex.core.util.LogUtil;
import com.agentsflex.llm.spark.SparkLlm;
import com.agentsflex.llm.spark.SparkLlmConfig;
import org.junit.Test;

import java.util.Scanner;

public class SparkLlmTest {

    private static SparkLlm getSparkLlm() {
        SparkLlmConfig config = new SparkLlmConfig();
        config.setAppId("****");
        config.setApiKey("****");
        config.setApiSecret("****");


        config.setDebug(true);
        return new SparkLlm(config);
    }

    @Test(expected = LlmException.class)
    public void testSimple() {
        Llm llm = getSparkLlm();
        String result = llm.chat("你好，请问你是谁？");
        System.out.println(result);
    }

    @Test
    public void testEmbedding() {
        Llm llm = getSparkLlm();
        VectorData vectorData = llm.embed(Document.of("你好，请问你是谁？"));
        System.out.println(vectorData);
    }


    @Test
    public void testFunctionCalling() throws InterruptedException {
        Llm llm = getSparkLlm();
        FunctionPrompt prompt = new FunctionPrompt("今天北京的天气怎么样", WeatherFunctions.class);
        AiMessageResponse response = llm.chat(prompt);

        System.out.println(response.callFunctions());
    }


    public static void main(String[] args) {
        Llm llm = getSparkLlm();

        HistoriesPrompt prompt = new HistoriesPrompt();

        LogUtil.println("您想问什么？");
        Scanner scanner = new Scanner(System.in);
        String userInput = scanner.nextLine();

        while (userInput != null) {

            prompt.addMessage(new HumanMessage(userInput));

            llm.chatStream(prompt, (context, response) -> {
                LogUtil.println(">>>> " + response.getMessage().getContent());
            });

            userInput = scanner.nextLine();
        }
    }
}
