| | |
| | | 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; |
| | |
| | | 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(); |
| | |
| | | */ |
| | | public void createChatSimilarTable(){ |
| | | Table chatSimilarTable = new Table(); |
| | | chatSimilarTable.setCollectionName("chat_similar"); |
| | | chatSimilarTable.setCollectionName(TABLE_CHAT_SIMILAR); |
| | | chatSimilarTable.setDescription("聊天提取工单摘要历史数据"); |
| | | chatSimilarTable.setShardsNum(1); |
| | | chatSimilarTable.setDimension(768); // 这个是根据使用向量模型维度定的 |
| | |
| | | |
| | | 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)); |
| | |
| | | } |
| | | |
| | | DataSet dataSet = new DataSet(); |
| | | dataSet.setTableName("chat_similar"); |
| | | dataSet.setTableName(TABLE_CHAT_SIMILAR); |
| | | |
| | | Map<String, List<?>> fieldMap = new HashMap(); |
| | | fieldMap.put("id", ids); |
| | |
| | | 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}; |