diff --git a/examples/301-discord-test-application/301-discord-test-application.csproj b/examples/301-discord-test-application/301-discord-test-application.csproj index 34b696e26..0a8d042b3 100644 --- a/examples/301-discord-test-application/301-discord-test-application.csproj +++ b/examples/301-discord-test-application/301-discord-test-application.csproj @@ -12,8 +12,8 @@ - - + + diff --git a/examples/301-discord-test-application/DiscordDbContext.cs b/examples/301-discord-test-application/DiscordDbContext.cs index ec0c007b4..1bbf5a454 100644 --- a/examples/301-discord-test-application/DiscordDbContext.cs +++ b/examples/301-discord-test-application/DiscordDbContext.cs @@ -8,6 +8,7 @@ public class DiscordDbContext : DbContext { public DbContextOptions Options { get; } + // Table to store Discord messages, table name is "Messages" public DbSet Messages { get; set; } public DiscordDbContext(DbContextOptions options) : base(options) diff --git a/examples/301-discord-test-application/DiscordMessageHandler.cs b/examples/301-discord-test-application/DiscordMessageHandler.cs index 3305d5c30..572e52b3d 100644 --- a/examples/301-discord-test-application/DiscordMessageHandler.cs +++ b/examples/301-discord-test-application/DiscordMessageHandler.cs @@ -12,7 +12,7 @@ namespace Microsoft.Discord.TestApplication; /// KM pipeline handler fetching discord data files from document storage /// and storing messages in Postgres. /// -public sealed class DiscordMessageHandler : IPipelineStepHandler, IDisposable, IAsyncDisposable +public sealed class DiscordMessageHandler : IPipelineStepHandler, IDisposable { // Name of the file where to store Discord data private readonly string _filename; @@ -23,8 +23,11 @@ public sealed class DiscordMessageHandler : IPipelineStepHandler, IDisposable, I // .NET service provider, used to get thread safe instances of EF DbContext private readonly IServiceProvider _serviceProvider; - // EF DbContext used to create the database - private DiscordDbContext? _firstInvokeDb; + // DB creation + private readonly object _dbCreation = new(); + private bool _dbCreated = false; + private bool _useScope = false; + private readonly IServiceScope _dbScope; // .NET logger private readonly ILogger _log; @@ -44,9 +47,16 @@ public DiscordMessageHandler( this._orchestrator = orchestrator; this._serviceProvider = serviceProvider; this._filename = config.FileName; + this._dbScope = this._serviceProvider.CreateScope(); - // This DbContext instance is used only to create the database - this._firstInvokeDb = serviceProvider.GetService() ?? throw new ConfigurationException("Discord DB Content is not defined"); + try + { + this.OnFirstInvoke(); + } + catch (Exception) + { + // ignore, will retry later + } } public async Task<(ReturnType returnType, DataPipeline updatedPipeline)> InvokeAsync(DataPipeline pipeline, CancellationToken cancellationToken = default) @@ -57,8 +67,7 @@ public DiscordMessageHandler( // exception: System.InvalidOperationException: a second operation was started on this context instance before a previous // operation completed. This is usually caused by different threads concurrently using the same instance of DbContext. // For more information on how to avoid threading issues with DbContext, see https://go.microsoft.com/fwlink/?linkid=2097913. - DiscordDbContext? db = this._serviceProvider.GetService(); - ArgumentNullExceptionEx.ThrowIfNull(db, nameof(db), "Discord DB context is NULL"); + var db = this.GetDb(); await using (db.ConfigureAwait(false)) { foreach (DataPipeline.FileDetails uploadedFile in pipeline.Files) @@ -95,27 +104,67 @@ public DiscordMessageHandler( public void Dispose() { - this._firstInvokeDb?.Dispose(); - this._firstInvokeDb = null; + this._dbScope.Dispose(); } - public async ValueTask DisposeAsync() + private void OnFirstInvoke() { - if (this._firstInvokeDb != null) { await this._firstInvokeDb.DisposeAsync(); } + if (this._dbCreated) { return; } - this._firstInvokeDb = null; + lock (this._dbCreation) + { + if (this._dbCreated) { return; } + + var db = this.GetDb(); + db.Database.EnsureCreated(); + db.Dispose(); + db = null; + + this._dbCreated = true; + + this._log.LogInformation("DB created"); + } } - private void OnFirstInvoke() + /// + /// Depending on the hosting type, the DB Context is retrieved in different ways. + /// Single host app: + /// db = _serviceProvider.GetService[DiscordDbContext](); // this throws an exception in multi-host mode + /// Multi host app: + /// db = serviceProvider.CreateScope().ServiceProvider.GetRequiredService[DiscordDbContext](); + /// + private DiscordDbContext GetDb() { - if (this._firstInvokeDb == null) { return; } + DiscordDbContext? db; - lock (this._firstInvokeDb) + if (this._useScope) { - // Create DB / Tables if needed - this._firstInvokeDb.Database.EnsureCreated(); - this._firstInvokeDb.Dispose(); - this._firstInvokeDb = null; + db = this._dbScope.ServiceProvider.GetRequiredService(); } + else + { + try + { + // Try the single app host first + this._log.LogTrace("Retrieving Discord DB context using service provider"); + db = this._serviceProvider.GetService(); + } + catch (InvalidOperationException) + { + // If the single app host fails, try the multi app host + this._log.LogInformation("Retrieving Discord DB context using scope"); + db = this._dbScope.ServiceProvider.GetRequiredService(); + + // If the multi app host succeeds, set a flag to remember to use the scope + if (db != null) + { + this._useScope = true; + } + } + } + + ArgumentNullExceptionEx.ThrowIfNull(db, nameof(db), "Discord DB context is NULL"); + + return db; } } diff --git a/examples/301-discord-test-application/Program.cs b/examples/301-discord-test-application/Program.cs index 3e81a500c..f22291195 100644 --- a/examples/301-discord-test-application/Program.cs +++ b/examples/301-discord-test-application/Program.cs @@ -2,11 +2,15 @@ using Microsoft.KernelMemory; using Microsoft.KernelMemory.DocumentStorage.DevTools; +using Microsoft.KernelMemory.MemoryStorage.DevTools; using Microsoft.KernelMemory.Sources.DiscordBot; namespace Microsoft.Discord.TestApplication; /* Example: Listen for new messages in Discord, and save them in a table in Postgres. + * + * Why this example: You can build on this example to populate a memory database with + * user messages, and then use the memory database to autogenerate answers. * * Use ASP.NET hosted services to host a Discord Bot. The discord bot logic is based * on DiscordConnector class. @@ -18,8 +22,15 @@ namespace Microsoft.Discord.TestApplication; * The call to KM.ImportDocument API asks to process the JSON file uploaded using * DiscordMessageHandler, included in this project. No other handlers are used. * - * DiscordMessageHandler, loads the uploaded file, deserializes its content, and - * save each Discord message into a table in Postgres, using Entity Framework. + * DiscordMessageHandler, loads the JSON file uploaded, deserializes its content, and + * saves each Discord message into a table in Postgres, using Entity Framework. + * + * Discord Server + * => Discord Bot + * => OnMessage Event + * => KM.ImportDocumentAsync(data, steps: ["store_discord_message"]) + * => DiscordMessageHandler + * => Postgres table */ internal static class Program @@ -51,8 +62,8 @@ public static void Main(string[] args) appBuilder.AddNpgsqlDbContext("postgresDb"); // Run Kernel Memory and DiscordMessageHandler - // var kmApp = BuildAsynchronousKernelMemoryApp(appBuilder, discordConfig); - var kmApp = BuildSynchronousKernelMemoryApp(appBuilder, discordCfg); + // var kmApp = BuildAsynchronousKernelMemoryApp(appBuilder, discordCfg); // run using queues and threads + var kmApp = BuildSynchronousKernelMemoryApp(appBuilder, discordCfg); // run everything in one thread Console.WriteLine("Starting KM application...\n"); kmApp.Run(); @@ -65,8 +76,9 @@ private static WebApplication BuildSynchronousKernelMemoryApp(WebApplicationBuil { // Note: there's no queue system, so the memory instance will be synchronous (ie MemoryServerless) - // Store files on disk + // Store files and vectors on disk kmb.WithSimpleFileStorage(SimpleFileStorageConfig.Persistent); + kmb.WithSimpleVectorDb(SimpleVectorDbConfig.Persistent); // Disable AI, not needed for this example kmb.WithoutEmbeddingGenerator(); @@ -76,29 +88,33 @@ private static WebApplication BuildSynchronousKernelMemoryApp(WebApplicationBuil WebApplication app = appBuilder.Build(); // In synchronous apps, handlers are added to the serverless memory orchestrator - (app.Services.GetService() as MemoryServerless)! - .Orchestrator - .AddHandler(discordConfig.Steps[0]); + var orchestrator = (app.Services.GetService() as MemoryServerless)!.Orchestrator; + orchestrator.AddHandler(discordConfig.Steps[0]); return app; } private static WebApplication BuildAsynchronousKernelMemoryApp(WebApplicationBuilder appBuilder, DiscordConnectorConfig discordConfig) { - appBuilder.Services.AddHandlerAsHostedService(discordConfig.Steps[0]); appBuilder.AddKernelMemory(kmb => { // Note: because of this the memory instance will be asynchronous (ie MemoryService) kmb.WithSimpleQueuesPipeline(); - // Store files on disk + // Store files and vectors on disk kmb.WithSimpleFileStorage(SimpleFileStorageConfig.Persistent); + kmb.WithSimpleVectorDb(SimpleVectorDbConfig.Persistent); // Disable AI, not needed for this example kmb.WithoutEmbeddingGenerator(); kmb.WithoutTextGenerator(); }); - return appBuilder.Build(); + // In asynchronous apps, handlers are added as hosted services to run on dedicated threads + appBuilder.Services.AddHandlerAsHostedService(discordConfig.Steps[0]); + + WebApplication app = appBuilder.Build(); + + return app; } } diff --git a/service/Abstractions/KernelMemoryBuilderBuildOptions.cs b/service/Abstractions/KernelMemoryBuilderBuildOptions.cs index e0739b6af..bf6c65403 100644 --- a/service/Abstractions/KernelMemoryBuilderBuildOptions.cs +++ b/service/Abstractions/KernelMemoryBuilderBuildOptions.cs @@ -4,5 +4,12 @@ namespace Microsoft.KernelMemory; public sealed class KernelMemoryBuilderBuildOptions { + public readonly static KernelMemoryBuilderBuildOptions Default = new(); + + public readonly static KernelMemoryBuilderBuildOptions WithVolatileAndPersistentData = new() + { + AllowMixingVolatileAndPersistentData = true + }; + public bool AllowMixingVolatileAndPersistentData { get; set; } = false; } diff --git a/service/Service.AspNetCore/WebApplicationBuilderExtensions.cs b/service/Service.AspNetCore/WebApplicationBuilderExtensions.cs index e253723ae..76de8c860 100644 --- a/service/Service.AspNetCore/WebApplicationBuilderExtensions.cs +++ b/service/Service.AspNetCore/WebApplicationBuilderExtensions.cs @@ -20,11 +20,13 @@ public static partial class WebApplicationBuilderExtensions /// Optional configuration steps for the memory builder /// Optional configuration steps for the memory instance /// Optional configuration for the internal dependencies + /// Optional options passed to Build() call public static WebApplicationBuilder AddKernelMemory( this WebApplicationBuilder appBuilder, Action? configureMemoryBuilder = null, Action? configureMemory = null, - Action? configureServices = null) + Action? configureServices = null, + KernelMemoryBuilderBuildOptions? buildOptions = null) { // Prepare memory builder, sharing the service collection used by the hosting service var memoryBuilder = new KernelMemoryBuilder(appBuilder.Services); @@ -35,7 +37,7 @@ public static WebApplicationBuilder AddKernelMemory( // Optional configuration provided by the user configureMemoryBuilder?.Invoke(memoryBuilder); - var memory = memoryBuilder.Build(); + var memory = memoryBuilder.Build(buildOptions); // Optional memory configuration provided by the user configureMemory?.Invoke(memory); diff --git a/service/tests/Core.UnitTests/KernelMemoryBuilderTest.cs b/service/tests/Core.UnitTests/KernelMemoryBuilderTest.cs index 6cc58e1cc..73cf1693b 100644 --- a/service/tests/Core.UnitTests/KernelMemoryBuilderTest.cs +++ b/service/tests/Core.UnitTests/KernelMemoryBuilderTest.cs @@ -118,10 +118,7 @@ public void ItDetectsMissingEmbeddingGenerator() public void ItCanMixPersistentAndVolatileStorageIfNeeded() { // Arrange - KernelMemoryBuilderBuildOptions kmbOptions = new() - { - AllowMixingVolatileAndPersistentData = true - }; + var kmbOptions = KernelMemoryBuilderBuildOptions.WithVolatileAndPersistentData; // Act - Assert no exception occurs new KernelMemoryBuilder() @@ -136,11 +133,7 @@ public void ItCanMixPersistentAndVolatileStorageIfNeeded() public void ItCanMixPersistentAndVolatileStorageIfNeeded2() { // Arrange - KernelMemoryBuilderBuildOptions kmbOptions = new() - { - AllowMixingVolatileAndPersistentData = true - }; - + var kmbOptions = KernelMemoryBuilderBuildOptions.WithVolatileAndPersistentData; var serviceCollection1 = new ServiceCollection(); var serviceCollection2 = new ServiceCollection();