package com.ai.config;

import com.ai.function.MyCalculator;
import com.ai.service.Assist;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
@RequiredArgsConstructor
public class AssistantInit {

    // All files in a directory, txt seems to be faster

    @Value("${langchain4j.api-key}")
    private String apiKey;

    @Value("${langchain4j.model}")
    private String model;

    @Bean
    public Assist init() {
        ChatLanguageModel qwenModel = QwenChatModel.builder()
                .apiKey(apiKey)
                .modelName(model)
                .build();
        List<Document> documents = FileSystemDocumentLoader.loadDocuments("E:\\ideaProject\\liang-ai");
        // for simplicity, we will use an in-memory one:
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        EmbeddingStoreIngestor.ingest(documents, embeddingStore);
        return AiServices.builder(Assist.class)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                .chatLanguageModel(qwenModel)
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
//                .tools(new MyCalculator())
                .build();
    }
}