package xyz.wbsite.ai;

import dev.langchain4j.data.document.*;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.SystemMessage;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.io.IOException;
import java.io.InputStream;
import java.util.List;

public class RagTest {

    public static void main(String[] args) {

        List<Document> documents = FileSystemDocumentLoader.loadDocuments("D:\\wbSource\\starter-ai\\src\\main\\resources");
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        EmbeddingStoreIngestor.ingest(documents, embeddingStore);

        EmbeddingStoreContentRetriever embeddingStoreContentRetriever = EmbeddingStoreContentRetriever.from(embeddingStore);

        List<Content> retrieve = embeddingStoreContentRetriever.retrieve(new Query("java中会存在内存泄漏吗，请简单描述"));

        System.out.println();

//        ChatLanguageModel model = OllamaChatModel.builder()
//                .baseUrl("http://36.138.207.178:11434")
//                .modelName("deepseek-r1:32B")
//                .build();
//
//        Assistant assistant = AiServices.create(Assistant.class,model);
//                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
//                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
//                .build();
//
//        String chat = assistant.chat("java中会存在内存泄漏吗，请简单描述。");
//        System.out.println(chat);
    }


    interface Assistant {

//        @SystemMessage("请扮演一名小学生,根据输入的文章题目写一篇100字以内的作文")
        String chat(String userMessage);
    }

}