From a154d27361b01b29df8752fce7270c16ef12b8b3 Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Sun, 1 Feb 2026 12:28:11 +0100 Subject: [PATCH 1/4] [MEVD] Fix byte/sbyte embeddings for CosmosNoSQL Fixes #12079 --- .../CosmosNoSql/ByteArrayJsonConverter.cs | 96 +++++++++++++++++++ .../CosmosNoSqlCollectionQueryBuilder.cs | 23 ++++- .../CosmosNoSql/CosmosNoSqlMapper.cs | 74 +++++++------- .../CosmosNoSql/CosmosNoSqlModelBuilder.cs | 9 ++ .../CosmosNoSqlEmbeddingTypeTests.cs | 12 +-- 5 files changed, 175 insertions(+), 39 deletions(-) create mode 100644 dotnet/src/VectorData/CosmosNoSql/ByteArrayJsonConverter.cs diff --git a/dotnet/src/VectorData/CosmosNoSql/ByteArrayJsonConverter.cs b/dotnet/src/VectorData/CosmosNoSql/ByteArrayJsonConverter.cs new file mode 100644 index 000000000000..bccc80bc8adf --- /dev/null +++ b/dotnet/src/VectorData/CosmosNoSql/ByteArrayJsonConverter.cs @@ -0,0 +1,96 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Text.Json; +using System.Text.Json.Serialization; + +namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql; + +/// +/// A JSON converter for byte arrays that serializes them as JSON arrays of numbers +/// instead of base64-encoded strings. +/// +/// +/// This is needed because Cosmos DB's VectorDistance function requires vectors to be arrays of numbers, +/// not base64-encoded strings. +/// +internal sealed class ByteArrayJsonConverter : JsonConverter +{ + /// + public override byte[] Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options) + { + // Support reading both base64 strings (for backward compatibility) and number arrays + if (reader.TokenType == JsonTokenType.String) + { + return reader.GetBytesFromBase64(); + } + + if (reader.TokenType != JsonTokenType.StartArray) + { + throw new JsonException($"Expected StartArray or String token, got {reader.TokenType}"); + } + + var list = new List(); + while (reader.Read() && reader.TokenType != JsonTokenType.EndArray) + { + list.Add(reader.GetByte()); + } + return list.ToArray(); + } + + /// + public override void Write(Utf8JsonWriter writer, byte[] value, JsonSerializerOptions options) + { + writer.WriteStartArray(); + foreach (var b in value) + { + writer.WriteNumberValue(b); + } + writer.WriteEndArray(); + } +} + +/// +/// A JSON converter for of byte that serializes as JSON arrays of numbers +/// instead of base64-encoded strings. +/// +/// +/// This is needed because Cosmos DB's VectorDistance function requires vectors to be arrays of numbers, +/// not base64-encoded strings. +/// +internal sealed class ReadOnlyMemoryByteJsonConverter : JsonConverter> +{ + /// + public override ReadOnlyMemory Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options) + { + // Support reading both base64 strings (for backward compatibility) and number arrays + if (reader.TokenType == JsonTokenType.String) + { + return new ReadOnlyMemory(reader.GetBytesFromBase64()); + } + + if (reader.TokenType != JsonTokenType.StartArray) + { + throw new JsonException($"Expected StartArray or String token, got {reader.TokenType}"); + } + + var list = new List(); + while (reader.Read() && reader.TokenType != JsonTokenType.EndArray) + { + list.Add(reader.GetByte()); + } + return new ReadOnlyMemory(list.ToArray()); + } + + /// + public override void Write(Utf8JsonWriter writer, ReadOnlyMemory value, JsonSerializerOptions options) + { + writer.WriteStartArray(); + foreach (var b in value.Span) + { + writer.WriteNumberValue(b); + } + writer.WriteEndArray(); + } +} diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs index 968dcc313bdf..57523ea42232 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs @@ -81,7 +81,13 @@ public static QueryDefinition BuildSearchQuery( var queryParameters = new Dictionary { - [VectorVariableName] = vector + // byte[] and ReadOnlyMemory 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 byteMemory => ConvertToIntArray(byteMemory.Span), + _ => vector + } }; // Add score threshold filter if specified. @@ -363,5 +369,20 @@ private static string EscapeJsonPropertyName(string propertyName) private static string GeneratePropertyAccess(char alias, string propertyName) => $"{alias}[\"{EscapeJsonPropertyName(propertyName)}\"]"; + /// + /// Converts a byte span to an int array. + /// This is needed because byte[] and ReadOnlyMemory<byte> are serialized as base64 by System.Text.Json, + /// which causes problems for the VectorDistance function. + /// + private static int[] ConvertToIntArray(ReadOnlySpan bytes) + { + var result = new int[bytes.Length]; + for (var i = 0; i < bytes.Length; i++) + { + result[i] = bytes[i]; + } + return result; + } + #endregion } diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlMapper.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlMapper.cs index fe37efd9b89a..7b176c339b50 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlMapper.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlMapper.cs @@ -15,15 +15,29 @@ namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql; /// Class for mapping between a json node stored in Azure CosmosDB NoSQL and the consumer data model. /// /// The consumer data model to map to or from. -internal sealed class CosmosNoSqlMapper(CollectionModel model, JsonSerializerOptions? jsonSerializerOptions) - : ICosmosNoSqlMapper +internal sealed class CosmosNoSqlMapper : ICosmosNoSqlMapper where TRecord : class { - private readonly KeyPropertyModel _keyProperty = model.KeyProperty; + private readonly CollectionModel _model; + private readonly KeyPropertyModel _keyProperty; + private readonly JsonSerializerOptions _jsonSerializerOptions; + + public CosmosNoSqlMapper(CollectionModel model, JsonSerializerOptions? jsonSerializerOptions) + { + this._model = model; + this._keyProperty = model.KeyProperty; + + // Add byte array converters to serialize byte[] and ReadOnlyMemory as JSON arrays of numbers + // instead of base64-encoded strings, which is required for Cosmos DB vector operations. + var newOptions = new JsonSerializerOptions(jsonSerializerOptions ?? JsonSerializerOptions.Default); + newOptions.Converters.Add(new ByteArrayJsonConverter()); + newOptions.Converters.Add(new ReadOnlyMemoryByteJsonConverter()); + this._jsonSerializerOptions = newOptions; + } public JsonObject MapFromDataToStorageModel(TRecord dataModel, int recordIndex, IReadOnlyList?[]? generatedEmbeddings) { - var jsonObject = JsonSerializer.SerializeToNode(dataModel, jsonSerializerOptions)!.AsObject(); + var jsonObject = JsonSerializer.SerializeToNode(dataModel, this._jsonSerializerOptions)!.AsObject(); // The key property in Azure CosmosDB NoSQL is always named 'id'. // But the external JSON serializer used just above isn't aware of that, and will produce a JSON object with another name, taking into @@ -32,9 +46,9 @@ public JsonObject MapFromDataToStorageModel(TRecord dataModel, int recordIndex, // Go over the vector properties; inject any generated embeddings to overwrite the JSON serialized above. // Also, for Embedding properties we also need to overwrite with a simple array (since Embedding gets serialized as a complex object). - for (var i = 0; i < model.VectorProperties.Count; i++) + for (var i = 0; i < this._model.VectorProperties.Count; i++) { - var property = model.VectorProperties[i]; + var property = this._model.VectorProperties[i]; Embedding? embedding = generatedEmbeddings?[i]?[recordIndex] is Embedding ge ? ge : null; @@ -42,28 +56,17 @@ public JsonObject MapFromDataToStorageModel(TRecord dataModel, int recordIndex, { switch (Nullable.GetUnderlyingType(property.Type) ?? property.Type) { + // For raw array/memory types, JsonSerializer already serialized them correctly above + // (including byte[] thanks to our custom converter). Nothing more to do. case var t when t == typeof(ReadOnlyMemory): case var t2 when t2 == typeof(float[]): - // The .NET vector property is a ReadOnlyMemory or float[] (not an Embedding), which means that JsonSerializer - // already serialized it correctly above. - // In addition, there's no generated embedding (which would be an Embedding which we'd need to handle manually). - // So there's nothing for us to do. - continue; - - // byte/sbyte is a special case, since it gets serialized as base64 by default; handle manually here. case var t3 when t3 == typeof(ReadOnlyMemory): - embedding = new Embedding((ReadOnlyMemory)property.GetValueAsObject(dataModel)!); - break; case var t4 when t4 == typeof(byte[]): - embedding = new Embedding((byte[])property.GetValueAsObject(dataModel)!); - break; case var t5 when t5 == typeof(ReadOnlyMemory): - embedding = new Embedding((ReadOnlyMemory)property.GetValueAsObject(dataModel)!); - break; case var t6 when t6 == typeof(sbyte[]): - embedding = new Embedding((sbyte[])property.GetValueAsObject(dataModel)!); - break; + continue; + // For Embedding types, we need to extract the vector and serialize it as a simple array case var t when t == typeof(Embedding): case var t1 when t1 == typeof(Embedding): case var t2 when t2 == typeof(Embedding): @@ -117,24 +120,31 @@ public TRecord MapFromStorageToDataModel(JsonObject storageModel, bool includeVe if (includeVectors) { - foreach (var vectorProperty in model.VectorProperties) + foreach (var vectorProperty in this._model.VectorProperties) { - if (vectorProperty.Type == typeof(Embedding)) + var arrayNode = storageModel[vectorProperty.StorageName]; + if (arrayNode is null) { - var arrayNode = storageModel[vectorProperty.StorageName]; - if (arrayNode is not null) + continue; + } + + // Embedding is stored as a simple JSON array, so convert it to the expected object shape for deserialization + if (vectorProperty.Type == typeof(Embedding) + || vectorProperty.Type == typeof(Embedding) + || vectorProperty.Type == typeof(Embedding)) + { + storageModel[vectorProperty.StorageName] = new JsonObject { - var embeddingNode = new JsonObject - { - [nameof(Embedding.Vector)] = arrayNode.DeepClone() - }; - storageModel[vectorProperty.StorageName] = embeddingNode; - } + [nameof(Embedding<>.Vector)] = arrayNode.DeepClone() + }; } + + // For byte[], ReadOnlyMemory, sbyte[], ReadOnlyMemory, float[], ReadOnlyMemory, + // the custom converters (for byte) and default converters (for others) handle deserialization correctly. } } - return storageModel.Deserialize(jsonSerializerOptions)!; + return storageModel.Deserialize(this._jsonSerializerOptions)!; } #region private diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs index 803d4aa4a2c9..feb9beaddf9a 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs @@ -58,6 +58,15 @@ static bool IsValid(Type type) protected override bool IsVectorPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) => IsVectorPropertyTypeValidCore(type, out supportedTypes); + protected override Type? ResolveEmbeddingType( + VectorPropertyModel vectorProperty, + IEmbeddingGenerator embeddingGenerator, + Type? userRequestedEmbeddingType) + // Resolve embedding type for float, byte, and sbyte embedding generators. + => vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType) + ?? vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType) + ?? vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType); + internal static bool IsVectorPropertyTypeValidCore(Type type, [NotNullWhen(false)] out string? supportedTypes) { supportedTypes = SupportedVectorTypes; diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/TypeTests/CosmosNoSqlEmbeddingTypeTests.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/TypeTests/CosmosNoSqlEmbeddingTypeTests.cs index 2601cfbe6c7d..1bb0e73d20fa 100644 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/TypeTests/CosmosNoSqlEmbeddingTypeTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/TypeTests/CosmosNoSqlEmbeddingTypeTests.cs @@ -14,41 +14,41 @@ namespace CosmosNoSql.ConformanceTests.TypeTests; public class CosmosNoSqlEmbeddingTypeTests(CosmosNoSqlEmbeddingTypeTests.Fixture fixture) : EmbeddingTypeTests(fixture), IClassFixture { - [ConditionalFact(Skip = "TODO: Figure out why int8/uint8 embeddings aren't working (float32 is working), #12079")] + [ConditionalFact] public virtual Task ReadOnlyMemory_of_byte() => this.Test>( new ReadOnlyMemory([1, 2, 3]), new ReadOnlyMemoryEmbeddingGenerator([1, 2, 3]), vectorEqualityAsserter: (e, a) => Assert.Equal(e.Span.ToArray(), a.Span.ToArray())); - [ConditionalFact(Skip = "TODO: Figure out why int8/uint8 embeddings aren't working (float32 is working), #12079")] + [ConditionalFact] public virtual Task Embedding_of_byte() => this.Test>( new Embedding(new ReadOnlyMemory([1, 2, 3])), new ReadOnlyMemoryEmbeddingGenerator([1, 2, 3]), vectorEqualityAsserter: (e, a) => Assert.Equal(e.Vector.Span.ToArray(), a.Vector.Span.ToArray())); - [ConditionalFact(Skip = "TODO: Figure out why int8/uint8 embeddings aren't working (float32 is working), #12079")] + [ConditionalFact] public virtual Task Array_of_byte() => this.Test( [1, 2, 3], new ReadOnlyMemoryEmbeddingGenerator([1, 2, 3])); - [ConditionalFact(Skip = "TODO: Figure out why int8/uint8 embeddings aren't working (float32 is working), #12079")] + [ConditionalFact] public virtual Task ReadOnlyMemory_of_sbyte() => this.Test>( new ReadOnlyMemory([1, 2, 3]), new ReadOnlyMemoryEmbeddingGenerator([1, 2, 3]), vectorEqualityAsserter: (e, a) => Assert.Equal(e.Span.ToArray(), a.Span.ToArray())); - [ConditionalFact(Skip = "TODO: Figure out why int8/uint8 embeddings aren't working (float32 is working), #12079")] + [ConditionalFact] public virtual Task Embedding_of_sbyte() => this.Test>( new Embedding(new ReadOnlyMemory([1, 2, 3])), new ReadOnlyMemoryEmbeddingGenerator([1, 2, 3]), vectorEqualityAsserter: (e, a) => Assert.Equal(e.Vector.Span.ToArray(), a.Vector.Span.ToArray())); - [ConditionalFact(Skip = "TODO: Figure out why int8/uint8 embeddings aren't working (float32 is working), #12079")] + [ConditionalFact] public virtual Task Array_of_sbyte() => this.Test( [1, 2, 3], From 489f4034a2b9f64b165c4cf6ea22f72ebacbc4e3 Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Sat, 14 Feb 2026 16:51:54 +0100 Subject: [PATCH 2/4] Redo Cosmos key/partition key handling Fixes #11352 Fixes #11353 --- .../CosmosNoSql/CosmosNoSqlCollection.cs | 178 +++++++++-------- .../CosmosNoSqlCollectionOptions.cs | 22 ++- .../CosmosNoSqlCollectionQueryBuilder.cs | 43 ++-- .../CosmosNoSql/CosmosNoSqlCompositeKey.cs | 19 -- .../CosmosNoSqlDynamicCollection.cs | 10 +- .../VectorData/CosmosNoSql/CosmosNoSqlKey.cs | 183 ++++++++++++++++++ .../CosmosNoSql/CosmosNoSqlModelBuilder.cs | 4 +- .../CosmosNoSqlServiceCollectionExtensions.cs | 12 +- .../CosmosNoSqlCollectionOptionsTests.cs | 23 +-- .../CosmosNoSqlDependencyInjectionTests.cs | 39 ++-- .../CosmosNoSqlEmbeddingGenerationTests.cs | 156 +++++++++++++++ .../Support/CosmosNoSqlCollectionAdapter.cs | 103 ++++++++++ .../CosmosNoSqlDynamicCollectionAdapter.cs | 91 +++++++++ .../Support/CosmosNoSqlTestStore.cs | 59 ++++++ .../CosmosNoSqlCollectionQueryBuilderTests.cs | 14 +- .../CosmosNoSqlCollectionTests.cs | 119 +++++++----- .../CosmosNoSqlVectorStoreTests.cs | 8 +- .../CollectionManagementTests.cs | 2 +- .../DistanceFunctionTests.cs | 2 +- .../EmbeddingGenerationTests.cs | 2 +- .../FilterTests.cs | 2 +- .../IndexKindTests.cs | 2 +- .../ModelTests/DynamicModelTests.cs | 2 +- .../Support/TestStore.cs | 20 ++ .../VectorStoreCollectionFixtureBase.cs | 2 +- .../Support/VectorStoreFixture.cs | 16 ++ .../TypeTests/DataTypeTests.cs | 2 +- .../TypeTests/EmbeddingTypeTests.cs | 10 +- .../TypeTests/KeyTypeTests.cs | 4 +- 29 files changed, 902 insertions(+), 247 deletions(-) delete mode 100644 dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCompositeKey.cs create mode 100644 dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlKey.cs create mode 100644 dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlCollectionAdapter.cs create mode 100644 dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlDynamicCollectionAdapter.cs diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs index 34fd0d12b3e0..680fb47a2fbc 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs @@ -26,13 +26,13 @@ namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql; /// /// Service for storing and retrieving vector records, that uses Azure CosmosDB NoSQL as the underlying storage. /// -/// The data type of the record key. Must be . +/// The data type of the record key. Must be . /// The data model to use for adding, updating and retrieving data from storage. #pragma warning disable CA1711 // Identifiers should not have incorrect suffix public class CosmosNoSqlCollection : VectorStoreCollection, IKeywordHybridSearchable where TKey : notnull where TRecord : class -#pragma warning restore CA1711 // Identifiers should not have incorrect +#pragma warning restore CA1711 // Identifiers should not have incorrect suffix { /// Metadata about vector store record collection. private readonly VectorStoreCollectionMetadata _collectionMetadata; @@ -42,6 +42,7 @@ public class CosmosNoSqlCollection : VectorStoreCollectionThe default options for hybrid vector search. private static readonly HybridSearchOptions s_defaultKeywordVectorizedHybridSearchOptions = new(); + private readonly ClientWrapper _clientWrapper; /// that can be used to manage the collections in Azure CosmosDB NoSQL. @@ -50,9 +51,9 @@ public class CosmosNoSqlCollection : VectorStoreCollectionThe model for this collection. private readonly CollectionModel _model; - // TODO: Refactor this into the model (Co) - /// The property to use as partition key. - private readonly PropertyModel _partitionKeyProperty; + // TODO: Refactor this into the model + /// The properties to use as partition key (supports hierarchical partition keys up to 3 levels). + private readonly List _partitionKeyProperties; /// The mapper to use when mapping between the consumer data model and the Azure CosmosDB NoSQL record. private readonly ICosmosNoSqlMapper _mapper; @@ -132,23 +133,21 @@ internal CosmosNoSqlCollection( Func modelFactory, CosmosNoSqlCollectionOptions? options) { - try + // Validate TKey: must be CosmosNoSqlKey, or object (used by CosmosNoSqlDynamicCollection). + if (typeof(TKey) != typeof(object) && typeof(TKey) != typeof(CosmosNoSqlKey)) { - if (typeof(TKey) != typeof(string) - && typeof(TKey) != typeof(Guid) - && typeof(TKey) != typeof(CosmosNoSqlCompositeKey) - && typeof(TKey) != typeof(object)) - { - throw new NotSupportedException($"Only string, Guid and {nameof(CosmosNoSqlCompositeKey)} keys are supported."); - } + throw new NotSupportedException($"The key type must be CosmosNoSqlKey. Received: {typeof(TKey).Name}"); + } + try + { this._database = databaseProvider(clientWrapper.Client); if (clientWrapper.Client.ClientOptions?.UseSystemTextJsonSerializerWithOptions is null) { throw new ArgumentException( $"Property {nameof(CosmosClientOptions.UseSystemTextJsonSerializerWithOptions)} in CosmosClient.ClientOptions " + - $"is required to be configured for {nameof(CosmosNoSqlCollection)}."); + $"is required to be configured for {this.GetType().Name}."); } options ??= CosmosNoSqlCollectionOptions.Default; @@ -165,25 +164,27 @@ internal CosmosNoSqlCollection( ? (new CosmosNoSqlDynamicMapper(this._model, jsonSerializerOptions) as ICosmosNoSqlMapper)! : new CosmosNoSqlMapper(this._model, options.JsonSerializerOptions); - // Setup partition key property - if (options.PartitionKeyPropertyName is not null) + // Setup partition key properties (supports hierarchical partition keys up to 3 levels) + if (options.PartitionKeyPropertyNames.Count > 3) { - if (!this._model.PropertyMap.TryGetValue(options.PartitionKeyPropertyName, out var property)) + throw new ArgumentException("Cosmos DB supports at most 3 levels of hierarchical partition keys."); + } + + this._partitionKeyProperties = new List(options.PartitionKeyPropertyNames.Count); + + foreach (var propertyName in options.PartitionKeyPropertyNames) + { + if (!this._model.PropertyMap.TryGetValue(propertyName, out var property)) { - throw new ArgumentException($"Partition key property '{options.PartitionKeyPropertyName}' is not part of the record definition."); + throw new ArgumentException($"Partition key property '{propertyName}' is not part of the record definition."); } - if (property.Type != typeof(string)) + if (property.Type != typeof(string) && property.Type != typeof(bool) && property.Type != typeof(double) && property.Type != typeof(Guid)) { - throw new ArgumentException("Partition key property must be string."); + throw new ArgumentException($"Partition key property '{propertyName}' must be string, bool, double, or Guid."); } - this._partitionKeyProperty = property; - } - else - { - // If partition key is not provided, use key property as a partition key. - this._partitionKeyProperty = this._model.KeyProperty; + this._partitionKeyProperties.Add(property); } this._collectionMetadata = new() @@ -295,12 +296,7 @@ await this._database /// public override Task DeleteAsync(TKey key, CancellationToken cancellationToken = default) { - Verify.NotNull(key); - - var compositeKey = GetCompositeKeys([key]).Single(); - - Verify.NotNullOrWhiteSpace(compositeKey.RecordKey); - Verify.NotNullOrWhiteSpace(compositeKey.PartitionKey); + var (documentId, partitionKey) = this.GetDocumentIdAndPartitionKey(key); return this.RunOperationAsync("DeleteItem", async () => { @@ -308,7 +304,7 @@ public override Task DeleteAsync(TKey key, CancellationToken cancellationToken = { await this._database .GetContainer(this.Name) - .DeleteItemAsync(compositeKey.RecordKey, new PartitionKey(compositeKey.PartitionKey), cancellationToken: cancellationToken) + .DeleteItemAsync(documentId, partitionKey, cancellationToken: cancellationToken) .ConfigureAwait(false); return 0; } @@ -348,17 +344,15 @@ public override async IAsyncEnumerable GetAsync( throw new NotSupportedException(VectorDataStrings.IncludeVectorsNotSupportedWithEmbeddingGeneration); } - var compositeKeys = GetCompositeKeys(keys).ToList(); - if (compositeKeys.Count == 0) + var documentIds = keys.Select(k => this.GetDocumentIdAndPartitionKey(k).DocumentId).ToList(); + if (documentIds.Count == 0) { yield break; } var queryDefinition = CosmosNoSqlCollectionQueryBuilder.BuildSelectQuery( this._model, - this._model.KeyProperty.StorageName, - this._partitionKeyProperty.StorageName, - compositeKeys, + documentIds, includeVectors); await foreach (var jsonObject in this.GetItemsAsync(queryDefinition, OperationName, cancellationToken).ConfigureAwait(false)) @@ -409,25 +403,21 @@ private async Task UpsertCoreAsync(TRecord record, int recordIndex, IReadOnlyLis keyProperty.SetValue(record, Guid.NewGuid()); } + var partitionKey = this.BuildPartitionKey(record); + var jsonObject = this._mapper.MapFromDataToStorageModel(record, recordIndex, generatedEmbeddings); var keyValue = jsonObject.TryGetPropertyValue(this._model.KeyProperty.StorageName!, out var jsonKey) ? jsonKey?.ToString() : null; - var partitionKeyValue = jsonObject.TryGetPropertyValue(this._partitionKeyProperty.StorageName, out var jsonPartitionKey) ? jsonPartitionKey?.ToString() : null; if (string.IsNullOrWhiteSpace(keyValue)) { throw new ArgumentException($"Key property {this._model.KeyProperty.ModelName} is not initialized."); } - if (string.IsNullOrWhiteSpace(partitionKeyValue)) - { - throw new ArgumentException($"Partition key property {this._partitionKeyProperty.ModelName} is not initialized."); - } - await this.RunOperationAsync(OperationName, () => this._database .GetContainer(this.Name) - .UpsertItemAsync(jsonObject, new PartitionKey(partitionKeyValue), cancellationToken: cancellationToken)) + .UpsertItemAsync(jsonObject, partitionKey, cancellationToken: cancellationToken)) .ConfigureAwait(false); } @@ -526,7 +516,7 @@ public override async IAsyncEnumerable> SearchAsync< this._model, vectorProperty.StorageName, vectorProperty.DistanceFunction, - null, + textPropertyName: null, ScorePropertyName, options.OldFilter, options.Filter, @@ -685,30 +675,45 @@ private Task RunOperationAsync(string operationName, Func> operati operationName, operation); + /// + /// Gets the document ID and partition key from a key. + /// + private (string DocumentId, PartitionKey PartitionKey) GetDocumentIdAndPartitionKey(TKey key) + => key is CosmosNoSqlKey cosmosKey + ? (cosmosKey.DocumentId, cosmosKey.PartitionKey) + : throw new InvalidOperationException( + $"Keys must be of type CosmosNoSqlKey. Received key of type '{key.GetType().FullName}'."); + /// /// Returns instance of with applied indexing policy. /// More information here: . /// private ContainerProperties GetContainerProperties() { - // Process Vector properties. - var embeddings = new Collection(); - var vectorIndexPaths = new Collection(); - var indexingPolicy = new IndexingPolicy { IndexingMode = this._indexingMode, Automatic = this._automatic }; + var containerProperties = new ContainerProperties( + this.Name, + partitionKeyPaths: this._partitionKeyProperties.Count > 0 + ? this._partitionKeyProperties.Select(p => $"/{p.StorageName}").ToList() + : ["/id"]) + { + IndexingPolicy = indexingPolicy + }; + if (this._indexingMode == IndexingMode.None) { - return new ContainerProperties(this.Name, partitionKeyPath: $"/{this._partitionKeyProperty.StorageName}") - { - IndexingPolicy = indexingPolicy - }; + return containerProperties; } + // Process Vector properties. + var embeddings = new Collection(); + var vectorIndexPaths = new Collection(); + foreach (var property in this._model.VectorProperties) { var path = $"/{property.StorageName}"; @@ -760,12 +765,10 @@ private ContainerProperties GetContainerProperties() indexingPolicy.ExcludedPaths.Add(new ExcludedPath { Path = $"{vectorIndexPath.Path}/*" }); } - return new ContainerProperties(this.Name, partitionKeyPath: $"/{this._partitionKeyProperty.StorageName}") - { - VectorEmbeddingPolicy = vectorEmbeddingPolicy, - IndexingPolicy = indexingPolicy, - FullTextPolicy = fullTextPolicy - }; + containerProperties.VectorEmbeddingPolicy = vectorEmbeddingPolicy; + containerProperties.FullTextPolicy = fullTextPolicy; + + return containerProperties; } /// @@ -855,30 +858,49 @@ private async IAsyncEnumerable> MapSearchResultsAsyn } } - private static IEnumerable GetCompositeKeys(IEnumerable keys) - => keys switch + /// + /// Builds a PartitionKey from the partition key property values in the record. + /// Supports single and hierarchical partition keys. + /// + private PartitionKey BuildPartitionKey(TRecord record) + { + switch (this._partitionKeyProperties) { - IEnumerable k => k, + case []: + return PartitionKey.None; - IEnumerable k => k.Select(key => new CosmosNoSqlCompositeKey(recordKey: key, partitionKey: key)), - - IEnumerable k => k.Select(key => + // Single partition key + case [var property]: { - var guidString = key.ToString(); - return new CosmosNoSqlCompositeKey(recordKey: guidString, partitionKey: guidString); - }), + return property.Type switch + { + Type t when t == typeof(string) => new PartitionKey(property.GetValue(record)), + Type t when t == typeof(bool) => new PartitionKey(property.GetValue(record)), + Type t when t == typeof(double) => new PartitionKey(property.GetValue(record)), + Type t when t == typeof(Guid) => new PartitionKey(property.GetValue(record).ToString()), + _ => throw new InvalidOperationException($"Unsupported partition key type: {property.Type.Name}") + }; + } - IEnumerable k => k.Select(key => key switch - { - string s => new CosmosNoSqlCompositeKey(recordKey: s, partitionKey: s), - Guid g when g.ToString() is var guidString => new CosmosNoSqlCompositeKey(recordKey: guidString, partitionKey: guidString), - CosmosNoSqlCompositeKey ck => ck, + // Hierarchical partition key (2 or 3 levels) + default: + var builder = new PartitionKeyBuilder(); - _ => throw new ArgumentException($"Invalid key type '{key.GetType().Name}'.") - }), + foreach (var property in this._partitionKeyProperties) + { + _ = property.Type switch + { + Type t when t == typeof(string) => builder.Add(property.GetValue(record)), + Type t when t == typeof(bool) => builder.Add(property.GetValue(record)), + Type t when t == typeof(double) => builder.Add(property.GetValue(record)), + Type t when t == typeof(Guid) => builder.Add(property.GetValue(record).ToString()), + _ => throw new InvalidOperationException($"Unsupported partition key type: {property.Type.Name}") + }; + } - _ => throw new UnreachableException() - }; + return builder.Build(); + } + } #endregion } diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionOptions.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionOptions.cs index c3e1265aa788..7dbea2dd9f15 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionOptions.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionOptions.cs @@ -1,5 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Collections.Generic; using System.Text.Json; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.VectorData; @@ -23,7 +24,7 @@ public CosmosNoSqlCollectionOptions() internal CosmosNoSqlCollectionOptions(CosmosNoSqlCollectionOptions? source) : base(source) { this.JsonSerializerOptions = source?.JsonSerializerOptions; - this.PartitionKeyPropertyName = source?.PartitionKeyPropertyName; + this.PartitionKeyPropertyNames = source?.PartitionKeyPropertyNames is { Count: > 0 } names ? [.. names] : []; this.IndexingMode = source?.IndexingMode ?? Default.IndexingMode; this.Automatic = source?.Automatic ?? Default.Automatic; } @@ -34,9 +35,24 @@ internal CosmosNoSqlCollectionOptions(CosmosNoSqlCollectionOptions? source) : ba public JsonSerializerOptions? JsonSerializerOptions { get; set; } /// - /// The property name to use as partition key. + /// Gets or sets the property names to use as partition key components. /// - public string? PartitionKeyPropertyName { get; set; } + /// + /// + /// Cosmos DB supports up to 3 levels of hierarchical partition keys. Provide the property names in hierarchical order. + /// For a single partition key, provide a list with one element. For hierarchical partition keys, provide up to 3 elements. + /// + /// + /// Selecting a partition key is critical for performance and scalability. Choose properties with high cardinality + /// that evenly distribute requests. See for guidance. + /// + /// + /// If empty, you must use as the key type with an explicitly constructed + /// . If your scenario does not require partitioning, use , + /// though this is not recommended for production workloads. + /// + /// + public IReadOnlyList PartitionKeyPropertyNames { get; set; } = []; /// /// Specifies the indexing mode in the Azure Cosmos DB service. diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs index 57523ea42232..305292c2c7f8 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs @@ -16,10 +16,6 @@ namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql; /// internal static class CosmosNoSqlCollectionQueryBuilder { - private const string SelectClauseDelimiter = ","; - private const string AndConditionDelimiter = " AND "; - private const string OrConditionDelimiter = " OR "; - /// /// Builds to get items from Azure CosmosDB NoSQL using vector search. /// @@ -66,7 +62,7 @@ public static QueryDefinition BuildSearchQuery( var rankingArgument = fullTextScoreArgument is null ? vectorDistanceArgument : $"RANK RRF({vectorDistanceArgument}, {fullTextScoreArgument})"; - var selectClauseArguments = string.Join(SelectClauseDelimiter, [.. fieldsArgument, vectorDistanceArgumentWithAlias]); + var selectClauseArguments = string.Join(",", [.. fieldsArgument, vectorDistanceArgumentWithAlias]); #pragma warning disable CS0618 // VectorSearchFilter is obsolete // Build filter object. @@ -127,7 +123,7 @@ public static QueryDefinition BuildSearchQuery( builder.Append(filterClause); if (scoreThresholdClause is not null) { - builder.Append(AndConditionDelimiter); + builder.Append(" AND "); } } @@ -185,7 +181,7 @@ internal static QueryDefinition BuildSearchQuery( var fieldsArgument = projectionProperties.Select(field => GeneratePropertyAccess(tableVariableName, field.StorageName)); - var selectClauseArguments = string.Join(SelectClauseDelimiter, [.. fieldsArgument]); + var selectClauseArguments = string.Join(",", [.. fieldsArgument]); // If Offset is not configured, use Top parameter instead of Limit/Offset // since it's more optimized. @@ -233,15 +229,12 @@ internal static QueryDefinition BuildSearchQuery( /// public static QueryDefinition BuildSelectQuery( CollectionModel model, - string keyStoragePropertyName, - string partitionKeyStoragePropertyName, - List keys, + List documentIds, bool includeVectors) { - Verify.True(keys.Count > 0, "At least one key should be provided.", nameof(keys)); + Verify.True(documentIds.Count > 0, "At least one document ID should be provided.", nameof(documentIds)); const string RecordKeyVariableName = "@rk"; - const string PartitionKeyVariableName = "@pk"; var tableVariableName = CosmosNoSqlConstants.ContainerAlias; @@ -251,13 +244,16 @@ public static QueryDefinition BuildSelectQuery( projectionProperties = projectionProperties.Where(p => p is not VectorPropertyModel); } - var selectClauseArguments = string.Join(SelectClauseDelimiter, + var selectClauseArguments = string.Join( + ",", projectionProperties.Select(field => GeneratePropertyAccess(tableVariableName, field.StorageName))); - var whereClauseArguments = string.Join(OrConditionDelimiter, - keys.Select((key, index) => - $"({GeneratePropertyAccess(tableVariableName, keyStoragePropertyName)} = {RecordKeyVariableName}{index} {AndConditionDelimiter} " + - $"{GeneratePropertyAccess(tableVariableName, partitionKeyStoragePropertyName)} = {PartitionKeyVariableName}{index})")); + // Build WHERE clause using only the document id. + // The partition key is provided via RequestOptions for point reads, not in the query. + var whereClauseArguments = string.Join( + " OR ", + documentIds.Select((_, index) => + $"({GeneratePropertyAccess(tableVariableName, model.KeyProperty.StorageName)} = {RecordKeyVariableName}{index})")); var query = $""" SELECT {selectClauseArguments} @@ -267,16 +263,11 @@ public static QueryDefinition BuildSelectQuery( var queryDefinition = new QueryDefinition(query); - for (var i = 0; i < keys.Count; i++) + for (var i = 0; i < documentIds.Count; i++) { - var recordKey = keys[i].RecordKey; - var partitionKey = keys[i].PartitionKey; - - Verify.NotNullOrWhiteSpace(recordKey); - Verify.NotNullOrWhiteSpace(partitionKey); - - queryDefinition.WithParameter($"{RecordKeyVariableName}{i}", recordKey); - queryDefinition.WithParameter($"{PartitionKeyVariableName}{i}", partitionKey); + var documentIdString = documentIds[i]; + Verify.NotNullOrWhiteSpace(documentIdString); + queryDefinition.WithParameter($"{RecordKeyVariableName}{i}", documentIdString); } return queryDefinition; diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCompositeKey.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCompositeKey.cs deleted file mode 100644 index dd54ecb407ca..000000000000 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCompositeKey.cs +++ /dev/null @@ -1,19 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql; - -/// -/// Composite key for Azure CosmosDB NoSQL record with record and partition keys. -/// -public sealed class CosmosNoSqlCompositeKey(string recordKey, string partitionKey) -{ - /// - /// Value of record key. - /// - public string RecordKey { get; } = recordKey; - - /// - /// Value of partition key. - /// - public string PartitionKey { get; } = partitionKey; -} diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlDynamicCollection.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlDynamicCollection.cs index c5474f28dd8e..661c6ca7c481 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlDynamicCollection.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlDynamicCollection.cs @@ -11,6 +11,12 @@ namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql; /// /// Represents a collection of vector store records in a CosmosNoSql database, mapped to a dynamic Dictionary<string, object?>. /// +/// +/// +/// This collection accepts keys of type , but only instances +/// are supported at runtime. Passing any other key type will result in an . +/// +/// #pragma warning disable CA1711 // Identifiers should not have incorrect suffix public sealed class CosmosNoSqlDynamicCollection : CosmosNoSqlCollection> #pragma warning restore CA1711 // Identifiers should not have incorrect suffix @@ -33,11 +39,11 @@ public CosmosNoSqlDynamicCollection(Database database, string name, CosmosNoSqlC } /// - /// Initializes a new instance of the class. + /// Initializes a new instance of the class. /// /// Connection string required to connect to Azure CosmosDB NoSQL. /// Database name for Azure CosmosDB NoSQL. - /// The name of the collection that this will access. + /// The name of the collection that this will access. /// Optional configuration options for . /// Optional configuration options for this class. [RequiresUnreferencedCode("The Cosmos NoSQL provider is currently incompatible with trimming.")] diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlKey.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlKey.cs new file mode 100644 index 000000000000..355ff02bfac0 --- /dev/null +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlKey.cs @@ -0,0 +1,183 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using Microsoft.Azure.Cosmos; + +namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql; + +/// +/// Represents a key for Azure CosmosDB NoSQL, containing both the record key (document ID) and the partition key. +/// +/// +/// +/// Azure Cosmos DB requires both a record key (the unique document identifier within a partition) and a partition key +/// to identify a document. This struct encapsulates both values together with the Cosmos SDK's type. +/// +/// +/// For simple partition keys, use the convenience constructors that accept , , or +/// partition key values. For hierarchical partition keys (up to 3 levels), construct a using +/// and pass it to the primary constructor. +/// +/// +// This is conceptually a record, but we're targeting .NET Standard 2.0 too. +public readonly struct CosmosNoSqlKey : IEquatable +{ + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + public CosmosNoSqlKey(string documentId, string partitionKey) + { + this.DocumentId = documentId; + this.PartitionKey = new PartitionKey(partitionKey); + } + + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + public CosmosNoSqlKey(Guid documentId, string partitionKey) + { + this.DocumentId = documentId.ToString(); + this.PartitionKey = new PartitionKey(partitionKey); + } + + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + public CosmosNoSqlKey(string documentId, Guid partitionKey) + { + this.DocumentId = documentId; + this.PartitionKey = new PartitionKey(partitionKey.ToString()); + } + + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + public CosmosNoSqlKey(Guid documentId, Guid partitionKey) + { + this.DocumentId = documentId.ToString(); + this.PartitionKey = new PartitionKey(partitionKey.ToString()); + } + + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + public CosmosNoSqlKey(string documentId, bool partitionKey) + { + this.DocumentId = documentId; + this.PartitionKey = new PartitionKey(partitionKey); + } + + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + public CosmosNoSqlKey(Guid documentId, bool partitionKey) + { + this.DocumentId = documentId.ToString(); + this.PartitionKey = new PartitionKey(partitionKey); + } + + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + public CosmosNoSqlKey(string documentId, double partitionKey) + { + this.DocumentId = documentId; + this.PartitionKey = new PartitionKey(partitionKey); + } + + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + public CosmosNoSqlKey(Guid documentId, double partitionKey) + { + this.DocumentId = documentId.ToString(); + this.PartitionKey = new PartitionKey(partitionKey); + } + + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + /// + /// Use this constructor for hierarchical partition keys or special partition key values like or . + /// For hierarchical partition keys, construct the using . + /// + public CosmosNoSqlKey(string documentId, PartitionKey partitionKey) + { + this.DocumentId = documentId; + this.PartitionKey = partitionKey; + } + + /// + /// Initializes a new instance of the struct with a document ID and partition key. + /// + /// The document ID. + /// The Cosmos DB partition key. + /// + /// Use this constructor for hierarchical partition keys or special partition key values like or . + /// For hierarchical partition keys, construct the using . + /// + public CosmosNoSqlKey(Guid documentId, PartitionKey partitionKey) + { + this.DocumentId = documentId.ToString(); + this.PartitionKey = partitionKey; + } + + /// + /// Gets the document ID. + /// + /// + /// The document ID uniquely identifies a document within a partition. It is stored in the id property of the Cosmos DB document. + /// + public string DocumentId { get; } + + /// + /// Gets the partition key. + /// + /// + /// The partition key determines which logical partition the document belongs to. + /// See for guidance on choosing partition keys. + /// + public PartitionKey PartitionKey { get; } + + /// + public bool Equals(CosmosNoSqlKey other) + => Equals(this.DocumentId, other.DocumentId) && this.PartitionKey.Equals(other.PartitionKey); + + /// + public override bool Equals(object? obj) + => obj is CosmosNoSqlKey other && this.Equals(other); + + /// + public override int GetHashCode() + => HashCode.Combine(this.DocumentId, this.PartitionKey); + + /// + /// Determines whether two instances are equal. + /// + public static bool operator ==(CosmosNoSqlKey left, CosmosNoSqlKey right) + => left.Equals(right); + + /// + /// Determines whether two instances are not equal. + /// + public static bool operator !=(CosmosNoSqlKey left, CosmosNoSqlKey right) + => !left.Equals(right); +} diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs index feb9beaddf9a..fcbd60877571 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs @@ -23,12 +23,14 @@ internal class CosmosNoSqlModelBuilder() : CollectionJsonModelBuilder(s_modelBui protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { + // Note that the key property in Cosmos NoSQL refers to the document ID, not to the CosmosNoSqlKey structure which includes both + // the document ID and the partition key (and which is the generic TKey type parameter of the collection). var type = keyProperty.Type; if (type != typeof(string) && type != typeof(Guid)) { throw new NotSupportedException( - $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, Guid."); + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, Guid, int, long."); } } diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlServiceCollectionExtensions.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlServiceCollectionExtensions.cs index 839aea3f7ce0..826e0a56e73f 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlServiceCollectionExtensions.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlServiceCollectionExtensions.cs @@ -159,15 +159,15 @@ public static IServiceCollection AddKeyedCosmosNoSqlCollection( Verify.NotNull(services); Verify.NotNullOrWhiteSpace(name); - services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection), serviceKey, (sp, _) => + services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection), serviceKey, (sp, _) => { var database = sp.GetRequiredService(); options = GetCollectionOptions(sp, _ => options); - return new CosmosNoSqlCollection(database, name, options); + return new CosmosNoSqlCollection(database, name, options); }, lifetime)); - AddAbstractions(services, serviceKey, lifetime); + AddAbstractions(services, serviceKey, lifetime); return services; } @@ -267,12 +267,12 @@ public static IServiceCollection AddKeyedCosmosNoSqlCollection( Verify.NotNull(connectionStringProvider); Verify.NotNull(databaseNameProvider); - services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection), serviceKey, (sp, _) => + services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection), serviceKey, (sp, _) => { var options = GetCollectionOptions(sp, optionsProvider); var clientOptions = CreateClientOptions(options?.JsonSerializerOptions); - return new CosmosNoSqlCollection( + return new CosmosNoSqlCollection( connectionString: connectionStringProvider(sp), databaseName: databaseNameProvider(sp), name: name, @@ -280,7 +280,7 @@ public static IServiceCollection AddKeyedCosmosNoSqlCollection( options); }, lifetime)); - AddAbstractions(services, serviceKey, lifetime); + AddAbstractions(services, serviceKey, lifetime); return services; } diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlCollectionOptionsTests.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlCollectionOptionsTests.cs index 63a9487dc2bc..869510aa8df1 100644 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlCollectionOptionsTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlCollectionOptionsTests.cs @@ -20,11 +20,11 @@ public async Task Collection_supports_partition_key_composite_key() var store = (CosmosNoSqlTestStore)fixture.TestStore; var collectionName = fixture.TestStore.AdjustCollectionName("PartitionKeyCompositeKey"); - using VectorStoreCollection collection = - new CosmosNoSqlCollection( - store.Database, - collectionName, - new() { PartitionKeyPropertyName = nameof(PartitionedHotel.HotelName) }); + // The key type for operations is CosmosNoSqlKey (containing both document ID and partition key). + using var collection = new CosmosNoSqlCollection( + store.Database, + collectionName, + new() { PartitionKeyPropertyNames = [nameof(PartitionedHotel.HotelName)] }); await collection.EnsureCollectionExistsAsync(); @@ -39,7 +39,7 @@ public async Task Collection_supports_partition_key_composite_key() }; await collection.UpsertAsync(record); - var key = new CosmosNoSqlCompositeKey(record.HotelId, record.HotelName); + var key = new CosmosNoSqlKey(record.HotelId, record.HotelName); var fetched = await collection.GetAsync(key, new() { IncludeVectors = true }); Assert.NotNull(fetched); @@ -62,10 +62,10 @@ public async Task Collection_supports_indexing_mode(IndexingMode indexingMode) var store = (CosmosNoSqlTestStore)fixture.TestStore; var collectionName = fixture.TestStore.AdjustCollectionName($"IndexingMode_{indexingMode}"); - using var collection = new CosmosNoSqlCollection( + using var collection = new CosmosNoSqlCollection( store.Database, collectionName, - new() { IndexingMode = indexingMode, Automatic = indexingMode != IndexingMode.None }); + new() { IndexingMode = indexingMode, Automatic = indexingMode != IndexingMode.None, PartitionKeyPropertyNames = ["HotelId"] }); await collection.EnsureCollectionExistsAsync(); @@ -79,12 +79,13 @@ public async Task Collection_supports_indexing_mode(IndexingMode indexingMode) }; await collection.UpsertAsync(record); - var fetched = await collection.GetAsync(record.HotelId); + var key = new CosmosNoSqlKey(record.HotelId, record.HotelId); + var fetched = await collection.GetAsync(key); Assert.NotNull(fetched); - await collection.DeleteAsync(record.HotelId); - Assert.Null(await collection.GetAsync(record.HotelId)); + await collection.DeleteAsync(key); + Assert.Null(await collection.GetAsync(key)); } finally { diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlDependencyInjectionTests.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlDependencyInjectionTests.cs index 1c7811681655..be63abe03387 100644 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlDependencyInjectionTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlDependencyInjectionTests.cs @@ -7,11 +7,12 @@ using Microsoft.SemanticKernel.Connectors.CosmosNoSql; using VectorData.ConformanceTests; using Xunit; +using TestRecord = VectorData.ConformanceTests.DependencyInjectionTests.Record; namespace CosmosNoSql.ConformanceTests.DependencyInjection; public class CosmosNoSqlDependencyInjectionTests - : DependencyInjectionTests.Record>, string, DependencyInjectionTests.Record> + : DependencyInjectionTests, CosmosNoSqlKey, TestRecord> { protected const string ConnectionString = "AccountEndpoint=https://test.documents.azure.com:443/;AccountKey=mock;"; protected const string DatabaseName = "dbName"; @@ -48,22 +49,22 @@ private static string DatabaseNameProvider(IServiceProvider sp, object serviceKe ? services .AddSingleton(sp => new CosmosClient(ConnectionString, s_clientOptions)) .AddSingleton(sp => sp.GetRequiredService().GetDatabase(DatabaseName)) - .AddCosmosNoSqlCollection(name, lifetime: lifetime) + .AddCosmosNoSqlCollection(name, lifetime: lifetime) : services .AddSingleton(sp => new CosmosClient(ConnectionString, s_clientOptions)) .AddSingleton(sp => sp.GetRequiredService().GetDatabase(DatabaseName)) - .AddKeyedCosmosNoSqlCollection(serviceKey, name, lifetime: lifetime); + .AddKeyedCosmosNoSqlCollection(serviceKey, name, lifetime: lifetime); yield return (services, serviceKey, name, lifetime) => serviceKey is null - ? services.AddCosmosNoSqlCollection( + ? services.AddCosmosNoSqlCollection( name, ConnectionString, DatabaseName, lifetime: lifetime) - : services.AddKeyedCosmosNoSqlCollection( + : services.AddKeyedCosmosNoSqlCollection( serviceKey, name, ConnectionString, DatabaseName, lifetime: lifetime); yield return (services, serviceKey, name, lifetime) => serviceKey is null - ? services.AddCosmosNoSqlCollection( + ? services.AddCosmosNoSqlCollection( name, ConnectionStringProvider, DatabaseNameProvider, lifetime: lifetime) - : services.AddKeyedCosmosNoSqlCollection( + : services.AddKeyedCosmosNoSqlCollection( serviceKey, name, sp => ConnectionStringProvider(sp, serviceKey), sp => DatabaseNameProvider(sp, serviceKey), lifetime: lifetime); } } @@ -95,9 +96,9 @@ public void ConnectionStringProviderCantBeNull() { IServiceCollection services = new ServiceCollection(); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", connectionStringProvider: null!, databaseNameProvider: DatabaseNameProvider)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", connectionStringProvider: null!, databaseNameProvider: DatabaseNameProvider)); } @@ -106,9 +107,9 @@ public void DatabaseNameProviderCantBeNull() { IServiceCollection services = new ServiceCollection(); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", connectionStringProvider: ConnectionStringProvider, databaseNameProvider: null!)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", connectionStringProvider: ConnectionStringProvider, databaseNameProvider: null!)); } @@ -119,13 +120,13 @@ public void ConnectionStringCantBeNullOrEmpty() Assert.Throws(() => services.AddCosmosNoSqlVectorStore(connectionString: null!, DatabaseName)); Assert.Throws(() => services.AddKeyedCosmosNoSqlVectorStore(serviceKey: "notNull", connectionString: null!, DatabaseName)); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", connectionString: null!, DatabaseName)); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", connectionString: "", DatabaseName)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", connectionString: null!, DatabaseName)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", connectionString: "", DatabaseName)); } @@ -136,13 +137,13 @@ public void DatabaseNameCantBeNullOrEmpty() Assert.Throws(() => services.AddCosmosNoSqlVectorStore(ConnectionString, databaseName: null!)); Assert.Throws(() => services.AddKeyedCosmosNoSqlVectorStore(serviceKey: "notNull", ConnectionString, databaseName: null!)); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", ConnectionString, databaseName: null!)); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", ConnectionString, databaseName: "")); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", ConnectionString, databaseName: null!)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", ConnectionString, databaseName: "")); } } diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlEmbeddingGenerationTests.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlEmbeddingGenerationTests.cs index a88c9254a8e8..5536025c72fb 100644 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlEmbeddingGenerationTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlEmbeddingGenerationTests.cs @@ -1,18 +1,66 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Reflection; using CosmosNoSql.ConformanceTests.Support; using Microsoft.Extensions.AI; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.VectorData; +using Microsoft.SemanticKernel.Connectors.CosmosNoSql; using VectorData.ConformanceTests; using VectorData.ConformanceTests.Support; using Xunit; namespace CosmosNoSql.ConformanceTests; +#pragma warning disable CA2000 // Dispose objects before losing scope + public class CosmosNoSqlEmbeddingGenerationTests(CosmosNoSqlEmbeddingGenerationTests.StringVectorFixture stringVectorFixture, CosmosNoSqlEmbeddingGenerationTests.RomOfFloatVectorFixture romOfFloatVectorFixture) : EmbeddingGenerationTests(stringVectorFixture, romOfFloatVectorFixture), IClassFixture, IClassFixture { + // Override the DI tests because the base tests call vectorStore.GetCollection() + // but CosmosNoSqlVectorStore.GetCollection requires CosmosNoSqlKey, not plain string. + public override async Task SearchAsync_with_store_dependency_injection() + { + foreach (var registrationDelegate in stringVectorFixture.DependencyInjectionStoreRegistrationDelegates) + { + IServiceCollection serviceCollection = new ServiceCollection(); + + serviceCollection.AddSingleton(new FakeEmbeddingGenerator(replaceLast: 1)); + registrationDelegate(serviceCollection); + + await using var serviceProvider = serviceCollection.BuildServiceProvider(); + + var vectorStore = serviceProvider.GetRequiredService(); + // Use the fixture's GetCollection override which routes through TestStore adapter + var collection = stringVectorFixture.GetCollection(vectorStore, stringVectorFixture.CollectionName, stringVectorFixture.CreateRecordDefinition()); + + var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync(); + + Assert.Equal("Store ([1, 1, 1])", result.Record.Text); + } + } + + public override async Task SearchAsync_with_collection_dependency_injection() + { + foreach (var registrationDelegate in stringVectorFixture.DependencyInjectionCollectionRegistrationDelegates) + { + IServiceCollection serviceCollection = new ServiceCollection(); + + serviceCollection.AddSingleton(new FakeEmbeddingGenerator(replaceLast: 1)); + registrationDelegate(serviceCollection); + + await using var serviceProvider = serviceCollection.BuildServiceProvider(); + + // Resolve with CosmosNoSqlKey since that's how AddAbstractions registers it, + // then wrap with the adapter for plain string key usage. + var innerCollection = serviceProvider.GetRequiredService>(); + var collection = new CosmosNoSqlCollectionAdapter(innerCollection); + + var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync(); + + Assert.Equal("Store ([1, 1, 1])", result.Record.Text); + } + } public new class StringVectorFixture : EmbeddingGenerationTests.StringVectorFixture { public override TestStore TestStore => CosmosNoSqlTestStore.Instance; @@ -20,6 +68,47 @@ public class CosmosNoSqlEmbeddingGenerationTests(CosmosNoSqlEmbeddingGenerationT public override VectorStore CreateVectorStore(IEmbeddingGenerator? embeddingGenerator) => CosmosNoSqlTestStore.Instance.GetVectorStore(new() { EmbeddingGenerator = embeddingGenerator }); + public override VectorStoreCollection GetCollection( + VectorStore vectorStore, + string collectionName, + VectorStoreCollectionDefinition? recordDefinition = null) + { + // If no definition was provided, build one from TRecord's attributes. + recordDefinition ??= BuildDefinitionFromType(typeof(TRecord)); + + // If the VectorStore has a store-level embedding generator, copy it onto the definition + // so that TestStore.CreateCollection can pass it to the collection options. + if (recordDefinition.EmbeddingGenerator is null + && vectorStore is CosmosNoSqlVectorStore cosmosStore) + { + var field = typeof(CosmosNoSqlVectorStore).GetField("_embeddingGenerator", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); + if (field?.GetValue(cosmosStore) is IEmbeddingGenerator generator) + { + recordDefinition.EmbeddingGenerator = generator; + } + } + + return this.CreateCollection(collectionName, recordDefinition); + } + + public override VectorStoreCollection> GetDynamicCollection( + VectorStore vectorStore, + string collectionName, + VectorStoreCollectionDefinition recordDefinition) + { + if (recordDefinition.EmbeddingGenerator is null + && vectorStore is CosmosNoSqlVectorStore cosmosStore) + { + var field = typeof(CosmosNoSqlVectorStore).GetField("_embeddingGenerator", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); + if (field?.GetValue(cosmosStore) is IEmbeddingGenerator generator) + { + recordDefinition.EmbeddingGenerator = generator; + } + } + + return this.CreateDynamicCollection(collectionName, recordDefinition); + } + public override Func[] DependencyInjectionStoreRegistrationDelegates => [ services => services @@ -42,6 +131,47 @@ public override VectorStore CreateVectorStore(IEmbeddingGenerator? embeddingGene public override VectorStore CreateVectorStore(IEmbeddingGenerator? embeddingGenerator) => CosmosNoSqlTestStore.Instance.GetVectorStore(new() { EmbeddingGenerator = embeddingGenerator }); + public override VectorStoreCollection GetCollection( + VectorStore vectorStore, + string collectionName, + VectorStoreCollectionDefinition? recordDefinition = null) + { + // If no definition was provided, build one from TRecord's attributes. + recordDefinition ??= BuildDefinitionFromType(typeof(TRecord)); + + // If the VectorStore has a store-level embedding generator, copy it onto the definition + // so that TestStore.CreateCollection can pass it to the collection options. + if (recordDefinition.EmbeddingGenerator is null + && vectorStore is CosmosNoSqlVectorStore cosmosStore) + { + var field = typeof(CosmosNoSqlVectorStore).GetField("_embeddingGenerator", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); + if (field?.GetValue(cosmosStore) is IEmbeddingGenerator generator) + { + recordDefinition.EmbeddingGenerator = generator; + } + } + + return this.CreateCollection(collectionName, recordDefinition); + } + + public override VectorStoreCollection> GetDynamicCollection( + VectorStore vectorStore, + string collectionName, + VectorStoreCollectionDefinition recordDefinition) + { + if (recordDefinition.EmbeddingGenerator is null + && vectorStore is CosmosNoSqlVectorStore cosmosStore) + { + var field = typeof(CosmosNoSqlVectorStore).GetField("_embeddingGenerator", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); + if (field?.GetValue(cosmosStore) is IEmbeddingGenerator generator) + { + recordDefinition.EmbeddingGenerator = generator; + } + } + + return this.CreateDynamicCollection(collectionName, recordDefinition); + } + public override Func[] DependencyInjectionStoreRegistrationDelegates => [ services => services @@ -56,4 +186,30 @@ public override VectorStore CreateVectorStore(IEmbeddingGenerator? embeddingGene .AddCosmosNoSqlCollection(this.CollectionName) ]; } + + /// + /// Builds a VectorStoreCollectionDefinition from a record type's VectorStore attributes. + /// Used when GetCollection is called without an explicit definition. + /// + private static VectorStoreCollectionDefinition BuildDefinitionFromType(Type recordType) + { + var properties = new List(); + foreach (var prop in recordType.GetProperties()) + { + if (prop.GetCustomAttribute() is not null) + { + properties.Add(new VectorStoreKeyProperty(prop.Name, prop.PropertyType)); + } + else if (prop.GetCustomAttribute() is VectorStoreVectorAttribute vecAttr) + { + properties.Add(new VectorStoreVectorProperty(prop.Name, prop.PropertyType, vecAttr.Dimensions)); + } + else if (prop.GetCustomAttribute() is not null) + { + properties.Add(new VectorStoreDataProperty(prop.Name, prop.PropertyType)); + } + } + + return new VectorStoreCollectionDefinition { Properties = properties }; + } } diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlCollectionAdapter.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlCollectionAdapter.cs new file mode 100644 index 000000000000..40e929f99081 --- /dev/null +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlCollectionAdapter.cs @@ -0,0 +1,103 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Linq.Expressions; +using Microsoft.Extensions.VectorData; +using Microsoft.SemanticKernel.Connectors.CosmosNoSql; + +namespace CosmosNoSql.ConformanceTests.Support; + +#pragma warning disable CA1812 // Internal class that is apparently never instantiated + +/// +/// Adapts a with keys +/// to a with simple string keys. +/// This allows Cosmos conformance tests to use simple key types (string, etc.) while still +/// using the underlying Cosmos collection with composite keys. +/// +internal sealed class CosmosNoSqlCollectionAdapter( + VectorStoreCollection inner) + : VectorStoreCollection, + IKeywordHybridSearchable + where TDocumentId : notnull + where TRecord : class +{ + public override string Name => inner.Name; + + public override Task CollectionExistsAsync(CancellationToken cancellationToken = default) + => inner.CollectionExistsAsync(cancellationToken); + + public override Task EnsureCollectionExistsAsync(CancellationToken cancellationToken = default) + => inner.EnsureCollectionExistsAsync(cancellationToken); + + public override Task EnsureCollectionDeletedAsync(CancellationToken cancellationToken = default) + => inner.EnsureCollectionDeletedAsync(cancellationToken); + + public override Task GetAsync(TDocumentId key, RecordRetrievalOptions? options = null, CancellationToken cancellationToken = default) + => key is null + ? throw new ArgumentNullException(nameof(key)) + : inner.GetAsync(ConvertKey(key), options, cancellationToken); + + public override IAsyncEnumerable GetAsync(IEnumerable keys, RecordRetrievalOptions? options = null, CancellationToken cancellationToken = default) + => keys is null + ? throw new ArgumentNullException(nameof(keys)) + : inner.GetAsync(ConvertKeys(keys), options, cancellationToken); + + public override IAsyncEnumerable GetAsync( + Expression> filter, + int top, + FilteredRecordRetrievalOptions? options = null, + CancellationToken cancellationToken = default) + => inner.GetAsync(filter, top, options, cancellationToken); + + public override Task UpsertAsync(TRecord record, CancellationToken cancellationToken = default) + => inner.UpsertAsync(record, cancellationToken); + + public override Task UpsertAsync(IEnumerable records, CancellationToken cancellationToken = default) + => inner.UpsertAsync(records, cancellationToken); + + public override Task DeleteAsync(TDocumentId key, CancellationToken cancellationToken = default) + => key is null + ? throw new ArgumentNullException(nameof(key)) + : inner.DeleteAsync(ConvertKey(key), cancellationToken); + + public override Task DeleteAsync(IEnumerable keys, CancellationToken cancellationToken = default) + => keys is null + ? throw new ArgumentNullException(nameof(keys)) + : inner.DeleteAsync(ConvertKeys(keys), cancellationToken); + + public override IAsyncEnumerable> SearchAsync( + TInput searchValue, + int top, + VectorSearchOptions? options = null, + CancellationToken cancellationToken = default) + => inner.SearchAsync(searchValue, top, options, cancellationToken); + + public IAsyncEnumerable> HybridSearchAsync( + TInput searchValue, + ICollection keywords, + int top, + HybridSearchOptions? options = default, + CancellationToken cancellationToken = default) + where TInput : notnull + => ((IKeywordHybridSearchable)inner).HybridSearchAsync(searchValue, keywords, top, options, cancellationToken); + + public override object? GetService(Type serviceType, object? serviceKey = null) + => inner.GetService(serviceType, serviceKey); + + // For tests, use the document ID as the partition key + private static CosmosNoSqlKey ConvertKey(TDocumentId documentId) + => documentId switch + { + string s => new(s, s), + Guid g => new(g, g), + _ => throw new InvalidOperationException() + }; + + private static IEnumerable ConvertKeys(IEnumerable keys) + { + foreach (var key in keys) + { + yield return ConvertKey(key); + } + } +} diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlDynamicCollectionAdapter.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlDynamicCollectionAdapter.cs new file mode 100644 index 000000000000..d8538e7541a9 --- /dev/null +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlDynamicCollectionAdapter.cs @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Linq.Expressions; +using Microsoft.Extensions.VectorData; +using Microsoft.SemanticKernel.Connectors.CosmosNoSql; + +namespace CosmosNoSql.ConformanceTests.Support; + +#pragma warning disable CA1812 // Internal class that is apparently never instantiated + +/// +/// Adapts a (which requires keys) +/// to accept plain keys (string, Guid, etc.) by wrapping them in . +/// Both the adapter and inner collection use object as the key type; the adapter intercepts Get/Delete calls +/// to convert plain keys to before forwarding. +/// +internal sealed class CosmosNoSqlDynamicCollectionAdapter( + VectorStoreCollection> inner) + : VectorStoreCollection> +{ + public override string Name => inner.Name; + + public override Task CollectionExistsAsync(CancellationToken cancellationToken = default) + => inner.CollectionExistsAsync(cancellationToken); + + public override Task EnsureCollectionExistsAsync(CancellationToken cancellationToken = default) + => inner.EnsureCollectionExistsAsync(cancellationToken); + + public override Task EnsureCollectionDeletedAsync(CancellationToken cancellationToken = default) + => inner.EnsureCollectionDeletedAsync(cancellationToken); + + public override Task?> GetAsync(object key, RecordRetrievalOptions? options = null, CancellationToken cancellationToken = default) + => key is null + ? throw new ArgumentNullException(nameof(key)) + : inner.GetAsync(ConvertKey(key), options, cancellationToken); + + public override IAsyncEnumerable> GetAsync(IEnumerable keys, RecordRetrievalOptions? options = null, CancellationToken cancellationToken = default) + => keys is null + ? throw new ArgumentNullException(nameof(keys)) + : inner.GetAsync(keys.Select(ConvertKey), options, cancellationToken); + + public override IAsyncEnumerable> GetAsync( + Expression, bool>> filter, + int top, + FilteredRecordRetrievalOptions>? options = null, + CancellationToken cancellationToken = default) + => inner.GetAsync(filter, top, options, cancellationToken); + + public override Task UpsertAsync(Dictionary record, CancellationToken cancellationToken = default) + => inner.UpsertAsync(record, cancellationToken); + + public override Task UpsertAsync(IEnumerable> records, CancellationToken cancellationToken = default) + => inner.UpsertAsync(records, cancellationToken); + + public override Task DeleteAsync(object key, CancellationToken cancellationToken = default) + => key is null + ? throw new ArgumentNullException(nameof(key)) + : inner.DeleteAsync(ConvertKey(key), cancellationToken); + + public override Task DeleteAsync(IEnumerable keys, CancellationToken cancellationToken = default) + => keys is null + ? throw new ArgumentNullException(nameof(keys)) + : inner.DeleteAsync(keys.Select(ConvertKey), cancellationToken); + + public override IAsyncEnumerable>> SearchAsync( + TInput searchValue, + int top, + VectorSearchOptions>? options = null, + CancellationToken cancellationToken = default) + => inner.SearchAsync(searchValue, top, options, cancellationToken); + + public override object? GetService(Type serviceType, object? serviceKey = null) + => inner.GetService(serviceType, serviceKey); + + /// + /// Wraps a plain key (e.g. string "1") in a using the key's string + /// representation as the partition key. + /// If the key is already a , returns it unchanged (boxed). + /// + private static object ConvertKey(object key) + { + if (key is CosmosNoSqlKey) + { + return key; + } + + // Wrap keys using their string representation + var keyString = key.ToString() ?? string.Empty; + return new CosmosNoSqlKey(keyString, keyString); + } +} diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlTestStore.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlTestStore.cs index 42db4212f9c9..f88a96bf74c1 100644 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlTestStore.cs +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlTestStore.cs @@ -3,14 +3,17 @@ #if NET472 using System.Net.Http; #endif +// using System.Globalization; using System.Text.Json; using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.VectorData; using Microsoft.SemanticKernel.Connectors.CosmosNoSql; using VectorData.ConformanceTests.Support; namespace CosmosNoSql.ConformanceTests.Support; #pragma warning disable CA1001 // Type owns disposable fields (_connection) but is not disposable +#pragma warning disable CA2000 // Dispose objects before losing scope internal sealed class CosmosNoSqlTestStore : TestStore { @@ -34,6 +37,62 @@ private CosmosNoSqlTestStore() { } + public override VectorStoreCollection CreateCollection( + string name, + VectorStoreCollectionDefinition definition) + { + // Cosmos NoSQL requires specifying a partition key for container creation. + // For the tests, use the key property as the partition key. + var keyPropertyName = definition.Properties.OfType().FirstOrDefault()?.Name + ?? throw new InvalidOperationException("Definition must have a key property"); + + var options = new CosmosNoSqlCollectionOptions + { + Definition = definition, + PartitionKeyPropertyNames = [keyPropertyName], + EmbeddingGenerator = definition.EmbeddingGenerator + }; + + // Also, in Cosmos NoSQL there's a discrepancy between: + // 1. The key property in the model and on the record type (representing the Cosmos document ID - string or Guid) + // 2. The TKey generic type variable on the collection (must be CosmosNoSqlKey, wrapping both the document ID and the partition key) + // To bridge this gap in the tests, if the test expects a simpler key type (string, Guid), we wrap it in an adapter that translates between the two. + var collection = new CosmosNoSqlCollection(this.Database, name, options); + + return typeof(TKey) switch + { + var t when t == typeof(CosmosNoSqlKey) + => (VectorStoreCollection)(object)collection, + + var t when t == typeof(string) + => (VectorStoreCollection)(object)new CosmosNoSqlCollectionAdapter(collection), + var t when t == typeof(Guid) + => (VectorStoreCollection)(object)new CosmosNoSqlCollectionAdapter(collection), + + _ => throw new InvalidOperationException( + $"Cosmos NoSQL tests must use string, Guid or CosmosNoSqlKey as the key type. Got: {typeof(TKey).Name}") + }; + } + + public override VectorStoreCollection> CreateDynamicCollection( + string name, + VectorStoreCollectionDefinition definition) + { + // See notes in CreateCollection above. + var keyPropertyName = definition.Properties.OfType().FirstOrDefault()?.Name + ?? throw new InvalidOperationException("Definition must have a key property"); + + return new CosmosNoSqlDynamicCollectionAdapter( + new CosmosNoSqlDynamicCollection( + this.Database, + name, + new CosmosNoSqlCollectionOptions + { + Definition = definition, + PartitionKeyPropertyNames = [keyPropertyName] + })); + } + #pragma warning disable CA5400 // HttpClient may be created without enabling CheckCertificateRevocationList protected override async Task StartAsync() { diff --git a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionQueryBuilderTests.cs b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionQueryBuilderTests.cs index 825e4015c192..ce3f579d8370 100644 --- a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionQueryBuilderTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionQueryBuilderTests.cs @@ -193,12 +193,9 @@ public void BuildSelectQueryByDefaultReturnsValidQueryDefinition() const string ExpectedQueryText = """ SELECT x["id"],x["TestProperty1"],x["TestProperty2"] FROM x - WHERE (x["id"] = @rk0 AND x["TestProperty1"] = @pk0) + WHERE (x["id"] = @rk0) """; - const string KeyStoragePropertyName = "id"; - const string PartitionKeyPropertyName = "TestProperty1"; - var model = new CosmosNoSqlModelBuilder().BuildDynamic( new() { @@ -210,14 +207,12 @@ FROM x ] }, defaultEmbeddingGenerator: null); - var keys = new List { new("id", "TestProperty1") }; + List documentIds = ["id"]; // Act var queryDefinition = CosmosNoSqlCollectionQueryBuilder.BuildSelectQuery( model, - KeyStoragePropertyName, - PartitionKeyPropertyName, - keys, + documentIds, includeVectors: true); var queryText = queryDefinition.QueryText; @@ -228,9 +223,6 @@ FROM x Assert.Equal("@rk0", queryParameters[0].Name); Assert.Equal("id", queryParameters[0].Value); - - Assert.Equal("@pk0", queryParameters[1].Name); - Assert.Equal("TestProperty1", queryParameters[1].Value); } [Fact] diff --git a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs index 4372c3de9561..7cc8cc3b2453 100644 --- a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs @@ -42,7 +42,7 @@ public CosmosNoSqlCollectionTests() public void ConstructorForModelWithoutKeyThrowsException() { // Act & Assert - var exception = Assert.Throws(() => new CosmosNoSqlCollection(this._mockDatabase.Object, "collection")); + var exception = Assert.Throws(() => new CosmosNoSqlCollection(this._mockDatabase.Object, "collection")); Assert.Contains("No key property found", exception.Message); } @@ -56,7 +56,7 @@ public void ConstructorWithoutSystemTextJsonSerializerOptionsThrowsArgumentExcep mockDatabase.Setup(l => l.Client).Returns(mockClient.Object); // Act & Assert - var exception = Assert.Throws(() => new CosmosNoSqlCollection(mockDatabase.Object, "collection")); + var exception = Assert.Throws(() => new CosmosNoSqlCollection(mockDatabase.Object, "collection")); Assert.Contains(nameof(CosmosClientOptions.UseSystemTextJsonSerializerWithOptions), exception.Message); } @@ -64,9 +64,10 @@ public void ConstructorWithoutSystemTextJsonSerializerOptionsThrowsArgumentExcep public void ConstructorWithDeclarativeModelInitializesCollection() { // Act & Assert - using var collection = new CosmosNoSqlCollection( + using var collection = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); Assert.NotNull(collection); } @@ -81,10 +82,10 @@ public void ConstructorWithImperativeModelInitializesCollection() }; // Act - using var collection = new CosmosNoSqlCollection( + using var collection = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection", - new() { Definition = definition }); + new() { Definition = definition, PartitionKeyPropertyNames = ["Id"] }); // Assert Assert.NotNull(collection); @@ -117,9 +118,10 @@ public async Task CollectionExistsReturnsValidResultAsync(List collectio It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - collectionName); + collectionName, + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act var actualResult = await sut.CollectionExistsAsync(); @@ -159,10 +161,15 @@ public async Task EnsureCollectionExistsUsesValidContainerPropertiesAsync(Indexi It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, CollectionName, - new() { IndexingMode = indexingMode, Automatic = indexingMode != IndexingMode.None }); + new() + { + IndexingMode = indexingMode, + Automatic = indexingMode != IndexingMode.None, + PartitionKeyPropertyNames = ["Id"] + }); var expectedVectorEmbeddingPolicy = new VectorEmbeddingPolicy( [ @@ -259,9 +266,10 @@ public async Task EnsureCollectionExistsInvokesValidMethodsAsync(List co It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - CollectionName); + CollectionName, + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act await sut.EnsureCollectionExistsAsync(); @@ -280,31 +288,23 @@ public async Task EnsureCollectionExistsInvokesValidMethodsAsync(List co [InlineData("partitionKey", true)] public async Task DeleteInvokesValidMethodsAsync( string expectedPartitionKey, - bool useCompositeKeyCollection) + bool useExplicitPartitionKey) { // Arrange const string RecordKey = "recordKey"; const string PartitionKey = "partitionKey"; - // Act - if (useCompositeKeyCollection) - { - using var sut = new CosmosNoSqlCollection( - this._mockDatabase.Object, - "collection"); + using var sut = new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection", + useExplicitPartitionKey ? null : new() { PartitionKeyPropertyNames = ["HotelId"] }); - await ((VectorStoreCollection)sut).DeleteAsync( - new CosmosNoSqlCompositeKey(RecordKey, PartitionKey)); - } - else - { - using var sut = new CosmosNoSqlCollection( - this._mockDatabase.Object, - "collection"); + // Act + var key = useExplicitPartitionKey + ? new CosmosNoSqlKey(RecordKey, PartitionKey) + : new CosmosNoSqlKey(RecordKey, RecordKey); - await ((VectorStoreCollection)sut).DeleteAsync( - RecordKey); - } + await sut.DeleteAsync(key); // Assert this._mockContainer.Verify(l => l.DeleteItemAsync( @@ -321,12 +321,15 @@ public async Task DeleteBatchInvokesValidMethodsAsync() // Arrange List recordKeys = ["key1", "key2"]; - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act - await sut.DeleteAsync(recordKeys); + // With the new design, keys must be CosmosNoSqlKey + var cosmosKeys = recordKeys.Select(k => new CosmosNoSqlKey(k, k)).ToList(); + await sut.DeleteAsync(cosmosKeys); // Assert foreach (var key in recordKeys) @@ -344,9 +347,10 @@ public async Task DeleteBatchInvokesValidMethodsAsync() public async Task DeleteCollectionInvokesValidMethodsAsync() { // Arrange - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act await sut.EnsureCollectionDeletedAsync(); @@ -388,12 +392,13 @@ public async Task GetReturnsValidRecordAsync() It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act - var result = await sut.GetAsync(RecordKey); + var result = await sut.GetAsync(new CosmosNoSqlKey(RecordKey, RecordKey)); // Assert Assert.NotNull(result); @@ -431,12 +436,19 @@ public async Task GetBatchReturnsValidRecordAsync() It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act - var results = await sut.GetAsync(["key1", "key2", "key3"]).ToListAsync(); + var keys = new List + { + new("key1", "key1"), + new("key2", "key2"), + new("key3", "key3") + }; + var results = await sut.GetAsync(keys).ToListAsync(); // Assert Assert.NotNull(results[0]); @@ -458,9 +470,10 @@ public async Task CanUpsertRecordAsync() // Arrange var hotel = new CosmosNoSqlHotel("key") { HotelName = "Test Name" }; - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act await sut.UpsertAsync(hotel); @@ -484,9 +497,10 @@ public async Task CanUpsertManyRecordsAsync() var hotel2 = new CosmosNoSqlHotel("key2") { HotelName = "Test Name 2" }; var hotel3 = new CosmosNoSqlHotel("key3") { HotelName = "Test Name 3" }; - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act await sut.UpsertAsync([hotel1, hotel2, hotel3]); @@ -528,9 +542,10 @@ public async Task VectorizedSearchReturnsValidRecordAsync() It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act var results = await sut.SearchAsync(new ReadOnlyMemory([1f, 2f, 3f]), top: 3).ToListAsync(); @@ -547,9 +562,10 @@ public async Task VectorizedSearchReturnsValidRecordAsync() public async Task VectorizedSearchWithUnsupportedVectorTypeThrowsExceptionAsync() { // Arrange - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); // Act & Assert await Assert.ThrowsAsync(async () => @@ -560,9 +576,10 @@ await Assert.ThrowsAsync(async () => public async Task VectorizedSearchWithNonExistentVectorPropertyNameThrowsExceptionAsync() { // Arrange - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection"); + "collection", + new() { PartitionKeyPropertyNames = ["HotelId"] }); var searchOptions = new VectorSearchOptions { VectorProperty = r => "non-existent-property" }; diff --git a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlVectorStoreTests.cs b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlVectorStoreTests.cs index 5f8f20cf865c..e829dd1cdf5a 100644 --- a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlVectorStoreTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlVectorStoreTests.cs @@ -48,12 +48,10 @@ public void GetCollectionWithSupportedKeyReturnsCollection() using var sut = new Microsoft.SemanticKernel.Connectors.CosmosNoSql.CosmosNoSqlVectorStore(this._mockDatabase.Object); // Act - var collectionWithStringKey = sut.GetCollection("collection1"); - var collectionWithCompositeKey = sut.GetCollection("collection1"); + var collection = sut.GetCollection("collection1"); // Assert - Assert.NotNull(collectionWithStringKey); - Assert.NotNull(collectionWithCompositeKey); + Assert.NotNull(collection); } [Fact] @@ -63,7 +61,7 @@ public void GetCollectionWithoutFactoryReturnsDefaultCollection() using var sut = new Microsoft.SemanticKernel.Connectors.CosmosNoSql.CosmosNoSqlVectorStore(this._mockDatabase.Object); // Act - var collection = sut.GetCollection("collection"); + var collection = sut.GetCollection("collection"); // Assert Assert.NotNull(collection); diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/CollectionManagementTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/CollectionManagementTests.cs index 9d85578c5ba8..1e32e41f37a6 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/CollectionManagementTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/CollectionManagementTests.cs @@ -97,7 +97,7 @@ public sealed class Record : TestRecord } public virtual VectorStoreCollection GetCollection() - => fixture.TestStore.DefaultVectorStore.GetCollection(this.CollectionName, this.CreateRecordDefinition()); + => fixture.TestStore.CreateCollection(this.CollectionName, this.CreateRecordDefinition()); public virtual VectorStoreCollectionDefinition CreateRecordDefinition() => new() diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/DistanceFunctionTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/DistanceFunctionTests.cs index 0ca4fb12c9a9..182954d9d985 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/DistanceFunctionTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/DistanceFunctionTests.cs @@ -185,7 +185,7 @@ public virtual VectorStoreCollection CreateCollection(string ] }; - return this.TestStore.DefaultVectorStore.GetCollection(this.CollectionName, definition); + return this.TestStore.CreateCollection(this.CollectionName, definition); } } diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/EmbeddingGenerationTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/EmbeddingGenerationTests.cs index f46d388146bc..c3b13dc8fde5 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/EmbeddingGenerationTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/EmbeddingGenerationTests.cs @@ -574,7 +574,7 @@ public virtual int GenerateNextCounter() => Interlocked.Increment(ref this._counter); } - private sealed class FakeEmbeddingGenerator(int? replaceLast = null) : IEmbeddingGenerator> + protected sealed class FakeEmbeddingGenerator(int? replaceLast = null) : IEmbeddingGenerator> { public Task>> GenerateAsync( IEnumerable values, diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/FilterTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/FilterTests.cs index aada9925ab24..1112df1056d0 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/FilterTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/FilterTests.cs @@ -640,7 +640,7 @@ public override async Task InitializeAsync() if (this.TestDynamic) { - this.DynamicCollection = this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition()); + this.DynamicCollection = this.TestStore.CreateDynamicCollection(this.CollectionName, this.CreateRecordDefinition()); } } diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/IndexKindTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/IndexKindTests.cs index ac87366edb8e..26738dbeba66 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/IndexKindTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/IndexKindTests.cs @@ -75,7 +75,7 @@ public virtual VectorStoreCollection CreateCollection(string ] }; - return this.TestStore.DefaultVectorStore.GetCollection(this.CollectionName, definition); + return this.TestStore.CreateCollection(this.CollectionName, definition); } } diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/ModelTests/DynamicModelTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/ModelTests/DynamicModelTests.cs index b63dcbe29abb..387f3ad05fe3 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/ModelTests/DynamicModelTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/ModelTests/DynamicModelTests.cs @@ -396,7 +396,7 @@ public abstract class Fixture : DynamicVectorStoreCollectionFixture protected override string KeyPropertyName => DynamicModelTests.KeyPropertyName; protected override VectorStoreCollection> GetCollection() - => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition()); + => this.TestStore.CreateDynamicCollection(this.CollectionName, this.CreateRecordDefinition()); public override VectorStoreCollectionDefinition CreateRecordDefinition() => new() diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/Support/TestStore.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/Support/TestStore.cs index f5a3924e464c..ced2b255d15a 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/Support/TestStore.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/Support/TestStore.cs @@ -92,6 +92,26 @@ _ when typeof(TKey) == typeof(Guid) => (TKey)(object)new Guid($"00000000-0000-00 public virtual string AdjustCollectionName(string baseName) => baseName; + /// + /// Creates a collection for the given name and definition. + /// Override this to provide provider-specific collection options (e.g., partition key configuration). + /// + public virtual VectorStoreCollection CreateCollection( + string name, + VectorStoreCollectionDefinition definition) + where TKey : notnull + where TRecord : class + => this.DefaultVectorStore.GetCollection(name, definition); + + /// + /// Creates a dynamic collection for the given name and definition. + /// Override this to provide provider-specific collection options (e.g., partition key configuration). + /// + public virtual VectorStoreCollection> CreateDynamicCollection( + string name, + VectorStoreCollectionDefinition definition) + => this.DefaultVectorStore.GetDynamicCollection(name, definition); + /// Loops until the expected number of records is visible in the given collection. /// Some databases upsert asynchronously, meaning that our seed data may not be visible immediately to tests. public virtual async Task WaitForDataAsync( diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/Support/VectorStoreCollectionFixtureBase.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/Support/VectorStoreCollectionFixtureBase.cs index 81877ffa5655..10ae77bc8a7d 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/Support/VectorStoreCollectionFixtureBase.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/Support/VectorStoreCollectionFixtureBase.cs @@ -34,7 +34,7 @@ public abstract class VectorStoreCollectionFixtureBase : VectorSt protected virtual string IndexKind => this.TestStore.DefaultIndexKind; protected virtual VectorStoreCollection GetCollection() - => this.TestStore.DefaultVectorStore.GetCollection(this.CollectionName, this.CreateRecordDefinition()); + => this.TestStore.CreateCollection(this.CollectionName, this.CreateRecordDefinition()); public override async Task InitializeAsync() { diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/Support/VectorStoreFixture.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/Support/VectorStoreFixture.cs index e88ba33ee78f..091de0bcee69 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/Support/VectorStoreFixture.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/Support/VectorStoreFixture.cs @@ -23,4 +23,20 @@ public virtual Task DisposeAsync() public virtual TKey GenerateNextKey() => this.TestStore.GenerateKey(Interlocked.Increment(ref this._nextKeyValue)); + + /// + /// Creates a collection for the given name and definition. + /// Delegates to which can be overridden for provider-specific options. + /// + public virtual VectorStoreCollection CreateCollection(string name, VectorStoreCollectionDefinition definition) + where TKey : notnull + where TRecord : class + => this.TestStore.CreateCollection(name, definition); + + /// + /// Creates a dynamic collection for the given name and definition. + /// Delegates to which can be overridden for provider-specific options. + /// + public virtual VectorStoreCollection> CreateDynamicCollection(string name, VectorStoreCollectionDefinition definition) + => this.TestStore.CreateDynamicCollection(name, definition); } diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/DataTypeTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/DataTypeTests.cs index 1a3342e80d79..09c2db775e79 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/DataTypeTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/DataTypeTests.cs @@ -192,7 +192,7 @@ protected virtual async Task Test( await fixture.TestStore.WaitForDataAsync(fixture.Collection, recordCount: 0); } - var dynamicCollection = fixture.VectorStore.GetDynamicCollection(fixture.CollectionName, fixture.CreateRecordDefinition()); + var dynamicCollection = fixture.CreateDynamicCollection(fixture.CollectionName, fixture.CreateRecordDefinition()); if (fixture.RecreateCollection) { diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/EmbeddingTypeTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/EmbeddingTypeTests.cs index 5bca4e5fd146..5fafde527ced 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/EmbeddingTypeTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/EmbeddingTypeTests.cs @@ -50,7 +50,7 @@ protected virtual async Task Test( await fixture.VectorStore.EnsureCollectionDeletedAsync(fixture.CollectionName); - var collection = fixture.VectorStore.GetCollection>(fixture.CollectionName, fixture.CreateRecordDefinition(embeddingGenerator: null, distanceFunction, dimensions)); + var collection = fixture.CreateCollection>(fixture.CollectionName, fixture.CreateRecordDefinition(embeddingGenerator: null, distanceFunction, dimensions)); await collection.EnsureCollectionExistsAsync(); var key = fixture.GenerateNextKey(); @@ -93,7 +93,7 @@ protected virtual async Task Test( await collection.DeleteAsync(key); } - var dynamicCollection = fixture.VectorStore.GetDynamicCollection(fixture.CollectionName, fixture.CreateRecordDefinition(embeddingGenerator, distanceFunction, dimensions)); + var dynamicCollection = fixture.CreateDynamicCollection(fixture.CollectionName, fixture.CreateRecordDefinition(embeddingGenerator, distanceFunction, dimensions)); if (fixture.RecreateCollection) { @@ -142,14 +142,14 @@ protected virtual async Task Test( await collection.DeleteAsync(key); } - var collection2 = fixture.VectorStore.GetCollection(fixture.CollectionName, fixture.CreateRecordDefinition(embeddingGenerator, distanceFunction, dimensions)); + var collection2 = fixture.CreateCollection(fixture.CollectionName, fixture.CreateRecordDefinition(embeddingGenerator, distanceFunction, dimensions)); if (fixture.RecreateCollection) { await collection2.EnsureCollectionExistsAsync(); } - var key2 = fixture.GenerateNextKey(); + key = fixture.GenerateNextKey(); var record2 = new RecordWithString { Key = key, @@ -217,7 +217,7 @@ public class Record public abstract class Fixture : VectorStoreFixture { - protected virtual string CollectionNameBase => nameof(EmbeddingTypeTests); + protected virtual string CollectionNameBase => nameof(EmbeddingTypeTests<>); public virtual string CollectionName => this.TestStore.AdjustCollectionName(this.CollectionNameBase); public virtual VectorStoreCollectionDefinition CreateRecordDefinition(IEmbeddingGenerator? embeddingGenerator, string? distanceFunction, int dimensions) diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/KeyTypeTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/KeyTypeTests.cs index fb9db07dac62..a10c469591f0 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/KeyTypeTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/KeyTypeTests.cs @@ -209,11 +209,11 @@ public abstract class Fixture : VectorStoreFixture public virtual VectorStoreCollection> CreateCollection(bool? withAutoGeneration) where TKey : notnull - => this.TestStore.DefaultVectorStore.GetCollection>(this.CollectionName, this.CreateRecordDefinition(withAutoGeneration)); + => this.TestStore.CreateCollection>(this.CollectionName, this.CreateRecordDefinition(withAutoGeneration)); public virtual VectorStoreCollection> CreateDynamicCollection(bool withAutoGeneration) where TKey : notnull - => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition(withAutoGeneration)); + => this.TestStore.CreateDynamicCollection(this.CollectionName, this.CreateRecordDefinition(withAutoGeneration)); public virtual VectorStoreCollectionDefinition CreateRecordDefinition(bool? withAutoGeneration) where TKey : notnull From 479169b3de530960ece925243b5eff86b3469113 Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Sat, 14 Feb 2026 18:53:20 +0100 Subject: [PATCH 3/4] Properly use ReadItem and partition keys in Cosmos NoSQL Closes #13549 --- .../CosmosNoSql/CosmosNoSqlCollection.cs | 49 ++++++++++++++----- .../CosmosNoSqlCollectionQueryBuilder.cs | 49 ------------------- .../CosmosNoSql/CosmosNoSqlMapper.cs | 42 +++++++++------- .../CosmosNoSqlCollectionQueryBuilderTests.cs | 40 --------------- .../CosmosNoSqlCollectionTests.cs | 47 ++++++------------ 5 files changed, 75 insertions(+), 152 deletions(-) diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs index 680fb47a2fbc..849db36df729 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs @@ -7,6 +7,7 @@ using System.Diagnostics.CodeAnalysis; using System.Linq; using System.Linq.Expressions; +using System.Net; using System.Runtime.CompilerServices; using System.Text.Json; using System.Text.Json.Nodes; @@ -323,9 +324,33 @@ await this._database { Verify.NotNull(key); - return await this.GetAsync([key], options, cancellationToken) - .FirstOrDefaultAsync(cancellationToken) - .ConfigureAwait(false); + const string OperationName = "ReadItem"; + + var includeVectors = options?.IncludeVectors ?? false; + if (includeVectors && this._model.EmbeddingGenerationRequired) + { + throw new NotSupportedException(VectorDataStrings.IncludeVectorsNotSupportedWithEmbeddingGeneration); + } + + var (documentId, partitionKey) = this.GetDocumentIdAndPartitionKey(key); + + var jsonObject = await this.RunOperationAsync(OperationName, async () => + { + try + { + var response = await this._database + .GetContainer(this.Name) + .ReadItemAsync(documentId, partitionKey, cancellationToken: cancellationToken) + .ConfigureAwait(false); + return response.Resource; + } + catch (CosmosException e) when (e.StatusCode == HttpStatusCode.NotFound) + { + return null; + } + }).ConfigureAwait(false); + + return jsonObject is null ? default : this._mapper.MapFromStorageToDataModel(jsonObject, includeVectors); } /// @@ -336,7 +361,7 @@ public override async IAsyncEnumerable GetAsync( { Verify.NotNull(keys); - const string OperationName = "GetItemQueryIterator"; + const string OperationName = "ReadManyItems"; var includeVectors = options?.IncludeVectors ?? false; if (includeVectors && this._model.EmbeddingGenerationRequired) @@ -344,18 +369,20 @@ public override async IAsyncEnumerable GetAsync( throw new NotSupportedException(VectorDataStrings.IncludeVectorsNotSupportedWithEmbeddingGeneration); } - var documentIds = keys.Select(k => this.GetDocumentIdAndPartitionKey(k).DocumentId).ToList(); - if (documentIds.Count == 0) + var items = keys.Select(this.GetDocumentIdAndPartitionKey).ToList(); + + if (items.Count == 0) { yield break; } - var queryDefinition = CosmosNoSqlCollectionQueryBuilder.BuildSelectQuery( - this._model, - documentIds, - includeVectors); + var response = await this.RunOperationAsync(OperationName, () => + this._database + .GetContainer(this.Name) + .ReadManyItemsAsync(items, cancellationToken: cancellationToken)) + .ConfigureAwait(false); - await foreach (var jsonObject in this.GetItemsAsync(queryDefinition, OperationName, cancellationToken).ConfigureAwait(false)) + foreach (var jsonObject in response.Resource) { var record = this._mapper.MapFromStorageToDataModel(jsonObject, includeVectors); diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs index 305292c2c7f8..e83b6cf577cb 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionQueryBuilder.cs @@ -224,55 +224,6 @@ internal static QueryDefinition BuildSearchQuery( return queryDefinition; } - /// - /// Builds to get items from Azure CosmosDB NoSQL. - /// - public static QueryDefinition BuildSelectQuery( - CollectionModel model, - List documentIds, - bool includeVectors) - { - Verify.True(documentIds.Count > 0, "At least one document ID should be provided.", nameof(documentIds)); - - const string RecordKeyVariableName = "@rk"; - - var tableVariableName = CosmosNoSqlConstants.ContainerAlias; - - IEnumerable projectionProperties = model.Properties; - if (!includeVectors) - { - projectionProperties = projectionProperties.Where(p => p is not VectorPropertyModel); - } - - var selectClauseArguments = string.Join( - ",", - projectionProperties.Select(field => GeneratePropertyAccess(tableVariableName, field.StorageName))); - - // Build WHERE clause using only the document id. - // The partition key is provided via RequestOptions for point reads, not in the query. - var whereClauseArguments = string.Join( - " OR ", - documentIds.Select((_, index) => - $"({GeneratePropertyAccess(tableVariableName, model.KeyProperty.StorageName)} = {RecordKeyVariableName}{index})")); - - var query = $""" - SELECT {selectClauseArguments} - FROM {tableVariableName} - WHERE {whereClauseArguments} - """; - - var queryDefinition = new QueryDefinition(query); - - for (var i = 0; i < documentIds.Count; i++) - { - var documentIdString = documentIds[i]; - Verify.NotNullOrWhiteSpace(documentIdString); - queryDefinition.WithParameter($"{RecordKeyVariableName}{i}", documentIdString); - } - - return queryDefinition; - } - #region private #pragma warning disable CS0618 // VectorSearchFilter is obsolete diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlMapper.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlMapper.cs index 7b176c339b50..8a8644b80198 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlMapper.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlMapper.cs @@ -118,30 +118,34 @@ public TRecord MapFromStorageToDataModel(JsonObject storageModel, bool includeVe // See above comment. RenameJsonProperty(storageModel, CosmosNoSqlConstants.ReservedKeyPropertyName, this._keyProperty.TemporaryStorageName!); - if (includeVectors) + foreach (var vectorProperty in this._model.VectorProperties) { - foreach (var vectorProperty in this._model.VectorProperties) + if (!includeVectors) { - var arrayNode = storageModel[vectorProperty.StorageName]; - if (arrayNode is null) - { - continue; - } + // Remove vector properties so they deserialize as default (e.g. empty ReadOnlyMemory). + storageModel.Remove(vectorProperty.StorageName); + continue; + } - // Embedding is stored as a simple JSON array, so convert it to the expected object shape for deserialization - if (vectorProperty.Type == typeof(Embedding) - || vectorProperty.Type == typeof(Embedding) - || vectorProperty.Type == typeof(Embedding)) - { - storageModel[vectorProperty.StorageName] = new JsonObject - { - [nameof(Embedding<>.Vector)] = arrayNode.DeepClone() - }; - } + var arrayNode = storageModel[vectorProperty.StorageName]; + if (arrayNode is null) + { + continue; + } - // For byte[], ReadOnlyMemory, sbyte[], ReadOnlyMemory, float[], ReadOnlyMemory, - // the custom converters (for byte) and default converters (for others) handle deserialization correctly. + // Embedding is stored as a simple JSON array, so convert it to the expected object shape for deserialization + if (vectorProperty.Type == typeof(Embedding) + || vectorProperty.Type == typeof(Embedding) + || vectorProperty.Type == typeof(Embedding)) + { + storageModel[vectorProperty.StorageName] = new JsonObject + { + [nameof(Embedding<>.Vector)] = arrayNode.DeepClone() + }; } + + // For byte[], ReadOnlyMemory, sbyte[], ReadOnlyMemory, float[], ReadOnlyMemory, + // the custom converters (for byte) and default converters (for others) handle deserialization correctly. } return storageModel.Deserialize(this._jsonSerializerOptions)!; diff --git a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionQueryBuilderTests.cs b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionQueryBuilderTests.cs index ce3f579d8370..a8556d2865b2 100644 --- a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionQueryBuilderTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionQueryBuilderTests.cs @@ -1,7 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. using System; -using System.Collections.Generic; using Microsoft.Extensions.VectorData; using Microsoft.Extensions.VectorData.ProviderServices; using Microsoft.SemanticKernel.Connectors.CosmosNoSql; @@ -186,45 +185,6 @@ public void BuildSearchQueryWithoutFilterDoesNotContainWhereClause() Assert.Equal(vector, queryParameters[0].Value); } - [Fact] - public void BuildSelectQueryByDefaultReturnsValidQueryDefinition() - { - // Arrange - const string ExpectedQueryText = """ - SELECT x["id"],x["TestProperty1"],x["TestProperty2"] - FROM x - WHERE (x["id"] = @rk0) - """; - - var model = new CosmosNoSqlModelBuilder().BuildDynamic( - new() - { - Properties = - [ - new VectorStoreKeyProperty("Key", typeof(string)), - new VectorStoreDataProperty("TestProperty1", typeof(string)), - new VectorStoreDataProperty("TestProperty2", typeof(string)) - ] - }, - defaultEmbeddingGenerator: null); - List documentIds = ["id"]; - - // Act - var queryDefinition = CosmosNoSqlCollectionQueryBuilder.BuildSelectQuery( - model, - documentIds, - includeVectors: true); - - var queryText = queryDefinition.QueryText; - var queryParameters = queryDefinition.GetQueryParameters(); - - // Assert - Assert.Equal(ExpectedQueryText, queryText); - - Assert.Equal("@rk0", queryParameters[0].Name); - Assert.Equal("id", queryParameters[0].Value); - } - [Fact] public void BuildSearchQueryWithHybridFieldsReturnsValidHybridQueryDefinition() { diff --git a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs index 7cc8cc3b2453..556fcc0cd60d 100644 --- a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs @@ -370,27 +370,18 @@ public async Task GetReturnsValidRecordAsync() var jsonObject = new JsonObject { ["id"] = RecordKey, ["HotelName"] = "Test Name" }; - var mockFeedResponse = new Mock>(); - mockFeedResponse + var mockItemResponse = new Mock>(); + mockItemResponse .Setup(l => l.Resource) - .Returns([jsonObject]); - - var mockFeedIterator = new Mock>(); - mockFeedIterator - .SetupSequence(l => l.HasMoreResults) - .Returns(true) - .Returns(false); - - mockFeedIterator - .Setup(l => l.ReadNextAsync(It.IsAny())) - .ReturnsAsync(mockFeedResponse.Object); + .Returns(jsonObject); this._mockContainer - .Setup(l => l.GetItemQueryIterator( - It.IsAny(), - It.IsAny(), - It.IsAny())) - .Returns(mockFeedIterator.Object); + .Setup(l => l.ReadItemAsync( + RecordKey, + new PartitionKey(RecordKey), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(mockItemResponse.Object); using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, @@ -419,22 +410,12 @@ public async Task GetBatchReturnsValidRecordAsync() .Setup(l => l.Resource) .Returns([jsonObject1, jsonObject2, jsonObject3]); - var mockFeedIterator = new Mock>(); - mockFeedIterator - .SetupSequence(l => l.HasMoreResults) - .Returns(true) - .Returns(false); - - mockFeedIterator - .Setup(l => l.ReadNextAsync(It.IsAny())) - .ReturnsAsync(mockFeedResponse.Object); - this._mockContainer - .Setup(l => l.GetItemQueryIterator( - It.IsAny(), - It.IsAny(), - It.IsAny())) - .Returns(mockFeedIterator.Object); + .Setup(l => l.ReadManyItemsAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(mockFeedResponse.Object); using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, From 55695ca26e9ac917d14212dc4ce1b2233da8756f Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Mon, 16 Feb 2026 17:11:59 +0100 Subject: [PATCH 4/4] Allow string/Guid keys with partition key = document ID --- .../samples/Demos/VectorStoreRAG/Program.cs | 2 +- .../CosmosNoSql/CosmosNoSqlCollection.cs | 65 +++-- .../CosmosNoSqlCollectionOptions.cs | 15 +- .../CosmosNoSqlServiceCollectionExtensions.cs | 45 +-- .../CosmosNoSqlCollectionOptionsTests.cs | 13 +- .../CosmosNoSqlDependencyInjectionTests.cs | 39 ++- .../CosmosNoSqlEmbeddingGenerationTests.cs | 160 +---------- .../Support/CosmosNoSqlCollectionAdapter.cs | 103 ------- .../CosmosNoSqlDynamicCollectionAdapter.cs | 91 ------ .../Support/CosmosNoSqlTestStore.cs | 58 ---- .../CosmosNoSqlCollectionTests.cs | 260 +++++++++++++----- .../CosmosNoSqlVectorStoreTests.cs | 37 +++ 12 files changed, 341 insertions(+), 547 deletions(-) delete mode 100644 dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlCollectionAdapter.cs delete mode 100644 dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlDynamicCollectionAdapter.cs diff --git a/dotnet/samples/Demos/VectorStoreRAG/Program.cs b/dotnet/samples/Demos/VectorStoreRAG/Program.cs index 71a26b541706..ff3b2e934e42 100644 --- a/dotnet/samples/Demos/VectorStoreRAG/Program.cs +++ b/dotnet/samples/Demos/VectorStoreRAG/Program.cs @@ -83,7 +83,7 @@ appConfig.CosmosMongoConfig.DatabaseName); break; case "CosmosNoSql": - kernelBuilder.Services.AddCosmosNoSqlCollection>( + kernelBuilder.Services.AddCosmosNoSqlCollection>( appConfig.RagConfig.CollectionName, appConfig.CosmosNoSqlConfig.ConnectionString, appConfig.CosmosNoSqlConfig.DatabaseName); diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs index 849db36df729..df0eac3a766b 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs @@ -27,7 +27,9 @@ namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql; /// /// Service for storing and retrieving vector records, that uses Azure CosmosDB NoSQL as the underlying storage. /// -/// The data type of the record key. Must be . +/// +/// The data type of the record key. Supported types are and (in which case the partition key +/// is the document ID), and (for other partition key configurations). /// The data model to use for adding, updating and retrieving data from storage. #pragma warning disable CA1711 // Identifiers should not have incorrect suffix public class CosmosNoSqlCollection : VectorStoreCollection, IKeywordHybridSearchable @@ -134,10 +136,10 @@ internal CosmosNoSqlCollection( Func modelFactory, CosmosNoSqlCollectionOptions? options) { - // Validate TKey: must be CosmosNoSqlKey, or object (used by CosmosNoSqlDynamicCollection). - if (typeof(TKey) != typeof(object) && typeof(TKey) != typeof(CosmosNoSqlKey)) + // Validate TKey: must be string or Guid (partition key = document ID), CosmosNoSqlKey (other parittion key cases), or object (used by CosmosNoSqlDynamicCollection). + if (typeof(TKey) != typeof(string) && typeof(TKey) != typeof(Guid) && typeof(TKey) != typeof(CosmosNoSqlKey) && typeof(TKey) != typeof(object)) { - throw new NotSupportedException($"The key type must be CosmosNoSqlKey. Received: {typeof(TKey).Name}"); + throw new NotSupportedException($"The key type must be string, Guid, or CosmosNoSqlKey. Received: {typeof(TKey).Name}"); } try @@ -166,26 +168,47 @@ internal CosmosNoSqlCollection( : new CosmosNoSqlMapper(this._model, options.JsonSerializerOptions); // Setup partition key properties (supports hierarchical partition keys up to 3 levels) - if (options.PartitionKeyPropertyNames.Count > 3) + if (options.PartitionKeyProperties is null) { - throw new ArgumentException("Cosmos DB supports at most 3 levels of hierarchical partition keys."); + // No partition key property names configured: default to using the key property as the partition key. + // This is the most common Cosmos DB strategy where the document ID and partition key are the same value. + this._partitionKeyProperties = [this._model.KeyProperty]; } - - this._partitionKeyProperties = new List(options.PartitionKeyPropertyNames.Count); - - foreach (var propertyName in options.PartitionKeyPropertyNames) + else { - if (!this._model.PropertyMap.TryGetValue(propertyName, out var property)) + if (options.PartitionKeyProperties.Count > 3) { - throw new ArgumentException($"Partition key property '{propertyName}' is not part of the record definition."); + throw new ArgumentException("Cosmos DB supports at most 3 levels of hierarchical partition keys."); } - if (property.Type != typeof(string) && property.Type != typeof(bool) && property.Type != typeof(double) && property.Type != typeof(Guid)) + this._partitionKeyProperties = new List(options.PartitionKeyProperties.Count); + + foreach (var propertyName in options.PartitionKeyProperties) { - throw new ArgumentException($"Partition key property '{propertyName}' must be string, bool, double, or Guid."); + if (!this._model.PropertyMap.TryGetValue(propertyName, out var property)) + { + throw new ArgumentException($"Partition key property '{propertyName}' is not part of the record definition."); + } + + if (property.Type != typeof(string) && property.Type != typeof(bool) && property.Type != typeof(double) && property.Type != typeof(Guid)) + { + throw new ArgumentException($"Partition key property '{propertyName}' must be string, bool, double, or Guid."); + } + + this._partitionKeyProperties.Add(property); } + } - this._partitionKeyProperties.Add(property); + // When using simple key types (string/Guid), the partition key must be the key property only, + // since the partition key value cannot be independently specified via a simple key. + // Users who need hierarchical partition keys, a non-key partition key, or no partition key at all + // must use CosmosNoSqlKey as the key type. + if ((typeof(TKey) == typeof(string) || typeof(TKey) == typeof(Guid)) + && (this._partitionKeyProperties is not [var singlePkProperty] || singlePkProperty != this._model.KeyProperty)) + { + throw new ArgumentException( + $"When using {typeof(TKey).Name} as the key type, the partition key must be the key property " + + $"('{this._model.KeyProperty.ModelName}'). To use a different partition key configuration, use CosmosNoSqlKey as the key type."); } this._collectionMetadata = new() @@ -706,10 +729,14 @@ private Task RunOperationAsync(string operationName, Func> operati /// Gets the document ID and partition key from a key. /// private (string DocumentId, PartitionKey PartitionKey) GetDocumentIdAndPartitionKey(TKey key) - => key is CosmosNoSqlKey cosmosKey - ? (cosmosKey.DocumentId, cosmosKey.PartitionKey) - : throw new InvalidOperationException( - $"Keys must be of type CosmosNoSqlKey. Received key of type '{key.GetType().FullName}'."); + => key switch + { + CosmosNoSqlKey cosmosKey => (cosmosKey.DocumentId, cosmosKey.PartitionKey), + string s => (s, new PartitionKey(s)), + Guid g => (g.ToString(), new PartitionKey(g.ToString())), + _ => throw new InvalidOperationException( + $"Keys must be of type string, Guid, or CosmosNoSqlKey. Received key of type '{key.GetType().FullName}'.") + }; /// /// Returns instance of with applied indexing policy. diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionOptions.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionOptions.cs index 7dbea2dd9f15..d4b55f33c449 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionOptions.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollectionOptions.cs @@ -24,7 +24,7 @@ public CosmosNoSqlCollectionOptions() internal CosmosNoSqlCollectionOptions(CosmosNoSqlCollectionOptions? source) : base(source) { this.JsonSerializerOptions = source?.JsonSerializerOptions; - this.PartitionKeyPropertyNames = source?.PartitionKeyPropertyNames is { Count: > 0 } names ? [.. names] : []; + this.PartitionKeyProperties = source?.PartitionKeyProperties is null ? null : [.. source.PartitionKeyProperties]; this.IndexingMode = source?.IndexingMode ?? Default.IndexingMode; this.Automatic = source?.Automatic ?? Default.Automatic; } @@ -39,20 +39,17 @@ internal CosmosNoSqlCollectionOptions(CosmosNoSqlCollectionOptions? source) : ba /// /// /// - /// Cosmos DB supports up to 3 levels of hierarchical partition keys. Provide the property names in hierarchical order. - /// For a single partition key, provide a list with one element. For hierarchical partition keys, provide up to 3 elements. - /// - /// /// Selecting a partition key is critical for performance and scalability. Choose properties with high cardinality /// that evenly distribute requests. See for guidance. /// /// - /// If empty, you must use as the key type with an explicitly constructed - /// . If your scenario does not require partitioning, use , - /// though this is not recommended for production workloads. + /// When (the default), the key property (document ID) is automatically used as the partition key - a common + /// Cosmos DB strategy; in this mode, the collection key type must be or . + /// To use a different partition key (or hierarchical partition keys), specify the key properties here and use + /// as the key type. /// /// - public IReadOnlyList PartitionKeyPropertyNames { get; set; } = []; + public IReadOnlyList? PartitionKeyProperties { get; set; } /// /// Specifies the indexing mode in the Azure Cosmos DB service. diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlServiceCollectionExtensions.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlServiceCollectionExtensions.cs index 826e0a56e73f..fe55f440536a 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlServiceCollectionExtensions.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlServiceCollectionExtensions.cs @@ -127,18 +127,20 @@ public static IServiceCollection AddKeyedCosmosNoSqlVectorStore( /// [RequiresUnreferencedCode(DynamicCodeMessage)] [RequiresDynamicCode(UnreferencedCodeMessage)] - public static IServiceCollection AddCosmosNoSqlCollection( + public static IServiceCollection AddCosmosNoSqlCollection( this IServiceCollection services, string name, CosmosNoSqlCollectionOptions? options = default, ServiceLifetime lifetime = ServiceLifetime.Singleton) + where TKey : notnull where TRecord : class - => AddKeyedCosmosNoSqlCollection(services, serviceKey: null, name, options, lifetime); + => AddKeyedCosmosNoSqlCollection(services, serviceKey: null, name, options, lifetime); /// /// Registers a keyed as /// with retrieved from the dependency injection container. /// + /// The type of the key. /// The type of the record. /// The to register the on. /// The key with which to associate the collection. @@ -148,26 +150,27 @@ public static IServiceCollection AddCosmosNoSqlCollection( /// Service collection. [RequiresUnreferencedCode(DynamicCodeMessage)] [RequiresDynamicCode(UnreferencedCodeMessage)] - public static IServiceCollection AddKeyedCosmosNoSqlCollection( + public static IServiceCollection AddKeyedCosmosNoSqlCollection( this IServiceCollection services, object? serviceKey, string name, CosmosNoSqlCollectionOptions? options = default, ServiceLifetime lifetime = ServiceLifetime.Singleton) + where TKey : notnull where TRecord : class { Verify.NotNull(services); Verify.NotNullOrWhiteSpace(name); - services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection), serviceKey, (sp, _) => + services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection), serviceKey, (sp, _) => { var database = sp.GetRequiredService(); options = GetCollectionOptions(sp, _ => options); - return new CosmosNoSqlCollection(database, name, options); + return new CosmosNoSqlCollection(database, name, options); }, lifetime)); - AddAbstractions(services, serviceKey, lifetime); + AddAbstractions(services, serviceKey, lifetime); return services; } @@ -176,23 +179,25 @@ public static IServiceCollection AddKeyedCosmosNoSqlCollection( /// Registers a as /// using the provided and . /// - /// + /// [RequiresUnreferencedCode(UnreferencedCodeMessage)] [RequiresDynamicCode(DynamicCodeMessage)] - public static IServiceCollection AddCosmosNoSqlCollection( + public static IServiceCollection AddCosmosNoSqlCollection( this IServiceCollection services, string name, string connectionString, string databaseName, CosmosNoSqlCollectionOptions? options = default, ServiceLifetime lifetime = ServiceLifetime.Singleton) + where TKey : notnull where TRecord : class - => AddKeyedCosmosNoSqlCollection(services, serviceKey: null, name, connectionString, databaseName, options, lifetime); + => AddKeyedCosmosNoSqlCollection(services, serviceKey: null, name, connectionString, databaseName, options, lifetime); /// /// Registers a keyed as /// using the provided and . /// + /// The type of the key. /// The type of the record. /// The to register the on. /// The key with which to associate the collection. @@ -204,7 +209,7 @@ public static IServiceCollection AddCosmosNoSqlCollection( /// Service collection. [RequiresUnreferencedCode(UnreferencedCodeMessage)] [RequiresDynamicCode(DynamicCodeMessage)] - public static IServiceCollection AddKeyedCosmosNoSqlCollection( + public static IServiceCollection AddKeyedCosmosNoSqlCollection( this IServiceCollection services, object? serviceKey, string name, @@ -212,35 +217,38 @@ public static IServiceCollection AddKeyedCosmosNoSqlCollection( string databaseName, CosmosNoSqlCollectionOptions? options = default, ServiceLifetime lifetime = ServiceLifetime.Singleton) + where TKey : notnull where TRecord : class { Verify.NotNullOrWhiteSpace(connectionString); Verify.NotNullOrWhiteSpace(databaseName); - return AddKeyedCosmosNoSqlCollection(services, serviceKey, name, _ => connectionString, _ => databaseName, _ => options!, lifetime); + return AddKeyedCosmosNoSqlCollection(services, serviceKey, name, _ => connectionString, _ => databaseName, _ => options!, lifetime); } /// /// Registers a as /// using the provided and . /// - /// + /// [RequiresUnreferencedCode(UnreferencedCodeMessage)] [RequiresDynamicCode(DynamicCodeMessage)] - public static IServiceCollection AddCosmosNoSqlCollection( + public static IServiceCollection AddCosmosNoSqlCollection( this IServiceCollection services, string name, Func connectionStringProvider, Func databaseNameProvider, Func? optionsProvider = default, ServiceLifetime lifetime = ServiceLifetime.Singleton) + where TKey : notnull where TRecord : class - => AddKeyedCosmosNoSqlCollection(services, serviceKey: null, name, connectionStringProvider, databaseNameProvider, optionsProvider, lifetime); + => AddKeyedCosmosNoSqlCollection(services, serviceKey: null, name, connectionStringProvider, databaseNameProvider, optionsProvider, lifetime); /// /// Registers a keyed as /// using the provided and . /// + /// The type of the key. /// The type of the record. /// The to register the on. /// The key with which to associate the collection. @@ -252,7 +260,7 @@ public static IServiceCollection AddCosmosNoSqlCollection( /// Service collection. [RequiresUnreferencedCode(UnreferencedCodeMessage)] [RequiresDynamicCode(DynamicCodeMessage)] - public static IServiceCollection AddKeyedCosmosNoSqlCollection( + public static IServiceCollection AddKeyedCosmosNoSqlCollection( this IServiceCollection services, object? serviceKey, string name, @@ -260,6 +268,7 @@ public static IServiceCollection AddKeyedCosmosNoSqlCollection( Func databaseNameProvider, Func? optionsProvider = default, ServiceLifetime lifetime = ServiceLifetime.Singleton) + where TKey : notnull where TRecord : class { Verify.NotNull(services); @@ -267,12 +276,12 @@ public static IServiceCollection AddKeyedCosmosNoSqlCollection( Verify.NotNull(connectionStringProvider); Verify.NotNull(databaseNameProvider); - services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection), serviceKey, (sp, _) => + services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection), serviceKey, (sp, _) => { var options = GetCollectionOptions(sp, optionsProvider); var clientOptions = CreateClientOptions(options?.JsonSerializerOptions); - return new CosmosNoSqlCollection( + return new CosmosNoSqlCollection( connectionString: connectionStringProvider(sp), databaseName: databaseNameProvider(sp), name: name, @@ -280,7 +289,7 @@ public static IServiceCollection AddKeyedCosmosNoSqlCollection( options); }, lifetime)); - AddAbstractions(services, serviceKey, lifetime); + AddAbstractions(services, serviceKey, lifetime); return services; } diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlCollectionOptionsTests.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlCollectionOptionsTests.cs index 869510aa8df1..b7fae39008be 100644 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlCollectionOptionsTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlCollectionOptionsTests.cs @@ -24,7 +24,7 @@ public async Task Collection_supports_partition_key_composite_key() using var collection = new CosmosNoSqlCollection( store.Database, collectionName, - new() { PartitionKeyPropertyNames = [nameof(PartitionedHotel.HotelName)] }); + new() { PartitionKeyProperties = [nameof(PartitionedHotel.HotelName)] }); await collection.EnsureCollectionExistsAsync(); @@ -62,10 +62,10 @@ public async Task Collection_supports_indexing_mode(IndexingMode indexingMode) var store = (CosmosNoSqlTestStore)fixture.TestStore; var collectionName = fixture.TestStore.AdjustCollectionName($"IndexingMode_{indexingMode}"); - using var collection = new CosmosNoSqlCollection( + using var collection = new CosmosNoSqlCollection( store.Database, collectionName, - new() { IndexingMode = indexingMode, Automatic = indexingMode != IndexingMode.None, PartitionKeyPropertyNames = ["HotelId"] }); + new() { IndexingMode = indexingMode, Automatic = indexingMode != IndexingMode.None }); await collection.EnsureCollectionExistsAsync(); @@ -79,13 +79,12 @@ public async Task Collection_supports_indexing_mode(IndexingMode indexingMode) }; await collection.UpsertAsync(record); - var key = new CosmosNoSqlKey(record.HotelId, record.HotelId); - var fetched = await collection.GetAsync(key); + var fetched = await collection.GetAsync(record.HotelId); Assert.NotNull(fetched); - await collection.DeleteAsync(key); - Assert.Null(await collection.GetAsync(key)); + await collection.DeleteAsync(record.HotelId); + Assert.Null(await collection.GetAsync(record.HotelId)); } finally { diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlDependencyInjectionTests.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlDependencyInjectionTests.cs index be63abe03387..c8004376dc66 100644 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlDependencyInjectionTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlDependencyInjectionTests.cs @@ -7,12 +7,11 @@ using Microsoft.SemanticKernel.Connectors.CosmosNoSql; using VectorData.ConformanceTests; using Xunit; -using TestRecord = VectorData.ConformanceTests.DependencyInjectionTests.Record; namespace CosmosNoSql.ConformanceTests.DependencyInjection; public class CosmosNoSqlDependencyInjectionTests - : DependencyInjectionTests, CosmosNoSqlKey, TestRecord> + : DependencyInjectionTests.Record>, string, DependencyInjectionTests.Record> { protected const string ConnectionString = "AccountEndpoint=https://test.documents.azure.com:443/;AccountKey=mock;"; protected const string DatabaseName = "dbName"; @@ -49,22 +48,22 @@ private static string DatabaseNameProvider(IServiceProvider sp, object serviceKe ? services .AddSingleton(sp => new CosmosClient(ConnectionString, s_clientOptions)) .AddSingleton(sp => sp.GetRequiredService().GetDatabase(DatabaseName)) - .AddCosmosNoSqlCollection(name, lifetime: lifetime) + .AddCosmosNoSqlCollection(name, lifetime: lifetime) : services .AddSingleton(sp => new CosmosClient(ConnectionString, s_clientOptions)) .AddSingleton(sp => sp.GetRequiredService().GetDatabase(DatabaseName)) - .AddKeyedCosmosNoSqlCollection(serviceKey, name, lifetime: lifetime); + .AddKeyedCosmosNoSqlCollection(serviceKey, name, lifetime: lifetime); yield return (services, serviceKey, name, lifetime) => serviceKey is null - ? services.AddCosmosNoSqlCollection( + ? services.AddCosmosNoSqlCollection( name, ConnectionString, DatabaseName, lifetime: lifetime) - : services.AddKeyedCosmosNoSqlCollection( + : services.AddKeyedCosmosNoSqlCollection( serviceKey, name, ConnectionString, DatabaseName, lifetime: lifetime); yield return (services, serviceKey, name, lifetime) => serviceKey is null - ? services.AddCosmosNoSqlCollection( + ? services.AddCosmosNoSqlCollection( name, ConnectionStringProvider, DatabaseNameProvider, lifetime: lifetime) - : services.AddKeyedCosmosNoSqlCollection( + : services.AddKeyedCosmosNoSqlCollection( serviceKey, name, sp => ConnectionStringProvider(sp, serviceKey), sp => DatabaseNameProvider(sp, serviceKey), lifetime: lifetime); } } @@ -96,9 +95,9 @@ public void ConnectionStringProviderCantBeNull() { IServiceCollection services = new ServiceCollection(); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", connectionStringProvider: null!, databaseNameProvider: DatabaseNameProvider)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", connectionStringProvider: null!, databaseNameProvider: DatabaseNameProvider)); } @@ -107,9 +106,9 @@ public void DatabaseNameProviderCantBeNull() { IServiceCollection services = new ServiceCollection(); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", connectionStringProvider: ConnectionStringProvider, databaseNameProvider: null!)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", connectionStringProvider: ConnectionStringProvider, databaseNameProvider: null!)); } @@ -120,13 +119,13 @@ public void ConnectionStringCantBeNullOrEmpty() Assert.Throws(() => services.AddCosmosNoSqlVectorStore(connectionString: null!, DatabaseName)); Assert.Throws(() => services.AddKeyedCosmosNoSqlVectorStore(serviceKey: "notNull", connectionString: null!, DatabaseName)); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", connectionString: null!, DatabaseName)); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", connectionString: "", DatabaseName)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", connectionString: null!, DatabaseName)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", connectionString: "", DatabaseName)); } @@ -137,13 +136,13 @@ public void DatabaseNameCantBeNullOrEmpty() Assert.Throws(() => services.AddCosmosNoSqlVectorStore(ConnectionString, databaseName: null!)); Assert.Throws(() => services.AddKeyedCosmosNoSqlVectorStore(serviceKey: "notNull", ConnectionString, databaseName: null!)); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", ConnectionString, databaseName: null!)); - Assert.Throws(() => services.AddCosmosNoSqlCollection( + Assert.Throws(() => services.AddCosmosNoSqlCollection( name: "notNull", ConnectionString, databaseName: "")); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", ConnectionString, databaseName: null!)); - Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( + Assert.Throws(() => services.AddKeyedCosmosNoSqlCollection( serviceKey: "notNull", name: "notNull", ConnectionString, databaseName: "")); } } diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlEmbeddingGenerationTests.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlEmbeddingGenerationTests.cs index 5536025c72fb..633aad44d577 100644 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlEmbeddingGenerationTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/CosmosNoSqlEmbeddingGenerationTests.cs @@ -1,66 +1,18 @@ // Copyright (c) Microsoft. All rights reserved. -using System.Reflection; using CosmosNoSql.ConformanceTests.Support; using Microsoft.Extensions.AI; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.VectorData; -using Microsoft.SemanticKernel.Connectors.CosmosNoSql; using VectorData.ConformanceTests; using VectorData.ConformanceTests.Support; using Xunit; namespace CosmosNoSql.ConformanceTests; -#pragma warning disable CA2000 // Dispose objects before losing scope - public class CosmosNoSqlEmbeddingGenerationTests(CosmosNoSqlEmbeddingGenerationTests.StringVectorFixture stringVectorFixture, CosmosNoSqlEmbeddingGenerationTests.RomOfFloatVectorFixture romOfFloatVectorFixture) : EmbeddingGenerationTests(stringVectorFixture, romOfFloatVectorFixture), IClassFixture, IClassFixture { - // Override the DI tests because the base tests call vectorStore.GetCollection() - // but CosmosNoSqlVectorStore.GetCollection requires CosmosNoSqlKey, not plain string. - public override async Task SearchAsync_with_store_dependency_injection() - { - foreach (var registrationDelegate in stringVectorFixture.DependencyInjectionStoreRegistrationDelegates) - { - IServiceCollection serviceCollection = new ServiceCollection(); - - serviceCollection.AddSingleton(new FakeEmbeddingGenerator(replaceLast: 1)); - registrationDelegate(serviceCollection); - - await using var serviceProvider = serviceCollection.BuildServiceProvider(); - - var vectorStore = serviceProvider.GetRequiredService(); - // Use the fixture's GetCollection override which routes through TestStore adapter - var collection = stringVectorFixture.GetCollection(vectorStore, stringVectorFixture.CollectionName, stringVectorFixture.CreateRecordDefinition()); - - var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync(); - - Assert.Equal("Store ([1, 1, 1])", result.Record.Text); - } - } - - public override async Task SearchAsync_with_collection_dependency_injection() - { - foreach (var registrationDelegate in stringVectorFixture.DependencyInjectionCollectionRegistrationDelegates) - { - IServiceCollection serviceCollection = new ServiceCollection(); - - serviceCollection.AddSingleton(new FakeEmbeddingGenerator(replaceLast: 1)); - registrationDelegate(serviceCollection); - - await using var serviceProvider = serviceCollection.BuildServiceProvider(); - - // Resolve with CosmosNoSqlKey since that's how AddAbstractions registers it, - // then wrap with the adapter for plain string key usage. - var innerCollection = serviceProvider.GetRequiredService>(); - var collection = new CosmosNoSqlCollectionAdapter(innerCollection); - - var result = await collection.SearchAsync("[1, 1, 0]", top: 1).SingleAsync(); - - Assert.Equal("Store ([1, 1, 1])", result.Record.Text); - } - } public new class StringVectorFixture : EmbeddingGenerationTests.StringVectorFixture { public override TestStore TestStore => CosmosNoSqlTestStore.Instance; @@ -68,47 +20,6 @@ public override async Task SearchAsync_with_collection_dependency_injection() public override VectorStore CreateVectorStore(IEmbeddingGenerator? embeddingGenerator) => CosmosNoSqlTestStore.Instance.GetVectorStore(new() { EmbeddingGenerator = embeddingGenerator }); - public override VectorStoreCollection GetCollection( - VectorStore vectorStore, - string collectionName, - VectorStoreCollectionDefinition? recordDefinition = null) - { - // If no definition was provided, build one from TRecord's attributes. - recordDefinition ??= BuildDefinitionFromType(typeof(TRecord)); - - // If the VectorStore has a store-level embedding generator, copy it onto the definition - // so that TestStore.CreateCollection can pass it to the collection options. - if (recordDefinition.EmbeddingGenerator is null - && vectorStore is CosmosNoSqlVectorStore cosmosStore) - { - var field = typeof(CosmosNoSqlVectorStore).GetField("_embeddingGenerator", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); - if (field?.GetValue(cosmosStore) is IEmbeddingGenerator generator) - { - recordDefinition.EmbeddingGenerator = generator; - } - } - - return this.CreateCollection(collectionName, recordDefinition); - } - - public override VectorStoreCollection> GetDynamicCollection( - VectorStore vectorStore, - string collectionName, - VectorStoreCollectionDefinition recordDefinition) - { - if (recordDefinition.EmbeddingGenerator is null - && vectorStore is CosmosNoSqlVectorStore cosmosStore) - { - var field = typeof(CosmosNoSqlVectorStore).GetField("_embeddingGenerator", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); - if (field?.GetValue(cosmosStore) is IEmbeddingGenerator generator) - { - recordDefinition.EmbeddingGenerator = generator; - } - } - - return this.CreateDynamicCollection(collectionName, recordDefinition); - } - public override Func[] DependencyInjectionStoreRegistrationDelegates => [ services => services @@ -120,7 +31,7 @@ public override VectorStoreCollection GetCollection( [ services => services .AddSingleton(CosmosNoSqlTestStore.Instance.Database) - .AddCosmosNoSqlCollection(this.CollectionName) + .AddCosmosNoSqlCollection(this.CollectionName) ]; } @@ -131,47 +42,6 @@ public override VectorStoreCollection GetCollection( public override VectorStore CreateVectorStore(IEmbeddingGenerator? embeddingGenerator) => CosmosNoSqlTestStore.Instance.GetVectorStore(new() { EmbeddingGenerator = embeddingGenerator }); - public override VectorStoreCollection GetCollection( - VectorStore vectorStore, - string collectionName, - VectorStoreCollectionDefinition? recordDefinition = null) - { - // If no definition was provided, build one from TRecord's attributes. - recordDefinition ??= BuildDefinitionFromType(typeof(TRecord)); - - // If the VectorStore has a store-level embedding generator, copy it onto the definition - // so that TestStore.CreateCollection can pass it to the collection options. - if (recordDefinition.EmbeddingGenerator is null - && vectorStore is CosmosNoSqlVectorStore cosmosStore) - { - var field = typeof(CosmosNoSqlVectorStore).GetField("_embeddingGenerator", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); - if (field?.GetValue(cosmosStore) is IEmbeddingGenerator generator) - { - recordDefinition.EmbeddingGenerator = generator; - } - } - - return this.CreateCollection(collectionName, recordDefinition); - } - - public override VectorStoreCollection> GetDynamicCollection( - VectorStore vectorStore, - string collectionName, - VectorStoreCollectionDefinition recordDefinition) - { - if (recordDefinition.EmbeddingGenerator is null - && vectorStore is CosmosNoSqlVectorStore cosmosStore) - { - var field = typeof(CosmosNoSqlVectorStore).GetField("_embeddingGenerator", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance); - if (field?.GetValue(cosmosStore) is IEmbeddingGenerator generator) - { - recordDefinition.EmbeddingGenerator = generator; - } - } - - return this.CreateDynamicCollection(collectionName, recordDefinition); - } - public override Func[] DependencyInjectionStoreRegistrationDelegates => [ services => services @@ -183,33 +53,7 @@ public override VectorStoreCollection GetCollection( [ services => services .AddSingleton(CosmosNoSqlTestStore.Instance.Database) - .AddCosmosNoSqlCollection(this.CollectionName) + .AddCosmosNoSqlCollection(this.CollectionName) ]; } - - /// - /// Builds a VectorStoreCollectionDefinition from a record type's VectorStore attributes. - /// Used when GetCollection is called without an explicit definition. - /// - private static VectorStoreCollectionDefinition BuildDefinitionFromType(Type recordType) - { - var properties = new List(); - foreach (var prop in recordType.GetProperties()) - { - if (prop.GetCustomAttribute() is not null) - { - properties.Add(new VectorStoreKeyProperty(prop.Name, prop.PropertyType)); - } - else if (prop.GetCustomAttribute() is VectorStoreVectorAttribute vecAttr) - { - properties.Add(new VectorStoreVectorProperty(prop.Name, prop.PropertyType, vecAttr.Dimensions)); - } - else if (prop.GetCustomAttribute() is not null) - { - properties.Add(new VectorStoreDataProperty(prop.Name, prop.PropertyType)); - } - } - - return new VectorStoreCollectionDefinition { Properties = properties }; - } } diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlCollectionAdapter.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlCollectionAdapter.cs deleted file mode 100644 index 40e929f99081..000000000000 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlCollectionAdapter.cs +++ /dev/null @@ -1,103 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System.Linq.Expressions; -using Microsoft.Extensions.VectorData; -using Microsoft.SemanticKernel.Connectors.CosmosNoSql; - -namespace CosmosNoSql.ConformanceTests.Support; - -#pragma warning disable CA1812 // Internal class that is apparently never instantiated - -/// -/// Adapts a with keys -/// to a with simple string keys. -/// This allows Cosmos conformance tests to use simple key types (string, etc.) while still -/// using the underlying Cosmos collection with composite keys. -/// -internal sealed class CosmosNoSqlCollectionAdapter( - VectorStoreCollection inner) - : VectorStoreCollection, - IKeywordHybridSearchable - where TDocumentId : notnull - where TRecord : class -{ - public override string Name => inner.Name; - - public override Task CollectionExistsAsync(CancellationToken cancellationToken = default) - => inner.CollectionExistsAsync(cancellationToken); - - public override Task EnsureCollectionExistsAsync(CancellationToken cancellationToken = default) - => inner.EnsureCollectionExistsAsync(cancellationToken); - - public override Task EnsureCollectionDeletedAsync(CancellationToken cancellationToken = default) - => inner.EnsureCollectionDeletedAsync(cancellationToken); - - public override Task GetAsync(TDocumentId key, RecordRetrievalOptions? options = null, CancellationToken cancellationToken = default) - => key is null - ? throw new ArgumentNullException(nameof(key)) - : inner.GetAsync(ConvertKey(key), options, cancellationToken); - - public override IAsyncEnumerable GetAsync(IEnumerable keys, RecordRetrievalOptions? options = null, CancellationToken cancellationToken = default) - => keys is null - ? throw new ArgumentNullException(nameof(keys)) - : inner.GetAsync(ConvertKeys(keys), options, cancellationToken); - - public override IAsyncEnumerable GetAsync( - Expression> filter, - int top, - FilteredRecordRetrievalOptions? options = null, - CancellationToken cancellationToken = default) - => inner.GetAsync(filter, top, options, cancellationToken); - - public override Task UpsertAsync(TRecord record, CancellationToken cancellationToken = default) - => inner.UpsertAsync(record, cancellationToken); - - public override Task UpsertAsync(IEnumerable records, CancellationToken cancellationToken = default) - => inner.UpsertAsync(records, cancellationToken); - - public override Task DeleteAsync(TDocumentId key, CancellationToken cancellationToken = default) - => key is null - ? throw new ArgumentNullException(nameof(key)) - : inner.DeleteAsync(ConvertKey(key), cancellationToken); - - public override Task DeleteAsync(IEnumerable keys, CancellationToken cancellationToken = default) - => keys is null - ? throw new ArgumentNullException(nameof(keys)) - : inner.DeleteAsync(ConvertKeys(keys), cancellationToken); - - public override IAsyncEnumerable> SearchAsync( - TInput searchValue, - int top, - VectorSearchOptions? options = null, - CancellationToken cancellationToken = default) - => inner.SearchAsync(searchValue, top, options, cancellationToken); - - public IAsyncEnumerable> HybridSearchAsync( - TInput searchValue, - ICollection keywords, - int top, - HybridSearchOptions? options = default, - CancellationToken cancellationToken = default) - where TInput : notnull - => ((IKeywordHybridSearchable)inner).HybridSearchAsync(searchValue, keywords, top, options, cancellationToken); - - public override object? GetService(Type serviceType, object? serviceKey = null) - => inner.GetService(serviceType, serviceKey); - - // For tests, use the document ID as the partition key - private static CosmosNoSqlKey ConvertKey(TDocumentId documentId) - => documentId switch - { - string s => new(s, s), - Guid g => new(g, g), - _ => throw new InvalidOperationException() - }; - - private static IEnumerable ConvertKeys(IEnumerable keys) - { - foreach (var key in keys) - { - yield return ConvertKey(key); - } - } -} diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlDynamicCollectionAdapter.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlDynamicCollectionAdapter.cs deleted file mode 100644 index d8538e7541a9..000000000000 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlDynamicCollectionAdapter.cs +++ /dev/null @@ -1,91 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System.Linq.Expressions; -using Microsoft.Extensions.VectorData; -using Microsoft.SemanticKernel.Connectors.CosmosNoSql; - -namespace CosmosNoSql.ConformanceTests.Support; - -#pragma warning disable CA1812 // Internal class that is apparently never instantiated - -/// -/// Adapts a (which requires keys) -/// to accept plain keys (string, Guid, etc.) by wrapping them in . -/// Both the adapter and inner collection use object as the key type; the adapter intercepts Get/Delete calls -/// to convert plain keys to before forwarding. -/// -internal sealed class CosmosNoSqlDynamicCollectionAdapter( - VectorStoreCollection> inner) - : VectorStoreCollection> -{ - public override string Name => inner.Name; - - public override Task CollectionExistsAsync(CancellationToken cancellationToken = default) - => inner.CollectionExistsAsync(cancellationToken); - - public override Task EnsureCollectionExistsAsync(CancellationToken cancellationToken = default) - => inner.EnsureCollectionExistsAsync(cancellationToken); - - public override Task EnsureCollectionDeletedAsync(CancellationToken cancellationToken = default) - => inner.EnsureCollectionDeletedAsync(cancellationToken); - - public override Task?> GetAsync(object key, RecordRetrievalOptions? options = null, CancellationToken cancellationToken = default) - => key is null - ? throw new ArgumentNullException(nameof(key)) - : inner.GetAsync(ConvertKey(key), options, cancellationToken); - - public override IAsyncEnumerable> GetAsync(IEnumerable keys, RecordRetrievalOptions? options = null, CancellationToken cancellationToken = default) - => keys is null - ? throw new ArgumentNullException(nameof(keys)) - : inner.GetAsync(keys.Select(ConvertKey), options, cancellationToken); - - public override IAsyncEnumerable> GetAsync( - Expression, bool>> filter, - int top, - FilteredRecordRetrievalOptions>? options = null, - CancellationToken cancellationToken = default) - => inner.GetAsync(filter, top, options, cancellationToken); - - public override Task UpsertAsync(Dictionary record, CancellationToken cancellationToken = default) - => inner.UpsertAsync(record, cancellationToken); - - public override Task UpsertAsync(IEnumerable> records, CancellationToken cancellationToken = default) - => inner.UpsertAsync(records, cancellationToken); - - public override Task DeleteAsync(object key, CancellationToken cancellationToken = default) - => key is null - ? throw new ArgumentNullException(nameof(key)) - : inner.DeleteAsync(ConvertKey(key), cancellationToken); - - public override Task DeleteAsync(IEnumerable keys, CancellationToken cancellationToken = default) - => keys is null - ? throw new ArgumentNullException(nameof(keys)) - : inner.DeleteAsync(keys.Select(ConvertKey), cancellationToken); - - public override IAsyncEnumerable>> SearchAsync( - TInput searchValue, - int top, - VectorSearchOptions>? options = null, - CancellationToken cancellationToken = default) - => inner.SearchAsync(searchValue, top, options, cancellationToken); - - public override object? GetService(Type serviceType, object? serviceKey = null) - => inner.GetService(serviceType, serviceKey); - - /// - /// Wraps a plain key (e.g. string "1") in a using the key's string - /// representation as the partition key. - /// If the key is already a , returns it unchanged (boxed). - /// - private static object ConvertKey(object key) - { - if (key is CosmosNoSqlKey) - { - return key; - } - - // Wrap keys using their string representation - var keyString = key.ToString() ?? string.Empty; - return new CosmosNoSqlKey(keyString, keyString); - } -} diff --git a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlTestStore.cs b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlTestStore.cs index f88a96bf74c1..40e800d893a0 100644 --- a/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlTestStore.cs +++ b/dotnet/test/VectorData/CosmosNoSql.ConformanceTests/Support/CosmosNoSqlTestStore.cs @@ -3,10 +3,8 @@ #if NET472 using System.Net.Http; #endif -// using System.Globalization; using System.Text.Json; using Microsoft.Azure.Cosmos; -using Microsoft.Extensions.VectorData; using Microsoft.SemanticKernel.Connectors.CosmosNoSql; using VectorData.ConformanceTests.Support; @@ -37,62 +35,6 @@ private CosmosNoSqlTestStore() { } - public override VectorStoreCollection CreateCollection( - string name, - VectorStoreCollectionDefinition definition) - { - // Cosmos NoSQL requires specifying a partition key for container creation. - // For the tests, use the key property as the partition key. - var keyPropertyName = definition.Properties.OfType().FirstOrDefault()?.Name - ?? throw new InvalidOperationException("Definition must have a key property"); - - var options = new CosmosNoSqlCollectionOptions - { - Definition = definition, - PartitionKeyPropertyNames = [keyPropertyName], - EmbeddingGenerator = definition.EmbeddingGenerator - }; - - // Also, in Cosmos NoSQL there's a discrepancy between: - // 1. The key property in the model and on the record type (representing the Cosmos document ID - string or Guid) - // 2. The TKey generic type variable on the collection (must be CosmosNoSqlKey, wrapping both the document ID and the partition key) - // To bridge this gap in the tests, if the test expects a simpler key type (string, Guid), we wrap it in an adapter that translates between the two. - var collection = new CosmosNoSqlCollection(this.Database, name, options); - - return typeof(TKey) switch - { - var t when t == typeof(CosmosNoSqlKey) - => (VectorStoreCollection)(object)collection, - - var t when t == typeof(string) - => (VectorStoreCollection)(object)new CosmosNoSqlCollectionAdapter(collection), - var t when t == typeof(Guid) - => (VectorStoreCollection)(object)new CosmosNoSqlCollectionAdapter(collection), - - _ => throw new InvalidOperationException( - $"Cosmos NoSQL tests must use string, Guid or CosmosNoSqlKey as the key type. Got: {typeof(TKey).Name}") - }; - } - - public override VectorStoreCollection> CreateDynamicCollection( - string name, - VectorStoreCollectionDefinition definition) - { - // See notes in CreateCollection above. - var keyPropertyName = definition.Properties.OfType().FirstOrDefault()?.Name - ?? throw new InvalidOperationException("Definition must have a key property"); - - return new CosmosNoSqlDynamicCollectionAdapter( - new CosmosNoSqlDynamicCollection( - this.Database, - name, - new CosmosNoSqlCollectionOptions - { - Definition = definition, - PartitionKeyPropertyNames = [keyPropertyName] - })); - } - #pragma warning disable CA5400 // HttpClient may be created without enabling CheckCertificateRevocationList protected override async Task StartAsync() { diff --git a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs index 556fcc0cd60d..7efb13bab864 100644 --- a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlCollectionTests.cs @@ -42,7 +42,7 @@ public CosmosNoSqlCollectionTests() public void ConstructorForModelWithoutKeyThrowsException() { // Act & Assert - var exception = Assert.Throws(() => new CosmosNoSqlCollection(this._mockDatabase.Object, "collection")); + var exception = Assert.Throws(() => new CosmosNoSqlCollection(this._mockDatabase.Object, "collection")); Assert.Contains("No key property found", exception.Message); } @@ -56,7 +56,7 @@ public void ConstructorWithoutSystemTextJsonSerializerOptionsThrowsArgumentExcep mockDatabase.Setup(l => l.Client).Returns(mockClient.Object); // Act & Assert - var exception = Assert.Throws(() => new CosmosNoSqlCollection(mockDatabase.Object, "collection")); + var exception = Assert.Throws(() => new CosmosNoSqlCollection(mockDatabase.Object, "collection")); Assert.Contains(nameof(CosmosClientOptions.UseSystemTextJsonSerializerWithOptions), exception.Message); } @@ -64,10 +64,9 @@ public void ConstructorWithoutSystemTextJsonSerializerOptionsThrowsArgumentExcep public void ConstructorWithDeclarativeModelInitializesCollection() { // Act & Assert - using var collection = new CosmosNoSqlCollection( + using var collection = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); Assert.NotNull(collection); } @@ -82,10 +81,10 @@ public void ConstructorWithImperativeModelInitializesCollection() }; // Act - using var collection = new CosmosNoSqlCollection( + using var collection = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection", - new() { Definition = definition, PartitionKeyPropertyNames = ["Id"] }); + new() { Definition = definition }); // Assert Assert.NotNull(collection); @@ -118,10 +117,9 @@ public async Task CollectionExistsReturnsValidResultAsync(List collectio It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - collectionName, - new() { PartitionKeyPropertyNames = ["HotelId"] }); + collectionName); // Act var actualResult = await sut.CollectionExistsAsync(); @@ -161,14 +159,13 @@ public async Task EnsureCollectionExistsUsesValidContainerPropertiesAsync(Indexi It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, CollectionName, new() { IndexingMode = indexingMode, - Automatic = indexingMode != IndexingMode.None, - PartitionKeyPropertyNames = ["Id"] + Automatic = indexingMode != IndexingMode.None }); var expectedVectorEmbeddingPolicy = new VectorEmbeddingPolicy( @@ -266,10 +263,9 @@ public async Task EnsureCollectionExistsInvokesValidMethodsAsync(List co It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - CollectionName, - new() { PartitionKeyPropertyNames = ["HotelId"] }); + CollectionName); // Act await sut.EnsureCollectionExistsAsync(); @@ -294,17 +290,25 @@ public async Task DeleteInvokesValidMethodsAsync( const string RecordKey = "recordKey"; const string PartitionKey = "partitionKey"; - using var sut = new CosmosNoSqlCollection( - this._mockDatabase.Object, - "collection", - useExplicitPartitionKey ? null : new() { PartitionKeyPropertyNames = ["HotelId"] }); - // Act - var key = useExplicitPartitionKey - ? new CosmosNoSqlKey(RecordKey, PartitionKey) - : new CosmosNoSqlKey(RecordKey, RecordKey); + if (useExplicitPartitionKey) + { + using var sut = new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection"); + + await ((VectorStoreCollection)sut).DeleteAsync( + new CosmosNoSqlKey(RecordKey, PartitionKey)); + } + else + { + using var sut = new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection"); - await sut.DeleteAsync(key); + await ((VectorStoreCollection)sut).DeleteAsync( + RecordKey); + } // Assert this._mockContainer.Verify(l => l.DeleteItemAsync( @@ -321,15 +325,12 @@ public async Task DeleteBatchInvokesValidMethodsAsync() // Arrange List recordKeys = ["key1", "key2"]; - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); // Act - // With the new design, keys must be CosmosNoSqlKey - var cosmosKeys = recordKeys.Select(k => new CosmosNoSqlKey(k, k)).ToList(); - await sut.DeleteAsync(cosmosKeys); + await sut.DeleteAsync(recordKeys); // Assert foreach (var key in recordKeys) @@ -347,10 +348,9 @@ public async Task DeleteBatchInvokesValidMethodsAsync() public async Task DeleteCollectionInvokesValidMethodsAsync() { // Arrange - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); // Act await sut.EnsureCollectionDeletedAsync(); @@ -383,13 +383,12 @@ public async Task GetReturnsValidRecordAsync() It.IsAny())) .ReturnsAsync(mockItemResponse.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); // Act - var result = await sut.GetAsync(new CosmosNoSqlKey(RecordKey, RecordKey)); + var result = await sut.GetAsync(RecordKey); // Assert Assert.NotNull(result); @@ -417,19 +416,12 @@ public async Task GetBatchReturnsValidRecordAsync() It.IsAny())) .ReturnsAsync(mockFeedResponse.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); // Act - var keys = new List - { - new("key1", "key1"), - new("key2", "key2"), - new("key3", "key3") - }; - var results = await sut.GetAsync(keys).ToListAsync(); + var results = await sut.GetAsync(["key1", "key2", "key3"]).ToListAsync(); // Assert Assert.NotNull(results[0]); @@ -451,10 +443,9 @@ public async Task CanUpsertRecordAsync() // Arrange var hotel = new CosmosNoSqlHotel("key") { HotelName = "Test Name" }; - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); // Act await sut.UpsertAsync(hotel); @@ -478,10 +469,9 @@ public async Task CanUpsertManyRecordsAsync() var hotel2 = new CosmosNoSqlHotel("key2") { HotelName = "Test Name 2" }; var hotel3 = new CosmosNoSqlHotel("key3") { HotelName = "Test Name 3" }; - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); // Act await sut.UpsertAsync([hotel1, hotel2, hotel3]); @@ -523,10 +513,9 @@ public async Task VectorizedSearchReturnsValidRecordAsync() It.IsAny())) .Returns(mockFeedIterator.Object); - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); // Act var results = await sut.SearchAsync(new ReadOnlyMemory([1f, 2f, 3f]), top: 3).ToListAsync(); @@ -543,10 +532,9 @@ public async Task VectorizedSearchReturnsValidRecordAsync() public async Task VectorizedSearchWithUnsupportedVectorTypeThrowsExceptionAsync() { // Arrange - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); // Act & Assert await Assert.ThrowsAsync(async () => @@ -557,10 +545,9 @@ await Assert.ThrowsAsync(async () => public async Task VectorizedSearchWithNonExistentVectorPropertyNameThrowsExceptionAsync() { // Arrange - using var sut = new CosmosNoSqlCollection( + using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, - "collection", - new() { PartitionKeyPropertyNames = ["HotelId"] }); + "collection"); var searchOptions = new VectorSearchOptions { VectorProperty = r => "non-existent-property" }; @@ -665,5 +652,152 @@ private sealed class TestIndexingModel } #pragma warning restore CA1812 + [Fact] + public void ConstructorWithStringKeyInitializesCollection() + { + // String TKey should work; partition key auto-defaults to the key property. + using var collection = new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection"); + + Assert.NotNull(collection); + } + + [Fact] + public void ConstructorWithStringKeyAndExplicitPartitionKeyInitializesCollection() + { + // String TKey should work when partition key explicitly points to the key property. + using var collection = new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection", + new() { PartitionKeyProperties = ["HotelId"] }); + + Assert.NotNull(collection); + } + + [Fact] + public void ConstructorWithStringKeyRejectsNonKeyPartitionKey() + { + // String TKey requires the partition key to be the key property. + var exception = Assert.Throws(() => new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection", + new() { PartitionKeyProperties = ["HotelName"] })); + + Assert.Contains("partition key must be the key property", exception.Message); + } + + [Fact] + public void ConstructorWithStringKeyRejectsEmptyPartitionKey() + { + // String TKey cannot use empty partition key (PartitionKey.None) - that requires CosmosNoSqlKey. + var exception = Assert.Throws(() => new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection", + new() { PartitionKeyProperties = [] })); + + Assert.Contains("partition key must be the key property", exception.Message); + } + + [Fact] + public void ConstructorWithGuidKeyInitializesCollection() + { + // Guid TKey should work; partition key auto-defaults to the key property. + using var collection = new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection"); + + Assert.NotNull(collection); + } + + [Fact] + public void ConstructorWithGuidKeyRejectsNonKeyPartitionKey() + { + // Guid TKey requires the partition key to be the key property. + var exception = Assert.Throws(() => new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection", + new() { PartitionKeyProperties = ["Name"] })); + + Assert.Contains("partition key must be the key property", exception.Message); + } + + [Fact] + public void ConstructorWithUnsupportedKeyTypeThrowsException() + { + var exception = Assert.Throws(() => new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection")); + + Assert.Contains("string, Guid, or CosmosNoSqlKey", exception.Message); + } + + [Fact] + public async Task DeleteWithStringKeyInvokesValidMethodsAsync() + { + // Arrange + const string RecordKey = "recordKey"; + + using var sut = new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection"); + + // Act + await sut.DeleteAsync(RecordKey); + + // Assert: with string key, partition key = document id + this._mockContainer.Verify(l => l.DeleteItemAsync( + RecordKey, + new PartitionKey(RecordKey), + It.IsAny(), + It.IsAny()), + Times.Once()); + } + + [Fact] + public async Task GetWithStringKeyReturnsValidRecordAsync() + { + // Arrange + const string RecordKey = "key"; + + var jsonObject = new JsonObject { ["id"] = RecordKey, ["HotelName"] = "Test Name" }; + + var mockItemResponse = new Mock>(); + mockItemResponse + .Setup(l => l.Resource) + .Returns(jsonObject); + + this._mockContainer + .Setup(l => l.ReadItemAsync( + RecordKey, + new PartitionKey(RecordKey), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(mockItemResponse.Object); + + using var sut = new CosmosNoSqlCollection( + this._mockDatabase.Object, + "collection"); + + // Act + var result = await sut.GetAsync(RecordKey); + + // Assert + Assert.NotNull(result); + Assert.Equal(RecordKey, result.HotelId); + Assert.Equal("Test Name", result.HotelName); + } + +#pragma warning disable CA1812 + private sealed class GuidKeyModel + { + [VectorStoreKey] + public Guid Id { get; set; } + + [VectorStoreData] + public string? Name { get; set; } + } +#pragma warning restore CA1812 + #endregion } diff --git a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlVectorStoreTests.cs b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlVectorStoreTests.cs index e829dd1cdf5a..21ada9f8710a 100644 --- a/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlVectorStoreTests.cs +++ b/dotnet/test/VectorData/CosmosNoSql.UnitTests/CosmosNoSqlVectorStoreTests.cs @@ -54,6 +54,32 @@ public void GetCollectionWithSupportedKeyReturnsCollection() Assert.NotNull(collection); } + [Fact] + public void GetCollectionWithStringKeyReturnsCollection() + { + // Arrange + using var sut = new Microsoft.SemanticKernel.Connectors.CosmosNoSql.CosmosNoSqlVectorStore(this._mockDatabase.Object); + + // Act + var collection = sut.GetCollection("collection1"); + + // Assert + Assert.NotNull(collection); + } + + [Fact] + public void GetCollectionWithGuidKeyReturnsCollection() + { + // Arrange + using var sut = new Microsoft.SemanticKernel.Connectors.CosmosNoSql.CosmosNoSqlVectorStore(this._mockDatabase.Object); + + // Act + var collection = sut.GetCollection("collection1"); + + // Assert + Assert.NotNull(collection); + } + [Fact] public void GetCollectionWithoutFactoryReturnsDefaultCollection() { @@ -103,4 +129,15 @@ public async Task ListCollectionNamesReturnsCollectionNamesAsync() // Assert Assert.Equal(expectedCollectionNames, actualCollectionNames); } + +#pragma warning disable CA1812 + private sealed class GuidKeyHotel + { + [Microsoft.Extensions.VectorData.VectorStoreKey] + public Guid Id { get; set; } + + [Microsoft.Extensions.VectorData.VectorStoreData] + public string? Name { get; set; } + } +#pragma warning restore CA1812 }