package com.egineering.ai.llmjavademo;

import com.egineering.ai.llmjavademo.agents.FaqAgent;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.BertTokenizer;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.junit.jupiter.api.Test;
import org.springframework.util.ResourceUtils;

import java.io.File;
import java.io.IOException;
import java.lang.reflect.Field;
import java.lang.reflect.Proxy;

import static dev.langchain4j.data.document.loader.FileSystemDocumentLoader.loadDocument;

public class Tests {

    @Test
    public void test() throws IOException {
        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        File fileResource = ResourceUtils.getFile("classpath:jackson_lottery.pdf");
        Document document = loadDocument(fileResource.toPath(), new ApachePdfBoxDocumentParser());

        DocumentSplitter documentSplitter = DocumentSplitters.recursive(100, 2, new BertTokenizer());
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(documentSplitter)
                .embeddingModel(new AllMiniLmL6V2EmbeddingModel())
                .embeddingStore(embeddingStore)
                .build();
        ingestor.ingest(document);
    }

    @Test
    public void test2() throws NoSuchFieldException, IllegalAccessException {
        StreamingChatLanguageModel model = OllamaStreamingChatModel.builder()
                .baseUrl("http://localhost:11434")
                .modelName("llama2")
                .temperature(0.0)
                .build();

        FaqAgent faqAgent = AiServices.builder(FaqAgent.class)
                .streamingChatLanguageModel(model)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(20))
                .build();

        Field defaultAiServiceField = Proxy.getInvocationHandler(faqAgent).getClass().getDeclaredField("context");
        defaultAiServiceField.setAccessible(true);
        Object defaultAiServices = defaultAiServiceField.get(AiServices.class);
        Proxy.getInvocationHandler(faqAgent);
    }
}

dev.langchain4j.service.AiServices.builder dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder dev.langchain4j.model.ollama.OllamaStreamingChatModel.builder
[((1713, 1937), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((1713, 1912), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((1713, 1864), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((1713, 1797), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((2112, 2293), 'dev.langchain4j.model.ollama.OllamaStreamingChatModel.builder'), ((2112, 2268), 'dev.langchain4j.model.ollama.OllamaStreamingChatModel.builder'), ((2112, 2234), 'dev.langchain4j.model.ollama.OllamaStreamingChatModel.builder'), ((2112, 2197), 'dev.langchain4j.model.ollama.OllamaStreamingChatModel.builder'), ((2324, 2507), 'dev.langchain4j.service.AiServices.builder'), ((2324, 2482), 'dev.langchain4j.service.AiServices.builder'), ((2324, 2409), 'dev.langchain4j.service.AiServices.builder'), ((2548, 2623), 'java.lang.reflect.Proxy.getInvocationHandler'), ((2548, 2595), 'java.lang.reflect.Proxy.getInvocationHandler')]