package com.example.application;

import com.example.application.services.BookingTools;
import com.example.application.services.CustomerSupportAgent;
import com.vaadin.flow.component.page.AppShellConfigurator;
import com.vaadin.flow.theme.Theme;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.Tokenizer;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;
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.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;

import java.io.IOException;

import static dev.langchain4j.data.document.loader.FileSystemDocumentLoader.loadDocument;
import static dev.langchain4j.model.openai.OpenAiModelName.GPT_3_5_TURBO;
import static dev.langchain4j.model.openai.OpenAiModelName.GPT_4;


@SpringBootApplication
@Theme(value = "customer-service-chatbot")
public class Application implements AppShellConfigurator {

    public static void main(String[] args) {
        SpringApplication.run(Application.class, args);
    }

    @Bean
    EmbeddingModel embeddingModel() {
        return new AllMiniLmL6V2EmbeddingModel();
    }

    @Bean
    EmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    @Bean
    Tokenizer tokenizer() {
        return new OpenAiTokenizer(GPT_3_5_TURBO);
    }


    // In the real world, ingesting documents would often happen separately, on a CI server or similar
    @Bean
    CommandLineRunner docsToEmbeddings(
            EmbeddingModel embeddingModel,
            EmbeddingStore<TextSegment> embeddingStore,
            Tokenizer tokenizer,
            ResourceLoader resourceLoader
    ) throws IOException {
        return args -> {
            Resource resource =
                    resourceLoader.getResource("classpath:terms-of-service.txt");
            var termsOfUse = loadDocument(resource.getFile().toPath(), new TextDocumentParser());

            DocumentSplitter documentSplitter = DocumentSplitters.recursive(200, 0,
                    tokenizer);

            EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                    .documentSplitter(documentSplitter)
                    .embeddingModel(embeddingModel)
                    .embeddingStore(embeddingStore)
                    .build();
            ingestor.ingest(termsOfUse);
        };
    }

    @Bean
    StreamingChatLanguageModel chatLanguageModel() {
        return OpenAiStreamingChatModel.builder()
                .apiKey(ApiKeys.OPENAI_API_KEY)
                .modelName(GPT_3_5_TURBO)
                .build();
    }


    @Bean
    ContentRetriever retriever(
            EmbeddingStore<TextSegment> embeddingStore,
            EmbeddingModel embeddingModel
    ) {
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(2)
                .minScore(0.6)
                .build();
    }


    @Bean
    CustomerSupportAgent customerSupportAgent(
            StreamingChatLanguageModel chatLanguageModel,
            Tokenizer tokenizer,
            ContentRetriever retriever,
            BookingTools tools
    ) {

        return AiServices.builder(CustomerSupportAgent.class)
                .streamingChatLanguageModel(chatLanguageModel)
                .chatMemoryProvider(chatId -> TokenWindowChatMemory.builder()
                        .id(chatId)
                        .maxTokens(1000, tokenizer)
                        .build())
                .contentRetriever(retriever)
                .tools(tools)
                .build();
    }
}
dev.langchain4j.service.AiServices.builder dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder dev.langchain4j.memory.chat.TokenWindowChatMemory.builder dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder
[((3196, 3417), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((3196, 3388), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((3196, 3336), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((3196, 3284), 'dev.langchain4j.store.embedding.EmbeddingStoreIngestor.builder'), ((3556, 3705), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((3556, 3680), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((3556, 3638), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((3878, 4101), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3878, 4076), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3878, 4045), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3878, 4014), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((3878, 3966), 'dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever.builder'), ((4354, 4763), 'dev.langchain4j.service.AiServices.builder'), ((4354, 4738), 'dev.langchain4j.service.AiServices.builder'), ((4354, 4708), 'dev.langchain4j.service.AiServices.builder'), ((4354, 4663), 'dev.langchain4j.service.AiServices.builder'), ((4354, 4463), 'dev.langchain4j.service.AiServices.builder'), ((4510, 4662), 'dev.langchain4j.memory.chat.TokenWindowChatMemory.builder'), ((4510, 4629), 'dev.langchain4j.memory.chat.TokenWindowChatMemory.builder'), ((4510, 4577), 'dev.langchain4j.memory.chat.TokenWindowChatMemory.builder')]