shikeying
2024-05-08 8924870a053f0b882ada86421c062cbdb9cff093
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,11 @@
public class MilvusEngine {
    protected final transient Logger logger = LoggerFactory.getLogger(this.getClass());
    public static final String TABLE_CHAT_SIMILAR = "chat_similar";
    //    private static final int VECTOR_DIMENSION = 768;
    private static final int VECTOR_DIMENSION = 512;
    public MilvusEngine(String ip, int port){
        DefaultOperateService service = new DefaultOperateService();
@@ -44,10 +51,11 @@
     */
    public void createChatSimilarTable(){
        Table chatSimilarTable = new Table();
        chatSimilarTable.setCollectionName("chat_similar");
        chatSimilarTable.setCollectionName(TABLE_CHAT_SIMILAR);
        chatSimilarTable.setDescription("聊天提取工单摘要历史数据");
        chatSimilarTable.setShardsNum(1);
        chatSimilarTable.setDimension(768); // 这个是根据使用向量模型维度定的
//        chatSimilarTable.setDimension(768); // 这个是根据使用向量模型维度定的
        chatSimilarTable.setDimension(VECTOR_DIMENSION); // 这个是根据使用向量模型维度定的
        // 设置字段
        FieldType id = FieldType.newBuilder()
@@ -59,7 +67,7 @@
        FieldType answer = FieldType.newBuilder()
                .withName("answer").withPrimaryKey(false).withMaxLength(255).withDataType(DataType.VarChar).build();
        FieldType embedding = FieldType.newBuilder()
                .withName("embedding").withPrimaryKey(false).withDataType(DataType.FloatVector).withDimension(768).build();
                .withName("embedding").withPrimaryKey(false).withDataType(DataType.FloatVector).withDimension(VECTOR_DIMENSION).build();
        List<FieldType> fieldTypeList = new ArrayList<>(8);
        fieldTypeList.add(id);
@@ -80,11 +88,12 @@
    public void dropChatSimilarTable(){
        this.operateService.dropTable("chat_similar");
        this.operateService.dropIndex("chat_similar", "chat_similar_index");
    }
    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 +141,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 +154,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};