package com.example.springaidemo.service;

import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
public class VectorService {

    @Autowired VectorStore vectorStore;

    public void storeVector(){
        List<Document> documents = List.of(
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!"),
                new Document("The World is Big and Salvation Lurks Around the Corner"),
                new Document("You walk forward facing the past and you turn back toward the future."));

// Add the documents to Elasticsearch 需要注意的是elasticSearch的版本需要大于8.10.0
        vectorStore.add(documents);

// Retrieve documents similar to a query
        SearchRequest searchRequest = SearchRequest.builder().query("Spring")
                .topK(5)
                .build();
        List<Document> results = this.vectorStore.similaritySearch(searchRequest);
        System.out.println(results);
    }

    public void modelVector() {
    // 创建包含自定义字段的元数据
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("modelName", "tb_vip_user");
        metadata.put("source", "dataCommon");
        metadata.put("timestamp", System.currentTimeMillis());
        List<Document> documents = List.of(
                new Document("表 tb_vip_user VIP用户表 | 字段 user_name用户名", metadata),
                new Document("表 tb_vip_user VIP用户表 | 字段 user_id用户id", metadata),
                new Document("表 tb_vip_user VIP用户表 | 字段 user_sex用户性别", metadata),
                new Document("表 tb_vip_user VIP用户表 | 字段 user_age用户年龄", metadata),
                new Document("表 tb_vip_user VIP用户表 | 字段 user_phone用户手机号", metadata),
                new Document("表 tb_vip_user VIP用户表 | 字段 user_email用户邮箱", metadata),
                new Document("表 tb_vip_user VIP用户表 | 字段 user_address用户地址", metadata));
        vectorStore.add(documents);
    }

    public List<Document> searchModel(String query, int topK) {
        return vectorStore.similaritySearch(SearchRequest.builder()
                .query(query)
                .topK(topK)
//                .filterExpression("source == 'dataCommon'")
                .build());
    }
}