import dev.langchain4j.chain.ConversationalRetrievalChain;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.injector.ContentInjector;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import java.util.Scanner;

import static java.util.Arrays.asList;

public class _05_Advanced_RAG_with_Metadata {

    /**
     * Please refer to previous examples for basic context.
     * <p>
     * Advanced RAG in LangChain4j is described here: https://github.com/langchain4j/langchain4j/pull/538
     * <p>
     * This example illustrates how to include document source or other metadata into the LLM prompt.
     * <p>
     * We will continue using {@link AiServices} for this example,
     * but the same principles apply to {@link ConversationalRetrievalChain}, or you can develop your custom RAG flow.
     */

    public static void main(String[] args) {

        CustomerSupportAgent agent = createCustomerSupportAgent();

        try (Scanner scanner = new Scanner(System.in)) {
            while (true) {
                System.out.println("==================================================");
                System.out.print("User: ");
                String userQuery = scanner.nextLine();
                System.out.println("==================================================");

                if ("exit".equalsIgnoreCase(userQuery)) {
                    break;
                }

                String agentAnswer = agent.answer(userQuery);
                System.out.println("==================================================");
                System.out.println("Agent: " + agentAnswer);
            }
        }
    }

    private static CustomerSupportAgent createCustomerSupportAgent() {

        // Check _01_Naive_RAG if you need more details on what is going on here

        ChatLanguageModel chatModel = OpenAiChatModel.builder()
                .apiKey("demo")
                .build();

        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

        Path documentPath = toPath("miles-of-smiles-terms-of-use.txt");
        EmbeddingStore<TextSegment> embeddingStore = embed(documentPath, embeddingModel);

        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .build();

        // Each retrieved segment should include "file_name" and "index" metadata values in the prompt
        ContentInjector contentInjector = DefaultContentInjector.builder()
                // .promptTemplate(...) // Formatting can also be changed
                .metadataKeysToInclude(asList("file_name", "index"))
                .build();

        RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                .contentRetriever(contentRetriever)
                .contentInjector(contentInjector)
                .build();

        return AiServices.builder(CustomerSupportAgent.class)
                .chatLanguageModel(chatModel)
                .retrievalAugmentor(retrievalAugmentor)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .build();
    }

    private static EmbeddingStore<TextSegment> embed(Path documentPath, EmbeddingModel embeddingModel) {
        DocumentParser documentParser = new TextDocumentParser();
        Document document = FileSystemDocumentLoader.loadDocument(documentPath, documentParser);

        DocumentSplitter splitter = DocumentSplitters.recursive(300, 0);
        List<TextSegment> segments = splitter.split(document);

        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings, segments);
        return embeddingStore;
    }

    interface CustomerSupportAgent {

        String answer(String query);
    }

    private static Path toPath(String fileName) {
        try {
            URL fileUrl = _05_Advanced_RAG_with_Metadata.class.getResource(fileName);
            return Paths.get(fileUrl.toURI());
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
    }
}
[((3112, 3194), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((3112, 3169), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((3480, 3641), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3480, 3616), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3480, 3568), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3789, 3989), 'dev.langchain4j.rag.content.injector.DefaultContentInjector.builder'), ((3789, 3964), 'dev.langchain4j.rag.content.injector.DefaultContentInjector.builder'), ((4040, 4202), 'dev.langchain4j.rag.DefaultRetrievalAugmentor.builder'), ((4040, 4177), 'dev.langchain4j.rag.DefaultRetrievalAugmentor.builder'), ((4040, 4127), 'dev.langchain4j.rag.DefaultRetrievalAugmentor.builder'), ((4220, 4466), 'dev.langchain4j.service.AiServices.builder'), ((4220, 4441), 'dev.langchain4j.service.AiServices.builder'), ((4220, 4368), 'dev.langchain4j.service.AiServices.builder'), ((4220, 4312), 'dev.langchain4j.service.AiServices.builder')]