diff --git a/Directory.Packages.props b/Directory.Packages.props index ffc08d0c3..2ea226435 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -32,6 +32,7 @@ + diff --git a/KernelMemory.sln b/KernelMemory.sln index 79a5db85a..698912d4f 100644 --- a/KernelMemory.sln +++ b/KernelMemory.sln @@ -227,6 +227,8 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "MongoDbAtlas", "extensions\ EndProject Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "MongoDbAtlas.FunctionalTests", "extensions\MongoDbAtlas\MongoDbAtlas.FunctionalTests\MongoDbAtlas.FunctionalTests.csproj", "{8A602227-B291-4F1B-ACB8-237F49501B6A}" EndProject +Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AzureCosmosDBMongoDB", "extensions\AzureCosmosDBMongoDB\AzureCosmosDBMongoDB\AzureCosmosDBMongoDB.csproj", "{8b62c632-9d70-4dc1-aeab-82d057a09a19}" +EndProject Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "107-dotnet-SemanticKernel-TextCompletion", "examples\107-dotnet-SemanticKernel-TextCompletion\107-dotnet-SemanticKernel-TextCompletion.csproj", "{494B8590-F0B2-4D40-A895-F9D7BDF26250}" EndProject Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "208-dotnet-lmstudio", "examples\208-dotnet-lmstudio\208-dotnet-lmstudio.csproj", "{BC8057DA-CB40-4308-96FB-EF0100822BAD}" @@ -523,6 +525,10 @@ Global {EE0D8645-2770-4E12-8E18-019B30970FE6}.Debug|Any CPU.ActiveCfg = Debug|Any CPU {EE0D8645-2770-4E12-8E18-019B30970FE6}.Debug|Any CPU.Build.0 = Debug|Any CPU {EE0D8645-2770-4E12-8E18-019B30970FE6}.Release|Any CPU.ActiveCfg = Release|Any CPU + {8b62c632-9d70-4dc1-aeab-82d057a09a19}.Debug|Any CPU.ActiveCfg = Debug|Any CPU + {8b62c632-9d70-4dc1-aeab-82d057a09a19}.Debug|Any CPU.Build.0 = Debug|Any CPU + {8b62c632-9d70-4dc1-aeab-82d057a09a19}.Release|Any CPU.ActiveCfg = Release|Any CPU + {8b62c632-9d70-4dc1-aeab-82d057a09a19}.Release|Any CPU.Build.0 = Release|Any CPU EndGlobalSection GlobalSection(SolutionProperties) = preSolution HideSolutionNode = FALSE @@ -607,8 +613,9 @@ Global {432AC1B4-8275-4284-9A44-44988A6F0C24} = {DBEA0A6B-474A-4E8C-BCC8-D5D43C063A54} {A0C81A29-715F-463E-A243-7E45DB8AE53F} = {155DA079-E267-49AF-973A-D1D44681970F} {EE0D8645-2770-4E12-8E18-019B30970FE6} = {0A43C65C-6007-4BB4-B3FE-8D439FC91841} + {8b62c632-9d70-4dc1-aeab-82d057a09a19} = {155DA079-E267-49AF-973A-D1D44681970F} EndGlobalSection GlobalSection(ExtensibilityGlobals) = postSolution SolutionGuid = {CC136C62-115C-41D1-B414-F9473EFF6EA8} EndGlobalSection -EndGlobal +EndGlobal \ No newline at end of file diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/AzureCosmosDBMongoDB.TestApplication.csproj b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/AzureCosmosDBMongoDB.TestApplication.csproj new file mode 100644 index 000000000..864ad3441 --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/AzureCosmosDBMongoDB.TestApplication.csproj @@ -0,0 +1,24 @@ + + + + Exe + net8.0 + LatestMajor + enable + enable + false + false + 5ee045b0-aea3-4f08-8d31-32d1a6f8fed0 + $(NoWarn);CA1050,CA2007,CA1826,CA1303,CA1307,SKEXP0001 + + + + + + + + + + + + diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/MockEmbeddingGenerator.cs b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/MockEmbeddingGenerator.cs new file mode 100644 index 000000000..ea486b6e6 --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/MockEmbeddingGenerator.cs @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.KernelMemory; +using Microsoft.KernelMemory.AI; + +namespace AzureCosmosDBMongoDB.TestApplication; + +internal sealed class MockEmbeddingGenerator : ITextEmbeddingGenerator +{ + private readonly Dictionary _embeddings = new(); + + internal void AddFakeEmbedding(string str, Embedding vector) + { + this._embeddings.Add(str, vector); + } + + /// + public int CountTokens(string text) => 0; + + /// + public int MaxTokens => 0; + + /// + public Task GenerateEmbeddingAsync(string text, CancellationToken cancellationToken = default) => + Task.FromResult(this._embeddings[text]); +} diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/Program.cs b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/Program.cs new file mode 100644 index 000000000..b652011bc --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/Program.cs @@ -0,0 +1,147 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using Microsoft.KernelMemory; +using Microsoft.KernelMemory.AI; +using Microsoft.KernelMemory.AI.OpenAI; +using Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB; +using Microsoft.KernelMemory.MemoryStorage; +using MongoDB.Driver; + +namespace AzureCosmosDBMongoDB.TestApplication; + +public static class Program +{ + private const string index = "default_index"; + + private const string Text1 = "this is test 1"; + private const string Text2 = "this is test 2"; + + public static async Task Main() + { + var (memory, embeddings) = await SetupAsync(); + + Console.WriteLine("++++ DELETE INDEX ++++"); + + // await memory.DeleteIndexAsync(index); + + Console.WriteLine("++++ CREATE INDEX ++++"); + + await memory.CreateIndexAsync(index, embeddings[0].Length); + + Console.WriteLine("++++ LIST INDEXES ++++"); + + IEnumerable indexes = await memory.GetIndexesAsync(); + foreach (var indexName in indexes) + { + Console.WriteLine(indexName); + } + + Console.WriteLine("===== INSERT RECORD 1 AND 2 ====="); + var memoryRecord1 = new MemoryRecord + { + Id = "memory 1", + Vector = embeddings[0], + Tags = new TagCollection { { "updated", "no" }, { "type", "email" } }, + Payload = new Dictionary() + }; + + var memoryRecord2 = new MemoryRecord + { + Id = "memory 2", + Vector = embeddings[0], + Tags = new TagCollection { { "updated", "no" }, { "type", "email" } }, + Payload = new Dictionary() + }; + + var id1 = await memory.UpsertAsync(index, memoryRecord1); + Console.WriteLine($"Insert 1: {id1} {memoryRecord1.Id}"); + + var id2 = await memory.UpsertAsync(index, memoryRecord2); + Console.WriteLine($"Insert 2: {id2} {memoryRecord2.Id}"); + + + Console.WriteLine("===== INSERT RECORD 3 ====="); + + var memoryRecord3 = new MemoryRecord + { + Id = "memory 3", + Vector = embeddings[1], + Tags = new TagCollection { { "type", "news" } }, + Payload = new Dictionary() + }; + + var id3 = await memory.UpsertAsync(index, memoryRecord3); + Console.WriteLine($"Insert 3: {id3} {memoryRecord3.Id}"); + + Console.WriteLine("===== UPDATE RECORD 3 ====="); + + memoryRecord3.Tags.Add("updated", "yes"); + id3 = await memory.UpsertAsync(index, memoryRecord3); + Console.WriteLine($"Update 3: {id3} {memoryRecord3.Id}"); + + Console.WriteLine("===== SEARCH 1 ====="); + + var similarList = memory.GetSimilarListAsync( + index, text: Text1, limit: 10, withEmbeddings: true, minRelevance: 0.7); + await foreach ((MemoryRecord, double) record in similarList) + { + Console.WriteLine(record.Item1.Id); + Console.WriteLine(" score: " + record.Item2); + Console.WriteLine(" tags: " + record.Item1.Tags.Count); + Console.WriteLine(" size: " + record.Item1.Vector.Length); + } + + Console.WriteLine("===== DELETE ====="); + + await memory.DeleteAsync("test", new MemoryRecord { Id = "memory 1" }); + await memory.DeleteAsync("test", new MemoryRecord { Id = "memory 2" }); + await memory.DeleteAsync("test", new MemoryRecord { Id = "memory 3" }); + + Console.WriteLine("== Done =="); + + } + + private static async Task<(AzureCosmosDBMongoDBMemory, Embedding[])> SetupAsync() + { + IConfiguration cfg = new ConfigurationBuilder() + .AddJsonFile("appsettings.json") + .AddJsonFile("appsettings.Development.json", optional: true) + .Build(); + + var config = cfg.GetSection("KernelMemory:Services:AzureCosmosDBMongoDB").Get() + ?? throw new ArgumentNullException(message: "AzureAISearch config not found", null); + var openAIConfig = cfg.GetSection("KernelMemory:Service:OpenAI").Get(); + var useRealEmbeddingGenerator = cfg.GetValue("UseRealEmbeddingGenerator"); + ITextEmbeddingGenerator embeddingGenerator; + + if (useRealEmbeddingGenerator) + { + embeddingGenerator = new OpenAITextEmbeddingGenerator(openAIConfig, log: null); + } + else + { + embeddingGenerator = new MockEmbeddingGenerator(); + } + + var memory = new AzureCosmosDBMongoDBMemory(config, embeddingGenerator); + + Embedding embedding1 = new[] { 0f, 0, 1, 0, 1 }; + Embedding embedding2 = new[] { 0, 0, 0.95f, 0.01f, 0.95f }; + if (useRealEmbeddingGenerator) + { + embedding1 = await embeddingGenerator.GenerateEmbeddingAsync(Text1); + embedding2 = await embeddingGenerator.GenerateEmbeddingAsync(Text2); + } + else + { + ((MockEmbeddingGenerator)embeddingGenerator).AddFakeEmbedding(Text1, embedding1); + ((MockEmbeddingGenerator)embeddingGenerator).AddFakeEmbedding(Text2, embedding2); + } + + return (memory, new[] { embedding1, embedding2 }); + + } + + +} diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/appsettings.json b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/appsettings.json new file mode 100644 index 000000000..29aab0aa7 --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.TestApplication/appsettings.json @@ -0,0 +1,70 @@ +{ + "Logging": { + "LogLevel": { + "Default": "Trace", + "Microsoft.AspNetCore": "Trace" + } + }, + "UseRealEmbeddingGenerator": false, + "KernelMemory": { + "Services": { + "AzureCosmosDBMongoDB": { + // Please refer here https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/vector-search for details about these configurations + "ConnectionString": "", + "DatabaseName": "cosmos_test_db", + "ContainerName": "cosmos_test_collection", + "ApplicationName": "DotNet_Kernel_Memory", + "IndexName": "default_index", + "NumLists": 1, + "NumberOfConnections": 16, + "EfConstruction": 64, + "EfSearch": 40 + }, + "AzureOpenAIText": { + // "ApiKey" or "AzureIdentity" + // AzureIdentity: use automatic AAD authentication mechanism. You can test locally + // using the env vars AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET. + "Auth": "AzureIdentity", + "Endpoint": "https://<...>.openai.azure.com/", + "APIKey": "", + "Deployment": "", + // The max number of tokens supported by model deployed + // See https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models + "MaxTokenTotal": 8191, + // "ChatCompletion" or "TextCompletion" + "APIType": "ChatCompletion", + "MaxRetries": 3 + }, + "AzureOpenAIEmbedding": { + // "ApiKey" or "AzureIdentity" + // AzureIdentity: use automatic AAD authentication mechanism. You can test locally + // using the env vars AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET. + "Auth": "AzureIdentity", + "Endpoint": "https://<...>.openai.azure.com/", + "APIKey": "", + "Deployment": "", + // The max number of tokens supported by model deployed + // See https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models + "MaxTokenTotal": 8191, + "MaxRetries": 3 + }, + "OpenAI": { + // Name of the model used to generate text (text completion or chat completion) + "TextModel": "gpt-3.5-turbo-16k", + // The max number of tokens supported by the text model. + "TextModelMaxTokenTotal": 16384, + // Name of the model used to generate text embeddings + "EmbeddingModel": "text-embedding-ada-002", + // The max number of tokens supported by the embedding model + // See https://platform.openai.com/docs/guides/embeddings/what-are-embeddings + "EmbeddingModelMaxTokenTotal": 8191, + // OpenAI API Key + "APIKey": "", + // OpenAI Organization ID (usually empty, unless you have multiple accounts on different orgs) + "OrgId": "", + // How many times to retry in case of throttling + "MaxRetries": 3 + } + } + } +} \ No newline at end of file diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.csproj b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.csproj new file mode 100644 index 000000000..6d2c1bdff --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB.csproj @@ -0,0 +1,37 @@ + + + + net8.0 + LatestMajor + Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB + Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB + $(NoWarn);CA1724;CS1591;CA1308; + + + + + + + + + + + + + + + + + + false + Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB + Azure Cosmos DB for MongoDB connector for Kernel Memory + Azure Cosmos DB for MongoDB connector for Microsoft Kernel Memory, to store and search memory using Azure AI Search vector indexing and semantic features. + Memory, RAG, Kernel Memory, Azure Cosmos DB for MongoDB, HNSW, AI, Artificial Intelligence, Embeddings, Vector DB, Vector Search, ETL + + + + + + + \ No newline at end of file diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBConfig.cs b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBConfig.cs new file mode 100644 index 000000000..e22a479b4 --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBConfig.cs @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB; + +#pragma warning disable IDE0130 // reduce number of "using" statements +// ReSharper disable once CheckNamespace - reduce number of "using" statements +namespace Microsoft.KernelMemory; + +/// +/// Get more details about Azure Cosmos DB for MongoDB and these configs +/// at https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/vector-search +/// +public class AzureCosmosDBMongoDBConfig +{ + /// + /// Connection string required to connect to Azure Cosmos DB for MongoDB + /// see https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/quickstart-portal + /// + public string ConnectionString { get; set; } = string.Empty; + + /// + /// Database name for the Mongo vCore DB + /// + public string DatabaseName { get; set; } = "default_KM_DB"; + + /// + /// Container name for the Mongo vCore DB + /// + public string CollectionName { get; set; } = "default_KM_Collection"; + + /// + /// Application name for the client for tracking and logging + /// + public string ApplicationName { get; set; } = "dotNet_Kernel_Memory"; + + /// + /// Index name for the Mongo vCore DB + /// Index name for the MongoDB + /// + public string IndexName { get; set; } = "default_index"; + + /// + /// Kind: Type of vector index to create. + /// Possible options are: + /// - vector-ivf + /// - vector-hnsw: available as a preview feature only, + /// to enable visit https://learn.microsoft.com/azure/azure-resource-manager/management/preview-features + /// + public AzureCosmosDBVectorSearchType Kind { get; set; } = AzureCosmosDBVectorSearchType.VectorHNSW; + + /// + /// NumLists: This integer is the number of clusters that the inverted file (IVF) index uses to group the vector data. + /// We recommend that numLists is set to documentCount/1000 for up to 1 million documents and to sqrt(documentCount) + /// for more than 1 million documents. Using a numLists value of 1 is akin to performing brute-force search, which has + /// limited performance. + /// + public int NumLists { get; set; } = 1; + + /// + /// Similarity: Similarity metric to use with the IVF index. + /// Possible options are: + /// - COS (cosine distance), + /// - L2 (Euclidean distance), and + /// - IP (inner product). + /// + public AzureCosmosDBSimilarityTypes Similarity { get; set; } = AzureCosmosDBSimilarityTypes.Cosine; + + /// + /// NumberOfConnections: The max number of connections per layer (16 by default, minimum value is 2, maximum value is + /// 100). Higher m is suitable for datasets with high dimensionality and/or high accuracy requirements. + /// + public int NumberOfConnections { get; set; } = 16; + + /// + /// EfConstruction: the size of the dynamic candidate list for constructing the graph (64 by default, minimum value is 4, + /// maximum value is 1000). Higher ef_construction will result in better index quality and higher accuracy, but it will + /// also increase the time required to build the index. EfConstruction has to be at least 2 * m + /// + public int EfConstruction { get; set; } = 64; + + /// + /// EfSearch: The size of the dynamic candidate list for search (40 by default). A higher value provides better recall at + /// the cost of speed. + /// + public int EfSearch { get; set; } = 40; +} diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemory.cs b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemory.cs new file mode 100644 index 000000000..2b92df5c7 --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemory.cs @@ -0,0 +1,401 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.Logging; +using Microsoft.KernelMemory.AI; +using Microsoft.KernelMemory.Diagnostics; +using Microsoft.KernelMemory.MemoryStorage; +using MongoDB.Bson; +using MongoDB.Bson.Serialization; +using MongoDB.Driver; + +namespace Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB; + +/// +/// Azure Cosmos DB for MongoDB connector for Kernel Memory, +/// Read more about Azure Cosmos DB for MongoDB and vector search here: +/// https://learn.microsoft.com/en-us/azure/cosmos-db/mongodb/vcore/ +/// https://learn.microsoft.com/en-us/azure/cosmos-db/mongodb/vcore/vector-search +/// +public class AzureCosmosDBMongoDBMemory : IMemoryDb +{ + private readonly ITextEmbeddingGenerator _embeddingGenerator; + private readonly ILogger _log; + private readonly AzureCosmosDBMongoDBConfig _config; + private readonly MongoClient _cosmosDBMongoClient; + private readonly IMongoDatabase _cosmosMongoDatabase; + private readonly IMongoCollection _cosmosMongoCollection; + + /// + /// Create a new instance + /// + /// Azure Cosmos DB for MongoDB configuration + /// Text embedding generator + /// Application logger + public AzureCosmosDBMongoDBMemory( + AzureCosmosDBMongoDBConfig config, + ITextEmbeddingGenerator embeddingGenerator, + ILogger? log = null) + { + this._embeddingGenerator = embeddingGenerator; + this._log = log ?? DefaultLogger.Instance; + this._config = config; + + if (string.IsNullOrEmpty(config.ConnectionString)) + { + this._log.LogCritical("Azure Cosmos DB for MongoDB connection string is empty."); + throw new AzureCosmosDBMongoDBMemoryException("Azure Cosmos DB for MongoDB connection string is empty."); + } + + if (this._embeddingGenerator == null) + { + throw new AzureCosmosDBMongoDBMemoryException("Embedding generator not configured"); + } + + MongoClientSettings settings = MongoClientSettings.FromConnectionString(config.ConnectionString); + settings.ApplicationName = config.ApplicationName; + this._cosmosDBMongoClient = new MongoClient(settings); + this._cosmosMongoDatabase = this._cosmosDBMongoClient.GetDatabase(config.DatabaseName); + + this._cosmosMongoDatabase.CreateCollectionAsync(config.CollectionName); + this._cosmosMongoCollection = this._cosmosMongoDatabase.GetCollection(config.CollectionName); + } + + /// + public async Task CreateIndexAsync(string index, int vectorSize, CancellationToken cancellationToken = default) + { + var indexes = await this._cosmosMongoCollection.Indexes.ListAsync(cancellationToken).ConfigureAwait(false); + if (!indexes.ToList(cancellationToken).Any(index => index["name"] == index)) + { + var command = new BsonDocument(); + switch (this._config.Kind) + { + case AzureCosmosDBVectorSearchType.VectorIVF: + command = this.GetIndexDefinitionVectorIVF(index, vectorSize); + break; + case AzureCosmosDBVectorSearchType.VectorHNSW: + command = this.GetIndexDefinitionVectorHNSW(index, vectorSize); + break; + } + + await this._cosmosMongoDatabase.RunCommandAsync(command, cancellationToken: cancellationToken).ConfigureAwait(false); + } + } + + /// + public async Task> GetIndexesAsync(CancellationToken cancellationToken = default) + { + var indexesCursor = await this._cosmosMongoCollection.Indexes.ListAsync(cancellationToken).ConfigureAwait(false); + var indexes = await indexesCursor.ToListAsync(cancellationToken).ConfigureAwait(false); + + return indexes.Select(index => index["name"].AsString); + } + + /// + public async Task DeleteIndexAsync(string index, CancellationToken cancellationToken = default) + { + var indexesCursor = await this._cosmosMongoCollection.Indexes.ListAsync(cancellationToken).ConfigureAwait(false); + var indexes = await indexesCursor.ToListAsync(cancellationToken).ConfigureAwait(false); + foreach (var indexDoc in indexes) + { + if (indexDoc["name"].AsString == index) + { + await this._cosmosMongoCollection.Indexes.DropOneAsync(index, cancellationToken).ConfigureAwait(false); + } + } + } + + /// + public async Task UpsertAsync(string index, MemoryRecord record, CancellationToken cancellationToken = default) + { + AzureCosmosDBMongoDBMemoryRecord localRecord = AzureCosmosDBMongoDBMemoryRecord.FromMemoryRecord(record); + var filter = Builders.Filter.Eq(r => r.Id, localRecord.Id); + + var replaceOptions = new ReplaceOptions() { IsUpsert = true }; + + await this._cosmosMongoCollection.ReplaceOneAsync(filter, localRecord, replaceOptions, cancellationToken).ConfigureAwait(false); + + return record.Id; + } + + /// + public async IAsyncEnumerable<(MemoryRecord, double)> GetSimilarListAsync( + string index, + string text, + ICollection? filters = null, + double minRelevance = 0, + int limit = 1, + bool withEmbeddings = false, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + if (limit <= 0) { limit = int.MaxValue; } + + Embedding embedding = await this._embeddingGenerator.GenerateEmbeddingAsync(text, cancellationToken).ConfigureAwait(false); + + BsonDocument[]? pipeline = null; + switch (this._config.Kind) + { + case AzureCosmosDBVectorSearchType.VectorIVF: + pipeline = this.GetVectorIVFSearchPipeline(embedding, limit); + break; + case AzureCosmosDBVectorSearchType.VectorHNSW: + pipeline = this.GetVectorHNSWSearchPipeline(embedding, limit); + break; + } + + using var cursor = await this._cosmosMongoCollection.AggregateAsync(pipeline, cancellationToken: cancellationToken).ConfigureAwait(false); + while (await cursor.MoveNextAsync(cancellationToken).ConfigureAwait(false)) + { + foreach (var doc in cursor.Current) + { + // Access the similarityScore from the BSON document + var similarityScore = doc.GetValue("similarityScore").AsDouble; + if (similarityScore < minRelevance) { continue; } + + MemoryRecord memoryRecord = AzureCosmosDBMongoDBMemoryRecord.ToMemoryRecord(doc, withEmbeddings); + + yield return (memoryRecord, similarityScore); + } + } + } + + /// + public async IAsyncEnumerable GetListAsync( + string index, + ICollection? filters = null, + int limit = 1, + bool withEmbeddings = false, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + var finalFilter = this.TranslateFilters(filters, index); + + // We need to perform a simple query without using vector search + var cursor = await this._cosmosMongoCollection + .FindAsync(finalFilter, + new FindOptions() + { + Limit = limit + }, + cancellationToken: cancellationToken).ConfigureAwait(false); + var documents = await cursor.ToListAsync(cancellationToken).ConfigureAwait(false); + + foreach (var document in documents) + { + var memoryRecord = FromAzureCosmosMemoryRecord(document, withEmbeddings); + yield return memoryRecord; + } + } + + private static MemoryRecord FromAzureCosmosMemoryRecord(AzureCosmosDBMongoDBMemoryRecord doc, bool withEmbeddings) + { + var record = new MemoryRecord + { + Id = doc.Id, + Payload = BsonSerializer.Deserialize>(doc.Payload) + ?? new Dictionary(), + Vector = withEmbeddings ? doc.Embedding ?? Array.Empty() : Array.Empty() + }; + + foreach (string[] keyValue in doc.Tags.Select(tag => tag.Split(Constants.ReservedEqualsChar, 2))) + { + string key = keyValue[0]; + string? value = keyValue.Length == 1 ? null : keyValue[1]; + record.Tags.Add(key, value); + } + + return record; + } + + /// + public async Task DeleteAsync( + string index, + MemoryRecord record, + CancellationToken cancellationToken = default) + { + AzureCosmosDBMongoDBMemoryRecord localRecord = AzureCosmosDBMongoDBMemoryRecord.FromMemoryRecord(record); + var filter = Builders.Filter.Eq(r => r.Id, localRecord.Id); + await this._cosmosMongoCollection.DeleteOneAsync(filter, cancellationToken).ConfigureAwait(false); + } + + private FilterDefinition? TranslateFilters(ICollection? filters, string index) + { + List> outerFiltersArray = new(); + foreach (var filter in filters ?? Array.Empty()) + { + var thisFilter = filter.GetFilters().ToArray(); + List> filtersArray = new(); + foreach (var singleFilter in thisFilter) + { + var condition = Builders.Filter.And( + Builders.Filter.Eq("Tags.Key", singleFilter.Key), + Builders.Filter.Eq("Tags.Values", singleFilter.Value) + ); + filtersArray.Add(condition); + } + + // if we have more than one condition, we need to compose all conditions with AND + // but if we have only a single filter we can directly use the filter. + if (filtersArray.Count > 1) + { + // More than one condition we need to create the condition + var andFilter = Builders.Filter.And(filtersArray); + outerFiltersArray.Add(andFilter); + } + else if (filtersArray.Count == 1) + { + // We do not need to include an and filter because we have only one condition + outerFiltersArray.Add(filtersArray[0]); + } + } + + FilterDefinition? finalFilter = null; + // Outer filters must be composed in or + if (outerFiltersArray.Count > 1) + { + finalFilter = Builders.Filter.Or(outerFiltersArray); + } + else if (outerFiltersArray.Count == 1) + { + // We do not need to include an or filter because we have only one condition + finalFilter = outerFiltersArray[0]; + } + + var indexFilter = Builders.Filter.Eq("Index", index); + if (finalFilter == null) + { + finalFilter = indexFilter; + } + else + { + // Compose in and + finalFilter = Builders.Filter.And(indexFilter, finalFilter); + } + return finalFilter; + } + + private BsonDocument GetIndexDefinitionVectorIVF(string index, int vectorSize) + { + return new BsonDocument + { + { "createIndexes", this._config.CollectionName }, + { + "indexes", + new BsonArray + { + new BsonDocument + { + { "name", index }, + { "key", new BsonDocument { { "embedding", "cosmosSearch" } } }, + { + "cosmosSearchOptions", new BsonDocument + { + { "kind", this._config.Kind.GetCustomName() }, + { "numLists", this._config.NumLists }, + { "similarity", this._config.Similarity.GetCustomName() }, + { "dimensions", vectorSize } + } + } + } + } + } + }; + } + + private BsonDocument GetIndexDefinitionVectorHNSW(string index, int vectorSize) + { + return new BsonDocument + { + { "createIndexes", this._config.CollectionName }, + { + "indexes", + new BsonArray + { + new BsonDocument + { + { "name", index }, + { "key", new BsonDocument { { "embedding", "cosmosSearch" } } }, + { + "cosmosSearchOptions", new BsonDocument + { + { "kind", this._config.Kind.GetCustomName() }, + { "m", this._config.NumberOfConnections }, + { "efConstruction", this._config.EfConstruction }, + { "similarity", this._config.Similarity.GetCustomName() }, + { "dimensions", vectorSize } + } + } + } + } + } + }; + } + + private BsonDocument[] GetVectorIVFSearchPipeline(Embedding embedding, int limit) + { +#pragma warning disable CA1305 // Specify IFormatProvider + string searchStage = @" + { + ""$search"": { + ""cosmosSearch"": { + ""vector"": [" + string.Join(",", embedding.Data.ToArray().Select(f => f.ToString())) + @"], + ""path"": ""embedding"", + ""k"": " + limit + @" + }, + ""returnStoredSource"": true + } + }"; +#pragma warning restore CA1305 // Specify IFormatProvider + + string projectStage = @" + { + ""$project"": { + ""similarityScore"": { ""$meta"": ""searchScore"" }, + ""document"": ""$$ROOT"" + } + }"; + + BsonDocument searchBson = BsonDocument.Parse(searchStage); + BsonDocument projectBson = BsonDocument.Parse(projectStage); + return new BsonDocument[] + { + searchBson, projectBson + }; + } + + private BsonDocument[] GetVectorHNSWSearchPipeline(Embedding embedding, int limit) + { +#pragma warning disable CA1305 // Specify IFormatProvider + string searchStage = @" + { + ""$search"": { + ""cosmosSearch"": { + ""vector"": [" + string.Join(",", embedding.Data.ToArray().Select(f => f.ToString())) + @"], + ""path"": ""embedding"", + ""k"": " + limit + @", + ""efSearch"": " + this._config.EfSearch + @" + } + } + }"; +#pragma warning restore CA1305 // Specify IFormatProvider + + string projectStage = @" + { + ""$project"": { + ""similarityScore"": { ""$meta"": ""searchScore"" }, + ""document"": ""$$ROOT"" + } + }"; + + BsonDocument searchBson = BsonDocument.Parse(searchStage); + BsonDocument projectBson = BsonDocument.Parse(projectStage); + return new BsonDocument[] + { + searchBson, projectBson + }; + } +} diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryException.cs b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryException.cs new file mode 100644 index 000000000..178148231 --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryException.cs @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; + +namespace Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB; + +public class AzureCosmosDBMongoDBMemoryException : KernelMemoryException +{ + /// + public AzureCosmosDBMongoDBMemoryException() + { + } + + /// + public AzureCosmosDBMongoDBMemoryException(string? message) : base(message) + { + } + + /// + public AzureCosmosDBMongoDBMemoryException(string? message, Exception? innerException) : base(message, innerException) + { + } +} diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryRecord.cs b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryRecord.cs new file mode 100644 index 000000000..eb44c8e58 --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryRecord.cs @@ -0,0 +1,108 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text; +using System.Text.Json; +using Microsoft.KernelMemory.MemoryStorage; +using MongoDB.Bson; +using MongoDB.Bson.Serialization; +using MongoDB.Bson.Serialization.Attributes; + +namespace Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB; + +public sealed class AzureCosmosDBMongoDBMemoryRecord +{ + [BsonId] + public string Id { get; set; } = null!; + + [BsonElement("embedding")] +#pragma warning disable CA1819 // Properties should not return arrays + public float[]? Embedding { get; set; } = null!; +#pragma warning restore CA1819 // Properties should not return arrays + + [BsonElement("tags")] + public List Tags { get; set; } = new(); + + [BsonElement("payload")] + public string Payload { get; set; } = string.Empty; + + [BsonElement("timestamp")] + [BsonDateTimeOptions(Kind = DateTimeKind.Utc, Representation = BsonType.DateTime)] + public DateTime? Timestamp { get; set; } + + [BsonElement("similarityScore")] + [BsonIgnoreIfDefault] + public double SimilarityScore { get; set; } + + public static MemoryRecord ToMemoryRecord(BsonDocument document, bool withEmbedding = true) + { + var doc = document["document"].AsBsonDocument; + MemoryRecord result = new() + { + Id = DecodeId(doc["_id"].AsString), + Payload = BsonSerializer.Deserialize>(doc["payload"].AsString) + ?? new Dictionary() + }; + + var timestamp = doc["timestamp"]; + if (timestamp != null) + { + result.Payload.Add("timeStamp", timestamp.ToUniversalTime()); + } + + if (withEmbedding) + { + result.Vector = doc["embedding"].AsBsonArray.Select(x => (float)x.AsDouble).ToArray(); + } + + foreach (string[] keyValue in doc["tags"].AsBsonArray.Select(tag => tag.AsString.Split(Constants.ReservedEqualsChar, 2))) + { + string key = keyValue[0]; + string? value = keyValue.Length == 1 ? null : keyValue[1]; + result.Tags.Add(key, value); + } + + return result; + } + + public static AzureCosmosDBMongoDBMemoryRecord FromMemoryRecord(MemoryRecord record) + { + AzureCosmosDBMongoDBMemoryRecord result = new() + { + Id = EncodeId(record.Id), + Embedding = record.Vector.Data.ToArray(), + Payload = JsonSerializer.Serialize(record.Payload, s_jsonOptions), + Timestamp = DateTime.UtcNow + }; + + foreach (var tag in record.Tags.Pairs) + { + result.Tags.Add($"{tag.Key}{Constants.ReservedEqualsChar}{tag.Value}"); + } + + return result; + } + + private static string EncodeId(string realId) + { + var bytes = Encoding.UTF8.GetBytes(realId); + return Convert.ToBase64String(bytes).Replace('=', '_'); + } + + private static string DecodeId(string encodedId) + { + var bytes = Convert.FromBase64String(encodedId.Replace('_', '=')); + return Encoding.UTF8.GetString(bytes); + } + + private static readonly JsonSerializerOptions s_jsonOptions = new() + { + AllowTrailingCommas = true, + MaxDepth = 10, + PropertyNameCaseInsensitive = true, + ReadCommentHandling = JsonCommentHandling.Disallow, + WriteIndented = false + }; +} diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBSimilarityTypes.cs b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBSimilarityTypes.cs new file mode 100644 index 000000000..d8c4ab422 --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBSimilarityTypes.cs @@ -0,0 +1,50 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Reflection; +using MongoDB.Bson; +using MongoDB.Bson.Serialization.Attributes; + +// ReSharper disable InconsistentNaming +namespace Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB; + +/// +/// Similarity metric to use with the index. Possible options are COS (cosine distance), L2 (Euclidean distance), and IP (inner product). +/// +public enum AzureCosmosDBSimilarityTypes +{ + /// + /// Cosine similarity + /// + [BsonElement("COS")] + Cosine, + + /// + /// Inner Product similarity + /// + [BsonElement("IP")] + InnerProduct, + + /// + /// Euclidean similarity + /// + [BsonElement("L2")] + Euclidean +} + +internal static class AzureCosmosDBSimilarityTypesExtensions +{ + public static string GetCustomName(this AzureCosmosDBSimilarityTypes type) + { + // Retrieve the FieldInfo object for the enum value, and check if it is null before accessing it. + var fieldInfo = type.GetType().GetField(type.ToString()); + if (fieldInfo == null) + { + // Optionally handle the situation when the field is not found, such as logging a warning or throwing an exception. + return type.ToString(); // Or handle differently as needed. + } + + // Retrieve the BsonElementAttribute from the field, if present. + var attribute = fieldInfo.GetCustomAttribute(); + return attribute?.ElementName ?? type.ToString(); + } +} diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBVectorSearchType.cs b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBVectorSearchType.cs new file mode 100644 index 000000000..98207121f --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/AzureCosmosDBVectorSearchType.cs @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Reflection; +using MongoDB.Bson.Serialization.Attributes; + +// ReSharper disable InconsistentNaming +namespace Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB; + +/// +/// Type of vector index to create. The options are vector-ivf and vector-hnsw. +/// +public enum AzureCosmosDBVectorSearchType +{ + /// + /// vector-ivf is available on all cluster tiers + /// + [BsonElement("vector-ivf")] + VectorIVF, + + /// + /// vector-hnsw is available on M40 cluster tiers and higher. + /// + [BsonElement("vector-hnsw")] + VectorHNSW +} + +internal static class AzureCosmosDBVectorSearchTypeExtensions +{ + public static string GetCustomName(this AzureCosmosDBVectorSearchType type) + { + // Retrieve the FieldInfo object for the enum value, and check if it is null before accessing it. + var fieldInfo = type.GetType().GetField(type.ToString()); + if (fieldInfo == null) + { + // Optionally handle the situation when the field is not found, such as logging a warning or throwing an exception. + return type.ToString(); // Or handle differently as needed. + } + + // Retrieve the BsonElementAttribute from the field, if present. + var attribute = fieldInfo.GetCustomAttribute(); + return attribute?.ElementName ?? type.ToString(); + } +} diff --git a/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/DependencyInjection.cs b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/DependencyInjection.cs new file mode 100644 index 000000000..9aa14a4d0 --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/AzureCosmosDBMongoDB/DependencyInjection.cs @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Extensions.DependencyInjection; +using Microsoft.KernelMemory.MemoryDb.AzureCosmosDBMongoDB; +using Microsoft.KernelMemory.MemoryStorage; + +#pragma warning disable IDE0130 // reduce number of "using" statements +// ReSharper disable once CheckNamespace - reduce number of "using" statements +namespace Microsoft.KernelMemory; + +public static partial class KernelMemoryBuilderExtensions +{ + public static IKernelMemoryBuilder WithAzureCosmosDBMongoDBMemoryDb(this IKernelMemoryBuilder builder, AzureCosmosDBMongoDBConfig config) + { + builder.Services.AddAzureCosmosDBMongoDBMemoryDb(config); + return builder; + } +} + +public static partial class DependencyInjection +{ + public static IServiceCollection AddAzureCosmosDBMongoDBMemoryDb(this IServiceCollection services, AzureCosmosDBMongoDBConfig config) + { + return services + .AddSingleton(config) + .AddSingleton(); + } +} diff --git a/extensions/AzureCosmosDBMongoDB/README.md b/extensions/AzureCosmosDBMongoDB/README.md new file mode 100644 index 000000000..e910c5ead --- /dev/null +++ b/extensions/AzureCosmosDBMongoDB/README.md @@ -0,0 +1,11 @@ +# Developing with Azure Cosmos DB for MongoDB + +Integrated Vector Database in Azure Cosmos DB for MongoDB can be used to seamlessly connect your AI-based applications +with your data that's stored in Azure Cosmos DB. This integration can include apps that you built by using Azure OpenAI +embeddings. The natively integrated vector database enables you to efficiently store, index, and query high-dimensional +vector data that's stored directly in Azure Cosmos DB for MongoDB vCore, along with the original data from which the +vector data is created. It eliminates the need to transfer your data to alternative vector stores and incur additional +costs. + +Please refer to this link for setting up your Azure Cosmos DB for MongoDB cluster, and more information on vector +search using [Azure CosmosDB MongoDB vCore](https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/vector-search). diff --git a/service/Service/appsettings.json b/service/Service/appsettings.json index 7d1ea8acd..ee58e3efe 100644 --- a/service/Service/appsettings.json +++ b/service/Service/appsettings.json @@ -137,8 +137,7 @@ // Multiple generators can be used, e.g. for data migration, A/B testing, etc. // None of these are used for `ITextEmbeddingGeneration` dependency injection, // see Retrieval settings. - "EmbeddingGeneratorTypes": [ - ], + "EmbeddingGeneratorTypes": [], // Vectors can be written to multiple storages, e.g. for data migration, A/B testing, etc. // "AzureAISearch", "Qdrant", "Postgres", "Redis", "SimpleVectorDb", "SqlServer", etc. "MemoryDbTypes": [ @@ -269,6 +268,16 @@ // Setting used only for country clouds "EndpointSuffix": "core.windows.net" }, + "AzureCosmosDBMongoDB": { + // Please refer here https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/vector-search for details about these configurations + "ConnectionString": "", + "ApplicationName": "DotNet_Kernel_Memory", + "IndexName": "default_index", + "NumLists": 1, + "NumberOfConnections": 16, + "EfConstruction": 64, + "EfSearch": 40 + }, "AzureOpenAIEmbedding": { // "ApiKey" or "AzureIdentity" // AzureIdentity: use automatic AAD authentication mechanism. You can test locally diff --git a/service/tests/TestHelpers/BaseFunctionalTestCase.cs b/service/tests/TestHelpers/BaseFunctionalTestCase.cs index c85852a1b..f81613c0f 100644 --- a/service/tests/TestHelpers/BaseFunctionalTestCase.cs +++ b/service/tests/TestHelpers/BaseFunctionalTestCase.cs @@ -28,6 +28,7 @@ public abstract class BaseFunctionalTestCase : IDisposable protected readonly MongoDbAtlasConfig MongoDbAtlasConfig; protected readonly SimpleVectorDbConfig SimpleVectorDbConfig; protected readonly LlamaSharpConfig LlamaSharpConfig; + protected readonly AzureCosmosDBMongoDBConfig AzureCosmosDBMongoDBConfig; protected readonly ElasticsearchConfig ElasticsearchConfig; // IMPORTANT: install Xunit.DependencyInjection package @@ -47,6 +48,7 @@ protected BaseFunctionalTestCase(IConfiguration cfg, ITestOutputHelper output) this.MongoDbAtlasConfig = cfg.GetSection("KernelMemory:Services:MongoDbAtlas").Get() ?? new(); this.SimpleVectorDbConfig = cfg.GetSection("KernelMemory:Services:SimpleVectorDb").Get() ?? new(); this.LlamaSharpConfig = cfg.GetSection("KernelMemory:Services:LlamaSharp").Get() ?? new(); + this.AzureCosmosDBMongoDBConfig = cfg.GetSection("KernelMemory:Services:AzureCosmosDBMongoDB").Get() ?? new(); this.ElasticsearchConfig = cfg.GetSection("KernelMemory:Services:Elasticsearch").Get() ?? new(); } diff --git a/service/tests/TestHelpers/TestHelpers.csproj b/service/tests/TestHelpers/TestHelpers.csproj index 51b4b7ac1..aef09ab18 100644 --- a/service/tests/TestHelpers/TestHelpers.csproj +++ b/service/tests/TestHelpers/TestHelpers.csproj @@ -24,6 +24,8 @@ + +