shikeying
2024-04-07 52c8ee3120e4a44e48e217bced6f6c799ff7980c
deploy-jar-template/src/main/java/com/iplatform/milvus/MilvusEngine.java
@@ -6,6 +6,8 @@
import com.walker.support.milvus.FieldType;
import com.walker.support.milvus.MetricType;
import com.walker.support.milvus.OperateService;
import com.walker.support.milvus.OutData;
import com.walker.support.milvus.Query;
import com.walker.support.milvus.Table;
import com.walker.support.milvus.engine.DefaultOperateService;
import org.slf4j.Logger;
@@ -20,6 +22,8 @@
public class MilvusEngine {
    protected final transient Logger logger = LoggerFactory.getLogger(this.getClass());
    public static final String TABLE_CHAT_SIMILAR = "chat_similar";
    public MilvusEngine(String ip, int port){
        DefaultOperateService service = new DefaultOperateService();
@@ -44,7 +48,7 @@
     */
    public void createChatSimilarTable(){
        Table chatSimilarTable = new Table();
        chatSimilarTable.setCollectionName("chat_similar");
        chatSimilarTable.setCollectionName(TABLE_CHAT_SIMILAR);
        chatSimilarTable.setDescription("聊天提取工单摘要历史数据");
        chatSimilarTable.setShardsNum(1);
        chatSimilarTable.setDimension(768); // 这个是根据使用向量模型维度定的
@@ -84,7 +88,7 @@
    public void insertTestData(){
        DataSet dataSet = new DataSet();
        dataSet.setTableName("chat_similar");
        dataSet.setTableName(TABLE_CHAT_SIMILAR);
        List<List<Float>> vectorList = new ArrayList<>();
        vectorList.add(Arrays.asList(mockVector));
@@ -132,7 +136,7 @@
        }
        DataSet dataSet = new DataSet();
        dataSet.setTableName("chat_similar");
        dataSet.setTableName(TABLE_CHAT_SIMILAR);
        Map<String, List<?>> fieldMap = new HashMap();
        fieldMap.put("id", ids);
@@ -145,6 +149,26 @@
        logger.info("写入了: {}", ids);
    }
    public OutData searchChatSimilar(List<List<Float>> vectors){
        Query query = new Query();
        query.setMetricType(MetricType.NLP.getIndex());
        query.setTableName(TABLE_CHAT_SIMILAR);
        query.setTopK(4);
        query.setVectorName("embedding");
        query.setOutFieldList(Arrays.asList(new String[]{"id","title","content"}));
        query.setFieldPrimaryKey("id");
        query.setSearchVectors(vectors);
        return this.operateService.searchVector(query);
    }
    /**
     * 必须在查询之前,加载数据到内存中。
     * @date 2024-03-31
     */
    public void loadChatSimilar4Search(){
        this.operateService.prepareSearch(TABLE_CHAT_SIMILAR);
    }
    private OperateService operateService;
//    private Double[] mockVector = new Double[]{-0.051114246249198914, 0.889954432};