Nexus.Reader/src/NexusReader.Application/Queries/Library/SearchLibrarySemanticallyQuery.cs

using FluentResults;
using Mapster;
using MediatR;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.AI;
using NexusReader.Application.DTOs.AI;
using NexusReader.Application.Abstractions.Persistence;
using Pgvector;
using Pgvector.EntityFrameworkCore;
using System.Text.Json;

namespace NexusReader.Application.Queries.Library;

public record SearchLibrarySemanticallyQuery(string QueryText, string TenantId, int Limit = 5) 
    : IRequest<Result<List<SemanticSearchResultDto>>>;

public class SearchLibrarySemanticallyQueryHandler : IRequestHandler<SearchLibrarySemanticallyQuery, Result<List<SemanticSearchResultDto>>>
{
    private readonly IApplicationDbContext _dbContext;
    private readonly IEmbeddingGenerator<string, Embedding<float>> _embeddingGenerator;

    public SearchLibrarySemanticallyQueryHandler(
        IApplicationDbContext dbContext,
        IEmbeddingGenerator<string, Embedding<float>> embeddingGenerator)
    {
        _dbContext = dbContext;
        _embeddingGenerator = embeddingGenerator;
    }

    public async Task<Result<List<SemanticSearchResultDto>>> Handle(SearchLibrarySemanticallyQuery request, CancellationToken cancellationToken)
    {
        if (string.IsNullOrWhiteSpace(request.QueryText))
        {
            return Result.Fail("Query text cannot be empty.");
        }

        try
        {
            // 1. Generate embedding for user query
            var embeddingResponse = await _embeddingGenerator.GenerateAsync(new[] { request.QueryText }, cancellationToken: cancellationToken);
            var queryVector = new Vector(embeddingResponse.First().Vector.ToArray());

            // 2. Perform Cosine Similarity Search on Knowledge Units
            var candidates = await _dbContext.KnowledgeUnits
                .AsNoTracking()
                .Where(x => (x.TenantId == request.TenantId || x.TenantId == "global") && x.Vector != null)
                .OrderBy(x => x.Vector!.CosineDistance(queryVector))
                .Take(request.Limit)
                .ToListAsync(cancellationToken);

            if (!candidates.Any())
            {
                // Fallback to legacy cache if no granular units found
                var legacyResults = await _dbContext.SemanticKnowledgeCache
                    .AsNoTracking()
                    .Where(x => x.TenantId == request.TenantId && x.Vector != null)
                    .OrderBy(x => x.Vector!.CosineDistance(queryVector))
                    .Take(request.Limit)
                    .ToListAsync(cancellationToken);

                return Result.Ok(legacyResults.Select(r => new SemanticSearchResultDto
                {
                    ContentHash = r.ContentHash,
                    Snippet = r.OriginalText,
                    RelevanceScore = (float)(1 - r.Vector!.CosineDistance(queryVector))
                }).ToList());
            }

            // 3. Graph Expansion: Pull related units (e.g. Definitions, Next steps)
            var candidateIds = candidates.Select(c => c.Id).ToList();
            var links = await _dbContext.KnowledgeUnitLinks
                .AsNoTracking()
                .Where(l => candidateIds.Contains(l.SourceUnitId) && (l.RelationType == "Defines" || l.RelationType == "Next"))
                .ToListAsync(cancellationToken);

            var relatedIds = links.Select(l => l.TargetUnitId).Distinct().ToList();
            var relatedUnits = await _dbContext.KnowledgeUnits
                .AsNoTracking()
                .Where(u => relatedIds.Contains(u.Id))
                .ToDictionaryAsync(u => u.Id, cancellationToken);

            // 4. Mapping with Context Enrichment
            var dtos = candidates.Select(c => 
            {
                var dto = new SemanticSearchResultDto
                {
                    ContentHash = c.Id,
                    Snippet = c.Content,
                    UnitType = c.Type.ToString(),
                    RelevanceScore = (float)(1 - c.Vector!.CosineDistance(queryVector)),
                    Metadata = string.IsNullOrEmpty(c.MetadataJson) 
                        ? null 
                        : JsonSerializer.Deserialize<Dictionary<string, object>>(c.MetadataJson)
                };

                // Enrich snippet with definitions if present
                var unitLinks = links.Where(l => l.SourceUnitId == c.Id && l.RelationType == "Defines").ToList();
                if (unitLinks.Any())
                {
                    var definitions = unitLinks
                        .Where(l => relatedUnits.ContainsKey(l.TargetUnitId))
                        .Select(l => relatedUnits[l.TargetUnitId].Content);
                    dto.Snippet = $"[Context: {string.Join("; ", definitions)}]\n{dto.Snippet}";
                }

                return dto;
            }).ToList();

            return Result.Ok(dtos);
        }
        catch (Exception ex)
        {
            return Result.Fail(new Error("Failed to perform semantic search").CausedBy(ex));
        }
    }
}