<dependencies>
    <dependency>
        <groupId>io.milvus</groupId>
        <artifactId>milvus-sdk-java</artifactId>
        <version>2.4.3</version>
    </dependency>
    <dependency>
        <groupId>org.apache.commons</groupId>
        <artifactId>commons-lang3</artifactId>
        <version>3.14.0</version>
    </dependency>
    <dependency>
        <groupId>org.slf4j</groupId>
        <artifactId>slf4j-simple</artifactId>
        <version>2.0.9</version>
    </dependency>
</dependencies>

import io.milvus.client.MilvusClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.MilvusClientParam;
import io.milvus.param.R;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.InsertResponse;
import io.milvus.response.SearchResponse;
import org.apache.commons.lang3.RandomUtils;

import java.util.ArrayList;
import java.util.List;

public class VectorDbDemo {
    // Milvus Lite 數(shù)據(jù)庫連接地址（本地文件路徑，自定義）
    private static final String MILVUS_URL = "localhost:19530";
    // 集合名稱（類似表名，自定義）
    private static final String COLLECTION_NAME = "java_vector_demo";
    // 向量維度（需與生成的向量維度一致，這里選128維）
    private static final int VECTOR_DIM = 128;

    public static void main(String[] args) {
        // 1. 構建連接參數(shù)
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withHost("localhost")
                .withPort(19530)
                .build();

        // 2. 創(chuàng)建 Milvus 客戶端
        MilvusClient client = new MilvusClient(connectParam);

        // 測試連接是否成功
        R<Boolean> healthCheck = client.checkHealth();
        if (healthCheck.getData()) {
            System.out.println("? 數(shù)據(jù)庫連接成功！");
        } else {
            System.out.println("? 數(shù)據(jù)庫連接失敗：" + healthCheck.getMessage());
            return;
        }

        // 后續(xù)操作（創(chuàng)建集合、插數(shù)、查詢）將在這里補充
    }
}

// 3. 創(chuàng)建集合（先判斷是否存在，存在則刪除）
if (client.hasCollection(COLLECTION_NAME).getData()) {
    client.dropCollection(COLLECTION_NAME);
    System.out.println("??  集合已存在，已刪除");
}

// 定義字段：主鍵字段（自增ID）
FieldType idField = FieldType.newBuilder()
        .withName("id")
        .withDataType(FieldType.DataType.Int64)
        .withPrimaryKey(true)
        .withAutoID(true) // 自增
        .build();

// 定義字段：向量字段（核心）
FieldType vectorField = FieldType.newBuilder()
        .withName("content_vector")
        .withDataType(FieldType.DataType.FloatVector)
        .withDimension(VECTOR_DIM) // 向量維度與生成的一致
        .build();

// 定義字段：普通屬性字段（存儲文本原文，方便查詢后展示）
FieldType textField = FieldType.newBuilder()
        .withName("text_content")
        .withDataType(FieldType.DataType.VarChar)
        .withMaxLength(512) // 文本最大長度
        .build();

// 構建創(chuàng)建集合的參數(shù)
CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
        .withCollectionName(COLLECTION_NAME)
        .addFieldType(idField)
        .addFieldType(vectorField)
        .addFieldType(textField)
        .withConsistencyLevel(CreateCollectionParam.ConsistencyLevelEnum.STRONG) // 一致性級別
        .build();

