使用 Milvus 与 Ollama 进行文本向量存储与检索

在这篇博客中,我们将展示如何将 MilvusOllama 结合使用,完成文本的向量化并存储在 Milvus 向量数据库中,同时演示如何进行相似度搜索。Milvus 是一个高效的向量数据库,适用于大规模的相似度搜索任务,而 Ollama 提供了强大的文本嵌入功能,能将文本转换为向量。通过结合这两个工具,我们能够快速实现高效的文本搜索系统。

项目概述

  1. Milvus 用于存储向量数据和执行快速相似度搜索。
  2. Ollama 提供文本嵌入 API,将文本转化为高维向量,作为 Milvus 中存储的向量数据。
  3. 我们将首先在 Milvus 中创建一个集合,并定义存储结构,然后将文本数据和对应的向量插入到该集合中,最后实现向量的相似度查询。

环境准备

  • Milvus 服务器:用于存储向量数据,支持高效的相似度搜索,安装 这里
  • Ollama 服务器:提供文本到向量的转换 API, 安装 这里
  • Java 开发环境:我们使用 Java 编写本示例代码,调用 Milvus 的客户端 API 和 Ollama 的 API。

步骤 1:连接到 Milvus 服务器

首先,我们需要连接到 Milvus 服务器。配置连接时,我们需要指定 Milvus 的服务地址和访问凭证。代码如下:

String CLUSTER_ENDPOINT = "http://localhost:19530";
String TOKEN = "root:Milvus";

// Milvus 连接配置
ConnectConfig connectConfig = ConnectConfig.builder()
        .uri(CLUSTER_ENDPOINT)
        .token(TOKEN)
        .build();

MilvusClientV2 client = new MilvusClientV2(connectConfig);

步骤 2:创建集合及其 Schema

Milvus 中,我们首先需要创建一个集合,并定义该集合的字段结构(Schema)。我们的集合将包含以下三个字段:

  1. id:主键字段,类型为 Int64
  2. text:文本字段,类型为 VarChar,用于存储原始文本。
  3. vector:向量字段,类型为 FloatVector,存储文本的向量表示。
CreateCollectionReq.CollectionSchema schema = client.createSchema();

schema.addField(AddFieldReq.builder()
        .fieldName("id")
        .dataType(DataType.Int64)
        .isPrimaryKey(true)
        .autoID(false)
        .build());

schema.addField(AddFieldReq.builder()
        .fieldName("text")
        .dataType(DataType.VarChar)
        .maxLength(512)
        .build());

schema.addField(AddFieldReq.builder()
        .fieldName("vector")
        .dataType(DataType.FloatVector)
        .dimension(768)
        .build());

步骤 3:创建索引

为了提高查询效率,我们需要为 idvector 字段创建索引。Milvus 支持多种索引类型,这里我们使用 Cosine 相似度作为度量方式。

IndexParam indexParamForIdField = IndexParam.builder()
        .fieldName("id")
        .indexType(IndexParam.IndexType.STL_SORT)
        .build();

IndexParam indexParamForVectorField = IndexParam.builder()
        .fieldName("vector")
        .indexType(IndexParam.IndexType.AUTOINDEX)
        .metricType(IndexParam.MetricType.COSINE)
        .build();

List<IndexParam> indexParams = new ArrayList<>();
indexParams.add(indexParamForIdField);
indexParams.add(indexParamForVectorField);

步骤 4:插入向量数据

接下来,我们使用 Ollama 提供的 API 将文本转换为向量,并将这些向量插入到 Milvus 中。

String[] textArray = {"Hello, this is a test.", "Milvus is great for vector search.", "Ollama helps in text embedding."};

List<float[]> vectors = new ArrayList<>();
for (String text : textArray) {
    // 调用 Ollama API 获取文本向量
    float[] vector = getTextEmbeddingFromOllama(OLLAMA_API_URL, text);
    vectors.add(vector);
}

// 将文本和向量组合成 JSON 格式的插入数据
List<JsonObject> data = new ArrayList<>();
for (int i = 0; i < vectors.size(); i++) {
    JsonObject jsonObject = new JsonObject();
    jsonObject.addProperty("id", i);
    jsonObject.add("vector", gson.toJsonTree(vectors.get(i)));
    jsonObject.addProperty("text", textArray[i]);
    data.add(jsonObject);
}

InsertReq insertReq = InsertReq.builder()
        .collectionName("milvus_example")
        .data(data)
        .build();

client.insert(insertReq);
System.out.println("Vectors inserted successfully into Milvus!");

步骤 5:进行相似度查询

最后,我们通过向量进行相似度查询。我们向 Milvus 提交一个查询向量,查询与之最相似的文本。

float[] queryVector = getTextEmbeddingFromOllama(OLLAMA_API_URL, "what about Milvus?");
FloatVec queryVectorFloatVec = new FloatVec(queryVector);

