MindWorkAI/AI-Studio
MindWorkAI/AI-Studio
1
0
Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2026-09-27 00:13:38 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Share data source descriptions and filter the ones ERI servers send

Browse Source
This commit is contained in:
Thorsten Sommer 2026-09-24 16:24:47 +02:00
parent 31410e52f5
commit abb4a41a71
Signed by untrusted user who does not match committer: tsommer
GPG Key ID: 371BBA77A02C0108
6 changed files with 131 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
app/MindWork AI Studio/Agents/AgentDataSourceSelection.cs
View File

@ -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<AgentDataSourceSelection> logger, ILogger<AgentBase> baseLogger, SettingsManager settingsManager, DataSourceService dataSourceService, ThreadSafeRandom rng) : AgentBase(baseLogger, settingsManager, dataSourceService, rng)
public sealed class AgentDataSourceSelection (ILogger<AgentDataSourceSelection> logger, ILogger<AgentBase> baseLogger, SettingsManager settingsManager, DataSourceService dataSourceService, DataSourceDescriptionService descriptionService, ThreadSafeRandom rng) : AgentBase(baseLogger, settingsManager, dataSourceService, rng)
{
private readonly List<ContentBlock> answers = new();
@ -187,75 +186,24 @@ public sealed class AgentDataSourceSelection (ILogger<AgentDataSourceSelection>
var additionalData = new Dictionary<string, string>();
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<RustService>()!;
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;
}
}

1
app/MindWork AI Studio/Program.cs
View File

@ -201,6 +201,7 @@ internal sealed class Program
builder.Services.AddSingleton<UpdatePolicy>();
builder.Services.AddSingleton<AssistantPluginGenerationService>();
builder.Services.AddSingleton<DataSourceService>();
builder.Services.AddSingleton<DataSourceDescriptionService>();
builder.Services.AddSingleton<DataSourceEmbeddingService>();
builder.Services.AddSingleton<DataSourceLocalRetrievalService>();
builder.Services.AddSingleton<DirectChatService>();

2
app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs
View File

@ -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);
}

1
app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKind.cs
View File

@ -7,4 +7,5 @@ public enum PromptInjectionSourceKind
FILE_CONTENT,
CHAT_ATTACHMENT,
RETRIEVAL_CONTEXT,
DATA_SOURCE_DESCRIPTION,
}

1
app/MindWork AI Studio/Tools/Security/PromptInjectionSourceKindExtensions.cs
View File

@ -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"),
};
}

120
app/MindWork AI Studio/Tools/Services/DataSourceDescriptionService.cs Normal file
View File

@ -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;
/// <summary>
/// Tells a model what a data source holds, so that it can decide where to search.
/// </summary>
/// <remarks>
/// 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.
/// </remarks>
public sealed class DataSourceDescriptionService(RustService rustService, PromptInjectionGuardService guardService, ILogger<DataSourceDescriptionService> 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);
/// <summary>
/// A description as the server sent it, filtered, together with the configuration it was asked with.
/// </summary>
private readonly record struct ServerDescription(IERIDataSource DataSource, string Description, DateTimeOffset ValidUntil);
private readonly ConcurrentDictionary<string, ServerDescription> serverDescriptions = new(StringComparer.Ordinal);
/// <summary>
/// What the data source holds, in a single line.
/// </summary>
/// <param name="dataSource">The data source to describe.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The description, or an empty string when there is none or its server could not be asked.</returns>
public async Task<string> 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<string> 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;
}
/// <returns>The filtered description, or null when the server could not be asked.</returns>
private async Task<string?> 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;
}
}
}