package ch.phoenixtechnologies.lc4j.workshop;

import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.enterprise.inject.Produces;

public class RagRetriever {
    @Produces
    @ApplicationScoped
    public RetrievalAugmentor create(EmbeddingStore store, EmbeddingModel model) {
        var contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingModel(model)
                .embeddingStore(store)
                .maxResults(3)
                .build();

        return DefaultRetrievalAugmentor.builder()
                .contentRetriever(contentRetriever)
                .contentInjector((contents, userMessage) -> {
                    StringBuffer prompt = new StringBuffer(userMessage.singleText());
                    prompt.append("\nPlease, only use the following information:\n");
                    prompt.append(
                            "- If you are asked for booking information and you aren't given a name or booking id, ask for it. DON'T try with randomly names or ids.\n");
                    prompt.append(
                            "- Don't inform about the function details you need to invoke, just invoke it when it's appropriate\n");
                    contents.forEach(content -> prompt.append("- ")
                            .append(content.textSegment().text())
                            .append("\n"));
                    return new UserMessage(prompt.toString());
                })
                .build();
    }
}
