package com.iwe3.langchain4j;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;

@SpringBootTest
class PineConeApplicationTest {

    @Autowired
    private EmbeddingModel embeddingModel;
    @Autowired
    private EmbeddingStore embeddingStore;


    /**
     * Pinecone - 相似度匹配
     */
    @Test
    public void embeddingSearch() {
        // 提问，并将问题转成向量数据
        var queryEmbedding = embeddingModel.embed("你最喜欢的运动是什么？").content();

        // 创建搜索请求对象
        var searchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)  // 匹配最相似的一条记录
                .minScore(0.8)
                .build();

        // 根据搜索请求 searchRequest 在向量存储中进行相似度搜索
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(searchRequest);

        // searchResult.matches(): 获取搜索结果中的匹配项列表。
        // .get(0): 从匹配项列表中获取第一个匹配项
        EmbeddingMatch<TextSegment> embeddingMatch = searchResult.matches().get(0);

        // 获取匹配项的相似度得分
        System.out.println(embeddingMatch.score()); // 输出：0.8144288515898701

        // 返回文本结果
        System.out.println(embeddingMatch.embedded().text());
    }



    /**
     * 将文本转换成向量，然后存储到 pinecone 中
     *
     * 参考:
     * https://docs.langchain4j.dev/tutorials/embedding-stores
     */
    @Test
    public void testPineconeEmbedded() {
        // 将文本转换成向量
        var seg1 = TextSegment.from("我喜欢羽毛球");
        var em1 = embeddingModel.embed(seg1).content();
        embeddingStore.add(em1, seg1);// 存入向量数据库

        var seg2 = TextSegment.from("今天天气很好");
        var em2 = embeddingModel.embed(seg2).content();
        embeddingStore.add(em2, seg2);// 存入向量数据库
    }


    @Test
    public void testEmbeddingModel() {
        var embed = embeddingModel.embed("你好");

        System.out.println("向量维度：" + embed.content().vector().length);
        System.out.println("向量输出：" + embed.toString());
    }
}