refactored confidence levels and ensured checks work

This commit is contained in:
PaulKoudelka committed 2026-08-14 12:03:16 +02:00
1 parent 786b106096
commit 88999d985c
44 files changed
+324 -335

No files matched your search

@@ -946,7 +946,7 @@ public sealed partial class DataSourceEmbeddingService
embeddingProvider.Hostname,
embeddingProvider.TokenizerPath,
embeddingProvider.EffectiveTokenLimit,
GetDataSourceComplianceLevel(dataSource).ToString(),
GetDataSourceConfidenceLevel(dataSource).ToString(),
dataSource is IInternalDataSource internalDataSource ? internalDataSource.MaxChunkTokenLength : 0,
dataSource is IInternalDataSource overlapDataSource ? overlapDataSource.ChunkOverlapTokenLength : 0,
chunkingOptions.MaxChunkTokenLength,
@@ -1050,7 +1050,7 @@ public sealed partial class DataSourceEmbeddingService
{
file.Refresh();
var absolutePath = Path.GetFullPath(file.FullName);
var complianceLevel = GetDataSourceComplianceLevel(dataSource);
var confidenceLevel = GetDataSourceConfidenceLevel(dataSource);
return new(
this.CreateParentFileId(dataSource.Id, absolutePath),
absolutePath,
@@ -1063,8 +1063,8 @@ public sealed partial class DataSourceEmbeddingService
file.Exists ? new DateTimeOffset(file.LastWriteTimeUtc) : DateTimeOffset.UnixEpoch,
embeddedAtUtc,
chunkCount,
complianceLevel.ToString(),
(int)complianceLevel);
confidenceLevel.ToString(),
(int)confidenceLevel);
}
private IReadOnlyList<EmbeddingStateChunk> CreateEmbeddingStateChunks(EmbeddingStateFile parentFile, IReadOnlyList<EmbeddingChunkDraft> batch, DateTimeOffset embeddedAtUtc)
@@ -1080,10 +1080,10 @@ public sealed partial class DataSourceEmbeddingService
.ToList();
}
private static ConfidenceLevel GetDataSourceComplianceLevel(IDataSource dataSource) =>
dataSource.ComplianceLevel is ConfidenceLevel.NONE
private static ConfidenceLevel GetDataSourceConfidenceLevel(IDataSource dataSource) =>
dataSource is not IInternalDataSource internalDataSource || internalDataSource.ConfidenceLevel is ConfidenceLevel.NONE
? ConfidenceLevel.UNKNOWN
: dataSource.ComplianceLevel;
: internalDataSource.ConfidenceLevel;
private static string GetFileType(FileInfo file)
{
@@ -393,6 +393,15 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
private async Task ProcessDataSourceAsync(IDataSource dataSource, DataSourceEmbeddingRefreshMode refreshMode, CancellationToken token)
{
if (dataSource is not IInternalDataSource internalDataSource)
{
logger.LogWarning(
"Skipping background embeddings for non-internal data source '{DataSourceName}' ({DataSourceId}).",
dataSource.Name,
dataSource.Id);
return;
}
logger.LogInformation(
"Starting background embedding hash check for data source '{DataSourceName}' ({DataSourceId}). RefreshMode={RefreshMode}.",
dataSource.Name,
@@ -451,17 +460,16 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
return;
}
if (!embeddingProvider.AllowsDataSourceAccess(settingsManager, dataSource.SecurityPolicy, dataSource.ComplianceLevel))
if (!embeddingProvider.GetConfidenceLevel(settingsManager).AllowsDataSourceConfidenceLevel(internalDataSource.ConfidenceLevel))
{
var errorMessage = $"The selected embedding provider is not allowed to embed this data source. The data source requires provider confidence '{dataSource.ComplianceLevel.GetName()}'. The embedding provider has confidence '{embeddingProvider.GetConfidenceLevel(settingsManager).GetName()}'.";
var errorMessage = $"The selected embedding provider is not allowed to embed this data source. The data source requires provider confidence '{internalDataSource.ConfidenceLevel.GetName()}'. The embedding provider has confidence '{embeddingProvider.GetConfidenceLevel(settingsManager).GetName()}'.";
logger.LogWarning(
"Skipping background embeddings for data source '{DataSourceName}' ({DataSourceId}) because embedding provider '{EmbeddingProviderName}' ({EmbeddingProviderId}) is not allowed. RequiredDataSecurity={RequiredDataSecurity}, RequiredCompliance={RequiredCompliance}, EmbeddingProviderConfidence={EmbeddingProviderConfidence}.",
"Skipping background embeddings for data source '{DataSourceName}' ({DataSourceId}) because embedding provider '{EmbeddingProviderName}' ({EmbeddingProviderId}) does not meet the required confidence. RequiredConfidence={RequiredConfidence}, EmbeddingProviderConfidence={EmbeddingProviderConfidence}.",
dataSource.Name,
dataSource.Id,
embeddingProvider.Name,
embeddingProvider.Id,
dataSource.SecurityPolicy,
dataSource.ComplianceLevel.GetName(),
internalDataSource.ConfidenceLevel.GetName(),
embeddingProvider.GetConfidenceLevel(settingsManager).GetName());
token.ThrowIfCancellationRequested();
@@ -877,8 +885,8 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
parentFile.CreationUtc,
parentFile.LastWriteUtc,
embeddedAtUtc,
parentFile.ComplianceLevel,
parentFile.ComplianceLevelRank)).ToList();
parentFile.ConfidenceLevel,
parentFile.ConfidenceLevelRank)).ToList();
await vectorStore.InsertEmbedding(collectionName, points, token);
}
@@ -38,8 +38,8 @@ public sealed class DataSourceLocalRetrievalService(
string Text,
double Score,
int Rank,
string ComplianceLevel,
int ComplianceLevelRank);
string ConfidenceLevel,
int ConfidenceLevelRank);
public Task<IReadOnlyList<IRetrievalContext>> RetrieveDataAsync(DataSourceLocalFile dataSource, IContent lastUserPrompt, ChatThread thread, CancellationToken token = default) =>
this.RetrieveDataAsync((IInternalDataSource)dataSource, lastUserPrompt, token);
@@ -296,8 +296,8 @@ public sealed class DataSourceLocalRetrievalService(
result.Text,
result.Score,
rank,
result.ComplianceLevel,
result.ComplianceLevelRank);
result.ConfidenceLevel,
result.ConfidenceLevelRank);
private static LocalRetrievalHit FromBm25Result(EmbeddingStateSearchResult result, int rank) =>
new(
@@ -316,8 +316,8 @@ public sealed class DataSourceLocalRetrievalService(
result.ChunkText,
result.Score,
rank,
result.ComplianceLevel,
result.ComplianceLevelRank);
result.ConfidenceLevel,
result.ConfidenceLevelRank);
private static RetrievalTextContext ToRetrievalContext(LocalRetrievalHit hit)
{
@@ -9,6 +9,8 @@ namespace AIStudio.Tools.Services;
public sealed class DataSourceService
{
private readonly record struct ParticipatingProvider(string Role, bool IsTrusted, ConfidenceLevel ConfidenceLevel);
private readonly RustService rustService;
private readonly SettingsManager settingsManager;
private readonly ILogger<DataSourceService> logger;
@@ -27,9 +29,10 @@ public sealed class DataSourceService
/// It also returns the data sources selected before when they are still allowed.
/// </summary>
/// <param name="selectedLLMProvider">The selected LLM provider.</param>
/// <param name="dataSourceOptions">The active data source options, which determine which agent providers participate.</param>
/// <param name="previousSelectedDataSources">The data sources selected before.</param>
/// <returns>The allowed data sources and the data sources selected before -- when they are still allowed.</returns>
public async Task<AllowedSelectedDataSources> GetDataSources(AIStudio.Settings.Provider selectedLLMProvider, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
public async Task<AllowedSelectedDataSources> GetDataSources(AIStudio.Settings.Provider selectedLLMProvider, DataSourceOptions dataSourceOptions, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
{
//
// Case: Somehow the selected LLM provider was not set. The default provider
@@ -42,10 +45,10 @@ public sealed class DataSourceService
return new([], []);
}
return await this.GetDataSources(
selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager),
selectedLLMProvider.GetConfidenceLevel(this.settingsManager),
previousSelectedDataSources);
var usingTrustedProvider = selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager);
var participatingProviders = this.GetParticipatingProviders(selectedLLMProvider.Id, dataSourceOptions,
new("chat provider", usingTrustedProvider, selectedLLMProvider.GetConfidenceLevel(this.settingsManager)));
return await this.GetDataSources(usingTrustedProvider, participatingProviders, previousSelectedDataSources);
}
/// <summary>
@@ -53,9 +56,10 @@ public sealed class DataSourceService
/// It also returns the data sources selected before when they are still allowed.
/// </summary>
/// <param name="selectedLLMProvider">The selected LLM provider.</param>
/// <param name="dataSourceOptions">The active data source options, which determine which agent providers participate.</param>
/// <param name="previousSelectedDataSources">The data sources selected before.</param>
/// <returns>The allowed data sources and the data sources selected before -- when they are still allowed.</returns>
public async Task<AllowedSelectedDataSources> GetDataSources(IProvider selectedLLMProvider, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
public async Task<AllowedSelectedDataSources> GetDataSources(IProvider selectedLLMProvider, DataSourceOptions dataSourceOptions, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
{
//
// Case: Somehow the selected LLM provider was not set. The default provider
@@ -68,13 +72,42 @@ public sealed class DataSourceService
return new([], []);
}
return await this.GetDataSources(
selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager),
selectedLLMProvider.GetConfidenceLevel(this.settingsManager),
previousSelectedDataSources);
var usingTrustedProvider = selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager);
var participatingProviders = this.GetParticipatingProviders(selectedLLMProvider.ConfiguredProviderId, dataSourceOptions,
new("chat provider", usingTrustedProvider, selectedLLMProvider.GetConfidenceLevel(this.settingsManager)));
return await this.GetDataSources(usingTrustedProvider, participatingProviders, previousSelectedDataSources);
}
private async Task<AllowedSelectedDataSources> GetDataSources(bool usingTrustedProvider, ConfidenceLevel providerConfidenceLevel, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
private IReadOnlyList<ParticipatingProvider> GetParticipatingProviders(string currentProviderId, DataSourceOptions dataSourceOptions, ParticipatingProvider currentProvider)
{
var providers = new List<ParticipatingProvider> { currentProvider };
if (dataSourceOptions.AutomaticDataSourceSelection)
this.AddAgentProvider(providers, Components.AGENT_DATA_SOURCE_SELECTION, currentProviderId, "data source selection agent");
if (dataSourceOptions.AutomaticValidation && this.settingsManager.ConfigurationData.AgentRetrievalContextValidation.EnableRetrievalContextValidation)
this.AddAgentProvider(providers, Components.AGENT_RETRIEVAL_CONTEXT_VALIDATION, currentProviderId, "retrieval context validation agent");
return providers;
}
private void AddAgentProvider(List<ParticipatingProvider> providers, Components component, string currentProviderId, string role)
{
var provider = this.settingsManager.GetPreselectedProvider(component, currentProviderId, true);
if (provider == Settings.Provider.NONE)
{
this.logger.LogWarning($"No provider is available for the {role}. Data sources cannot be made available while this agent is enabled.");
providers.Add(new(role, false, ConfidenceLevel.NONE));
return;
}
providers.Add(new(
role,
provider.IsTrustedForDataSourceSecurityChecks(this.settingsManager),
provider.GetConfidenceLevel(this.settingsManager)));
}
private async Task<AllowedSelectedDataSources> GetDataSources(bool usingTrustedProvider, IReadOnlyList<ParticipatingProvider> participatingProviders, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
{
var allDataSources = this.settingsManager.ConfigurationData.DataSources.ToList();
var previousSelectedDataSourceIds = previousSelectedDataSources?.Select(source => source.Id).ToHashSet(StringComparer.Ordinal) ?? [];
@@ -84,7 +117,7 @@ public sealed class DataSourceService
// Start all checks in parallel:
foreach (var source in allDataSources)
tasks.Add(this.CheckOneDataSource(source, usingTrustedProvider, providerConfidenceLevel));
tasks.Add(this.CheckOneDataSource(source, usingTrustedProvider, participatingProviders));
// Wait for all checks and collect the results:
foreach (var task in tasks)
@@ -101,16 +134,34 @@ public sealed class DataSourceService
return new(filteredDataSources, filteredSelectedDataSources);
}
private async Task<IDataSource?> CheckOneDataSource(IDataSource source, bool usingTrustedProvider, ConfidenceLevel providerConfidenceLevel)
private async Task<IDataSource?> CheckOneDataSource(IDataSource source, bool usingTrustedProvider, IReadOnlyList<ParticipatingProvider> participatingProviders)
{
if (!providerConfidenceLevel.AllowsDataSourceComplianceLevel(source.ComplianceLevel))
if (source is IInternalDataSource internalSource)
{
this.logger.LogWarning($"The data source '{source.Name}' (id={source.Id}) requires provider confidence '{source.ComplianceLevel.GetName()}'. The selected provider only has confidence '{providerConfidenceLevel.GetName()}'. We skip this source.");
return null;
}
foreach (var provider in participatingProviders)
{
if (!provider.ConfidenceLevel.AllowsDataSourceConfidenceLevel(internalSource.ConfidenceLevel))
{
this.logger.LogWarning($"The internal data source '{source.Name}' (id={source.Id}) requires provider confidence '{internalSource.ConfidenceLevel.GetName()}'. The {provider.Role} only has confidence '{provider.ConfidenceLevel.GetName()}'. We skip this source.");
return null;
}
}
if (!DataSourceEmbeddingProviders.TryResolve(this.settingsManager, source, out var embeddingProvider))
{
this.logger.LogWarning($"The internal data source '{source.Name}' (id={source.Id}) has no usable embedding provider. We skip this source.");
return null;
}
var embeddingProviderConfidence = embeddingProvider.GetConfidenceLevel(this.settingsManager);
if (!embeddingProviderConfidence.AllowsDataSourceConfidenceLevel(internalSource.ConfidenceLevel))
{
this.logger.LogWarning($"The internal data source '{source.Name}' (id={source.Id}) requires provider confidence '{internalSource.ConfidenceLevel.GetName()}'. Its embedding provider '{embeddingProvider.Name}' only has confidence '{embeddingProviderConfidence.GetName()}'. We skip this source.");
return null;
}
if (source is IInternalDataSource)
return source;
}
//
// Unfortunately, we have to live-check any ERI source for its security requirements.
@@ -146,8 +197,11 @@ public sealed class DataSourceService
eriSourceRequirements = securityRequest.Data;
this.logger.LogInformation($"Security requirements for ERI source '{source.Name}' (id={source.Id}) retrieved successfully.");
}
switch (source.SecurityPolicy)
if (source is not IExternalDataSource externalSource)
return source;
switch (externalSource.SecurityPolicy)
{
case DataSourceSecurity.ALLOW_ANY: