Browse Source

feat: Encrypt the Workspace ApiKey

pull/1421/head
colin 2 months ago
parent
commit
024d904722
  1. 7
      aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/DynamicWorkspaceDefinitionStoreInMemoryCache.cs
  2. 2
      aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/IWorkspaceDefinitionRecordRepository.cs
  3. 23
      aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/WorkspaceDefinitionRecord.cs
  4. 13
      aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/WorkspaceDefinitionSerializer.cs
  5. 45
      aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/WorkspaceDynamicInitializer.cs

7
aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/DynamicWorkspaceDefinitionStoreInMemoryCache.cs

@ -6,6 +6,7 @@ using System.Threading;
using System.Threading.Tasks; using System.Threading.Tasks;
using Volo.Abp.DependencyInjection; using Volo.Abp.DependencyInjection;
using Volo.Abp.Localization; using Volo.Abp.Localization;
using Volo.Abp.Security.Encryption;
using Volo.Abp.SimpleStateChecking; using Volo.Abp.SimpleStateChecking;
namespace LINGYUN.Abp.AIManagement.Workspaces; namespace LINGYUN.Abp.AIManagement.Workspaces;
@ -13,6 +14,7 @@ public class DynamicWorkspaceDefinitionStoreInMemoryCache : IDynamicWorkspaceDef
{ {
public string CacheStamp { get; set; } public string CacheStamp { get; set; }
protected IDictionary<string, WorkspaceDefinition> WorkspaceDefinitions { get; } protected IDictionary<string, WorkspaceDefinition> WorkspaceDefinitions { get; }
protected IStringEncryptionService StringEncryptionService { get; }
protected ISimpleStateCheckerSerializer StateCheckerSerializer { get; } protected ISimpleStateCheckerSerializer StateCheckerSerializer { get; }
protected ILocalizableStringSerializer LocalizableStringSerializer { get; } protected ILocalizableStringSerializer LocalizableStringSerializer { get; }
@ -21,9 +23,11 @@ public class DynamicWorkspaceDefinitionStoreInMemoryCache : IDynamicWorkspaceDef
public DateTime? LastCheckTime { get; set; } public DateTime? LastCheckTime { get; set; }
public DynamicWorkspaceDefinitionStoreInMemoryCache( public DynamicWorkspaceDefinitionStoreInMemoryCache(
IStringEncryptionService stringEncryptionService,
ISimpleStateCheckerSerializer stateCheckerSerializer, ISimpleStateCheckerSerializer stateCheckerSerializer,
ILocalizableStringSerializer localizableStringSerializer) ILocalizableStringSerializer localizableStringSerializer)
{ {
StringEncryptionService = stringEncryptionService;
StateCheckerSerializer = stateCheckerSerializer; StateCheckerSerializer = stateCheckerSerializer;
LocalizableStringSerializer = localizableStringSerializer; LocalizableStringSerializer = localizableStringSerializer;
@ -51,7 +55,8 @@ public class DynamicWorkspaceDefinitionStoreInMemoryCache : IDynamicWorkspaceDef
if (!workspace.ApiKey.IsNullOrWhiteSpace()) if (!workspace.ApiKey.IsNullOrWhiteSpace())
{ {
workspaceDef.WithApiKey(workspace.ApiKey); var decryptApiKey = StringEncryptionService.Decrypt(workspace.ApiKey);
workspaceDef.WithApiKey(decryptApiKey!);
} }
if (!workspace.ApiBaseUrl.IsNullOrWhiteSpace()) if (!workspace.ApiBaseUrl.IsNullOrWhiteSpace())
{ {

2
aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/IWorkspaceDefinitionRecordRepository.cs

@ -4,7 +4,7 @@ using System.Threading.Tasks;
using Volo.Abp.Domain.Repositories; using Volo.Abp.Domain.Repositories;
namespace LINGYUN.Abp.AIManagement.Workspaces; namespace LINGYUN.Abp.AIManagement.Workspaces;
public interface IWorkspaceDefinitionRecordRepository : IBasicRepository<WorkspaceDefinitionRecord, Guid> public interface IWorkspaceDefinitionRecordRepository : IRepository<WorkspaceDefinitionRecord, Guid>
{ {
Task<WorkspaceDefinitionRecord?> FindByNameAsync( Task<WorkspaceDefinitionRecord?> FindByNameAsync(
string name, string name,

23
aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/WorkspaceDefinitionRecord.cs

@ -14,7 +14,7 @@ public class WorkspaceDefinitionRecord : AuditedAggregateRoot<Guid>
public string DisplayName { get; private set; } public string DisplayName { get; private set; }
public string? Description { get; private set; } public string? Description { get; set; }
public string? ApiKey { get; private set; } public string? ApiKey { get; private set; }
@ -49,8 +49,6 @@ public class WorkspaceDefinitionRecord : AuditedAggregateRoot<Guid>
string modelName, string modelName,
string displayName, string displayName,
string? description = null, string? description = null,
string? apiKey = null,
string? apiBaseUrl = null,
string? systemPrompt = null, string? systemPrompt = null,
string? instructions = null, string? instructions = null,
float? temperature = null, float? temperature = null,
@ -65,8 +63,6 @@ public class WorkspaceDefinitionRecord : AuditedAggregateRoot<Guid>
ModelName = Check.NotNullOrWhiteSpace(modelName, nameof(modelName), WorkspaceDefinitionRecordConsts.MaxModelNameLength); ModelName = Check.NotNullOrWhiteSpace(modelName, nameof(modelName), WorkspaceDefinitionRecordConsts.MaxModelNameLength);
DisplayName = Check.NotNullOrWhiteSpace(displayName, nameof(displayName), WorkspaceDefinitionRecordConsts.MaxDisplayNameLength); DisplayName = Check.NotNullOrWhiteSpace(displayName, nameof(displayName), WorkspaceDefinitionRecordConsts.MaxDisplayNameLength);
Description = Check.Length(description, nameof(description), WorkspaceDefinitionRecordConsts.MaxDescriptionLength); Description = Check.Length(description, nameof(description), WorkspaceDefinitionRecordConsts.MaxDescriptionLength);
ApiKey = Check.Length(apiKey, nameof(apiKey), WorkspaceDefinitionRecordConsts.MaxApiKeyLength);
ApiBaseUrl = Check.Length(apiBaseUrl, nameof(apiBaseUrl), WorkspaceDefinitionRecordConsts.MaxApiBaseUrlLength);
SystemPrompt = Check.Length(systemPrompt, nameof(systemPrompt), WorkspaceDefinitionRecordConsts.MaxSystemPromptLength); SystemPrompt = Check.Length(systemPrompt, nameof(systemPrompt), WorkspaceDefinitionRecordConsts.MaxSystemPromptLength);
Instructions = Check.Length(instructions, nameof(instructions), WorkspaceDefinitionRecordConsts.MaxInstructionsLength); Instructions = Check.Length(instructions, nameof(instructions), WorkspaceDefinitionRecordConsts.MaxInstructionsLength);
StateCheckers = Check.Length(stateCheckers, nameof(stateCheckers), WorkspaceDefinitionRecordConsts.MaxStateCheckersLength); StateCheckers = Check.Length(stateCheckers, nameof(stateCheckers), WorkspaceDefinitionRecordConsts.MaxStateCheckersLength);
@ -80,6 +76,23 @@ public class WorkspaceDefinitionRecord : AuditedAggregateRoot<Guid>
this.SetDefaultsForExtraProperties(); this.SetDefaultsForExtraProperties();
} }
public void SetDisplayName(string displayName)
{
DisplayName = Check.NotNullOrWhiteSpace(displayName, nameof(displayName), WorkspaceDefinitionRecordConsts.MaxDisplayNameLength);
}
public void SetModel(string provider, string modelName)
{
Provider = Check.NotNullOrWhiteSpace(provider, nameof(provider), WorkspaceDefinitionRecordConsts.MaxProviderLength);
ModelName = Check.NotNullOrWhiteSpace(modelName, nameof(modelName), WorkspaceDefinitionRecordConsts.MaxModelNameLength);
}
public void SetApiKey(string? apiKey = null, string? apiBaseUrl = null)
{
ApiKey = Check.Length(apiKey, nameof(apiKey), WorkspaceDefinitionRecordConsts.MaxApiKeyLength);
ApiBaseUrl = Check.Length(apiBaseUrl, nameof(apiBaseUrl), WorkspaceDefinitionRecordConsts.MaxApiBaseUrlLength);
}
public bool HasSameData(WorkspaceDefinitionRecord otherWorkspace) public bool HasSameData(WorkspaceDefinitionRecord otherWorkspace)
{ {
if (Name != otherWorkspace.Name) if (Name != otherWorkspace.Name)

13
aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/WorkspaceDefinitionSerializer.cs

@ -1,4 +1,5 @@
using LINGYUN.Abp.AI.Workspaces; using LINGYUN.Abp.AI.Workspaces;
using System;
using System.Collections.Generic; using System.Collections.Generic;
using System.Globalization; using System.Globalization;
using System.Threading.Tasks; using System.Threading.Tasks;
@ -6,21 +7,25 @@ using Volo.Abp.Data;
using Volo.Abp.DependencyInjection; using Volo.Abp.DependencyInjection;
using Volo.Abp.Guids; using Volo.Abp.Guids;
using Volo.Abp.Localization; using Volo.Abp.Localization;
using Volo.Abp.Security.Encryption;
using Volo.Abp.SimpleStateChecking; using Volo.Abp.SimpleStateChecking;
namespace LINGYUN.Abp.AIManagement.Workspaces; namespace LINGYUN.Abp.AIManagement.Workspaces;
public class WorkspaceDefinitionSerializer : IWorkspaceDefinitionSerializer, ITransientDependency public class WorkspaceDefinitionSerializer : IWorkspaceDefinitionSerializer, ITransientDependency
{ {
protected IGuidGenerator GuidGenerator { get; } protected IGuidGenerator GuidGenerator { get; }
protected IStringEncryptionService StringEncryptionService { get; }
protected ISimpleStateCheckerSerializer StateCheckerSerializer { get; } protected ISimpleStateCheckerSerializer StateCheckerSerializer { get; }
protected ILocalizableStringSerializer LocalizableStringSerializer { get; } protected ILocalizableStringSerializer LocalizableStringSerializer { get; }
public WorkspaceDefinitionSerializer( public WorkspaceDefinitionSerializer(
IGuidGenerator guidGenerator, IGuidGenerator guidGenerator,
IStringEncryptionService stringEncryptionService,
ISimpleStateCheckerSerializer stateCheckerSerializer, ISimpleStateCheckerSerializer stateCheckerSerializer,
ILocalizableStringSerializer localizableStringSerializer) ILocalizableStringSerializer localizableStringSerializer)
{ {
GuidGenerator = guidGenerator; GuidGenerator = guidGenerator;
StringEncryptionService = stringEncryptionService;
StateCheckerSerializer = stateCheckerSerializer; StateCheckerSerializer = stateCheckerSerializer;
LocalizableStringSerializer = localizableStringSerializer; LocalizableStringSerializer = localizableStringSerializer;
} }
@ -47,8 +52,6 @@ public class WorkspaceDefinitionSerializer : IWorkspaceDefinitionSerializer, ITr
definition.ModelName, definition.ModelName,
LocalizableStringSerializer.Serialize(definition.DisplayName)!, LocalizableStringSerializer.Serialize(definition.DisplayName)!,
definition.Description != null ? LocalizableStringSerializer.Serialize(definition.Description) : null, definition.Description != null ? LocalizableStringSerializer.Serialize(definition.Description) : null,
definition.ApiKey,
definition.ApiBaseUrl,
definition.SystemPrompt, definition.SystemPrompt,
definition.Instructions, definition.Instructions,
definition.Temperature, definition.Temperature,
@ -57,6 +60,12 @@ public class WorkspaceDefinitionSerializer : IWorkspaceDefinitionSerializer, ITr
definition.PresencePenalty, definition.PresencePenalty,
SerializeStateCheckers(definition.StateCheckers)); SerializeStateCheckers(definition.StateCheckers));
if (!definition.ApiKey.IsNullOrWhiteSpace())
{
var encryptApiKey = StringEncryptionService.Encrypt(definition.ApiKey);
workspace.SetApiKey(encryptApiKey, definition.ApiBaseUrl);
}
foreach (var property in definition.Properties) foreach (var property in definition.Properties)
{ {
workspace.SetProperty(property.Key, property.Value); workspace.SetProperty(property.Key, property.Value);

45
aspnet-core/modules/ai/LINGYUN.Abp.AIManagement.Domain/LINGYUN/Abp/AIManagement/Workspaces/WorkspaceDynamicInitializer.cs

@ -1,5 +1,4 @@
using JetBrains.Annotations; using LINGYUN.Abp.AI.Workspaces;
using LINGYUN.Abp.AI.Workspaces;
using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting; using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging; using Microsoft.Extensions.Logging;
@ -18,33 +17,17 @@ public class WorkspaceDynamicInitializer : ITransientDependency
public ILogger<WorkspaceDynamicInitializer> Logger { get; set; } public ILogger<WorkspaceDynamicInitializer> Logger { get; set; }
protected IServiceProvider ServiceProvider { get; } protected IServiceProvider ServiceProvider { get; }
protected IOptions<AIManagementOptions> Options { get; }
[CanBeNull] public WorkspaceDynamicInitializer(IServiceProvider serviceProvider)
protected IHostApplicationLifetime ApplicationLifetime { get; }
protected ICancellationTokenProvider CancellationTokenProvider { get; }
protected IDynamicWorkspaceDefinitionStore DynamicWorkspaceDefinitionStore { get; }
protected IStaticWorkspaceSaver StaticWorkspaceSaver { get; }
public WorkspaceDynamicInitializer(
IServiceProvider serviceProvider,
IOptions<AIManagementOptions> options,
ICancellationTokenProvider cancellationTokenProvider,
IDynamicWorkspaceDefinitionStore dynamicWorkspaceDefinitionStore,
IStaticWorkspaceSaver staticWorkspaceSaver)
{ {
Logger = NullLogger<WorkspaceDynamicInitializer>.Instance; Logger = NullLogger<WorkspaceDynamicInitializer>.Instance;
ServiceProvider = serviceProvider; ServiceProvider = serviceProvider;
Options = options;
CancellationTokenProvider = cancellationTokenProvider;
DynamicWorkspaceDefinitionStore = dynamicWorkspaceDefinitionStore;
StaticWorkspaceSaver = staticWorkspaceSaver;
ApplicationLifetime = ServiceProvider.GetRequiredService<IHostApplicationLifetime>();
} }
public virtual Task InitializeAsync(bool runInBackground, CancellationToken cancellationToken = default) public virtual Task InitializeAsync(bool runInBackground, CancellationToken cancellationToken = default)
{ {
var options = Options.Value; var options = ServiceProvider.GetRequiredService<IOptions<AIManagementOptions>>().Value;
if (!options.SaveStaticWorkspacesToDatabase && !options.IsDynamicWorkspaceStoreEnabled) if (!options.SaveStaticWorkspacesToDatabase && !options.IsDynamicWorkspaceStoreEnabled)
{ {
@ -53,11 +36,12 @@ public class WorkspaceDynamicInitializer : ITransientDependency
if (runInBackground) if (runInBackground)
{ {
var applicationLifetime = ServiceProvider.GetService<IHostApplicationLifetime>();
Task.Run(async () => Task.Run(async () =>
{ {
if (cancellationToken == default && ApplicationLifetime?.ApplicationStopping != null) if (cancellationToken == default && applicationLifetime?.ApplicationStopping != null)
{ {
cancellationToken = ApplicationLifetime.ApplicationStopping; cancellationToken = applicationLifetime.ApplicationStopping;
} }
await ExecuteInitializationAsync(options, cancellationToken); await ExecuteInitializationAsync(options, cancellationToken);
}, cancellationToken); }, cancellationToken);
@ -72,16 +56,17 @@ public class WorkspaceDynamicInitializer : ITransientDependency
{ {
try try
{ {
using (CancellationTokenProvider.Use(cancellationToken)) var cancellationTokenProvider = ServiceProvider.GetRequiredService<ICancellationTokenProvider>();
using (cancellationTokenProvider.Use(cancellationToken))
{ {
if (CancellationTokenProvider.Token.IsCancellationRequested) if (cancellationTokenProvider.Token.IsCancellationRequested)
{ {
return; return;
} }
await SaveStaticWorkspacesToDatabaseAsync(options, cancellationToken); await SaveStaticWorkspacesToDatabaseAsync(options, cancellationToken);
if (CancellationTokenProvider.Token.IsCancellationRequested) if (cancellationTokenProvider.Token.IsCancellationRequested)
{ {
return; return;
} }
@ -104,6 +89,8 @@ public class WorkspaceDynamicInitializer : ITransientDependency
return; return;
} }
var staticWorkspaceSaver = ServiceProvider.GetRequiredService<IStaticWorkspaceSaver>();
await Policy await Policy
.Handle<Exception>(ex => ex is not OperationCanceledException) .Handle<Exception>(ex => ex is not OperationCanceledException)
.WaitAndRetryAsync( .WaitAndRetryAsync(
@ -118,7 +105,7 @@ public class WorkspaceDynamicInitializer : ITransientDependency
{ {
try try
{ {
await StaticWorkspaceSaver.SaveAsync(); await staticWorkspaceSaver.SaveAsync();
} }
catch (Exception ex) catch (Exception ex)
{ {
@ -135,10 +122,12 @@ public class WorkspaceDynamicInitializer : ITransientDependency
return; return;
} }
var dynamicWorkspaceDefinitionStore = ServiceProvider.GetRequiredService<IDynamicWorkspaceDefinitionStore>();
try try
{ {
// Pre-cache Workspaces, so first request doesn't wait // Pre-cache Workspaces, so first request doesn't wait
await DynamicWorkspaceDefinitionStore.GetAllAsync(); await dynamicWorkspaceDefinitionStore.GetAllAsync();
} }
catch (Exception ex) catch (Exception ex)
{ {

Loading…
Cancel
Save