69 lines
2.5 KiB
C#
69 lines
2.5 KiB
C#
using Qdrant.Client;
|
|
using Qdrant.Client.Grpc;
|
|
using VectorSearchApp.Configuration;
|
|
using VectorSearchApp.Models;
|
|
|
|
namespace VectorSearchApp.Services;
|
|
|
|
public interface IQdrantService
|
|
{
|
|
Task InitializeCollectionAsync(CancellationToken cancellationToken = default);
|
|
Task StoreAddressAsync(Address address, float[] embedding, CancellationToken cancellationToken = default);
|
|
Task<List<AddressEmbedding>> SearchSimilarAddressesAsync(float[] queryEmbedding, int limit = 5, CancellationToken cancellationToken = default);
|
|
}
|
|
|
|
public class QdrantService : IQdrantService
|
|
{
|
|
private readonly QdrantClient _client;
|
|
private readonly string _collectionName;
|
|
private readonly int _vectorDimension;
|
|
|
|
public QdrantService(QdrantConfiguration config, int vectorDimension)
|
|
{
|
|
_client = new QdrantClient(config.Host, config.GrpcPort);
|
|
_collectionName = config.CollectionName;
|
|
_vectorDimension = vectorDimension;
|
|
}
|
|
|
|
public async Task InitializeCollectionAsync(CancellationToken cancellationToken = default)
|
|
{
|
|
var collections = await _client.ListCollectionsAsync(cancellationToken: cancellationToken);
|
|
|
|
if (!collections.Contains(_collectionName))
|
|
{
|
|
await _client.CreateCollectionAsync(_collectionName, new VectorParams
|
|
{
|
|
Size = (ulong)_vectorDimension,
|
|
Distance = Distance.Cosine
|
|
}, cancellationToken: cancellationToken);
|
|
}
|
|
}
|
|
|
|
public async Task StoreAddressAsync(Address address, float[] embedding, CancellationToken cancellationToken = default)
|
|
{
|
|
var point = new PointStruct
|
|
{
|
|
Id = new PointId { Uuid = address.Id.ToString() },
|
|
Vectors = embedding,
|
|
Payload =
|
|
{
|
|
["address"] = address.FullAddress,
|
|
["created_at"] = address.CreatedAt.ToString("O")
|
|
}
|
|
};
|
|
|
|
await _client.UpsertAsync(_collectionName, new[] { point }, cancellationToken: cancellationToken);
|
|
}
|
|
|
|
public async Task<List<AddressEmbedding>> SearchSimilarAddressesAsync(float[] queryEmbedding, int limit = 5, CancellationToken cancellationToken = default)
|
|
{
|
|
var results = await _client.SearchAsync(_collectionName, queryEmbedding, limit: (ulong)limit, cancellationToken: cancellationToken);
|
|
|
|
return results.Select(r => new AddressEmbedding
|
|
{
|
|
Id = Guid.Parse(r.Id.Uuid),
|
|
FullAddress = r.Payload["address"].StringValue,
|
|
Vector = Array.Empty<float>()
|
|
}).ToList();
|
|
}
|
|
} |