Help with AzureCosmosDBNoSQL Vector Store Use

[dccrain]

I am using Microsoft.SemanticKernel 1.26 and Microsoft.SemanticKernel.Connectors.AzureCosmosDBNoSQL 1.26.0-preview. I am trying to insert records into the vector store but when I run the following code I get an error. Here is the code and TestRecord class for reference.

// class for ref
public class TestRecord
 {
        [VectorStoreRecordKey]
        public string Id { get; set; }

        [VectorStoreRecordData(IsFilterable = true)]
        public string DocumentId { get; set; }

        [VectorStoreRecordVector(1536)]
        public ReadOnlyMemory<float> ChunkEmbedding { get; set; }
}

// 
CosmosClient client = new CosmosClient(_cosmosEndpoint, _cosmosKey);
Database database = client.GetDatabase(_cosmosdb);
AzureCosmosDBNoSQLVectorStoreOptions options = new()
{
     JsonSerializerOptions = new()
     {
       ReferenceHandler = ReferenceHandler.IgnoreCycles
      }
 };

var cosmosVectorStore = new AzureCosmosDBNoSQLVectorStore(database, options);
var collection = cosmosVectorStore.GetCollection<string, TestRecord>(request.Collection);
await collection.CreateCollectionIfNotExistsAsync();

var testRecords = new List<TestRecord>() { 
new() { Id = "test", DocumentId= "doc1", ChunkEmbedding = await _embeddingGenerationService.GenerateEmbeddingAsync("test") }, 
new() { Id = "test2", DocumentId = "doc2", ChunkEmbedding = await _embeddingGenerationService.GenerateEmbeddingAsync("test12") }
 };

await foreach (var record in collection.UpsertBatchAsync(testRecords))
{
    _logger.LogInformation($"Method DocumentProcessingService.GenerateAndStoreEmbeddingsForFileAsync | Successfully saved chunk:{record} for file {request.FileName}");
}

Here is the error that is caught when upserting

Microsoft.Extensions.VectorData.VectorStoreOperationException: 'Call to vector store failed.'
Inner Exception
JsonSerializationException: Self referencing loop detected for property 'Parent' with type 'System.Text.Json.Nodes.JsonObject'. Path 'DocumentId'.

Am I doing something wrong here? Any help would be appreciated.

[evchaki]

@westey-m - can you take a look at this?

[westey-m]

Hi @dccrain,
I'm still investigating the root cause of the issue, but in the mean time I was able to find a workaround. The error is actually happening on the Cosmos DB client side, therefore, if you pass the json serializer options with IgnoreCycles into the CosmosClient constructor, this is stopping the issue from reproducing.

var jsonSerializerOptions = new JsonSerializerOptions
{
    ReferenceHandler = ReferenceHandler.IgnoreCycles
};

CosmosClient client = new(
    _cosmosEndpoint,
    _cosmosKey,
    new CosmosClientOptions() { UseSystemTextJsonSerializerWithOptions = jsonSerializerOptions });
Database database = client.GetDatabase(_cosmosdb);
var cosmosVectorStore = new AzureCosmosDBNoSQLVectorStore(database);

[westey-m]

@dccrain, a further update. The default serializer used by the CosmosClient has a limitation in its ability to serialize certain types including JsonObject and ReadOnlyMemory<float>, so in order for the CosmosClient to be used with Vectors and by extension the VectorStores implementation, its serializer has to be overridden. Setting UseSystemTextJsonSerializerWithOptions does this overriding so that System.Text.Json is used instead of the built serializer.

It is not actually required to set the IgnoreCycles setting either. The following works just as well.

CosmosClient client = new(
    _cosmosEndpoint,
    _cosmosKey,
    new CosmosClientOptions() { UseSystemTextJsonSerializerWithOptions = JsonSerializerOptions.Default });

We missed adding this restriction to our documentation. Apologies for that. We'll be adding an update to the documentation within the next couple of days.

[dccrain]

Thank you for your answer, I was able to get it working. I will mark the response as the answer.

I have some other questions about best practices for using this connector to store vectors. My current experience is that the VectorStoreRecordKey attribute corresponds to the /id field in cosmos, and from what I can tell I can't set a different partition key because Upsert operations will fail with the error if I have any other field set as a partition key.

Response status code does not indicate success: BadRequest (400); Substatus: 1001; ActivityId: 95b5b72b-82d7-4412-93a3-da01c7e3ccf3; Reason: (Message: {"Errors":["PartitionKey extracted from document doesn't match the one specified in the header. Learn more: https://aka.ms/CosmosDB/sql/errors/wrong-pk-value"]}

This also is the case when attempting to use hierarchical partition keys.

Is best practice to only use the unique id of the record as the partition key for the container for your vector records, and that unique field should be the VectorStoreRecordKey attributed field?

In my particular use case, I was planning on having a document broken into chunks and vectorized, stored in Cosmos. The partition key would be the document id /documentId, with the /id field corresponding to the unique id for each chunk.

Thanks again for the quick response and the help.

[westey-m]

HI @dccrain,
Actually, we do support partition keys, but we don't support hierarchical partition keys at this time. If hierarchical partition keys are really important, let us know, and we can add it to the backlog. Choosing a good partition key in Cosmos DB is really important, since cross partition queries can be more expensive, while too few partition keys can stop you from scaling. I think this is a good article to read with regards to this topic.

To use a partition key, you can specify it as part of the collection options when creating the collection. This is most easily done by creating the collection directly, instead of via the IVectorStore.GetCollection method. If you have to do it via the IVectorStore.GetCollection method, you will need to implement a factory to pass to the vector store on construction, using the mechanism described here.

To use a partition key, you also need to use the alternative key type that we support for Cosmos DB NoSQL: AzureCosmosDBNoSQLCompositeKey. Note that the Cosmos DB NoSQL collection type supports two key types: string or AzureCosmosDBNoSQLCompositeKey.

So first, to construct the collection, you can do the following to specify the partition key name and the AzureCosmosDBNoSQLCompositeKey type:

var collection = new AzureCosmosDBNoSQLVectorStoreRecordCollection<TestRecord>(
    database,
    request.Collection,
    new()
    {
        PartitionKeyPropertyName = nameof(TestRecord.DocumentId)
    }) as IVectorStoreRecordCollection<AzureCosmosDBNoSQLCompositeKey, TestRecord>;

GetAsync and DeleteAsync can then be used like the following, where you can specify both a partition key and an id via the AzureCosmosDBNoSQLCompositeKey key type:

var testRecord = await collection.GetAsync(new("test", "doc1"));

The experience doesn't change for searches or upsert.

Apologies again for the gap in docs around this. We'll get the docs improved to explain partition key usage properly.

[westey-m]

The documentation is now updated:
https://learn.microsoft.com/en-us/semantic-kernel/concepts/vector-store-connectors/out-of-the-box-connectors/azure-cosmosdb-nosql-connector?pivots=programming-language-csharp