Files
VectorDBDemo/MilvusDemo/MilvusService.cs
2025-10-17 02:39:47 +08:00

144 lines
4.9 KiB
C#

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using Milvus.Client;
namespace MilvusDemo
{
/// <summary>
/// 提供与Milvus向量数据库交互的服务类
/// </summary>
public class MilvusService
{
private readonly MilvusClient _milvusClient;
private readonly string _collectionName;
private readonly int _dim;
/// <summary>
/// 初始化MilvusService的新实例
/// </summary>
/// <param name="host">Milvus服务器主机名</param>
/// <param name="port">Milvus服务器端口</param>
/// <param name="collectionName">集合名称</param>
/// <param name="dim">向量维度</param>
public MilvusService(string host, int port, string collectionName, int dim)
{
_milvusClient = new MilvusClient(host, port);
_collectionName = collectionName;
_dim = dim;
}
/// <summary>
/// 初始化集合,如果存在则删除并重新创建
/// </summary>
public async Task InitializeCollectionAsync()
{
// 检查集合是否存在,如果存在则删除
if (await _milvusClient.HasCollectionAsync(_collectionName))
{
var collectionToDrop = _milvusClient.GetCollection(_collectionName);
await collectionToDrop.DropAsync();
Console.WriteLine($"Collection '{_collectionName}' dropped.");
}
// 创建新的集合
await _milvusClient.CreateCollectionAsync(
_collectionName,
new[]
{
FieldSchema.Create<long>("id", isPrimaryKey: true),
FieldSchema.CreateVarchar("varchar", 256),
FieldSchema.CreateFloatVector("vector", _dim)
});
Console.WriteLine($"Collection '{_collectionName}' created.");
}
/// <summary>
/// 创建索引并加载集合
/// </summary>
public async Task CreateIndexAndLoadAsync()
{
var collection = _milvusClient.GetCollection(_collectionName);
// 创建索引
await collection.CreateIndexAsync("vector", IndexType.Flat, SimilarityMetricType.L2);
Console.WriteLine("Index created.");
// 加载集合
await collection.LoadAsync();
Console.WriteLine("Collection loaded.");
}
/// <summary>
/// 插入随机向量数据
/// </summary>
/// <param name="count">要插入的向量数量</param>
public async Task InsertRandomDataAsync(int count)
{
var collection = _milvusClient.GetCollection(_collectionName);
var random = new Random();
// 生成随机向量
var vectors = new List<ReadOnlyMemory<float>>();
for (int i = 0; i < count; i++)
{
var vector = new float[_dim];
for (int j = 0; j < _dim; j++)
{
vector[j] = (float)random.NextDouble();
}
vectors.Add(vector);
}
// 生成ID和字符串字段
var ids = Enumerable.Range(0, count).Select(i => (long)i).ToList();
var varchars = ids.Select(i => $"varchar_{i}").ToList();
// 打包数据
var insertData = new List<FieldData>
{
FieldData.Create("id", ids),
FieldData.Create("varchar", varchars),
FieldData.CreateFloatVector("vector", vectors)
};
// 执行插入
await collection.InsertAsync(insertData);
Console.WriteLine("Data inserted.");
}
/// <summary>
/// 执行向量相似度搜索
/// </summary>
/// <param name="topK">返回最相似的结果数量</param>
/// <returns>搜索结果</returns>
public async Task<SearchResults> SearchVectorAsync(int topK = 5)
{
var collection = _milvusClient.GetCollection(_collectionName);
var random = new Random();
// 生成随机目标向量
var targetVector = new float[_dim];
for (int j = 0; j < _dim; j++)
{
targetVector[j] = (float)random.NextDouble();
}
// 设置搜索参数
var searchParameters = new SearchParameters();
searchParameters.OutputFields.Add("id");
searchParameters.OutputFields.Add("varchar");
// 执行搜索
var searchResults = await collection.SearchAsync(
"vector",
new[] { (ReadOnlyMemory<float>)targetVector },
SimilarityMetricType.L2,
topK,
searchParameters);
return searchResults;
}
}
}