diff --git a/MEVD/src/AzureDocumentDB/AzureDocumentDB.csproj b/MEVD/src/AzureDocumentDB/AzureDocumentDB.csproj index d5dc7e5..244c119 100644 --- a/MEVD/src/AzureDocumentDB/AzureDocumentDB.csproj +++ b/MEVD/src/AzureDocumentDB/AzureDocumentDB.csproj @@ -1,7 +1,7 @@  - 1.0.2 + 1.1.0 CommunityToolkit.VectorData.AzureDocumentDB $(AssemblyName) net10.0;net8.0;netstandard2.1;net472 diff --git a/MEVD/src/AzureDocumentDB/DocumentDBCollection.cs b/MEVD/src/AzureDocumentDB/DocumentDBCollection.cs index 88902da..4dab057 100644 --- a/MEVD/src/AzureDocumentDB/DocumentDBCollection.cs +++ b/MEVD/src/AzureDocumentDB/DocumentDBCollection.cs @@ -415,7 +415,7 @@ _ when vectorProperty.EmbeddingGenerationDispatcher is not null () => _mongoCollection.AggregateAsync(pipeline, cancellationToken: cancellationToken)).ConfigureAwait(false); using var errorHandlingAsyncCursor = new ErrorHandlingAsyncCursor(cursor, _collectionMetadata, OperationName); - await foreach (var result in EnumerateAndMapSearchResultsAsync(errorHandlingAsyncCursor, options, cancellationToken).ConfigureAwait(false)) + await foreach (var result in EnumerateAndMapSearchResultsAsync(errorHandlingAsyncCursor, options, vectorProperty.DistanceFunction, cancellationToken).ConfigureAwait(false)) { yield return result; } @@ -562,6 +562,7 @@ private async Task> FindAsync( private async IAsyncEnumerable> EnumerateAndMapSearchResultsAsync( ErrorHandlingAsyncCursor cursor, MEVD.VectorSearchOptions searchOptions, + string? distanceFunction, [EnumeratorCancellation] CancellationToken cancellationToken) { var skipCounter = 0; @@ -572,7 +573,7 @@ private async IAsyncEnumerable> EnumerateAndMapSearc { if (skipCounter >= searchOptions.Skip) { - var score = response[ScorePropertyName].AsDouble; + var score = DocumentDBCollectionSearchMapping.GetScore(response[ScorePropertyName].AsDouble, distanceFunction); var record = _mapper.MapFromStorageToDataModel(response[DocumentPropertyName].AsBsonDocument, includeVectors: searchOptions.IncludeVectors); yield return new VectorSearchResult(record, score); diff --git a/MEVD/src/AzureDocumentDB/DocumentDBCollectionCreateMapping.cs b/MEVD/src/AzureDocumentDB/DocumentDBCollectionCreateMapping.cs index 467368b..06c1e91 100644 --- a/MEVD/src/AzureDocumentDB/DocumentDBCollectionCreateMapping.cs +++ b/MEVD/src/AzureDocumentDB/DocumentDBCollectionCreateMapping.cs @@ -120,7 +120,7 @@ private static string GetIndexKind(string? indexKind, string vectorPropertyName) private static string GetDistanceFunction(string? distanceFunction, string vectorPropertyName) => DocumentDBCollectionSearchMapping.GetVectorPropertyDistanceFunction(distanceFunction) switch { - DistanceFunction.CosineDistance => "COS", + DistanceFunction.CosineDistance or DistanceFunction.CosineSimilarity => "COS", DistanceFunction.DotProductSimilarity => "IP", DistanceFunction.EuclideanDistance => "L2", _ => throw new NotSupportedException($"Distance function '{distanceFunction}' for {nameof(VectorStoreVectorProperty)} '{vectorPropertyName}' is not supported by the Azure DocumentDB VectorStore.") diff --git a/MEVD/src/AzureDocumentDB/DocumentDBCollectionSearchMapping.cs b/MEVD/src/AzureDocumentDB/DocumentDBCollectionSearchMapping.cs index b343d4f..0fe6cf2 100644 --- a/MEVD/src/AzureDocumentDB/DocumentDBCollectionSearchMapping.cs +++ b/MEVD/src/AzureDocumentDB/DocumentDBCollectionSearchMapping.cs @@ -96,18 +96,27 @@ public static BsonDocument GetProjectionQuery(string scorePropertyName, string d }; } + /// Returns the score of a result from the search score that Azure DocumentDB returns. + /// + /// For COS Azure DocumentDB returns the cosine similarity: CosineSimilarity returns it as it is, and CosineDistance returns + /// 1 - similarity. IP and L2 return the dot product and the Euclidean distance as they are. + /// + public static double GetScore(double searchScore, string? distanceFunction) + => GetVectorPropertyDistanceFunction(distanceFunction) == DistanceFunction.CosineDistance ? 1 - searchScore : searchScore; + /// Returns a $match stage to filter results by score threshold. /// - /// For COS and IP Azure DocumentDB returns a similarity score where higher values mean more similar, - /// so we filter with $gte. For L2 it returns a distance where lower values mean more similar, - /// so we filter with $lte. + /// The stage filters the search score that Azure DocumentDB returns. For COS and IP higher values mean more similar, + /// so we filter with $gte; for CosineDistance a distance of at most the threshold is a similarity of at least 1 - threshold. + /// For L2 it returns a distance where lower values mean more similar, so we filter with $lte. /// public static BsonDocument GetScoreThresholdMatchQuery(string scorePropertyName, double scoreThreshold, string? distanceFunction) { - var comparisonOperator = GetVectorPropertyDistanceFunction(distanceFunction) switch + var (comparisonOperator, searchScore) = GetVectorPropertyDistanceFunction(distanceFunction) switch { - DistanceFunction.CosineDistance or DistanceFunction.DotProductSimilarity => "$gte", - DistanceFunction.EuclideanDistance => "$lte", + DistanceFunction.CosineDistance => ("$gte", 1 - scoreThreshold), + DistanceFunction.CosineSimilarity or DistanceFunction.DotProductSimilarity => ("$gte", scoreThreshold), + DistanceFunction.EuclideanDistance => ("$lte", scoreThreshold), _ => throw new NotSupportedException($"Score threshold is not supported for distance function '{distanceFunction}'.") }; @@ -116,7 +125,7 @@ public static BsonDocument GetScoreThresholdMatchQuery(string scorePropertyName, { "$match", new BsonDocument { - { scorePropertyName, new BsonDocument { { comparisonOperator, scoreThreshold } } } + { scorePropertyName, new BsonDocument { { comparisonOperator, searchScore } } } } } }; diff --git a/MEVD/test/AzureDocumentDB.ConformanceTests/DocumentDBDistanceFunctionTests.cs b/MEVD/test/AzureDocumentDB.ConformanceTests/DocumentDBDistanceFunctionTests.cs index 8cd94cd..fe2504b 100644 --- a/MEVD/test/AzureDocumentDB.ConformanceTests/DocumentDBDistanceFunctionTests.cs +++ b/MEVD/test/AzureDocumentDB.ConformanceTests/DocumentDBDistanceFunctionTests.cs @@ -11,7 +11,6 @@ namespace AzureDocumentDB.ConformanceTests; public class DocumentDBDistanceFunctionTests(DocumentDBDistanceFunctionTests.Fixture fixture) : DistanceFunctionTests(fixture), IClassFixture { - public override Task CosineSimilarity() => Assert.ThrowsAsync(base.CosineSimilarity); public override Task EuclideanSquaredDistance() => Assert.ThrowsAsync(base.EuclideanSquaredDistance); public override Task NegativeDotProductSimilarity() => Assert.ThrowsAsync(base.NegativeDotProductSimilarity); public override Task HammingDistance() => Assert.ThrowsAsync(base.HammingDistance); @@ -20,7 +19,5 @@ public class DocumentDBDistanceFunctionTests(DocumentDBDistanceFunctionTests.Fix public new class Fixture() : DistanceFunctionTests.Fixture { public override TestStore TestStore => DocumentDBTestStore.Instance; - - public override bool AssertScores { get; } = false; } } diff --git a/MEVD/test/AzureDocumentDB.UnitTests/DocumentDBCollectionDistanceFunctionTests.cs b/MEVD/test/AzureDocumentDB.UnitTests/DocumentDBCollectionDistanceFunctionTests.cs new file mode 100644 index 0000000..10efa7c --- /dev/null +++ b/MEVD/test/AzureDocumentDB.UnitTests/DocumentDBCollectionDistanceFunctionTests.cs @@ -0,0 +1,55 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System; +using System.Collections.Generic; +using CommunityToolkit.VectorData.AzureDocumentDB; +using Microsoft.Extensions.VectorData; +using Microsoft.Extensions.VectorData.ProviderServices; +using MongoDB.Bson; +using Xunit; + +namespace AzureDocumentDB.UnitTests; + +public sealed class DocumentDBCollectionDistanceFunctionTests +{ + [Theory] + [InlineData(DistanceFunction.CosineDistance, "COS")] + [InlineData(DistanceFunction.CosineSimilarity, "COS")] + [InlineData(DistanceFunction.DotProductSimilarity, "IP")] + [InlineData(DistanceFunction.EuclideanDistance, "L2")] + [InlineData(null, "COS")] + public void IndexUsesTheSimilarityOfTheDistanceFunction(string? distanceFunction, string expectedSimilarity) + { + var vectorProperty = new VectorPropertyModel("vector", typeof(ReadOnlyMemory)) { Dimensions = 4, DistanceFunction = distanceFunction }; + + var indexes = DocumentDBCollectionCreateMapping.GetVectorIndexes([vectorProperty], new HashSet(), numLists: 1, efConstruction: 64); + + Assert.Equal(expectedSimilarity, indexes[0]["cosmosSearchOptions"]["similarity"].AsString); + } + + // COS returns the cosine similarity: CosineDistance gives 1 - similarity, and its threshold keeps the distances up to it. + [Theory] + [InlineData(DistanceFunction.CosineDistance, 0.25, 0.75)] + [InlineData(null, 0.25, 0.75)] + [InlineData(DistanceFunction.CosineSimilarity, 0.75, 0.75)] + [InlineData(DistanceFunction.DotProductSimilarity, 0.75, 0.75)] + [InlineData(DistanceFunction.EuclideanDistance, 0.75, 0.75)] + public void ScoreFollowsTheDistanceFunction(string? distanceFunction, double expectedScore, double searchScore) + => Assert.Equal(expectedScore, DocumentDBCollectionSearchMapping.GetScore(searchScore, distanceFunction), precision: 10); + + [Theory] + [InlineData(DistanceFunction.CosineDistance, 0.25, "$gte", 0.75)] + [InlineData(null, 0.25, "$gte", 0.75)] + [InlineData(DistanceFunction.CosineSimilarity, 0.75, "$gte", 0.75)] + [InlineData(DistanceFunction.DotProductSimilarity, 0.75, "$gte", 0.75)] + [InlineData(DistanceFunction.EuclideanDistance, 0.75, "$lte", 0.75)] + public void ScoreThresholdFollowsTheDistanceFunction(string? distanceFunction, double scoreThreshold, string expectedOperator, double expectedSearchScore) + { + var match = DocumentDBCollectionSearchMapping.GetScoreThresholdMatchQuery("score", scoreThreshold, distanceFunction); + + var condition = match["$match"]["score"].AsBsonDocument; + Assert.Equal(expectedOperator, condition.GetElement(0).Name); + Assert.Equal(expectedSearchScore, condition[0].AsDouble, precision: 10); + } +} diff --git a/MEVD/test/AzureDocumentDB.UnitTests/DocumentDBCollectionTests.cs b/MEVD/test/AzureDocumentDB.UnitTests/DocumentDBCollectionTests.cs index 3d2f588..b3f406f 100644 --- a/MEVD/test/AzureDocumentDB.UnitTests/DocumentDBCollectionTests.cs +++ b/MEVD/test/AzureDocumentDB.UnitTests/DocumentDBCollectionTests.cs @@ -561,7 +561,9 @@ public async Task SearchReturnsRecordWithScoreAsync() Assert.NotNull(result); Assert.Equal("key", result.Record.HotelId); Assert.Equal("Test Name", result.Record.HotelName); - Assert.Equal(0.99f, result.Score); + + // The model uses CosineDistance: the score is 1 - the cosine similarity that COS returns. + Assert.Equal(1 - 0.99f, result.Score!.Value, precision: 6); } public static TheoryData, string, bool> CollectionExistsData => new()