ChatRAG/Data/SearchTextData.cs

using ChatRAG.Models;
using MongoDB.Driver;
using System.ComponentModel;

namespace ChatRAG.Data
{
    public class SearchTextData
    {
        private readonly IMongoCollection<TextoComEmbedding> _textsCollection;

        public SearchTextData(IMongoCollection<TextoComEmbedding> textsCollection)
        {
            _textsCollection = textsCollection;
        }

        // Busca otimizada com filtros pré-aplicados
        public async Task<List<ResultadoSimilaridade>> BuscarSimilaridadeOtimizada(
            double[] queryEmbedding,
            string projectId = null,
            string tipoDocumento = null,
            double similaridadeMinima = 0.3,
            int? limite = 3)
        {
            // 1. Aplica filtros primeiro para reduzir dataset
            var filterBuilder = Builders<TextoComEmbedding>.Filter;
            var filters = new List<FilterDefinition<TextoComEmbedding>>();

            if (!string.IsNullOrEmpty(projectId))
                filters.Add(filterBuilder.Eq(x => x.ProjetoId, projectId));

            if (!string.IsNullOrEmpty(tipoDocumento))
                filters.Add(filterBuilder.Eq(x => x.TipoDocumento, tipoDocumento));

            var finalFilter = filters.Any()
                ? filterBuilder.And(filters)
                : filterBuilder.Empty;

            // 2. Busca apenas os documentos filtrados
            var documentosFiltrados = await _textsCollection
                .Find(finalFilter)
                .ToListAsync();

            // 3. Calcula similaridade apenas nos documentos filtrados
            var resultados = documentosFiltrados
                .AsParallel() // Paralelização para performance
                .Select(doc => new ResultadoSimilaridade
                {
                    Documento = doc,
                    Similaridade = CalcularSimilaridadeCoseno(queryEmbedding, doc.Embedding)
                })
                .Where(r => r.Similaridade >= similaridadeMinima)
                .OrderByDescending(r => r.Similaridade)
                .ToList();

            // 4. Aplica limite se especificado
            if (limite.HasValue)
                resultados = resultados.Take(limite.Value).ToList();

            return resultados;
        }

        // Busca dinâmica baseada em threshold de similaridade
        public async Task<List<ResultadoSimilaridade>> BuscarPorSimilaridadeDinamica(
            double[] queryEmbedding,
            string projectId,
            double similaridadeMinima = 0.5,
            int limiteSuperior = 50) // Limite máximo para evitar sobrecarga
        {
            var resultados = await BuscarSimilaridadeOtimizada(
                queryEmbedding,
                projectId,
                similaridadeMinima: similaridadeMinima,
                limite: limiteSuperior
            );

            // Se não encontrar resultados suficientes, relaxa o threshold
            if (resultados.Count < 3)
            {
                resultados = await BuscarSimilaridadeOtimizada(
                    queryEmbedding,
                    projectId,
                    similaridadeMinima: similaridadeMinima * 0.7, // Reduz 30%
                    limite: limiteSuperior
                );
            }

            return resultados;
        }

        double CalcularSimilaridadeCoseno(double[] embedding1, double[] embedding2)
        {
            double dotProduct = 0.0;
            double normA = 0.0;
            double normB = 0.0;
            for (int i = 0; i < embedding1.Length; i++)
            {
                dotProduct += embedding1[i] * embedding2[i];
                normA += Math.Pow(embedding1[i], 2);
                normB += Math.Pow(embedding2[i], 2);
            }
            return dotProduct / (Math.Sqrt(normA) * Math.Sqrt(normB));
        }

    }

    public class ResultadoSimilaridade
    {
        public TextoComEmbedding Documento { get; set; }
        public double Similaridade { get; set; }
    }

    public class GrupoSimilaridade
    {
        public ResultadoSimilaridade Representante { get; set; }
        public List<ResultadoSimilaridade> Membros { get; set; }
        public int TotalMembros => Membros.Count;
    }
}