using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.Logging;
using Qdrant.Client;
using Qdrant.Client.Grpc;
using Polly;
using Polly.Registry;
using NexusReader.Application.Abstractions.Persistence;
namespace NexusReader.Infrastructure.Persistence;
///
/// Infrastructure implementation of utilizing
/// and to execute semantic vector queries.
///
internal sealed class VectorSearchStore : IVectorSearchStore
{
private readonly QdrantClient _qdrantClient;
private readonly IEmbeddingGenerator> _embeddingGenerator;
private readonly ResiliencePipeline _retryPipeline;
private readonly ILogger _logger;
public VectorSearchStore(
QdrantClient qdrantClient,
IEmbeddingGenerator> embeddingGenerator,
ResiliencePipelineProvider pipelineProvider,
ILogger logger)
{
_qdrantClient = qdrantClient;
_embeddingGenerator = embeddingGenerator;
_retryPipeline = pipelineProvider.GetPipeline("ai-retry");
_logger = logger;
}
///
public async Task> SearchGlobalAsync(string queryText, string tenantId, int limit, CancellationToken cancellationToken = default)
{
var queryVector = await GenerateEmbeddingAsync(queryText, cancellationToken);
var filter = BuildTenantFilter(tenantId);
return await ExecuteSearchAsync(queryVector, filter, limit, cancellationToken);
}
///
public async Task> SearchLocalAsync(string queryText, string tenantId, List whitelistedBookIds, int limit, CancellationToken cancellationToken = default)
{
if (whitelistedBookIds == null || !whitelistedBookIds.Any())
{
return new List();
}
var queryVector = await GenerateEmbeddingAsync(queryText, cancellationToken);
var filter = BuildTenantFilter(tenantId);
var whitelistFilter = new Qdrant.Client.Grpc.Filter();
foreach (var bookId in whitelistedBookIds)
{
whitelistFilter.Should.Add(new Qdrant.Client.Grpc.Condition
{
Field = new Qdrant.Client.Grpc.FieldCondition
{
Key = "ebookId",
Match = new Qdrant.Client.Grpc.Match { Text = bookId.ToString() }
}
});
}
filter.Must.Add(new Qdrant.Client.Grpc.Condition { Filter = whitelistFilter });
return await ExecuteSearchAsync(queryVector, filter, limit, cancellationToken);
}
///
public async Task> SearchGlobalExcludeAsync(string queryText, string tenantId, Guid excludeBookId, int limit, CancellationToken cancellationToken = default)
{
var queryVector = await GenerateEmbeddingAsync(queryText, cancellationToken);
var filter = BuildTenantFilter(tenantId);
// Exclude current book
filter.MustNot.Add(new Qdrant.Client.Grpc.Condition
{
Field = new Qdrant.Client.Grpc.FieldCondition
{
Key = "ebookId",
Match = new Qdrant.Client.Grpc.Match { Text = excludeBookId.ToString() }
}
});
return await ExecuteSearchAsync(queryVector, filter, limit, cancellationToken);
}
private async Task GenerateEmbeddingAsync(string text, CancellationToken cancellationToken)
{
if (string.IsNullOrWhiteSpace(text))
{
_logger.LogWarning("[VectorSearchStore] Attempted to generate embedding from empty text. Returning zero vector.");
return Array.Empty();
}
var sw = Stopwatch.StartNew();
var response = await _retryPipeline.ExecuteAsync(async ct =>
await _embeddingGenerator.GenerateAsync(
new[] { text },
new EmbeddingGenerationOptions { Dimensions = 768 },
cancellationToken: ct), cancellationToken);
sw.Stop();
_logger.LogDebug("[VectorSearchStore] Embedding generated in {ElapsedMs}ms for text of {Length} chars.", sw.ElapsedMilliseconds, text.Length);
return response.First().Vector.ToArray();
}
private Qdrant.Client.Grpc.Filter BuildTenantFilter(string tenantId)
{
var filter = new Qdrant.Client.Grpc.Filter();
var tenantFilter = new Qdrant.Client.Grpc.Filter();
tenantFilter.Should.Add(new Qdrant.Client.Grpc.Condition
{
Field = new Qdrant.Client.Grpc.FieldCondition
{
Key = "tenantId",
Match = new Qdrant.Client.Grpc.Match { Text = tenantId }
}
});
tenantFilter.Should.Add(new Qdrant.Client.Grpc.Condition
{
Field = new Qdrant.Client.Grpc.FieldCondition
{
Key = "tenantId",
Match = new Qdrant.Client.Grpc.Match { Text = "global" }
}
});
filter.Must.Add(new Qdrant.Client.Grpc.Condition { Filter = tenantFilter });
return filter;
}
private async Task> ExecuteSearchAsync(float[] queryVector, Qdrant.Client.Grpc.Filter filter, int limit, CancellationToken cancellationToken)
{
if (queryVector.Length == 0)
{
_logger.LogWarning("[VectorSearchStore] Empty query vector — skipping Qdrant search.");
return new List();
}
try
{
await EnsureCollectionExistsAsync("knowledge_units", cancellationToken);
var sw = Stopwatch.StartNew();
var response = await _qdrantClient.SearchAsync(
collectionName: "knowledge_units",
vector: queryVector,
filter: filter,
limit: (ulong)limit,
cancellationToken: cancellationToken
);
sw.Stop();
_logger.LogInformation("[VectorSearchStore] Qdrant search returned {Count} results in {ElapsedMs}ms.", response.Count, sw.ElapsedMilliseconds);
return response.Select(point =>
{
var content = point.Payload.TryGetValue("content", out var cv) ? cv.StringValue : string.Empty;
var ebookId = point.Payload.TryGetValue("ebookId", out var ev) ? ev.StringValue : string.Empty;
var metadataJson = point.Payload.TryGetValue("metadataJson", out var mv) ? mv.StringValue : string.Empty;
var bookTitle = point.Payload.TryGetValue("bookTitle", out var btv) ? btv.StringValue : string.Empty;
var chapterTitle = point.Payload.TryGetValue("chapterTitle", out var ctv) ? ctv.StringValue : string.Empty;
return new VectorChunk(content, ebookId, point.Score, metadataJson, bookTitle, chapterTitle);
}).ToList();
}
catch (Exception ex)
{
_logger.LogError(ex, "[VectorSearchStore] Qdrant search execution failed.");
throw;
}
}
private async Task EnsureCollectionExistsAsync(string collectionName, CancellationToken cancellationToken)
{
try
{
var exists = await _qdrantClient.CollectionExistsAsync(collectionName, cancellationToken);
if (!exists)
{
_logger.LogInformation("[VectorSearchStore] Collection '{CollectionName}' does not exist — creating.", collectionName);
await _qdrantClient.CreateCollectionAsync(
collectionName: collectionName,
vectorsConfig: new Qdrant.Client.Grpc.VectorParams
{
Size = 768,
Distance = Distance.Cosine
},
cancellationToken: cancellationToken
);
_logger.LogInformation("[VectorSearchStore] Collection '{CollectionName}' created successfully.", collectionName);
}
}
catch (Exception ex)
{
// Log concurrent creation conflicts (e.g., AlreadyExists gRPC status) but do not propagate.
_logger.LogWarning(ex, "[VectorSearchStore] Non-fatal error while ensuring collection '{CollectionName}' exists. Possible concurrent creation.", collectionName);
}
}
}