diff --git a/app/MindWork AI Studio/Agents/AgentDataSourceSelection.cs b/app/MindWork AI Studio/Agents/AgentDataSourceSelection.cs index dc7c295c..fa72fcdb 100644 --- a/app/MindWork AI Studio/Agents/AgentDataSourceSelection.cs +++ b/app/MindWork AI Studio/Agents/AgentDataSourceSelection.cs @@ -5,12 +5,11 @@ using AIStudio.Chat; using AIStudio.Provider; using AIStudio.Settings; using AIStudio.Settings.DataModel; -using AIStudio.Tools.ERIClient; using AIStudio.Tools.Services; namespace AIStudio.Agents; -public sealed class AgentDataSourceSelection (ILogger logger, ILogger baseLogger, SettingsManager settingsManager, DataSourceService dataSourceService, ThreadSafeRandom rng) : AgentBase(baseLogger, settingsManager, dataSourceService, rng) +public sealed class AgentDataSourceSelection (ILogger logger, ILogger baseLogger, SettingsManager settingsManager, DataSourceService dataSourceService, DataSourceDescriptionService descriptionService, ThreadSafeRandom rng) : AgentBase(baseLogger, settingsManager, dataSourceService, rng) { private readonly List answers = new(); @@ -187,75 +186,24 @@ public sealed class AgentDataSourceSelection (ILogger var additionalData = new Dictionary(); logger.LogInformation("Preparing the list of allowed data sources for the agent to choose from."); - // Notice: We do not dispose the Rust service here. The Rust service is a singleton - // and will be disposed when the application shuts down: - var rustService = Program.SERVICE_PROVIDER.GetService()!; - var sb = new StringBuilder(); sb.AppendLine("The following data sources are available for selection:"); foreach (var ds in dataSources.AllowedDataSources) { + var description = await descriptionService.GetDescriptionAsync(ds, token); + var descriptionPart = string.IsNullOrWhiteSpace(description) ? string.Empty : $", description='{description}'"; switch (ds) { case DataSourceLocalDirectory localDirectory: - if (string.IsNullOrWhiteSpace(localDirectory.Description)) - sb.AppendLine($"- Id={ds.Id}, name='{localDirectory.Name}', type=local directory, path='{localDirectory.Path}'"); - else - { - var description = localDirectory.Description.Replace("\n", " ").Replace("\r", " "); - sb.AppendLine($"- Id={ds.Id}, name='{localDirectory.Name}', type=local directory, path='{localDirectory.Path}', description='{description}'"); - } + sb.AppendLine($"- Id={ds.Id}, name='{localDirectory.Name}', type=local directory, path='{localDirectory.Path}'{descriptionPart}"); break; case DataSourceLocalFile localFile: - if (string.IsNullOrWhiteSpace(localFile.Description)) - sb.AppendLine($"- Id={ds.Id}, name='{localFile.Name}', type=local file, path='{localFile.FilePath}'"); - else - { - var description = localFile.Description.Replace("\n", " ").Replace("\r", " "); - sb.AppendLine($"- Id={ds.Id}, name='{localFile.Name}', type=local file, path='{localFile.FilePath}', description='{description}'"); - } + sb.AppendLine($"- Id={ds.Id}, name='{localFile.Name}', type=local file, path='{localFile.FilePath}'{descriptionPart}"); break; case IERIDataSource eriDataSource: - var eriServerDescription = string.Empty; - - try - { - // - // Call the ERI server to get the server description: - // - using var eriClient = ERIClientFactory.Get(eriDataSource.Version, eriDataSource)!; - var authResponse = await eriClient.AuthenticateAsync(rustService, cancellationToken: token); - if (authResponse.Successful) - { - var serverDescriptionResponse = await eriClient.GetDataSourceInfoAsync(token); - if (serverDescriptionResponse.Successful) - { - eriServerDescription = serverDescriptionResponse.Data.Description; - - // Remove all line breaks from the description: - eriServerDescription = eriServerDescription.Replace("\n", " ").Replace("\r", " "); - } - else - logger.LogWarning($"Was not able to retrieve the server description from the ERI data source '{eriDataSource.Name}'. Message: {serverDescriptionResponse.Message}"); - } - else - logger.LogWarning($"Was not able to authenticate with the ERI data source '{eriDataSource.Name}'. Message: {authResponse.Message}"); - } - catch (Exception e) - { - logger.LogWarning($"The ERI data source '{eriDataSource.Name}' is not available. Thus, we cannot retrieve the server description. Error: {e.Message}"); - } - - // - // Append the ERI data source to the list. Use the server description if available: - // - if (string.IsNullOrWhiteSpace(eriServerDescription)) - sb.AppendLine($"- Id={ds.Id}, name='{eriDataSource.Name}', type=external data source"); - else - sb.AppendLine($"- Id={ds.Id}, name='{eriDataSource.Name}', type=external data source, description='{eriServerDescription}'"); - + sb.AppendLine($"- Id={ds.Id}, name='{eriDataSource.Name}', type=external data source{descriptionPart}"); break; } } diff --git a/app/MindWork AI Studio/Program.cs b/app/MindWork AI Studio/Program.cs index 54c1608b..21524b0e 100644 --- a/app/MindWork AI Studio/Program.cs +++ b/app/MindWork AI Studio/Program.cs @@ -201,6 +201,7 @@ internal sealed class Program builder.Services.AddSingleton(); builder.Services.AddSingleton(); builder.Services.AddSingleton(); + builder.Services.AddSingleton(); builder.Services.AddSingleton(); builder.Services.AddSingleton(); builder.Services.AddSingleton(); diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs index d13b1dc4..b8a8fb5f 100644 --- a/app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs @@ -13,4 +13,6 @@ public readonly record struct PromptInjectionSource(PromptInjectionSourceKind Ki public static PromptInjectionSource ChatAttachment(string filePath) => new(PromptInjectionSourceKind.CHAT_ATTACHMENT, filePath); public static PromptInjectionSource RetrievalContext(string dataSourceName, string path) => new(PromptInjectionSourceKind.RETRIEVAL_CONTEXT, $"{dataSourceName}: {path}"); + + public static PromptInjectionSource DataSourceDescription(string dataSourceName) => new(PromptInjectionSourceKind.DATA_SOURCE_DESCRIPTION, dataSourceName); } \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKind.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKind.cs index 3df49619..81d30d65 100644 --- a/app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKind.cs +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKind.cs @@ -7,4 +7,5 @@ public enum PromptInjectionSourceKind FILE_CONTENT, CHAT_ATTACHMENT, RETRIEVAL_CONTEXT, + DATA_SOURCE_DESCRIPTION, } \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKindExtensions.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKindExtensions.cs index cf5511a3..6b14e283 100644 --- a/app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKindExtensions.cs +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKindExtensions.cs @@ -12,6 +12,7 @@ public static class PromptInjectionSourceKindExtensions PromptInjectionSourceKind.FILE_CONTENT => TB("File content"), PromptInjectionSourceKind.CHAT_ATTACHMENT => TB("Chat attachment"), PromptInjectionSourceKind.RETRIEVAL_CONTEXT => TB("Retrieved context"), + PromptInjectionSourceKind.DATA_SOURCE_DESCRIPTION => TB("Data source description"), _ => TB("Unknown"), }; } \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceDescriptionService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceDescriptionService.cs new file mode 100644 index 00000000..4d7857ae --- /dev/null +++ b/app/MindWork AI Studio/Tools/Services/DataSourceDescriptionService.cs @@ -0,0 +1,120 @@ +using System.Collections.Concurrent; + +using AIStudio.Settings; +using AIStudio.Settings.DataModel; +using AIStudio.Tools.ERIClient; +using AIStudio.Tools.Security; + +namespace AIStudio.Tools.Services; + +/// +/// Tells a model what a data source holds, so that it can decide where to search. +/// +/// +/// Both the agent which selects data sources and Semantic Search describe the data sources to a +/// model. The user describes a local data source. An ERI data source is described by its server, +/// which costs two requests and is written by somebody else: that description is filtered for +/// prompt injections like any other external content before it is kept. It is kept for a few +/// minutes, because Semantic Search describes the data sources with every request it makes. +/// +public sealed class DataSourceDescriptionService(RustService rustService, PromptInjectionGuardService guardService, ILogger logger) +{ + private static readonly TimeSpan SERVER_DESCRIPTION_LIFETIME = TimeSpan.FromMinutes(5); + + // As long as the check of its security requirements may take, cf. DataSourceService: + private static readonly TimeSpan SERVER_TIMEOUT = TimeSpan.FromSeconds(6); + + /// + /// A description as the server sent it, filtered, together with the configuration it was asked with. + /// + private readonly record struct ServerDescription(IERIDataSource DataSource, string Description, DateTimeOffset ValidUntil); + + private readonly ConcurrentDictionary serverDescriptions = new(StringComparer.Ordinal); + + /// + /// What the data source holds, in a single line. + /// + /// The data source to describe. + /// The cancellation token. + /// The description, or an empty string when there is none or its server could not be asked. + public async Task GetDescriptionAsync(IDataSource dataSource, CancellationToken token = default) + { + var description = dataSource switch + { + DataSourceLocalDirectory localDirectory => localDirectory.Description, + DataSourceLocalFile localFile => localFile.Description, + IERIDataSource eriDataSource => await this.GetServerDescriptionAsync(eriDataSource, token), + _ => string.Empty, + }; + + // A description is written into a list, one data source per line: + return description.Replace("\n", " ").Replace("\r", " "); + } + + private async Task GetServerDescriptionAsync(IERIDataSource dataSource, CancellationToken token) + { + // + // A changed configuration, e.g., another server or another account, asks anew rather than + // waiting for the old description to expire: + // + if (this.serverDescriptions.TryGetValue(dataSource.Id, out var known) && known.DataSource.Equals(dataSource) && known.ValidUntil > DateTimeOffset.UtcNow) + return known.Description; + + var description = await this.FetchServerDescriptionAsync(dataSource, token); + if (description is null) + return string.Empty; + + // + // Only an answer is kept. A server which gave none is asked again next time; it is rarely + // asked at all, since a data source whose server cannot be reached is not offered anyway. + // + this.serverDescriptions[dataSource.Id] = new(dataSource, description, DateTimeOffset.UtcNow + SERVER_DESCRIPTION_LIFETIME); + return description; + } + + /// The filtered description, or null when the server could not be asked. + private async Task FetchServerDescriptionAsync(IERIDataSource dataSource, CancellationToken token) + { + try + { + using var timeout = CancellationTokenSource.CreateLinkedTokenSource(token); + timeout.CancelAfter(SERVER_TIMEOUT); + + using var eriClient = ERIClientFactory.Get(dataSource.Version, dataSource); + if (eriClient is null) + { + logger.LogWarning($"Could not create an ERI client for the data source '{dataSource.Name}'. Thus, we cannot retrieve the server description."); + return null; + } + + var authResponse = await eriClient.AuthenticateAsync(rustService, cancellationToken: timeout.Token); + if (!authResponse.Successful) + { + logger.LogWarning($"Was not able to authenticate with the ERI data source '{dataSource.Name}'. Message: {authResponse.Message}"); + return null; + } + + var serverDescriptionResponse = await eriClient.GetDataSourceInfoAsync(timeout.Token); + if (!serverDescriptionResponse.Successful) + { + logger.LogWarning($"Was not able to retrieve the server description from the ERI data source '{dataSource.Name}'. Message: {serverDescriptionResponse.Message}"); + return null; + } + + // + // Whoever runs the server writes this, and a model reads it as the description of + // where to search -- a fine place to tell it what to do instead: + // + return await guardService.SanitizeAsync(serverDescriptionResponse.Data.Description, PromptInjectionSource.DataSourceDescription(dataSource.Name)); + } + catch (OperationCanceledException) when (token.IsCancellationRequested) + { + throw; + } + catch (Exception e) + { + logger.LogWarning($"The ERI data source '{dataSource.Name}' is not available. Thus, we cannot retrieve the server description. Error: {e.Message}"); + return null; + } + } +} \ No newline at end of file