Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add/Remove handlers without code changes #316

Merged
merged 2 commits into from
Feb 20, 2024
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
1 change: 1 addition & 0 deletions KernelMemory.sln.DotSettings
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,7 @@ public void It$SOMENAME$()
<s:String x:Key="/Default/CodeInspection/Highlighting/InspectionSeverities/=Xunit_002EXunitTestWithConsoleOutput/@EntryIndexedValue">DO_NOT_SHOW</s:String>
<s:Boolean x:Key="/Default/UserDictionary/Words/=smemory/@EntryIndexedValue">True</s:Boolean>
<s:Boolean x:Key="/Default/UserDictionary/Words/=SVCS/@EntryIndexedValue">True</s:Boolean>
<s:Boolean x:Key="/Default/UserDictionary/Words/=syntaxes/@EntryIndexedValue">True</s:Boolean>
<s:Boolean x:Key="/Default/UserDictionary/Words/=testsettings/@EntryIndexedValue">True</s:Boolean>
<s:Boolean x:Key="/Default/UserDictionary/Words/=tldr/@EntryIndexedValue">True</s:Boolean>
<s:Boolean x:Key="/Default/UserDictionary/Words/=typeparam/@EntryIndexedValue">True</s:Boolean>
Expand Down
17 changes: 17 additions & 0 deletions examples/004-dotnet-ServerlessCustomPipeline/Program.cs
Original file line number Diff line number Diff line change
@@ -1,5 +1,22 @@
// Copyright (c) Microsoft. All rights reserved.

/*
* The ingestion pipeline is composed by a set of default STEPS to process documents:
* - extract
* - partition
* - gen_embeddings
* - save_records
*
* Each step is managed by a HANDLER, see the Core/Handlers for a list of available handlers.
*
* You can create new handlers, and customize the pipeline in multiple ways:
*
* - If you are using the Memory Web Service, the list of handlers can be configured in appsettings.json and appsettings.<ENV>.json
* - You can create new handlers and load them in the configuration, passing the path to the assembly and the handler class.
* - If you are using the Serverless Memory, see the example below.
* - You can also remove the default handlers calling .WithoutDefaultHandlers() in the memory builder.
*/

using Microsoft.KernelMemory;
using Microsoft.KernelMemory.Handlers;

Expand Down
6 changes: 3 additions & 3 deletions examples/202-dotnet-CustomHandlerAsAService/Program.cs
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@

using Microsoft.KernelMemory;

/* The following code shows how to create a custom handler, attached
* to a queue and listening for work to do. You can also add multiple handlers
* the same way.
/* The following code shows how to create a custom handler and run it as a standalone service.
* The handler will automatically attach to a queue and listen for work to do.
* You can also add multiple handlers the same way.
*/

// Usual .NET web app builder
Expand Down
14 changes: 7 additions & 7 deletions extensions/AzureOpenAI/AzureOpenAIConfig.cs
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ public void SetCredential(TokenCredential credential)
public TokenCredential GetTokenCredential()
{
return this._tokenCredential
?? throw new ConfigurationException("TokenCredential not defined");
?? throw new ConfigurationException("Azure OpenAI TokenCredential not defined");
}

/// <summary>
Expand All @@ -94,33 +94,33 @@ public void Validate()
{
if (this.Auth == AuthTypes.Unknown)
{
throw new ArgumentOutOfRangeException(nameof(this.Auth), "The authentication type is not defined");
throw new ArgumentOutOfRangeException(nameof(this.Auth), "The Azure OpenAI Authentication Type is not defined");
}

if (this.Auth == AuthTypes.APIKey && string.IsNullOrWhiteSpace(this.APIKey))
{
throw new ArgumentOutOfRangeException(nameof(this.APIKey), "The API Key is empty");
throw new ArgumentOutOfRangeException(nameof(this.APIKey), "The Azure OpenAI API Key is empty");
}

if (string.IsNullOrWhiteSpace(this.Endpoint))
{
throw new ArgumentOutOfRangeException(nameof(this.Endpoint), "The endpoint value is empty");
throw new ArgumentOutOfRangeException(nameof(this.Endpoint), "The Azure OpenAI Endpoint value is empty");
}

if (!this.Endpoint.StartsWith("https://", StringComparison.OrdinalIgnoreCase))
{
throw new ArgumentOutOfRangeException(nameof(this.Endpoint), "The endpoint value must start with https://");
throw new ArgumentOutOfRangeException(nameof(this.Endpoint), "The Azure OpenAI Endpoint value must start with https://");
}

if (string.IsNullOrWhiteSpace(this.Deployment))
{
throw new ArgumentOutOfRangeException(nameof(this.Deployment), "The deployment value is empty");
throw new ArgumentOutOfRangeException(nameof(this.Deployment), "The Azure OpenAI Deployment Name is empty");
}

if (this.MaxTokenTotal < 1)
{
throw new ArgumentOutOfRangeException(nameof(this.MaxTokenTotal),
$"{nameof(this.MaxTokenTotal)} cannot be less than 1");
$"Azure OpenAI: {nameof(this.MaxTokenTotal)} cannot be less than 1");
}
}
}
61 changes: 60 additions & 1 deletion service/Core/AppBuilders/DependencyInjection.cs
Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@
// Copyright (c) Microsoft. All rights reserved.

