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.cohere.CohereScoringModel;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.scoring.ScoringModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.aggregator.ContentAggregator;
import dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator;
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;

public class _04_Advanced_RAG_with_ReRanking {

    /**
     * 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 the implementation of a more advanced RAG application
     * using a technique known as "re-ranking".
     * <p>
     * Frequently, not all results retrieved by {@link ContentRetriever} are truly relevant to the user query.
     * This is because, during the initial retrieval stage, it is often preferable to use faster
     * and more cost-effective models, particularly when dealing with a large volume of data.
     * The trade-off is that the retrieval quality may be lower.
     * Providing irrelevant information to the LLM can be costly and, in the worst case, lead to hallucinations.
     * Therefore, in the second stage, we can perform re-ranking of the results obtained in the first stage
     * and eliminate irrelevant results using a more advanced model (e.g., Cohere Rerank).
     * <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();

        // First, say "Hi". Observe how all segments retrieved in the first stage were filtered out.
        // Then, ask "Can I cancel my reservation?" and observe how all but one segment were filtered out.

        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")
                .modelName("gpt-3.5-turbo")
                .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)
                .maxResults(5) // let's get more results
                .build();

        // To register and get a free API key for Cohere, please visit the following link:
        // https://dashboard.cohere.com/welcome/register
        ScoringModel scoringModel = CohereScoringModel.withApiKey(System.getenv("COHERE_API_KEY"));

        ContentAggregator contentAggregator = ReRankingContentAggregator.builder()
                .scoringModel(scoringModel)
                .minScore(0.8) // we want to present the LLM with only the truly relevant segments for the user's query

                .build();

        RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                .contentRetriever(contentRetriever)
                .contentAggregator(contentAggregator)
                .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 = _04_Advanced_RAG_with_ReRanking.class.getResource(fileName);
            return Paths.get(fileUrl.toURI());
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
    }
}
dev.langchain4j.service.AiServices.builder dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder dev.langchain4j.rag.DefaultRetrievalAugmentor.builder dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator.builder dev.langchain4j.model.openai.OpenAiChatModel.builder
[((4121, 4247), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((4121, 4222), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((4121, 4178), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((4533, 4751), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((4533, 4700), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((4533, 4669), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((4533, 4621), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((5049, 5275), 'dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator.builder'), ((5049, 5160), 'dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator.builder'), ((5049, 5129), 'dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator.builder'), ((5326, 5492), 'dev.langchain4j.rag.DefaultRetrievalAugmentor.builder'), ((5326, 5467), 'dev.langchain4j.rag.DefaultRetrievalAugmentor.builder'), ((5326, 5413), 'dev.langchain4j.rag.DefaultRetrievalAugmentor.builder'), ((5510, 5756), 'dev.langchain4j.service.AiServices.builder'), ((5510, 5731), 'dev.langchain4j.service.AiServices.builder'), ((5510, 5658), 'dev.langchain4j.service.AiServices.builder'), ((5510, 5602), 'dev.langchain4j.service.AiServices.builder')]