Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,8 @@

<ItemGroup>
<PackageReference Include="Aspire.Npgsql.EntityFrameworkCore.PostgreSQL" Version="9.0.0" />
<PackageReference Include="Discord.Net" Version="3.17.0" />
<PackageReference Include="Microsoft.KernelMemory.Service.AspNetCore" Version="0.96.250116.1" />
<PackageReference Include="Discord.Net" Version="3.17.1" />
<PackageReference Include="Microsoft.KernelMemory.Service.AspNetCore" Version="0.96.250120.1" />
</ItemGroup>

</Project>
1 change: 1 addition & 0 deletions examples/301-discord-test-application/DiscordDbContext.cs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ public class DiscordDbContext : DbContext
{
public DbContextOptions<DiscordDbContext> Options { get; }

// Table to store Discord messages, table name is "Messages"
public DbSet<DiscordDbMessage> Messages { get; set; }

public DiscordDbContext(DbContextOptions<DiscordDbContext> options) : base(options)
Expand Down
87 changes: 68 additions & 19 deletions examples/301-discord-test-application/DiscordMessageHandler.cs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ namespace Microsoft.Discord.TestApplication;
/// KM pipeline handler fetching discord data files from document storage
/// and storing messages in Postgres.
/// </summary>
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;
Expand All @@ -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<DiscordMessageHandler> _log;
Expand All @@ -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<DiscordDbContext>() ?? 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)
Expand All @@ -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<DiscordDbContext>();
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)
Expand Down Expand Up @@ -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()
/// <summary>
/// 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]();
/// </summary>
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<DiscordDbContext>();
}
else
{
try
{
// Try the single app host first
this._log.LogTrace("Retrieving Discord DB context using service provider");
db = this._serviceProvider.GetService<DiscordDbContext>();
}
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<DiscordDbContext>();

// 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;
}
}
38 changes: 27 additions & 11 deletions examples/301-discord-test-application/Program.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -51,8 +62,8 @@ public static void Main(string[] args)
appBuilder.AddNpgsqlDbContext<DiscordDbContext>("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();
Expand All @@ -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();
Expand All @@ -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<IKernelMemory>() as MemoryServerless)!
.Orchestrator
.AddHandler<DiscordMessageHandler>(discordConfig.Steps[0]);
var orchestrator = (app.Services.GetService<IKernelMemory>() as MemoryServerless)!.Orchestrator;
orchestrator.AddHandler<DiscordMessageHandler>(discordConfig.Steps[0]);

return app;
}

private static WebApplication BuildAsynchronousKernelMemoryApp(WebApplicationBuilder appBuilder, DiscordConnectorConfig discordConfig)
{
appBuilder.Services.AddHandlerAsHostedService<DiscordMessageHandler>(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<DiscordMessageHandler>(discordConfig.Steps[0]);

WebApplication app = appBuilder.Build();

return app;
}
}
7 changes: 7 additions & 0 deletions service/Abstractions/KernelMemoryBuilderBuildOptions.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
6 changes: 4 additions & 2 deletions service/Service.AspNetCore/WebApplicationBuilderExtensions.cs
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,13 @@ public static partial class WebApplicationBuilderExtensions
/// <param name="configureMemoryBuilder">Optional configuration steps for the memory builder</param>
/// <param name="configureMemory">Optional configuration steps for the memory instance</param>
/// <param name="configureServices">Optional configuration for the internal dependencies</param>
/// <param name="buildOptions">Optional options passed to Build() call</param>
public static WebApplicationBuilder AddKernelMemory(
this WebApplicationBuilder appBuilder,
Action<IKernelMemoryBuilder>? configureMemoryBuilder = null,
Action<IKernelMemory>? configureMemory = null,
Action<IServiceCollection>? configureServices = null)
Action<IServiceCollection>? configureServices = null,
KernelMemoryBuilderBuildOptions? buildOptions = null)
{
// Prepare memory builder, sharing the service collection used by the hosting service
var memoryBuilder = new KernelMemoryBuilder(appBuilder.Services);
Expand All @@ -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);
Expand Down
11 changes: 2 additions & 9 deletions service/tests/Core.UnitTests/KernelMemoryBuilderTest.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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();

Expand Down
Loading