using System;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.KernelMemory.Configuration;
using Microsoft.KernelMemory.Handlers;
using Microsoft.KernelMemory.Pipeline;

namespace Microsoft.KernelMemory;
Expand All @@ -16,8 +20,63 @@ public static partial class DependencyInjection
public static void AddHandlerAsHostedService<THandler>(this IServiceCollection services, string stepName) where THandler : class, IPipelineStepHandler
{
services.AddTransient<THandler>(serviceProvider => ActivatorUtilities.CreateInstance<THandler>(serviceProvider, stepName));

services.AddHostedService<HandlerAsAHostedService<THandler>>(
serviceProvider => ActivatorUtilities.CreateInstance<HandlerAsAHostedService<THandler>>(serviceProvider, stepName));
}

/// <summary>
/// Register the handler as a hosted service, passing the step name to the handler ctor
/// </summary>
/// <param name="services">Application builder service collection</param>
/// <param name="tHandler">Handler class</param>
/// <param name="stepName">Pipeline step name</param>
public static void AddHandlerAsHostedService(this IServiceCollection services, Type tHandler, string stepName)
{
if (!typeof(IPipelineStepHandler).IsAssignableFrom(tHandler))
{
throw new ArgumentException($"'{tHandler.FullName}' doesn't implement interface '{nameof(IPipelineStepHandler)}'", nameof(tHandler));
}

if (tHandler == null)
{
throw new ArgumentNullException(nameof(tHandler), $"Handler type for '{stepName}' is NULL");
}

services.AddTransient(tHandler, serviceProvider => ActivatorUtilities.CreateInstance(serviceProvider, tHandler, stepName));

// Build generic type: HandlerAsAHostedService<THandler>
Type handlerAsAHostedServiceTHandler = typeof(HandlerAsAHostedService<>).MakeGenericType(tHandler);

Func<IServiceProvider, IHostedService> implementationFactory =
serviceProvider => (IHostedService)ActivatorUtilities.CreateInstance(serviceProvider, handlerAsAHostedServiceTHandler, stepName);

// See https://github.com/dotnet/runtime/issues/38751 for troubleshooting
services.Add(ServiceDescriptor.Singleton<IHostedService>(implementationFactory));
}

/// <summary>
/// Register the handler as a hosted service, passing the step name to the handler ctor
/// </summary>
/// <param name="services">Application builder service collection</param>
/// <param name="config">Handler type configuration</param>
/// <param name="stepName">Pipeline step name</param>
public static void AddHandlerAsHostedService(this IServiceCollection services, HandlerConfig config, string stepName)
{
if (HandlerTypeLoader.TryGetHandlerType(config, out var handlerType))
{
services.AddHandlerAsHostedService(handlerType, stepName);
}
}

/// <summary>
/// Register the handler as a hosted service, passing the step name to the handler ctor
/// </summary>
/// <param name="services">Application builder service collection</param>
/// <param name="assemblyFile">Path to assembly containing handler class</param>
/// <param name="typeFullName">Handler type, within the assembly</param>
/// <param name="stepName">Pipeline step name</param>
public static void AddHandlerAsHostedService(this IServiceCollection services, string assemblyFile, string typeFullName, string stepName)
{
services.AddHandlerAsHostedService(new HandlerConfig(assemblyFile, typeFullName), stepName);
}
}
60 changes: 60 additions & 0 deletions service/Core/AppBuilders/HandlerTypeLoader.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Reflection;
using Microsoft.KernelMemory.Configuration;
using Microsoft.KernelMemory.Pipeline;

namespace Microsoft.KernelMemory.Handlers;