// 執(zhí)行創(chuàng)建集合
R<Boolean> createResult = client.createCollection(createParam);
if (createResult.getData()) {
    System.out.println("? 集合創(chuàng)建成功！");
} else {
    System.out.println("? 集合創(chuàng)建失?。? + createResult.getMessage());
    client.close();
    return;
}

// 4. 生成測試數(shù)據(jù)（5條文本，模擬用戶問答場景）
List<String> textList = new ArrayList<>();
textList.add("Java 中如何創(chuàng)建線程？有幾種方式？");
textList.add("Java 線程的生命周期包括哪些狀態(tài)？");
textList.add("Spring Boot 如何配置數(shù)據(jù)庫連接池？");
textList.add("Spring Cloud 與 Spring Boot 的區(qū)別是什么？");
textList.add("Java 8 的 Lambda 表達式有什么用法？");

// 生成向量：將每條文本轉化為128維向量（模擬，實際用模型生成）
List<List<Float>> vectorList = new ArrayList<>();
for (int i = 0; i < textList.size(); i++) {
    List<Float> vector = new ArrayList<>();
    for (int j = 0; j < VECTOR_DIM; j++) {
        // 用隨機數(shù)模擬向量（真實場景替換為模型輸出）
        vector.add(RandomUtils.nextFloat(0.0f, 1.0f));
    }
    vectorList.add(vector);
}

// 構建插入?yún)?shù)
InsertParam insertParam = InsertParam.newBuilder()
        .withCollectionName(COLLECTION_NAME)
        .addField("text_content", textList) // 插入文本字段
        .addField("content_vector", vectorList) // 插入向量字段
        .build();

// 執(zhí)行插入
R<InsertResponse> insertResult = client.insert(insertParam);
if (insertResult.getData() != null) {
    System.out.println("? 數(shù)據(jù)插入成功！插入條數(shù)：" + textList.size());
    System.out.println("插入數(shù)據(jù)ID：" + insertResult.getData().getInsertIds());
} else {
    System.out.println("? 數(shù)據(jù)插入失敗：" + insertResult.getMessage());
    client.close();
    return;
}

// 5. 創(chuàng)建向量索引（提升查詢效率）
// 索引參數(shù)（IVF_FLAT 是基礎索引類型，適合小數(shù)據(jù)量）
String indexParam = String.format("{\"index_type\": \"IVF_FLAT\", \"params\": {\"nlist\": 1024}, \"metric_type\": \"L2\"}");
R<Boolean> createIndexResult = client.createIndex(
        COLLECTION_NAME,
        "content_vector", // 向量字段名
        indexParam
);
if (createIndexResult.getData()) {
    System.out.println("? 向量索引創(chuàng)建成功！");
} else {
    System.out.println("? 索引創(chuàng)建失敗：" + createIndexResult.getMessage());
    client.close();
    return;
}

// 6. 加載集合到內存（查詢前必須執(zhí)行）
R<Boolean> loadResult = client.loadCollection(COLLECTION_NAME);
if (loadResult.getData()) {
    System.out.println("? 集合加載到內存成功！");
} else {
    System.out.println("? 集合加載失敗：" + loadResult.getMessage());
    client.close();
    return;
}

// 7. 執(zhí)行相似度查詢（模擬用戶查詢："Java 線程相關問題"）
String queryText = "Java 線程相關問題";
// 生成查詢向量（同樣用隨機數(shù)模擬，真實場景與數(shù)據(jù)向量用同一模型生成）
List<Float> queryVector = new ArrayList<>();
for (int j = 0; j < VECTOR_DIM; j++) {
    queryVector.add(RandomUtils.nextFloat(0.0f, 1.0f));
}

// 構建查詢參數(shù)
SearchParam searchParam = SearchParam.newBuilder()
        .withCollectionName(COLLECTION_NAME)
        .withMetricType(SearchParam.MetricType.L2) // 距離計算方式（L2為歐氏距離）
        .withTopK(2) // 返回最相似的2條數(shù)據(jù)
        .withVectorFieldName("content_vector")
        .addVector(queryVector)
        .withParams("{\"nprobe\": 10}") // 索引查詢參數(shù)
        .addOutputField("text_content") // 要返回的字段
        .build();

// 執(zhí)行查詢
R<SearchResponse> searchResult = client.search(searchParam);
if (searchResult.getData() != null) {
    System.out.println("\n?? 查詢結果（與\"" + queryText + "\"最相似的2條數(shù)據(jù)）：");
    // 解析查詢結果
    SearchResponse.Data data = searchResult.getData().get(0);
    for (int i = 0; i < data.getScoreCount(); i++) {
        float score = data.getScores().get(i); // 相似度分數(shù)（L2距離越小越相似）
        String text = data.getFields().get("text_content").get(i).toString(); // 文本內容
        System.out.println((i+1) + ". 相似度：" + String.format("%.4f", score) + "，內容：" + text);
    }
} else {
    System.out.println("? 查詢失?。? + searchResult.getMessage());
}

// 關閉客戶端
client.close();