SearchReq searchReq = SearchReq.builder()
        .collectionName("milvus_example")
        .topK(1)  // 获取最相似的 1 个结果
        .metricType(IndexParam.MetricType.COSINE)  // 使用余弦相似度
        .data(Collections.singletonList(queryVectorFloatVec))
        .outputFields(Collections.singletonList("text"))
        .build();

SearchResp searchResp = client.search(searchReq);
if (searchResp != null
        && searchResp.getSearchResults() != null
        && !searchResp.getSearchResults().get(0).isEmpty()
        && searchResp.getSearchResults().get(0).get(0).getEntity() != null) {
    System.out.println("Search Results: " + gson.toJson(searchResp.getSearchResults().get(0).get(0).getEntity()));
} else {
    System.out.println("Search Results: null");
}

步骤 6:获取 Ollama 文本向量

为了获取文本的向量,我们调用 Ollama 的 Embedding API。通过向其发送请求,我们可以获得文本的嵌入向量。

public static float[] getTextEmbeddingFromOllama(String apiUrl, String text) {
    try {
        // 构造请求体
        JsonObject requestBody = new JsonObject();
        JsonArray inputArray = new JsonArray();
        inputArray.add(text);
        requestBody.add("input", inputArray);
        requestBody.addProperty("model", "nomic-embed-text");

        // 创建 OkHttp 请求
        RequestBody body = RequestBody.create(
                okhttp3.MediaType.parse("application/json; charset=utf-8"),
                gson.toJson(requestBody)
        );

        OkHttpClient client = new OkHttpClient();

        Request request = new Request.Builder()
                .url(apiUrl)
                .post(body)
                .build();

        // 执行请求并获取响应
        Response response = client.newCall(request).execute();
        if (!response.isSuccessful()) {
            System.out.println("Request failed: " + response.code());
            return new float[0];
        }

        // 解析响应数据
        String responseBody = response.body().string();
        JsonObject jsonResponse = gson.fromJson(responseBody, JsonObject.class);
        JsonArray embeddingsArray = jsonResponse.getAsJsonArray("data")
                .get(0).getAsJsonObject().getAsJsonArray("embedding");

        // 将嵌入向量转换为 float 数组
        float[] vector = new float[embeddingsArray.size()];
        for (int i = 0; i < embeddingsArray.size(); i++) {
            vector[i] = embeddingsArray.get(i).getAsFloat();
        }

        return vector;

    } catch (IOException e) {
        e.printStackTrace();
        return new float[0];
    }
}

完整代码实例