internal static class HandlerTypeLoader
{
internal static bool TryGetHandlerType(HandlerConfig config, [NotNullWhen(true)] out Type? handlerType)
{
handlerType = null;

// If part of the config is empty, the handler is disabled
if (string.IsNullOrEmpty(config.Class) || string.IsNullOrEmpty(config.Assembly))
{
return false;
}

// Search the assembly in a few directories
var path = string.Empty;
var assemblyFilePaths = new HashSet<string>
{
config.Assembly,
Path.Join(Environment.CurrentDirectory, config.Assembly),
Path.Join(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location), config.Assembly),
};

foreach (var p in assemblyFilePaths)
{
if (!File.Exists(p)) { continue; }

path = p;
break;
}

// Check if the assembly exists
if (string.IsNullOrEmpty(path))
{
throw new ConfigurationException($"Handler assembly not found: {config.Assembly}");
}

Assembly assembly = Assembly.LoadFrom(path);

// IPipelineStepHandler
handlerType = assembly.GetType(config.Class);

if (!typeof(IPipelineStepHandler).IsAssignableFrom(handlerType))
{
throw new ConfigurationException($"Invalid handler definition: `{config.Class}` class doesn't implement interface {nameof(IPipelineStepHandler)}");
}

return true;
}
}
28 changes: 28 additions & 0 deletions service/Core/Configuration/HandlerConfig.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
// Copyright (c) Microsoft. All rights reserved.

namespace Microsoft.KernelMemory.Configuration;

public class HandlerConfig
{
/// <summary>
/// .NET assembly containing the handler class
/// </summary>
public string Assembly { get; set; }

/// <summary>
/// .NET class in the assembly, containing the handler logic
/// </summary>
public string Class { get; set; }

public HandlerConfig()
{
this.Assembly = string.Empty;
this.Class = string.Empty;
}

public HandlerConfig(string assembly, string className)
{
this.Assembly = assembly;
this.Class = className;
}
}
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
// Copyright (c) Microsoft. All rights reserved.

using System.Collections.Generic;

namespace Microsoft.KernelMemory.Configuration;

public class ServiceConfig
Expand All @@ -20,4 +22,9 @@ public class ServiceConfig
/// Web service settings, e.g. whether to expose OpenAPI swagger docs.
/// </summary>
public bool OpenApiEnabled { get; set; } = false;

/// <summary>
/// List of handlers to enable
/// </summary>
public Dictionary<string, HandlerConfig> Handlers { get; set; } = new();
}
25 changes: 1 addition & 24 deletions service/Core/KernelMemoryBuilder.cs
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ private enum ClientTypes
// Services required to build the memory client class
private readonly IServiceCollection _memoryServiceCollection;

// Services of the host application, when hosting pipeline handlers, e.g. in service mode
// Services of the host application
private readonly IServiceCollection? _hostServiceCollection;

// List of all the embedding generators to use during ingestion
Expand Down Expand Up @@ -296,29 +296,6 @@ private MemoryService BuildAsyncClient()
var orchestrator = serviceProvider.GetService<DistributedPipelineOrchestrator>() ?? throw new ConfigurationException("Unable to build orchestrator");
var searchClient = serviceProvider.GetService<ISearchClient>() ?? throw new ConfigurationException("Unable to build search client");

if (this._useDefaultHandlers)
{
if (this._hostServiceCollection == null)
{
const string ClassName = nameof(KernelMemoryBuilder);
const string MethodName = nameof(this.WithoutDefaultHandlers);
throw new ConfigurationException("Host service collection not available: unable to register default handlers. " +
$"If you'd like using the default handlers use `new {ClassName}(<your service collection provider>)`, " +
$"otherwise use `new {ClassName}(...).{MethodName}()` to manage the list of handlers manually.");
}

// Handlers - Register these handlers to run as hosted services in the caller app.
// At start each hosted handler calls IPipelineOrchestrator.AddHandlerAsync() to register in the orchestrator.
this._hostServiceCollection.AddHandlerAsHostedService<TextExtractionHandler>(Constants.PipelineStepsExtract);
this._hostServiceCollection.AddHandlerAsHostedService<TextPartitioningHandler>(Constants.PipelineStepsPartition);
this._hostServiceCollection.AddHandlerAsHostedService<GenerateEmbeddingsHandler>(Constants.PipelineStepsGenEmbeddings);
this._hostServiceCollection.AddHandlerAsHostedService<SaveRecordsHandler>(Constants.PipelineStepsSaveRecords);
this._hostServiceCollection.AddHandlerAsHostedService<SummarizationHandler>(Constants.PipelineStepsSummarize);
this._hostServiceCollection.AddHandlerAsHostedService<DeleteDocumentHandler>(Constants.PipelineStepsDeleteDocument);
this._hostServiceCollection.AddHandlerAsHostedService<DeleteIndexHandler>(Constants.PipelineStepsDeleteIndex);
this._hostServiceCollection.AddHandlerAsHostedService<DeleteGeneratedFilesHandler>(Constants.PipelineStepsDeleteGeneratedFiles);
}

this.CheckForMissingDependencies();

return new MemoryService(orchestrator, searchClient);
Expand Down
1 change: 1 addition & 0 deletions service/Service/Auth/HttpAuthHandler.cs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ public class HttpAuthEndpointFilter : IEndpointFilter

public HttpAuthEndpointFilter(ServiceAuthorizationConfig config)
{
config.Validate();
this._config = config;
}

Expand Down
Loading
Loading