mirror of
https://github.com/microsoft/semantic-kernel.git
synced 2026-10-12 02:19:06 +00:00
.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:
@@ -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<byte> 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<string, object?></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;
|
||||
}
|
||||
|
||||
+6
-6
@@ -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);
|
||||
|
||||
+18
-18
@@ -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: ""));
|
||||
}
|
||||
}
|
||||
|
||||
+2
-2
@@ -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
|
||||
{
|
||||
|
||||
+6
-6
@@ -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],
|
||||
|
||||
-48
@@ -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>(
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user