import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.DescribeCollectionReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.SearchResp;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class MilvusOllamaExample {

    private static final Gson gson = new Gson();

    public static void main(String[] args) {

        // Milvus 集群的连接配置
        String CLUSTER_ENDPOINT = "http://localhost:19530";
        String TOKEN = "root:Milvus";

        // Ollama 的 Embedding API 端点
        String OLLAMA_API_URL = "http://localhost:11434/v1/embeddings";

        String CLLECTION_NAME = "milvus_example";

        // 1. 连接到 Milvus 服务器
        ConnectConfig connectConfig = ConnectConfig.builder()
                .uri(CLUSTER_ENDPOINT)
                .token(TOKEN)
                .build();

        MilvusClientV2 client = new MilvusClientV2(connectConfig);

        // 检查并删除现有集合(如果存在)
        if (client.hasCollection(HasCollectionReq.builder().collectionName(CLLECTION_NAME).build())) {
            client.describeCollection(DescribeCollectionReq.builder().collectionName(CLLECTION_NAME).build());
            System.out.println("Collection deleted successfully!");
        }

        try {
            // 2. 创建 schema
            CreateCollectionReq.CollectionSchema schema = client.createSchema();

            // 2.1 添加字段到 schema
            schema.addField(AddFieldReq.builder()
                    .fieldName("id")
                    .dataType(DataType.Int64)
                    .isPrimaryKey(true)
                    .autoID(false)
                    .build());

            schema.addField(AddFieldReq.builder()
                    .fieldName("text")
                    .dataType(DataType.VarChar)
                    .maxLength(512)  // 设置 varchar 字段最大长度
                    .build());

            schema.addField(AddFieldReq.builder()
                    .fieldName("vector")
                    .dataType(DataType.FloatVector)
                    .dimension(768)  // 设置向量维度
                    .build());

            // 3. 准备索引参数
            IndexParam indexParamForIdField = IndexParam.builder()
                    .fieldName("id")
                    .indexType(IndexParam.IndexType.STL_SORT)
                    .build();

            IndexParam indexParamForVectorField = IndexParam.builder()
                    .fieldName("vector")
                    .indexType(IndexParam.IndexType.AUTOINDEX)
                    .metricType(IndexParam.MetricType.COSINE)  // 使用 COSINE 相似度
                    .build();

            List<IndexParam> indexParams = new ArrayList<>();
            indexParams.add(indexParamForIdField);
            indexParams.add(indexParamForVectorField);

            // 4. 创建集合并设置索引
            CreateCollectionReq customizedSetupReq1 = CreateCollectionReq.builder()
                    .collectionName("milvus_example")
                    .collectionSchema(schema)
                    .indexParams(indexParams)
                    .build();

            client.createCollection(customizedSetupReq1);
            System.out.println("Collection created successfully!");

            // 5. 使用 Ollama 获取文本向量
            String[] textArray = {"Hello, this is a test.", "Milvus is great for vector search.", "Ollama helps in text embedding."};

            List<float[]> vectors = new ArrayList<>();
            for (String text : textArray) {
                // 调用 Ollama API 获取文本向量
                float[] vector = getTextEmbeddingFromOllama(OLLAMA_API_URL, text);
                vectors.add(vector);
            }

            // 6. 插入向量到 Milvus
            List<Long> ids = new ArrayList<>();
            for (int i = 0; i < textArray.length; i++) {
                ids.add((long) i);  // 简单地用索引作为 ID
            }

            // 将文本和向量组合成 JSON 格式的插入数据
            List<JsonObject> data = new ArrayList<>();
            for (int i = 0; i < vectors.size(); i++) {
                JsonObject jsonObject = new JsonObject();
                jsonObject.addProperty("id", i);
                jsonObject.add("vector", gson.toJsonTree(vectors.get(i)));
                jsonObject.addProperty("text", textArray[i]);
                data.add(jsonObject);
            }

            InsertReq insertReq = InsertReq.builder()
                    .collectionName("milvus_example")
                    .data(data)
                    .build();

            client.insert(insertReq);
            System.out.println("Vectors inserted successfully into Milvus!");

            // 7. 创建查询向量并执行搜索
            float[] queryVector = getTextEmbeddingFromOllama(OLLAMA_API_URL, "what about Milvus?");
            FloatVec queryVectorFloatVec = new FloatVec(queryVector);

            // 创建查询请求
            SearchReq searchReq = SearchReq.builder()
                    .collectionName("milvus_example")
                    .topK(1)  // 获取最相似的 1 个结果
                    .metricType(IndexParam.MetricType.COSINE)  // 使用余弦相似度
                    .data(Collections.singletonList(queryVectorFloatVec))
                    .outputFields(Collections.singletonList("text"))
                    .build();

            // 执行查询操作
            SearchResp searchResp = client.search(searchReq);
            if (searchResp != null
                    && searchResp.getSearchResults() != null
                    && !searchResp.getSearchResults().get(0).isEmpty()
                    && searchResp.getSearchResults().get(0).get(0).getEntity() != null) {
                System.out.println("Search Results: " + gson.toJson(searchResp.getSearchResults().get(0).get(0).getEntity()));
            } else {
                System.out.println("Search Results: null");
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            // 8. 关闭 Milvus 客户端连接
            client.close();
        }
    }

    // 调用 Ollama API 获取文本的向量
    public static float[] getTextEmbeddingFromOllama(String apiUrl, String text) {
        try {
            // 构造请求体
            JsonObject requestBody = new JsonObject();
            JsonArray inputArray = new JsonArray();
            inputArray.add(text);
            requestBody.add("input", inputArray);
            requestBody.addProperty("model", "nomic-embed-text");

            // 创建 OkHttp 请求
            RequestBody body = RequestBody.create(
                    okhttp3.MediaType.parse("application/json; charset=utf-8"),
                    gson.toJson(requestBody)
            );

            OkHttpClient client = new OkHttpClient();

            Request request = new Request.Builder()
                    .url(apiUrl)
                    .post(body)
                    .build();

            // 执行请求并获取响应
            Response response = client.newCall(request).execute();
            if (!response.isSuccessful()) {
                System.out.println("Request failed: " + response.code());
                return new float[0];
            }

            // 解析响应数据
            String responseBody = response.body().string();
            JsonObject jsonResponse = gson.fromJson(responseBody, JsonObject.class);
            JsonArray embeddingsArray = jsonResponse.getAsJsonArray("data")
                    .get(0).getAsJsonObject().getAsJsonArray("embedding");

            // 将嵌入向量转换为 float 数组
            float[] vector = new float[embeddingsArray.size()];
            for (int i = 0; i < embeddingsArray.size(); i++) {
                vector[i] = embeddingsArray.get(i).getAsFloat();
            }

            return vector;

        } catch (IOException e) {
            e.printStackTrace();
            return new float[0];
        }
    }
}

总结

通过本示例,我们展示了如何结合 MilvusOllama 完成文本向量的存储和相似度检索。Milvus 提供了强大的向量存储和检索功能,而 Ollama 则通过其 API 提供了高效的文本嵌入服务。两者的结合,使得构建文本搜索系统变得更加高效和简单。

希望这篇博客对你有所帮助,如果你有任何问题,欢迎在评论区留言!

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值