CosmosNoSqlCollectionQueryBuilder.cs

// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Linq.Expressions;
using System.Text;
using Microsoft.Azure.Cosmos;
using Microsoft.Extensions.VectorData;
using Microsoft.Extensions.VectorData.ProviderServices;

namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql;

/// <summary>
/// Contains helpers to build queries for Azure CosmosDB NoSQL.
/// </summary>
internal static class CosmosNoSqlCollectionQueryBuilder
{
    /// <summary>
    /// Builds <see cref="QueryDefinition"/> to get items from Azure CosmosDB NoSQL using vector search.
    /// </summary>
    public static QueryDefinition BuildSearchQuery<TRecord>(
        object vector,
        ICollection<string>? keywords,
        CollectionModel model,
        string vectorPropertyName,
        string? distanceFunction,
        string? textPropertyName,
        string scorePropertyName,
        Expression<Func<TRecord, bool>>? filter,
        double? scoreThreshold,
        int top,
        int skip,
        bool includeVectors)
    {
        Verify.NotNull(vector);

        const string VectorVariableName = "@vector";
        const string KeywordVariablePrefix = "@keyword";

        var tableVariableName = CosmosNoSqlConstants.ContainerAlias;

        IEnumerable<PropertyModel> projectionProperties = model.Properties;
        if (!includeVectors)
        {
            projectionProperties = projectionProperties.Where(p => p is not VectorPropertyModel);
        }
        var fieldsArgument = projectionProperties.Select(p => GeneratePropertyAccess(tableVariableName, p.StorageName));
        var vectorDistanceArgument = $"VectorDistance({GeneratePropertyAccess(tableVariableName, vectorPropertyName)}, {VectorVariableName})";
        var vectorDistanceArgumentWithAlias = $"{vectorDistanceArgument} AS {scorePropertyName}";

        var selectClauseArguments = string.Join(",", [.. fieldsArgument, vectorDistanceArgumentWithAlias]);

        // Build filter object.
        var (filterClause, filterParameters) = filter is not null
            ? new CosmosNoSqlFilterTranslator().Translate(filter, model)
            : ((string?)null, new Dictionary<string, object?>());

        var queryParameters = new Dictionary<string, object?>
        {
            // byte[] and ReadOnlyMemory<byte> are serialized as base64 by System.Text.Json, which causes problems for the VectorDistance function.
            // Convert to int[] which serializes as a JSON array of numbers.
            [VectorVariableName] = vector switch
            {
                ReadOnlyMemory<byte> byteMemory => ConvertToIntArray(byteMemory.Span),
                _ => vector
            }
        };

        string? fullTextScoreArgument = null;
        if (textPropertyName is not null && keywords is not null)
        {
            var fullTextScoreBuilder = new StringBuilder();
            fullTextScoreBuilder.Append($"FullTextScore({GeneratePropertyAccess(tableVariableName, textPropertyName)}");
            var i = 0;
            foreach (var keyword in keywords)
            {
                var paramName = $"{KeywordVariablePrefix}{i}";
                fullTextScoreBuilder.Append(", ").Append(paramName);
                queryParameters[paramName] = keyword;
                i++;
            }
            fullTextScoreBuilder.Append(')');
            fullTextScoreArgument = fullTextScoreBuilder.ToString();
        }

        var rankingArgument = fullTextScoreArgument is null ? vectorDistanceArgument : $"RANK RRF({vectorDistanceArgument}, {fullTextScoreArgument})";

        // Add score threshold filter if specified.
        // For similarity functions (CosineSimilarity, DotProductSimilarity), higher scores are better, so filter with >=.
        // For distance functions (EuclideanDistance), lower scores are better, so filter with <=.
        const string ScoreThresholdVariableName = "@scoreThreshold";
        string? scoreThresholdClause = null;
        if (scoreThreshold.HasValue)
        {
            var comparisonOperator = distanceFunction switch
            {
                Microsoft.Extensions.VectorData.DistanceFunction.CosineSimilarity => ">=",
                Microsoft.Extensions.VectorData.DistanceFunction.DotProductSimilarity => ">=",
                Microsoft.Extensions.VectorData.DistanceFunction.EuclideanDistance => "<=",
                _ => throw new NotSupportedException($"Score threshold is not supported for distance function '{distanceFunction}'.")
            };
            scoreThresholdClause = $"{vectorDistanceArgument} {comparisonOperator} {ScoreThresholdVariableName}";
            queryParameters[ScoreThresholdVariableName] = scoreThreshold.Value;
        }

        // If Offset is not configured, use Top parameter instead of Limit/Offset
        // since it's more optimized. Hybrid search doesn't allow top to be passed as a parameter
        // so directly add it to the query here.
        var topArgument = skip == 0 ? $"TOP {top} " : string.Empty;

        var builder = new StringBuilder();

        builder.AppendLine($"SELECT {topArgument}{selectClauseArguments}");
        builder.AppendLine($"FROM {tableVariableName}");

        if (filterClause is not null || scoreThresholdClause is not null)
        {
            builder.Append("WHERE ");

            if (filterClause is not null)
            {
                builder.Append(filterClause);
                if (scoreThresholdClause is not null)
                {
                    builder.Append(" AND ");
                }
            }

            if (scoreThresholdClause is not null)
            {
                builder.Append(scoreThresholdClause);
            }

            builder.AppendLine();
        }

        builder.AppendLine($"ORDER BY {rankingArgument}");

        if (string.IsNullOrEmpty(topArgument))
        {
            // Hybrid search doesn't allow offset and limit to be passed as parameters
            // so directly add it to the query here.
            builder.AppendLine($"OFFSET {skip} LIMIT {top}");
        }

        var queryDefinition = new QueryDefinition(builder.ToString());

        if (filterParameters is { Count: > 0 })
        {
            queryParameters = queryParameters.Union(filterParameters).ToDictionary(k => k.Key, v => v.Value);
        }

        foreach (var queryParameter in queryParameters)
        {
            queryDefinition.WithParameter(queryParameter.Key, queryParameter.Value);
        }

        return queryDefinition;
    }

