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); } } }