diff --git a/dotnet/src/Connectors/Connectors.Memory.AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStore.cs b/dotnet/src/Connectors/Connectors.Memory.AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStore.cs index be8a82165e9e..219889d8e3e1 100644 --- a/dotnet/src/Connectors/Connectors.Memory.AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStore.cs +++ b/dotnet/src/Connectors/Connectors.Memory.AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStore.cs @@ -22,6 +22,7 @@ public class AzureCosmosDBMongoDBMemoryStore : IMemoryStore, IDisposable private readonly MongoClient _mongoClient; private readonly IMongoDatabase _mongoDatabase; private readonly AzureCosmosDBMongoDBConfig _config; + private readonly bool _ownsMongoClient; /// /// Initiates a AzureCosmosDBMongoDBMemoryStore instance using a Azure CosmosDB Mongo vCore connection string @@ -41,6 +42,7 @@ AzureCosmosDBMongoDBConfig config settings.ApplicationName = this._config.ApplicationName; this._mongoClient = new MongoClient(settings); this._mongoDatabase = this._mongoClient.GetDatabase(databaseName); + this._ownsMongoClient = true; } /// @@ -48,15 +50,13 @@ AzureCosmosDBMongoDBConfig config /// and other properties required for vector search. /// public AzureCosmosDBMongoDBMemoryStore( - IMongoClient mongoClient, + MongoClient mongoClient, string databaseName, AzureCosmosDBMongoDBConfig config ) { - MongoClientSettings settings = mongoClient.Settings; this._config = config; - settings.ApplicationName = this._config.ApplicationName; - this._mongoClient = new MongoClient(settings); + this._mongoClient = mongoClient; this._mongoDatabase = this._mongoClient.GetDatabase(databaseName); } @@ -318,7 +318,10 @@ protected virtual void Dispose(bool disposing) { if (disposing) { - this._mongoClient.Cluster.Dispose(); + if (this._ownsMongoClient) + { + this._mongoClient.Cluster.Dispose(); + } } } diff --git a/dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTests.cs b/dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTests.cs index 080c486ddcf9..cc0d1238b95a 100644 --- a/dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTests.cs +++ b/dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTests.cs @@ -49,7 +49,6 @@ public async Task ItCanBatchUpsertGetRemoveAsync(bool withEmbeddings) var memoryStore = this._fixture.MemoryStore; var records = DataHelper.CreateBatchRecords(Count); - await memoryStore.CreateCollectionAsync(collectionName); var keys = await memoryStore.UpsertBatchAsync(collectionName, records).ToListAsync(); var actualRecords = await memoryStore .GetBatchAsync(collectionName, keys, withEmbeddings: withEmbeddings) @@ -86,6 +85,12 @@ public async Task ItCanGetNearestMatchesAsync(int limit, bool withEmbeddings) var memoryStore = this._fixture.MemoryStore; var searchEmbedding = DataHelper.VectorSearchTestEmbedding; var nearestMatchesExpected = DataHelper.VectorSearchExpectedResults; + var records = DataHelper.VectorSearchTestRecords; + + var keys = await memoryStore.UpsertBatchAsync(collectionName, records).ToListAsync(); + var actualRecords = await memoryStore + .GetBatchAsync(collectionName, keys, withEmbeddings: withEmbeddings) + .ToListAsync(); var nearestMatchesActual = await memoryStore .GetNearestMatchesAsync( diff --git a/dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTestsFixture.cs b/dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTestsFixture.cs index 1074765559a8..1b1255c46b68 100644 --- a/dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTestsFixture.cs +++ b/dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTestsFixture.cs @@ -28,7 +28,6 @@ public AzureCosmosDBMongoDBMemoryStoreTestsFixture() ) .AddEnvironmentVariables() .Build(); - var connectionString = GetSetting(configuration, "ConnectionString"); this.DatabaseName = "DotNetSKTestDB"; this.CollectionName = "DotNetSKTestCollection"; @@ -42,7 +41,6 @@ public AzureCosmosDBMongoDBMemoryStoreTestsFixture() public async Task InitializeAsync() { await this.MemoryStore.CreateCollectionAsync(this.CollectionName); - await this .MemoryStore.UpsertBatchAsync(this.CollectionName, DataHelper.VectorSearchTestRecords) .ToListAsync();