.Net: [MEVD] Cosmos NoSQL provider work on keys, partition keys and point reads (#13550)

This is a sizable cleanup/redo of large chunks of the Cosmos NoSQL
provider; I've broken it down to three commits for easier reviewing.

* Cosmos NoSQL collections can now have either string/Guid as their key,
or CosmosNoSqlKey. In the former case, this configures the partition key
to be the same as the string/Guid document key - this is a sensible
default that's common (and encouraged) in Cosmos. All other scenarios
require CosmosNoSqlKey:
    * A single partition key that isn't the document ID
    * A hierarchical partition key (more than one property)
    * No partition key (legacy/discouraged)
* Note that CosmosNoSqlKey is only passed to GetAsync and DeleteAsync;
on the user's .NET record type, the properties are there as usual (e.g.
a string document Id property, an int partition key property). This
introduces a discrepancy for the first time between the collection's
TKey (CosmosNoSqlKey) and its actual key property on the .NET type (a
string); this made some things trickier, but is the correct way to do
things for composite key cases (we'd do the same if we implement
composite keys for relational providers).
* Added support for Cosmos hierarchical partition keys (2-3 properties
instead of just one).
* GetAsync no longer does a SQL query with a WHERE clause, but rather an
efficient point read via ReadItem. The multiple-key overload does
ReadManyItemsAsync, which again is far more efficient.
* UpdateAsync and DeleteAsync now correctly supply the partition key in
the call's request options, again making it much more efficient.
* Various other cleanup and fixes all around.
This commit is contained in:
Shay Rojansky
2026-02-17 13:57:34 +00:00
committed by GitHub
parent 91f795605e
commit 8513c2ac00
31 changed files with 881 additions and 373 deletions
@@ -83,7 +83,7 @@ switch (appConfig.RagConfig.VectorStoreType)
appConfig.CosmosMongoConfig.DatabaseName);
break;
case "CosmosNoSql":
kernelBuilder.Services.AddCosmosNoSqlCollection<TextSnippet<string>>(
kernelBuilder.Services.AddCosmosNoSqlCollection<string, TextSnippet<string>>(
appConfig.RagConfig.CollectionName,
appConfig.CosmosNoSqlConfig.ConnectionString,
appConfig.CosmosNoSqlConfig.DatabaseName);
@@ -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;
/// <summary>
/// A JSON converter for byte arrays that serializes them as JSON arrays of numbers
/// instead of base64-encoded strings.
/// </summary>
/// <remarks>
/// This is needed because Cosmos DB's VectorDistance function requires vectors to be arrays of numbers,
/// not base64-encoded strings.
/// </remarks>
internal sealed class ByteArrayJsonConverter : JsonConverter<byte[]>
{
/// <inheritdoc/>
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<byte>();
while (reader.Read() && reader.TokenType != JsonTokenType.EndArray)
{
list.Add(reader.GetByte());
}
return list.ToArray();
}
/// <inheritdoc/>
public override void Write(Utf8JsonWriter writer, byte[] value, JsonSerializerOptions options)
{
writer.WriteStartArray();
foreach (var b in value)
{
writer.WriteNumberValue(b);
}
writer.WriteEndArray();
}
}
/// <summary>
/// A JSON converter for <see cref="ReadOnlyMemory{T}"/> of byte that serializes as JSON arrays of numbers
/// instead of base64-encoded strings.
/// </summary>
/// <remarks>
/// This is needed because Cosmos DB's VectorDistance function requires vectors to be arrays of numbers,
/// not base64-encoded strings.
/// </remarks>
internal sealed class ReadOnlyMemoryByteJsonConverter : JsonConverter<ReadOnlyMemory<byte>>
{
/// <inheritdoc/>
public override ReadOnlyMemory<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 new ReadOnlyMemory<byte>(reader.GetBytesFromBase64());
}
if (reader.TokenType != JsonTokenType.StartArray)
{
throw new JsonException($"Expected StartArray or String token, got {reader.TokenType}");
}
var list = new List<byte>();
while (reader.Read() && reader.TokenType != JsonTokenType.EndArray)
{
list.Add(reader.GetByte());
}
return new ReadOnlyMemory<byte>(list.ToArray());
}
/// <inheritdoc/>
public override void Write(Utf8JsonWriter writer, ReadOnlyMemory<byte> value, JsonSerializerOptions options)
{
writer.WriteStartArray();
foreach (var b in value.Span)
{
writer.WriteNumberValue(b);
}
writer.WriteEndArray();
}
}
@@ -7,6 +7,7 @@ using System.Diagnostics;
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;
@@ -26,13 +27,15 @@ namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql;
/// <summary>
/// Service for storing and retrieving vector records, that uses Azure CosmosDB NoSQL as the underlying storage.
/// </summary>
/// <typeparam name="TKey">The data type of the record key. Must be <see cref="string"/>.</typeparam>
/// <typeparam name="TKey">
/// The data type of the record key. Supported types are <see cref="string"/> and <see cref="Guid"/> (in which case the partition key
/// is the document ID), and <see cref="CosmosNoSqlKey"/> (for other partition key configurations).</typeparam>
/// <typeparam name="TRecord">The data model to use for adding, updating and retrieving data from storage.</typeparam>
#pragma warning disable CA1711 // Identifiers should not have incorrect suffix
public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey, TRecord>, IKeywordHybridSearchable<TRecord>
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
{
/// <summary>Metadata about vector store record collection.</summary>
private readonly VectorStoreCollectionMetadata _collectionMetadata;
@@ -42,6 +45,7 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
/// <summary>The default options for hybrid vector search.</summary>
private static readonly HybridSearchOptions<TRecord> s_defaultKeywordVectorizedHybridSearchOptions = new();
private readonly ClientWrapper _clientWrapper;
/// <summary><see cref="Database"/> that can be used to manage the collections in Azure CosmosDB NoSQL.</summary>
@@ -50,9 +54,9 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
/// <summary>The model for this collection.</summary>
private readonly CollectionModel _model;
// TODO: Refactor this into the model (Co)
/// <summary>The property to use as partition key.</summary>
private readonly PropertyModel _partitionKeyProperty;
// TODO: Refactor this into the model
/// <summary>The properties to use as partition key (supports hierarchical partition keys up to 3 levels).</summary>
private readonly List<PropertyModel> _partitionKeyProperties;
/// <summary>The mapper to use when mapping between the consumer data model and the Azure CosmosDB NoSQL record.</summary>
private readonly ICosmosNoSqlMapper<TRecord> _mapper;
@@ -132,23 +136,21 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
Func<CosmosNoSqlCollectionOptions, CollectionModel> modelFactory,
CosmosNoSqlCollectionOptions? options)
{
// 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 string, Guid, or CosmosNoSqlKey. Received: {typeof(TKey).Name}");
}
try
{
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.");
}
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<TKey, TRecord>)}.");
$"is required to be configured for {this.GetType().Name}.");
}
options ??= CosmosNoSqlCollectionOptions.Default;
@@ -165,25 +167,48 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
? (new CosmosNoSqlDynamicMapper(this._model, jsonSerializerOptions) as ICosmosNoSqlMapper<TRecord>)!
: new CosmosNoSqlMapper<TRecord>(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.PartitionKeyProperties is null)
{
if (!this._model.PropertyMap.TryGetValue(options.PartitionKeyPropertyName, out var property))
{
throw new ArgumentException($"Partition key property '{options.PartitionKeyPropertyName}' is not part of the record definition.");
}
if (property.Type != typeof(string))
{
throw new ArgumentException("Partition key property must be string.");
}
this._partitionKeyProperty = property;
// 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];
}
else
{
// If partition key is not provided, use key property as a partition key.
this._partitionKeyProperty = this._model.KeyProperty;
if (options.PartitionKeyProperties.Count > 3)
{
throw new ArgumentException("Cosmos DB supports at most 3 levels of hierarchical partition keys.");
}
this._partitionKeyProperties = new List<PropertyModel>(options.PartitionKeyProperties.Count);
foreach (var propertyName in options.PartitionKeyProperties)
{
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);
}
}
// 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()
@@ -295,12 +320,7 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
/// <inheritdoc />
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 +328,7 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
{
await this._database
.GetContainer(this.Name)
.DeleteItemAsync<JsonObject>(compositeKey.RecordKey, new PartitionKey(compositeKey.PartitionKey), cancellationToken: cancellationToken)
.DeleteItemAsync<JsonObject>(documentId, partitionKey, cancellationToken: cancellationToken)
.ConfigureAwait(false);
return 0;
}
@@ -327,9 +347,33 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
{
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<JsonObject>(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);
}
/// <inheritdoc />
@@ -340,7 +384,7 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
{
Verify.NotNull(keys);
const string OperationName = "GetItemQueryIterator";
const string OperationName = "ReadManyItems";
var includeVectors = options?.IncludeVectors ?? false;
if (includeVectors && this._model.EmbeddingGenerationRequired)
@@ -348,20 +392,20 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
throw new NotSupportedException(VectorDataStrings.IncludeVectorsNotSupportedWithEmbeddingGeneration);
}
var compositeKeys = GetCompositeKeys(keys).ToList();
if (compositeKeys.Count == 0)
var items = keys.Select(this.GetDocumentIdAndPartitionKey).ToList();
if (items.Count == 0)
{
yield break;
}
var queryDefinition = CosmosNoSqlCollectionQueryBuilder.BuildSelectQuery(
this._model,
this._model.KeyProperty.StorageName,
this._partitionKeyProperty.StorageName,
compositeKeys,
includeVectors);
var response = await this.RunOperationAsync(OperationName, () =>
this._database
.GetContainer(this.Name)
.ReadManyItemsAsync<JsonObject>(items, cancellationToken: cancellationToken))
.ConfigureAwait(false);
await foreach (var jsonObject in this.GetItemsAsync<JsonObject>(queryDefinition, OperationName, cancellationToken).ConfigureAwait(false))
foreach (var jsonObject in response.Resource)
{
var record = this._mapper.MapFromStorageToDataModel(jsonObject, includeVectors);
@@ -409,25 +453,21 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
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 +566,7 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
this._model,
vectorProperty.StorageName,
vectorProperty.DistanceFunction,
null,
textPropertyName: null,
ScorePropertyName,
options.OldFilter,
options.Filter,
@@ -685,30 +725,49 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
operationName,
operation);
/// <summary>
/// Gets the document ID and partition key from a key.
/// </summary>
private (string DocumentId, PartitionKey PartitionKey) GetDocumentIdAndPartitionKey(TKey key)
=> 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}'.")
};
/// <summary>
/// Returns instance of <see cref="ContainerProperties"/> with applied indexing policy.
/// More information here: <see href="https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-manage-indexing-policy"/>.
/// </summary>
private ContainerProperties GetContainerProperties()
{
// Process Vector properties.
var embeddings = new Collection<Azure.Cosmos.Embedding>();
var vectorIndexPaths = new Collection<VectorIndexPath>();
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<Azure.Cosmos.Embedding>();
var vectorIndexPaths = new Collection<VectorIndexPath>();
foreach (var property in this._model.VectorProperties)
{
var path = $"/{property.StorageName}";
@@ -760,12 +819,10 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
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;
}
/// <summary>
@@ -855,30 +912,49 @@ public class CosmosNoSqlCollection<TKey, TRecord> : VectorStoreCollection<TKey,
}
}
private static IEnumerable<CosmosNoSqlCompositeKey> GetCompositeKeys(IEnumerable<TKey> keys)
=> keys switch
/// <summary>
/// Builds a PartitionKey from the partition key property values in the record.
/// Supports single and hierarchical partition keys.
/// </summary>
private PartitionKey BuildPartitionKey(TRecord record)
{
switch (this._partitionKeyProperties)
{
IEnumerable<CosmosNoSqlCompositeKey> k => k,
case []:
return PartitionKey.None;
IEnumerable<string> k => k.Select(key => new CosmosNoSqlCompositeKey(recordKey: key, partitionKey: key)),
IEnumerable<Guid> 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<string>(record)),
Type t when t == typeof(bool) => new PartitionKey(property.GetValue<bool>(record)),
Type t when t == typeof(double) => new PartitionKey(property.GetValue<double>(record)),
Type t when t == typeof(Guid) => new PartitionKey(property.GetValue<Guid>(record).ToString()),
_ => throw new InvalidOperationException($"Unsupported partition key type: {property.Type.Name}")
};
}
IEnumerable<object> 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<string>(record)),
Type t when t == typeof(bool) => builder.Add(property.GetValue<bool>(record)),
Type t when t == typeof(double) => builder.Add(property.GetValue<double>(record)),
Type t when t == typeof(Guid) => builder.Add(property.GetValue<Guid>(record).ToString()),
_ => throw new InvalidOperationException($"Unsupported partition key type: {property.Type.Name}")
};
}
_ => throw new UnreachableException()
};
return builder.Build();
}
}
#endregion
}
@@ -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 sealed class CosmosNoSqlCollectionOptions : VectorStoreCollectionOptions
internal CosmosNoSqlCollectionOptions(CosmosNoSqlCollectionOptions? source) : base(source)
{
this.JsonSerializerOptions = source?.JsonSerializerOptions;
this.PartitionKeyPropertyName = source?.PartitionKeyPropertyName;
this.PartitionKeyProperties = source?.PartitionKeyProperties is null ? null : [.. source.PartitionKeyProperties];
this.IndexingMode = source?.IndexingMode ?? Default.IndexingMode;
this.Automatic = source?.Automatic ?? Default.Automatic;
}
@@ -34,9 +35,21 @@ public sealed class CosmosNoSqlCollectionOptions : VectorStoreCollectionOptions
public JsonSerializerOptions? JsonSerializerOptions { get; set; }
/// <summary>
/// The property name to use as partition key.
/// Gets or sets the property names to use as partition key components.
/// </summary>
public string? PartitionKeyPropertyName { get; set; }
/// <remarks>
/// <para>
/// Selecting a partition key is critical for performance and scalability. Choose properties with high cardinality
/// that evenly distribute requests. See <see href="https://learn.microsoft.com/azure/cosmos-db/partitioning-overview"/> for guidance.
/// </para>
/// <para>
/// When <see langword="null" /> (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 <see cref="string"/> or <see cref="System.Guid"/>.
/// To use a different partition key (or hierarchical partition keys), specify the key properties here and use
/// <see cref="CosmosNoSqlKey"/> as the key type.
/// </para>
/// </remarks>
public IReadOnlyList<string>? PartitionKeyProperties { get; set; }
/// <summary>
/// Specifies the indexing mode in the Azure Cosmos DB service.
@@ -16,10 +16,6 @@ namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql;
/// </summary>
internal static class CosmosNoSqlCollectionQueryBuilder
{
private const string SelectClauseDelimiter = ",";
private const string AndConditionDelimiter = " AND ";
private const string OrConditionDelimiter = " OR ";
/// <summary>
/// Builds <see cref="QueryDefinition"/> to get items from Azure CosmosDB NoSQL using vector search.
/// </summary>
@@ -66,7 +62,7 @@ internal static class CosmosNoSqlCollectionQueryBuilder
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.
@@ -81,7 +77,13 @@ internal static class CosmosNoSqlCollectionQueryBuilder
var queryParameters = new Dictionary<string, object?>
{
[VectorVariableName] = vector
// byte[] and ReadOnlyMemory<byte> are serialized as base64 by System.Text.Json, which causes problems for the VectorDistance function.
// Convert to int[] which serializes as a JSON array of numbers.
[VectorVariableName] = vector switch
{
ReadOnlyMemory<byte> byteMemory => ConvertToIntArray(byteMemory.Span),
_ => vector
}
};
// Add score threshold filter if specified.
@@ -121,7 +123,7 @@ internal static class CosmosNoSqlCollectionQueryBuilder
builder.Append(filterClause);
if (scoreThresholdClause is not null)
{
builder.Append(AndConditionDelimiter);
builder.Append(" AND ");
}
}
@@ -179,7 +181,7 @@ internal static class CosmosNoSqlCollectionQueryBuilder
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.
@@ -222,60 +224,6 @@ internal static class CosmosNoSqlCollectionQueryBuilder
return queryDefinition;
}
/// <summary>
/// Builds <see cref="QueryDefinition"/> to get items from Azure CosmosDB NoSQL.
/// </summary>
public static QueryDefinition BuildSelectQuery(
CollectionModel model,
string keyStoragePropertyName,
string partitionKeyStoragePropertyName,
List<CosmosNoSqlCompositeKey> keys,
bool includeVectors)
{
Verify.True(keys.Count > 0, "At least one key should be provided.", nameof(keys));
const string RecordKeyVariableName = "@rk";
const string PartitionKeyVariableName = "@pk";
var tableVariableName = CosmosNoSqlConstants.ContainerAlias;
IEnumerable<PropertyModel> projectionProperties = model.Properties;
if (!includeVectors)
{
projectionProperties = projectionProperties.Where(p => p is not VectorPropertyModel);
}
var selectClauseArguments = string.Join(SelectClauseDelimiter,
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})"));
var query = $"""
SELECT {selectClauseArguments}
FROM {tableVariableName}
WHERE {whereClauseArguments}
""";
var queryDefinition = new QueryDefinition(query);
for (var i = 0; i < keys.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);
}
return queryDefinition;
}
#region private
#pragma warning disable CS0618 // VectorSearchFilter is obsolete
@@ -363,5 +311,20 @@ internal static class CosmosNoSqlCollectionQueryBuilder
private static string GeneratePropertyAccess(char alias, string propertyName)
=> $"{alias}[\"{EscapeJsonPropertyName(propertyName)}\"]";
/// <summary>
/// Converts a byte span to an int array.
/// This is needed because byte[] and ReadOnlyMemory&lt;byte&gt; are serialized as base64 by System.Text.Json,
/// which causes problems for the VectorDistance function.
/// </summary>
private static int[] ConvertToIntArray(ReadOnlySpan<byte> bytes)
{
var result = new int[bytes.Length];
for (var i = 0; i < bytes.Length; i++)
{
result[i] = bytes[i];
}
return result;
}
#endregion
}
@@ -1,19 +0,0 @@
// Copyright (c) Microsoft. All rights reserved.
namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql;
/// <summary>
/// Composite key for Azure CosmosDB NoSQL record with record and partition keys.
/// </summary>
public sealed class CosmosNoSqlCompositeKey(string recordKey, string partitionKey)
{
/// <summary>
/// Value of record key.
/// </summary>
public string RecordKey { get; } = recordKey;
/// <summary>
/// Value of partition key.
/// </summary>
public string PartitionKey { get; } = partitionKey;
}
@@ -11,6 +11,12 @@ namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql;
/// <summary>
/// Represents a collection of vector store records in a CosmosNoSql database, mapped to a dynamic <c>Dictionary&lt;string, object?&gt;</c>.
/// </summary>
/// <remarks>
/// <para>
/// This collection accepts keys of type <see cref="object"/>, but only <see cref="CosmosNoSqlKey"/> instances
/// are supported at runtime. Passing any other key type will result in an <see cref="InvalidOperationException"/>.
/// </para>
/// </remarks>
#pragma warning disable CA1711 // Identifiers should not have incorrect suffix
public sealed class CosmosNoSqlDynamicCollection : CosmosNoSqlCollection<object, Dictionary<string, object?>>
#pragma warning restore CA1711 // Identifiers should not have incorrect suffix
@@ -33,11 +39,11 @@ public sealed class CosmosNoSqlDynamicCollection : CosmosNoSqlCollection<object,
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlCollection{TKey, TRecord}"/> class.
/// Initializes a new instance of the <see cref="CosmosNoSqlDynamicCollection"/> class.
/// </summary>
/// <param name="connectionString">Connection string required to connect to Azure CosmosDB NoSQL.</param>
/// <param name="databaseName">Database name for Azure CosmosDB NoSQL.</param>
/// <param name="collectionName">The name of the collection that this <see cref="CosmosNoSqlCollection{TKey, TRecord}"/> will access.</param>
/// <param name="collectionName">The name of the collection that this <see cref="CosmosNoSqlDynamicCollection"/> will access.</param>
/// <param name="clientOptions">Optional configuration options for <see cref="CosmosClient"/>.</param>
/// <param name="options">Optional configuration options for this class.</param>
[RequiresUnreferencedCode("The Cosmos NoSQL provider is currently incompatible with trimming.")]
@@ -0,0 +1,183 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using Microsoft.Azure.Cosmos;
namespace Microsoft.SemanticKernel.Connectors.CosmosNoSql;
/// <summary>
/// Represents a key for Azure CosmosDB NoSQL, containing both the record key (document ID) and the partition key.
/// </summary>
/// <remarks>
/// <para>
/// 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 <see cref="PartitionKey"/> type.
/// </para>
/// <para>
/// For simple partition keys, use the convenience constructors that accept <see cref="string"/>, <see cref="bool"/>, or <see cref="double"/>
/// partition key values. For hierarchical partition keys (up to 3 levels), construct a <see cref="PartitionKey"/> using
/// <see cref="PartitionKeyBuilder"/> and pass it to the primary constructor.
/// </para>
/// </remarks>
// This is conceptually a record, but we're targeting .NET Standard 2.0 too.
public readonly struct CosmosNoSqlKey : IEquatable<CosmosNoSqlKey>
{
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
public CosmosNoSqlKey(string documentId, string partitionKey)
{
this.DocumentId = documentId;
this.PartitionKey = new PartitionKey(partitionKey);
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
public CosmosNoSqlKey(Guid documentId, string partitionKey)
{
this.DocumentId = documentId.ToString();
this.PartitionKey = new PartitionKey(partitionKey);
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
public CosmosNoSqlKey(string documentId, Guid partitionKey)
{
this.DocumentId = documentId;
this.PartitionKey = new PartitionKey(partitionKey.ToString());
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
public CosmosNoSqlKey(Guid documentId, Guid partitionKey)
{
this.DocumentId = documentId.ToString();
this.PartitionKey = new PartitionKey(partitionKey.ToString());
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
public CosmosNoSqlKey(string documentId, bool partitionKey)
{
this.DocumentId = documentId;
this.PartitionKey = new PartitionKey(partitionKey);
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
public CosmosNoSqlKey(Guid documentId, bool partitionKey)
{
this.DocumentId = documentId.ToString();
this.PartitionKey = new PartitionKey(partitionKey);
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
public CosmosNoSqlKey(string documentId, double partitionKey)
{
this.DocumentId = documentId;
this.PartitionKey = new PartitionKey(partitionKey);
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
public CosmosNoSqlKey(Guid documentId, double partitionKey)
{
this.DocumentId = documentId.ToString();
this.PartitionKey = new PartitionKey(partitionKey);
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
/// <remarks>
/// Use this constructor for hierarchical partition keys or special partition key values like <see cref="PartitionKey.None"/> or <see cref="PartitionKey.Null"/>.
/// For hierarchical partition keys, construct the <see cref="PartitionKey"/> using <see cref="PartitionKeyBuilder"/>.
/// </remarks>
public CosmosNoSqlKey(string documentId, PartitionKey partitionKey)
{
this.DocumentId = documentId;
this.PartitionKey = partitionKey;
}
/// <summary>
/// Initializes a new instance of the <see cref="CosmosNoSqlKey"/> struct with a document ID and partition key.
/// </summary>
/// <param name="documentId">The document ID.</param>
/// <param name="partitionKey">The Cosmos DB partition key.</param>
/// <remarks>
/// Use this constructor for hierarchical partition keys or special partition key values like <see cref="PartitionKey.None"/> or <see cref="PartitionKey.Null"/>.
/// For hierarchical partition keys, construct the <see cref="PartitionKey"/> using <see cref="PartitionKeyBuilder"/>.
/// </remarks>
public CosmosNoSqlKey(Guid documentId, PartitionKey partitionKey)
{
this.DocumentId = documentId.ToString();
this.PartitionKey = partitionKey;
}
/// <summary>
/// Gets the document ID.
/// </summary>
/// <remarks>
/// The document ID uniquely identifies a document within a partition. It is stored in the <c>id</c> property of the Cosmos DB document.
/// </remarks>
public string DocumentId { get; }
/// <summary>
/// Gets the partition key.
/// </summary>
/// <remarks>
/// The partition key determines which logical partition the document belongs to.
/// See <see href="https://learn.microsoft.com/azure/cosmos-db/partitioning-overview"/> for guidance on choosing partition keys.
/// </remarks>
public PartitionKey PartitionKey { get; }
/// <inheritdoc/>
public bool Equals(CosmosNoSqlKey other)
=> Equals(this.DocumentId, other.DocumentId) && this.PartitionKey.Equals(other.PartitionKey);
/// <inheritdoc/>
public override bool Equals(object? obj)
=> obj is CosmosNoSqlKey other && this.Equals(other);
/// <inheritdoc/>
public override int GetHashCode()
=> HashCode.Combine(this.DocumentId, this.PartitionKey);
/// <summary>
/// Determines whether two <see cref="CosmosNoSqlKey"/> instances are equal.
/// </summary>
public static bool operator ==(CosmosNoSqlKey left, CosmosNoSqlKey right)
=> left.Equals(right);
/// <summary>
/// Determines whether two <see cref="CosmosNoSqlKey"/> instances are not equal.
/// </summary>
public static bool operator !=(CosmosNoSqlKey left, CosmosNoSqlKey right)
=> !left.Equals(right);
}
@@ -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.
/// </summary>
/// <typeparam name="TRecord">The consumer data model to map to or from.</typeparam>
internal sealed class CosmosNoSqlMapper<TRecord>(CollectionModel model, JsonSerializerOptions? jsonSerializerOptions)
: ICosmosNoSqlMapper<TRecord>
internal sealed class CosmosNoSqlMapper<TRecord> : ICosmosNoSqlMapper<TRecord>
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<byte> 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<MEAI.Embedding>?[]? 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 @@ internal sealed class CosmosNoSqlMapper<TRecord>(CollectionModel model, JsonSeri
// Go over the vector properties; inject any generated embeddings to overwrite the JSON serialized above.
// Also, for Embedding<T> properties we also need to overwrite with a simple array (since Embedding<T> 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 @@ internal sealed class CosmosNoSqlMapper<TRecord>(CollectionModel model, JsonSeri
{
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<float>):
case var t2 when t2 == typeof(float[]):
// The .NET vector property is a ReadOnlyMemory<float> 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.
case var t3 when t3 == typeof(ReadOnlyMemory<byte>):
case var t4 when t4 == typeof(byte[]):
case var t5 when t5 == typeof(ReadOnlyMemory<sbyte>):
case var t6 when t6 == typeof(sbyte[]):
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<byte>):
embedding = new Embedding<byte>((ReadOnlyMemory<byte>)property.GetValueAsObject(dataModel)!);
break;
case var t4 when t4 == typeof(byte[]):
embedding = new Embedding<byte>((byte[])property.GetValueAsObject(dataModel)!);
break;
case var t5 when t5 == typeof(ReadOnlyMemory<sbyte>):
embedding = new Embedding<sbyte>((ReadOnlyMemory<sbyte>)property.GetValueAsObject(dataModel)!);
break;
case var t6 when t6 == typeof(sbyte[]):
embedding = new Embedding<sbyte>((sbyte[])property.GetValueAsObject(dataModel)!);
break;
// For Embedding<T> types, we need to extract the vector and serialize it as a simple array
case var t when t == typeof(Embedding<float>):
case var t1 when t1 == typeof(Embedding<byte>):
case var t2 when t2 == typeof(Embedding<sbyte>):
@@ -115,26 +118,37 @@ internal sealed class CosmosNoSqlMapper<TRecord>(CollectionModel model, JsonSeri
// See above comment.
RenameJsonProperty(storageModel, CosmosNoSqlConstants.ReservedKeyPropertyName, this._keyProperty.TemporaryStorageName!);
if (includeVectors)
foreach (var vectorProperty in this._model.VectorProperties)
{
foreach (var vectorProperty in model.VectorProperties)
if (!includeVectors)
{
if (vectorProperty.Type == typeof(Embedding<float>))
{
var arrayNode = storageModel[vectorProperty.StorageName];
if (arrayNode is not null)
{
var embeddingNode = new JsonObject
{
[nameof(Embedding<float>.Vector)] = arrayNode.DeepClone()
};
storageModel[vectorProperty.StorageName] = embeddingNode;
}
}
// Remove vector properties so they deserialize as default (e.g. empty ReadOnlyMemory<float>).
storageModel.Remove(vectorProperty.StorageName);
continue;
}
var arrayNode = storageModel[vectorProperty.StorageName];
if (arrayNode is null)
{
continue;
}
// Embedding<T> is stored as a simple JSON array, so convert it to the expected object shape for deserialization
if (vectorProperty.Type == typeof(Embedding<float>)
|| vectorProperty.Type == typeof(Embedding<byte>)
|| vectorProperty.Type == typeof(Embedding<sbyte>))
{
storageModel[vectorProperty.StorageName] = new JsonObject
{
[nameof(Embedding<>.Vector)] = arrayNode.DeepClone()
};
}
// For byte[], ReadOnlyMemory<byte>, sbyte[], ReadOnlyMemory<sbyte>, float[], ReadOnlyMemory<float>,
// the custom converters (for byte) and default converters (for others) handle deserialization correctly.
}
return storageModel.Deserialize<TRecord>(jsonSerializerOptions)!;
return storageModel.Deserialize<TRecord>(this._jsonSerializerOptions)!;
}
#region private
@@ -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.");
}
}
@@ -58,6 +60,15 @@ internal class CosmosNoSqlModelBuilder() : CollectionJsonModelBuilder(s_modelBui
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<Embedding<float>>(embeddingGenerator, userRequestedEmbeddingType)
?? vectorProperty.ResolveEmbeddingType<Embedding<byte>>(embeddingGenerator, userRequestedEmbeddingType)
?? vectorProperty.ResolveEmbeddingType<Embedding<sbyte>>(embeddingGenerator, userRequestedEmbeddingType);
internal static bool IsVectorPropertyTypeValidCore(Type type, [NotNullWhen(false)] out string? supportedTypes)
{
supportedTypes = SupportedVectorTypes;
@@ -127,18 +127,20 @@ public static class CosmosNoSqlServiceCollectionExtensions
/// <inheritdoc cref="AddKeyedCosmosNoSqlVectorStore(IServiceCollection, object?, CosmosNoSqlVectorStoreOptions?, ServiceLifetime)"/>
[RequiresUnreferencedCode(DynamicCodeMessage)]
[RequiresDynamicCode(UnreferencedCodeMessage)]
public static IServiceCollection AddCosmosNoSqlCollection<TRecord>(
public static IServiceCollection AddCosmosNoSqlCollection<TKey, TRecord>(
this IServiceCollection services,
string name,
CosmosNoSqlCollectionOptions? options = default,
ServiceLifetime lifetime = ServiceLifetime.Singleton)
where TKey : notnull
where TRecord : class
=> AddKeyedCosmosNoSqlCollection<TRecord>(services, serviceKey: null, name, options, lifetime);
=> AddKeyedCosmosNoSqlCollection<TKey, TRecord>(services, serviceKey: null, name, options, lifetime);
/// <summary>
/// Registers a keyed <see cref="CosmosNoSqlCollection{TKey, TRecord}"/> as <see cref="VectorStoreCollection{TKey, TRecord}"/>
/// with <see cref="Database"/> retrieved from the dependency injection container.
/// </summary>
/// <typeparam name="TKey">The type of the key.</typeparam>
/// <typeparam name="TRecord">The type of the record.</typeparam>
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="CosmosNoSqlCollection{TKey, TRecord}"/> on.</param>
/// <param name="serviceKey">The key with which to associate the collection.</param>
@@ -148,26 +150,27 @@ public static class CosmosNoSqlServiceCollectionExtensions
/// <returns>Service collection.</returns>
[RequiresUnreferencedCode(DynamicCodeMessage)]
[RequiresDynamicCode(UnreferencedCodeMessage)]
public static IServiceCollection AddKeyedCosmosNoSqlCollection<TRecord>(
public static IServiceCollection AddKeyedCosmosNoSqlCollection<TKey, TRecord>(
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<string, TRecord>), serviceKey, (sp, _) =>
services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection<TKey, TRecord>), serviceKey, (sp, _) =>
{
var database = sp.GetRequiredService<Database>();
options = GetCollectionOptions(sp, _ => options);
return new CosmosNoSqlCollection<string, TRecord>(database, name, options);
return new CosmosNoSqlCollection<TKey, TRecord>(database, name, options);
}, lifetime));
AddAbstractions<string, TRecord>(services, serviceKey, lifetime);
AddAbstractions<TKey, TRecord>(services, serviceKey, lifetime);
return services;
}
@@ -176,23 +179,25 @@ public static class CosmosNoSqlServiceCollectionExtensions
/// Registers a <see cref="CosmosNoSqlCollection{TKey, TRecord}"/> as <see cref="VectorStoreCollection{TKey,TRecord}"/>
/// using the provided <paramref name="connectionString"/> and <paramref name="databaseName"/>.
/// </summary>
/// <inheritdoc cref="AddKeyedCosmosNoSqlCollection{TRecord}(IServiceCollection, object?, string, string, string, CosmosNoSqlCollectionOptions?, ServiceLifetime)"/>
/// <inheritdoc cref="AddKeyedCosmosNoSqlCollection{TKey, TRecord}(IServiceCollection, object?, string, string, string, CosmosNoSqlCollectionOptions?, ServiceLifetime)"/>
[RequiresUnreferencedCode(UnreferencedCodeMessage)]
[RequiresDynamicCode(DynamicCodeMessage)]
public static IServiceCollection AddCosmosNoSqlCollection<TRecord>(
public static IServiceCollection AddCosmosNoSqlCollection<TKey, TRecord>(
this IServiceCollection services,
string name,
string connectionString,
string databaseName,
CosmosNoSqlCollectionOptions? options = default,
ServiceLifetime lifetime = ServiceLifetime.Singleton)
where TKey : notnull
where TRecord : class
=> AddKeyedCosmosNoSqlCollection<TRecord>(services, serviceKey: null, name, connectionString, databaseName, options, lifetime);
=> AddKeyedCosmosNoSqlCollection<TKey, TRecord>(services, serviceKey: null, name, connectionString, databaseName, options, lifetime);
/// <summary>
/// Registers a keyed <see cref="CosmosNoSqlCollection{TKey, TRecord}"/> as <see cref="VectorStoreCollection{TKey,TRecord}"/>
/// using the provided <paramref name="connectionString"/> and <paramref name="databaseName"/>.
/// </summary>
/// <typeparam name="TKey">The type of the key.</typeparam>
/// <typeparam name="TRecord">The type of the record.</typeparam>
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="CosmosNoSqlCollection{TKey, TRecord}"/> on.</param>
/// <param name="serviceKey">The key with which to associate the collection.</param>
@@ -204,7 +209,7 @@ public static class CosmosNoSqlServiceCollectionExtensions
/// <returns>Service collection.</returns>
[RequiresUnreferencedCode(UnreferencedCodeMessage)]
[RequiresDynamicCode(DynamicCodeMessage)]
public static IServiceCollection AddKeyedCosmosNoSqlCollection<TRecord>(
public static IServiceCollection AddKeyedCosmosNoSqlCollection<TKey, TRecord>(
this IServiceCollection services,
object? serviceKey,
string name,
@@ -212,35 +217,38 @@ public static class CosmosNoSqlServiceCollectionExtensions
string databaseName,
CosmosNoSqlCollectionOptions? options = default,
ServiceLifetime lifetime = ServiceLifetime.Singleton)
where TKey : notnull
where TRecord : class
{
Verify.NotNullOrWhiteSpace(connectionString);
Verify.NotNullOrWhiteSpace(databaseName);
return AddKeyedCosmosNoSqlCollection<TRecord>(services, serviceKey, name, _ => connectionString, _ => databaseName, _ => options!, lifetime);
return AddKeyedCosmosNoSqlCollection<TKey, TRecord>(services, serviceKey, name, _ => connectionString, _ => databaseName, _ => options!, lifetime);
}
/// <summary>
/// Registers a <see cref="CosmosNoSqlCollection{TKey, TRecord}"/> as <see cref="VectorStoreCollection{TKey,TRecord}"/>
/// using the provided <paramref name="connectionStringProvider"/> and <paramref name="databaseNameProvider"/>.
/// </summary>
/// <inheritdoc cref="AddKeyedCosmosNoSqlCollection{TRecord}(IServiceCollection, object?, string, Func{IServiceProvider, string}, Func{IServiceProvider, string}, Func{IServiceProvider, CosmosNoSqlCollectionOptions}?, ServiceLifetime)"/>
/// <inheritdoc cref="AddKeyedCosmosNoSqlCollection{TKey, TRecord}(IServiceCollection, object?, string, Func{IServiceProvider, string}, Func{IServiceProvider, string}, Func{IServiceProvider, CosmosNoSqlCollectionOptions}?, ServiceLifetime)"/>
[RequiresUnreferencedCode(UnreferencedCodeMessage)]
[RequiresDynamicCode(DynamicCodeMessage)]
public static IServiceCollection AddCosmosNoSqlCollection<TRecord>(
public static IServiceCollection AddCosmosNoSqlCollection<TKey, TRecord>(
this IServiceCollection services,
string name,
Func<IServiceProvider, string> connectionStringProvider,
Func<IServiceProvider, string> databaseNameProvider,
Func<IServiceProvider, CosmosNoSqlCollectionOptions>? optionsProvider = default,
ServiceLifetime lifetime = ServiceLifetime.Singleton)
where TKey : notnull
where TRecord : class
=> AddKeyedCosmosNoSqlCollection<TRecord>(services, serviceKey: null, name, connectionStringProvider, databaseNameProvider, optionsProvider, lifetime);
=> AddKeyedCosmosNoSqlCollection<TKey, TRecord>(services, serviceKey: null, name, connectionStringProvider, databaseNameProvider, optionsProvider, lifetime);
/// <summary>
/// Registers a keyed <see cref="CosmosNoSqlCollection{TKey, TRecord}"/> as <see cref="VectorStoreCollection{TKey,TRecord}"/>
/// using the provided <paramref name="connectionStringProvider"/> and <paramref name="databaseNameProvider"/>.
/// </summary>
/// <typeparam name="TKey">The type of the key.</typeparam>
/// <typeparam name="TRecord">The type of the record.</typeparam>
/// <param name="services">The <see cref="IServiceCollection"/> to register the <see cref="VectorStoreCollection{TKey, TRecord}"/> on.</param>
/// <param name="serviceKey">The key with which to associate the collection.</param>
@@ -252,7 +260,7 @@ public static class CosmosNoSqlServiceCollectionExtensions
/// <returns>Service collection.</returns>
[RequiresUnreferencedCode(UnreferencedCodeMessage)]
[RequiresDynamicCode(DynamicCodeMessage)]
public static IServiceCollection AddKeyedCosmosNoSqlCollection<TRecord>(
public static IServiceCollection AddKeyedCosmosNoSqlCollection<TKey, TRecord>(
this IServiceCollection services,
object? serviceKey,
string name,
@@ -260,6 +268,7 @@ public static class CosmosNoSqlServiceCollectionExtensions
Func<IServiceProvider, string> databaseNameProvider,
Func<IServiceProvider, CosmosNoSqlCollectionOptions>? optionsProvider = default,
ServiceLifetime lifetime = ServiceLifetime.Singleton)
where TKey : notnull
where TRecord : class
{
Verify.NotNull(services);
@@ -267,12 +276,12 @@ public static class CosmosNoSqlServiceCollectionExtensions
Verify.NotNull(connectionStringProvider);
Verify.NotNull(databaseNameProvider);
services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection<string, TRecord>), serviceKey, (sp, _) =>
services.Add(new ServiceDescriptor(typeof(CosmosNoSqlCollection<TKey, TRecord>), serviceKey, (sp, _) =>
{
var options = GetCollectionOptions(sp, optionsProvider);
var clientOptions = CreateClientOptions(options?.JsonSerializerOptions);
return new CosmosNoSqlCollection<string, TRecord>(
return new CosmosNoSqlCollection<TKey, TRecord>(
connectionString: connectionStringProvider(sp),
databaseName: databaseNameProvider(sp),
name: name,
@@ -280,7 +289,7 @@ public static class CosmosNoSqlServiceCollectionExtensions
options);
}, lifetime));
AddAbstractions<string, TRecord>(services, serviceKey, lifetime);
AddAbstractions<TKey, TRecord>(services, serviceKey, lifetime);
return services;
}
@@ -20,11 +20,11 @@ public sealed class CosmosNoSqlCollectionOptionsTests(CosmosNoSqlCollectionOptio
var store = (CosmosNoSqlTestStore)fixture.TestStore;
var collectionName = fixture.TestStore.AdjustCollectionName("PartitionKeyCompositeKey");
using VectorStoreCollection<CosmosNoSqlCompositeKey, PartitionedHotel> collection =
new CosmosNoSqlCollection<CosmosNoSqlCompositeKey, PartitionedHotel>(
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<CosmosNoSqlKey, PartitionedHotel>(
store.Database,
collectionName,
new() { PartitionKeyProperties = [nameof(PartitionedHotel.HotelName)] });
await collection.EnsureCollectionExistsAsync();
@@ -39,7 +39,7 @@ public sealed class CosmosNoSqlCollectionOptionsTests(CosmosNoSqlCollectionOptio
};
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);
@@ -48,22 +48,22 @@ public class CosmosNoSqlDependencyInjectionTests
? services
.AddSingleton<CosmosClient>(sp => new CosmosClient(ConnectionString, s_clientOptions))
.AddSingleton<Database>(sp => sp.GetRequiredService<CosmosClient>().GetDatabase(DatabaseName))
.AddCosmosNoSqlCollection<Record>(name, lifetime: lifetime)
.AddCosmosNoSqlCollection<string, Record>(name, lifetime: lifetime)
: services
.AddSingleton<CosmosClient>(sp => new CosmosClient(ConnectionString, s_clientOptions))
.AddSingleton<Database>(sp => sp.GetRequiredService<CosmosClient>().GetDatabase(DatabaseName))
.AddKeyedCosmosNoSqlCollection<Record>(serviceKey, name, lifetime: lifetime);
.AddKeyedCosmosNoSqlCollection<string, Record>(serviceKey, name, lifetime: lifetime);
yield return (services, serviceKey, name, lifetime) => serviceKey is null
? services.AddCosmosNoSqlCollection<Record>(
? services.AddCosmosNoSqlCollection<string, Record>(
name, ConnectionString, DatabaseName, lifetime: lifetime)
: services.AddKeyedCosmosNoSqlCollection<Record>(
: services.AddKeyedCosmosNoSqlCollection<string, Record>(
serviceKey, name, ConnectionString, DatabaseName, lifetime: lifetime);
yield return (services, serviceKey, name, lifetime) => serviceKey is null
? services.AddCosmosNoSqlCollection<Record>(
? services.AddCosmosNoSqlCollection<string, Record>(
name, ConnectionStringProvider, DatabaseNameProvider, lifetime: lifetime)
: services.AddKeyedCosmosNoSqlCollection<Record>(
: services.AddKeyedCosmosNoSqlCollection<string, Record>(
serviceKey, name, sp => ConnectionStringProvider(sp, serviceKey), sp => DatabaseNameProvider(sp, serviceKey), lifetime: lifetime);
}
}
@@ -95,9 +95,9 @@ public class CosmosNoSqlDependencyInjectionTests
{
IServiceCollection services = new ServiceCollection();
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlCollection<string, Record>(
name: "notNull", connectionStringProvider: null!, databaseNameProvider: DatabaseNameProvider));
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlCollection<string, Record>(
serviceKey: "notNull", name: "notNull", connectionStringProvider: null!, databaseNameProvider: DatabaseNameProvider));
}
@@ -106,9 +106,9 @@ public class CosmosNoSqlDependencyInjectionTests
{
IServiceCollection services = new ServiceCollection();
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlCollection<string, Record>(
name: "notNull", connectionStringProvider: ConnectionStringProvider, databaseNameProvider: null!));
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlCollection<string, Record>(
serviceKey: "notNull", name: "notNull", connectionStringProvider: ConnectionStringProvider, databaseNameProvider: null!));
}
@@ -119,13 +119,13 @@ public class CosmosNoSqlDependencyInjectionTests
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlVectorStore(connectionString: null!, DatabaseName));
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlVectorStore(serviceKey: "notNull", connectionString: null!, DatabaseName));
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlCollection<string, Record>(
name: "notNull", connectionString: null!, DatabaseName));
Assert.Throws<ArgumentException>(() => services.AddCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentException>(() => services.AddCosmosNoSqlCollection<string, Record>(
name: "notNull", connectionString: "", DatabaseName));
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlCollection<string, Record>(
serviceKey: "notNull", name: "notNull", connectionString: null!, DatabaseName));
Assert.Throws<ArgumentException>(() => services.AddKeyedCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentException>(() => services.AddKeyedCosmosNoSqlCollection<string, Record>(
serviceKey: "notNull", name: "notNull", connectionString: "", DatabaseName));
}
@@ -136,13 +136,13 @@ public class CosmosNoSqlDependencyInjectionTests
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlVectorStore(ConnectionString, databaseName: null!));
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlVectorStore(serviceKey: "notNull", ConnectionString, databaseName: null!));
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentNullException>(() => services.AddCosmosNoSqlCollection<string, Record>(
name: "notNull", ConnectionString, databaseName: null!));
Assert.Throws<ArgumentException>(() => services.AddCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentException>(() => services.AddCosmosNoSqlCollection<string, Record>(
name: "notNull", ConnectionString, databaseName: ""));
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentNullException>(() => services.AddKeyedCosmosNoSqlCollection<string, Record>(
serviceKey: "notNull", name: "notNull", ConnectionString, databaseName: null!));
Assert.Throws<ArgumentException>(() => services.AddKeyedCosmosNoSqlCollection<Record>(
Assert.Throws<ArgumentException>(() => services.AddKeyedCosmosNoSqlCollection<string, Record>(
serviceKey: "notNull", name: "notNull", ConnectionString, databaseName: ""));
}
}
@@ -31,7 +31,7 @@ public class CosmosNoSqlEmbeddingGenerationTests(CosmosNoSqlEmbeddingGenerationT
[
services => services
.AddSingleton(CosmosNoSqlTestStore.Instance.Database)
.AddCosmosNoSqlCollection<RecordWithAttributes>(this.CollectionName)
.AddCosmosNoSqlCollection<string, RecordWithAttributes>(this.CollectionName)
];
}
@@ -53,7 +53,7 @@ public class CosmosNoSqlEmbeddingGenerationTests(CosmosNoSqlEmbeddingGenerationT
[
services => services
.AddSingleton(CosmosNoSqlTestStore.Instance.Database)
.AddCosmosNoSqlCollection<RecordWithAttributes>(this.CollectionName)
.AddCosmosNoSqlCollection<string, RecordWithAttributes>(this.CollectionName)
];
}
}
@@ -11,6 +11,7 @@ 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
{
@@ -14,41 +14,41 @@ namespace CosmosNoSql.ConformanceTests.TypeTests;
public class CosmosNoSqlEmbeddingTypeTests(CosmosNoSqlEmbeddingTypeTests.Fixture fixture)
: EmbeddingTypeTests<string>(fixture), IClassFixture<CosmosNoSqlEmbeddingTypeTests.Fixture>
{
[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<ReadOnlyMemory<byte>>(
new ReadOnlyMemory<byte>([1, 2, 3]),
new ReadOnlyMemoryEmbeddingGenerator<byte>([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<Embedding<byte>>(
new Embedding<byte>(new ReadOnlyMemory<byte>([1, 2, 3])),
new ReadOnlyMemoryEmbeddingGenerator<byte>([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<byte[]>(
[1, 2, 3],
new ReadOnlyMemoryEmbeddingGenerator<byte>([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<ReadOnlyMemory<sbyte>>(
new ReadOnlyMemory<sbyte>([1, 2, 3]),
new ReadOnlyMemoryEmbeddingGenerator<sbyte>([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<Embedding<sbyte>>(
new Embedding<sbyte>(new ReadOnlyMemory<sbyte>([1, 2, 3])),
new ReadOnlyMemoryEmbeddingGenerator<sbyte>([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<sbyte[]>(
[1, 2, 3],
@@ -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,53 +185,6 @@ public sealed class CosmosNoSqlCollectionQueryBuilderTests
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 AND x["TestProperty1"] = @pk0)
""";
const string KeyStoragePropertyName = "id";
const string PartitionKeyPropertyName = "TestProperty1";
var model = new CosmosNoSqlModelBuilder().BuildDynamic(
new()
{
Properties =
[
new VectorStoreKeyProperty("Key", typeof(string)),
new VectorStoreDataProperty("TestProperty1", typeof(string)),
new VectorStoreDataProperty("TestProperty2", typeof(string))
]
},
defaultEmbeddingGenerator: null);
var keys = new List<CosmosNoSqlCompositeKey> { new("id", "TestProperty1") };
// Act
var queryDefinition = CosmosNoSqlCollectionQueryBuilder.BuildSelectQuery(
model,
KeyStoragePropertyName,
PartitionKeyPropertyName,
keys,
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);
Assert.Equal("@pk0", queryParameters[1].Name);
Assert.Equal("TestProperty1", queryParameters[1].Value);
}
[Fact]
public void BuildSearchQueryWithHybridFieldsReturnsValidHybridQueryDefinition()
{
@@ -162,7 +162,11 @@ public sealed class CosmosNoSqlCollectionTests
using var sut = new CosmosNoSqlCollection<string, TestIndexingModel>(
this._mockDatabase.Object,
CollectionName,
new() { IndexingMode = indexingMode, Automatic = indexingMode != IndexingMode.None });
new()
{
IndexingMode = indexingMode,
Automatic = indexingMode != IndexingMode.None
});
var expectedVectorEmbeddingPolicy = new VectorEmbeddingPolicy(
[
@@ -280,21 +284,21 @@ public sealed class CosmosNoSqlCollectionTests
[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)
if (useExplicitPartitionKey)
{
using var sut = new CosmosNoSqlCollection<CosmosNoSqlCompositeKey, CosmosNoSqlHotel>(
using var sut = new CosmosNoSqlCollection<CosmosNoSqlKey, CosmosNoSqlHotel>(
this._mockDatabase.Object,
"collection");
await ((VectorStoreCollection<CosmosNoSqlCompositeKey, CosmosNoSqlHotel>)sut).DeleteAsync(
new CosmosNoSqlCompositeKey(RecordKey, PartitionKey));
await ((VectorStoreCollection<CosmosNoSqlKey, CosmosNoSqlHotel>)sut).DeleteAsync(
new CosmosNoSqlKey(RecordKey, PartitionKey));
}
else
{
@@ -366,27 +370,18 @@ public sealed class CosmosNoSqlCollectionTests
var jsonObject = new JsonObject { ["id"] = RecordKey, ["HotelName"] = "Test Name" };
var mockFeedResponse = new Mock<FeedResponse<JsonObject>>();
mockFeedResponse
var mockItemResponse = new Mock<ItemResponse<JsonObject>>();
mockItemResponse
.Setup(l => l.Resource)
.Returns([jsonObject]);
var mockFeedIterator = new Mock<FeedIterator<JsonObject>>();
mockFeedIterator
.SetupSequence(l => l.HasMoreResults)
.Returns(true)
.Returns(false);
mockFeedIterator
.Setup(l => l.ReadNextAsync(It.IsAny<CancellationToken>()))
.ReturnsAsync(mockFeedResponse.Object);
.Returns(jsonObject);
this._mockContainer
.Setup(l => l.GetItemQueryIterator<JsonObject>(
It.IsAny<QueryDefinition>(),
It.IsAny<string>(),
It.IsAny<QueryRequestOptions>()))
.Returns(mockFeedIterator.Object);
.Setup(l => l.ReadItemAsync<JsonObject>(
RecordKey,
new PartitionKey(RecordKey),
It.IsAny<ItemRequestOptions>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync(mockItemResponse.Object);
using var sut = new CosmosNoSqlCollection<string, CosmosNoSqlHotel>(
this._mockDatabase.Object,
@@ -414,22 +409,12 @@ public sealed class CosmosNoSqlCollectionTests
.Setup(l => l.Resource)
.Returns([jsonObject1, jsonObject2, jsonObject3]);
var mockFeedIterator = new Mock<FeedIterator<JsonObject>>();
mockFeedIterator
.SetupSequence(l => l.HasMoreResults)
.Returns(true)
.Returns(false);
mockFeedIterator
.Setup(l => l.ReadNextAsync(It.IsAny<CancellationToken>()))
.ReturnsAsync(mockFeedResponse.Object);
this._mockContainer
.Setup(l => l.GetItemQueryIterator<JsonObject>(
It.IsAny<QueryDefinition>(),
It.IsAny<string>(),
It.IsAny<QueryRequestOptions>()))
.Returns(mockFeedIterator.Object);
.Setup(l => l.ReadManyItemsAsync<JsonObject>(
It.IsAny<IReadOnlyList<(string id, PartitionKey partitionKey)>>(),
It.IsAny<ReadManyRequestOptions>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync(mockFeedResponse.Object);
using var sut = new CosmosNoSqlCollection<string, CosmosNoSqlHotel>(
this._mockDatabase.Object,
@@ -667,5 +652,152 @@ public sealed class CosmosNoSqlCollectionTests
}
#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<string, CosmosNoSqlHotel>(
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<string, CosmosNoSqlHotel>(
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<ArgumentException>(() => new CosmosNoSqlCollection<string, CosmosNoSqlHotel>(
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<ArgumentException>(() => new CosmosNoSqlCollection<string, CosmosNoSqlHotel>(
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<Guid, GuidKeyModel>(
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<ArgumentException>(() => new CosmosNoSqlCollection<Guid, GuidKeyModel>(
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<NotSupportedException>(() => new CosmosNoSqlCollection<int, CosmosNoSqlHotel>(
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<string, CosmosNoSqlHotel>(
this._mockDatabase.Object,
"collection");
// Act
await sut.DeleteAsync(RecordKey);
// Assert: with string key, partition key = document id
this._mockContainer.Verify(l => l.DeleteItemAsync<JsonObject>(
RecordKey,
new PartitionKey(RecordKey),
It.IsAny<ItemRequestOptions>(),
It.IsAny<CancellationToken>()),
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<ItemResponse<JsonObject>>();
mockItemResponse
.Setup(l => l.Resource)
.Returns(jsonObject);
this._mockContainer
.Setup(l => l.ReadItemAsync<JsonObject>(
RecordKey,
new PartitionKey(RecordKey),
It.IsAny<ItemRequestOptions>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync(mockItemResponse.Object);
using var sut = new CosmosNoSqlCollection<string, CosmosNoSqlHotel>(
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
}
@@ -48,12 +48,36 @@ public sealed class CosmosNoSqlVectorStoreTests
using var sut = new Microsoft.SemanticKernel.Connectors.CosmosNoSql.CosmosNoSqlVectorStore(this._mockDatabase.Object);
// Act
var collectionWithStringKey = sut.GetCollection<string, CosmosNoSqlHotel>("collection1");
var collectionWithCompositeKey = sut.GetCollection<CosmosNoSqlCompositeKey, CosmosNoSqlHotel>("collection1");
var collection = sut.GetCollection<CosmosNoSqlKey, CosmosNoSqlHotel>("collection1");
// Assert
Assert.NotNull(collectionWithStringKey);
Assert.NotNull(collectionWithCompositeKey);
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<string, CosmosNoSqlHotel>("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<Guid, GuidKeyHotel>("collection1");
// Assert
Assert.NotNull(collection);
}
[Fact]
@@ -63,7 +87,7 @@ public sealed class CosmosNoSqlVectorStoreTests
using var sut = new Microsoft.SemanticKernel.Connectors.CosmosNoSql.CosmosNoSqlVectorStore(this._mockDatabase.Object);
// Act
var collection = sut.GetCollection<string, CosmosNoSqlHotel>("collection");
var collection = sut.GetCollection<CosmosNoSqlKey, CosmosNoSqlHotel>("collection");
// Assert
Assert.NotNull(collection);
@@ -105,4 +129,15 @@ public sealed class CosmosNoSqlVectorStoreTests
// 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
}
@@ -97,7 +97,7 @@ public abstract class CollectionManagementTests<TKey>(VectorStoreFixture fixture
}
public virtual VectorStoreCollection<TKey, Record> GetCollection()
=> fixture.TestStore.DefaultVectorStore.GetCollection<TKey, Record>(this.CollectionName, this.CreateRecordDefinition());
=> fixture.TestStore.CreateCollection<TKey, Record>(this.CollectionName, this.CreateRecordDefinition());
public virtual VectorStoreCollectionDefinition CreateRecordDefinition()
=> new()
@@ -185,7 +185,7 @@ public abstract class DistanceFunctionTests<TKey>(DistanceFunctionTests<TKey>.Fi
]
};
return this.TestStore.DefaultVectorStore.GetCollection<TKey, SearchRecord>(this.CollectionName, definition);
return this.TestStore.CreateCollection<TKey, SearchRecord>(this.CollectionName, definition);
}
}
@@ -574,7 +574,7 @@ public abstract class EmbeddingGenerationTests<TKey>(EmbeddingGenerationTests<TK
=> Interlocked.Increment(ref this._counter);
}
private sealed class FakeEmbeddingGenerator(int? replaceLast = null) : IEmbeddingGenerator<string, Embedding<float>>
protected sealed class FakeEmbeddingGenerator(int? replaceLast = null) : IEmbeddingGenerator<string, Embedding<float>>
{
public Task<GeneratedEmbeddings<Embedding<float>>> GenerateAsync(
IEnumerable<string> values,
@@ -640,7 +640,7 @@ public abstract class FilterTests<TKey>(FilterTests<TKey>.Fixture fixture)
if (this.TestDynamic)
{
this.DynamicCollection = this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition());
this.DynamicCollection = this.TestStore.CreateDynamicCollection(this.CollectionName, this.CreateRecordDefinition());
}
}
@@ -75,7 +75,7 @@ public abstract class IndexKindTests<TKey>(IndexKindTests<TKey>.Fixture fixture)
]
};
return this.TestStore.DefaultVectorStore.GetCollection<TKey, SearchRecord>(this.CollectionName, definition);
return this.TestStore.CreateCollection<TKey, SearchRecord>(this.CollectionName, definition);
}
}
@@ -396,7 +396,7 @@ public abstract class DynamicModelTests<TKey>(DynamicModelTests<TKey>.Fixture fi
protected override string KeyPropertyName => DynamicModelTests<TKey>.KeyPropertyName;
protected override VectorStoreCollection<object, Dictionary<string, object?>> GetCollection()
=> this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition());
=> this.TestStore.CreateDynamicCollection(this.CollectionName, this.CreateRecordDefinition());
public override VectorStoreCollectionDefinition CreateRecordDefinition()
=> new()
@@ -92,6 +92,26 @@ public abstract class TestStore
public virtual string AdjustCollectionName(string baseName)
=> baseName;
/// <summary>
/// Creates a collection for the given name and definition.
/// Override this to provide provider-specific collection options (e.g., partition key configuration).
/// </summary>
public virtual VectorStoreCollection<TKey, TRecord> CreateCollection<TKey, TRecord>(
string name,
VectorStoreCollectionDefinition definition)
where TKey : notnull
where TRecord : class
=> this.DefaultVectorStore.GetCollection<TKey, TRecord>(name, definition);
/// <summary>
/// Creates a dynamic collection for the given name and definition.
/// Override this to provide provider-specific collection options (e.g., partition key configuration).
/// </summary>
public virtual VectorStoreCollection<object, Dictionary<string, object?>> CreateDynamicCollection(
string name,
VectorStoreCollectionDefinition definition)
=> this.DefaultVectorStore.GetDynamicCollection(name, definition);
/// <summary>Loops until the expected number of records is visible in the given collection.</summary>
/// <remarks>Some databases upsert asynchronously, meaning that our seed data may not be visible immediately to tests.</remarks>
public virtual async Task WaitForDataAsync<TKey, TRecord>(
@@ -34,7 +34,7 @@ public abstract class VectorStoreCollectionFixtureBase<TKey, TRecord> : VectorSt
protected virtual string IndexKind => this.TestStore.DefaultIndexKind;
protected virtual VectorStoreCollection<TKey, TRecord> GetCollection()
=> this.TestStore.DefaultVectorStore.GetCollection<TKey, TRecord>(this.CollectionName, this.CreateRecordDefinition());
=> this.TestStore.CreateCollection<TKey, TRecord>(this.CollectionName, this.CreateRecordDefinition());
public override async Task InitializeAsync()
{
@@ -23,4 +23,20 @@ public abstract class VectorStoreFixture : IAsyncLifetime
public virtual TKey GenerateNextKey<TKey>()
=> this.TestStore.GenerateKey<TKey>(Interlocked.Increment(ref this._nextKeyValue));
/// <summary>
/// Creates a collection for the given name and definition.
/// Delegates to <see cref="TestStore.CreateCollection{TKey, TRecord}"/> which can be overridden for provider-specific options.
/// </summary>
public virtual VectorStoreCollection<TKey, TRecord> CreateCollection<TKey, TRecord>(string name, VectorStoreCollectionDefinition definition)
where TKey : notnull
where TRecord : class
=> this.TestStore.CreateCollection<TKey, TRecord>(name, definition);
/// <summary>
/// Creates a dynamic collection for the given name and definition.
/// Delegates to <see cref="TestStore.CreateDynamicCollection"/> which can be overridden for provider-specific options.
/// </summary>
public virtual VectorStoreCollection<object, Dictionary<string, object?>> CreateDynamicCollection(string name, VectorStoreCollectionDefinition definition)
=> this.TestStore.CreateDynamicCollection(name, definition);
}
@@ -192,7 +192,7 @@ public abstract class DataTypeTests<TKey, TRecord>(DataTypeTests<TKey, TRecord>.
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)
{
@@ -50,7 +50,7 @@ public abstract class EmbeddingTypeTests<TKey>(EmbeddingTypeTests<TKey>.Fixture
await fixture.VectorStore.EnsureCollectionDeletedAsync(fixture.CollectionName);
var collection = fixture.VectorStore.GetCollection<TKey, Record<TVector>>(fixture.CollectionName, fixture.CreateRecordDefinition<TVector>(embeddingGenerator: null, distanceFunction, dimensions));
var collection = fixture.CreateCollection<TKey, Record<TVector>>(fixture.CollectionName, fixture.CreateRecordDefinition<TVector>(embeddingGenerator: null, distanceFunction, dimensions));
await collection.EnsureCollectionExistsAsync();
var key = fixture.GenerateNextKey<TKey>();
@@ -93,7 +93,7 @@ public abstract class EmbeddingTypeTests<TKey>(EmbeddingTypeTests<TKey>.Fixture
await collection.DeleteAsync(key);
}
var dynamicCollection = fixture.VectorStore.GetDynamicCollection(fixture.CollectionName, fixture.CreateRecordDefinition<TVector>(embeddingGenerator, distanceFunction, dimensions));
var dynamicCollection = fixture.CreateDynamicCollection(fixture.CollectionName, fixture.CreateRecordDefinition<TVector>(embeddingGenerator, distanceFunction, dimensions));
if (fixture.RecreateCollection)
{
@@ -142,14 +142,14 @@ public abstract class EmbeddingTypeTests<TKey>(EmbeddingTypeTests<TKey>.Fixture
await collection.DeleteAsync(key);
}
var collection2 = fixture.VectorStore.GetCollection<TKey, RecordWithString>(fixture.CollectionName, fixture.CreateRecordDefinition<string>(embeddingGenerator, distanceFunction, dimensions));
var collection2 = fixture.CreateCollection<TKey, RecordWithString>(fixture.CollectionName, fixture.CreateRecordDefinition<string>(embeddingGenerator, distanceFunction, dimensions));
if (fixture.RecreateCollection)
{
await collection2.EnsureCollectionExistsAsync();
}
var key2 = fixture.GenerateNextKey<TKey>();
key = fixture.GenerateNextKey<TKey>();
var record2 = new RecordWithString
{
Key = key,
@@ -217,7 +217,7 @@ public abstract class EmbeddingTypeTests<TKey>(EmbeddingTypeTests<TKey>.Fixture
public abstract class Fixture : VectorStoreFixture
{
protected virtual string CollectionNameBase => nameof(EmbeddingTypeTests<int>);
protected virtual string CollectionNameBase => nameof(EmbeddingTypeTests<>);
public virtual string CollectionName => this.TestStore.AdjustCollectionName(this.CollectionNameBase);
public virtual VectorStoreCollectionDefinition CreateRecordDefinition<TVectorProperty>(IEmbeddingGenerator? embeddingGenerator, string? distanceFunction, int dimensions)
@@ -209,11 +209,11 @@ public abstract class KeyTypeTests(KeyTypeTests.Fixture fixture)
public virtual VectorStoreCollection<TKey, Record<TKey>> CreateCollection<TKey>(bool? withAutoGeneration)
where TKey : notnull
=> this.TestStore.DefaultVectorStore.GetCollection<TKey, Record<TKey>>(this.CollectionName, this.CreateRecordDefinition<TKey>(withAutoGeneration));
=> this.TestStore.CreateCollection<TKey, Record<TKey>>(this.CollectionName, this.CreateRecordDefinition<TKey>(withAutoGeneration));
public virtual VectorStoreCollection<object, Dictionary<string, object?>> CreateDynamicCollection<TKey>(bool withAutoGeneration)
where TKey : notnull
=> this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition<TKey>(withAutoGeneration));
=> this.TestStore.CreateDynamicCollection(this.CollectionName, this.CreateRecordDefinition<TKey>(withAutoGeneration));
public virtual VectorStoreCollectionDefinition CreateRecordDefinition<TKey>(bool? withAutoGeneration)
where TKey : notnull