    internal static QueryDefinition BuildSearchQuery<TRecord>(
        CollectionModel model,
        string whereClause, Dictionary<string, object?> filterParameters,
        FilteredRecordRetrievalOptions<TRecord> filterOptions,
        int top)
    {
        var tableVariableName = CosmosNoSqlConstants.ContainerAlias;

        IEnumerable<PropertyModel> projectionProperties = model.Properties;
        if (!filterOptions.IncludeVectors)
        {
            projectionProperties = projectionProperties.Where(p => p is not VectorPropertyModel);
        }

        var fieldsArgument = projectionProperties.Select(field => GeneratePropertyAccess(tableVariableName, field.StorageName));

        var selectClauseArguments = string.Join(",", [.. fieldsArgument]);

        // If Offset is not configured, use Top parameter instead of Limit/Offset
        // since it's more optimized.
        var topArgument = filterOptions.Skip == 0 ? $"TOP {top} " : string.Empty;

        var builder = new StringBuilder();

        builder.AppendLine($"SELECT {topArgument}{selectClauseArguments}");
        builder.AppendLine($"FROM {tableVariableName}");
        builder.Append("WHERE ").AppendLine(whereClause);

        var orderBy = filterOptions.OrderBy?.Invoke(new()).Values;
        if (orderBy is { Count: > 0 })
        {
            builder.Append("ORDER BY ");

            foreach (var sortInfo in orderBy)
            {
                builder
                    .Append(GeneratePropertyAccess(tableVariableName, model.GetDataOrKeyProperty(sortInfo.PropertySelector).StorageName))
                    .Append(sortInfo.Ascending ? " ASC," : " DESC,");
            }

            builder.Length--; // remove the last comma
            builder.AppendLine();
        }

        if (string.IsNullOrEmpty(topArgument))
        {
            builder.AppendLine($"OFFSET {filterOptions.Skip} LIMIT {top}");
        }

        var queryDefinition = new QueryDefinition(builder.ToString());

        foreach (var queryParameter in filterParameters)
        {
            queryDefinition.WithParameter(queryParameter.Key, queryParameter.Value);
        }

        return queryDefinition;
    }

    #region private

    private static string GetStoragePropertyName(string propertyName, CollectionModel model)
    {
        if (!model.PropertyMap.TryGetValue(propertyName, out var property))
        {
            throw new InvalidOperationException($"Property name '{propertyName}' provided as part of the filter clause is not a valid property name.");
        }

        return property.StorageName;
    }

    /// <summary>
    /// Escapes a JSON property name for use in Cosmos NoSQL SQL queries.
    /// JSON property names within bracket notation need backslash and double-quote escaping.
    /// </summary>
    private static string EscapeJsonPropertyName(string propertyName)
        => propertyName.Replace(@"\", @"\\").Replace("\"", "\\\"");

    /// <summary>
    /// Generates a property access expression using bracket notation with proper escaping.
    /// </summary>
    private static string GeneratePropertyAccess(char alias, string propertyName)
        => $"{alias}[\"{EscapeJsonPropertyName(propertyName)}\"]";

    /// <summary>
    /// Converts a byte span to an int array.
    /// This is needed because byte[] and ReadOnlyMemory&lt;byte&gt; are serialized as base64 by System.Text.Json,
    /// which causes problems for the VectorDistance function.
    /// </summary>
    private static int[] ConvertToIntArray(ReadOnlySpan<byte> bytes)
    {
        var result = new int[bytes.Length];
        for (var i = 0; i < bytes.Length; i++)
        {
            result[i] = bytes[i];
        }
        return result;
    }

    #endregion
}