mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-18 02:57:59 +00:00
Added a semantic search tool (#1005)
Build and Release / Determine run mode (push) Waiting to run
Build and Release / Read metadata (push) Blocked by required conditions
Build and Release / Sync Flatpak repo (push) Blocked by required conditions
Build and Release / Collect Flatpak artifacts (push) Blocked by required conditions
Build and Release / Verify (push) Waiting to run
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-aarch64-apple-darwin, osx-arm64, macos-latest, aarch64-apple-darwin, dmg,app,updater, dmg) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-aarch64-pc-windows-msvc.exe, win-arm64, windows-latest, aarch64-pc-windows-msvc, nsis,updater, nsis) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-aarch64-unknown-linux-gnu, linux-arm64, ubuntu-22.04-arm, aarch64-unknown-linux-gnu, appimage,updater, appimage) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-x86_64-apple-darwin, osx-x64, macos-latest, x86_64-apple-darwin, dmg,app,updater, dmg) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-x86_64-pc-windows-msvc.exe, win-x64, windows-latest, x86_64-pc-windows-msvc, nsis,updater, nsis) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-x86_64-unknown-linux-gnu, linux-x64, ubuntu-22.04, x86_64-unknown-linux-gnu, appimage,updater, appimage) (push) Blocked by required conditions
Build and Release / Prepare & create release (push) Blocked by required conditions
Build and Release / Publish release (push) Blocked by required conditions
Build and Release / Determine run mode (push) Waiting to run
Build and Release / Read metadata (push) Blocked by required conditions
Build and Release / Sync Flatpak repo (push) Blocked by required conditions
Build and Release / Collect Flatpak artifacts (push) Blocked by required conditions
Build and Release / Verify (push) Waiting to run
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-aarch64-apple-darwin, osx-arm64, macos-latest, aarch64-apple-darwin, dmg,app,updater, dmg) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-aarch64-pc-windows-msvc.exe, win-arm64, windows-latest, aarch64-pc-windows-msvc, nsis,updater, nsis) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-aarch64-unknown-linux-gnu, linux-arm64, ubuntu-22.04-arm, aarch64-unknown-linux-gnu, appimage,updater, appimage) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-x86_64-apple-darwin, osx-x64, macos-latest, x86_64-apple-darwin, dmg,app,updater, dmg) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-x86_64-pc-windows-msvc.exe, win-x64, windows-latest, x86_64-pc-windows-msvc, nsis,updater, nsis) (push) Blocked by required conditions
Build and Release / Build app (${{ matrix.dotnet_runtime }}) (-x86_64-unknown-linux-gnu, linux-x64, ubuntu-22.04, x86_64-unknown-linux-gnu, appimage,updater, appimage) (push) Blocked by required conditions
Build and Release / Prepare & create release (push) Blocked by required conditions
Build and Release / Publish release (push) Blocked by required conditions
This commit is contained in:
1 parent
4baf21656a
commit
1eaca9b12f
91 files changed
+4248
-487
No files matched your search
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -54,24 +54,62 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
int Rank);
|
||||
// ReSharper restore NotAccessedPositionalProperty.Local
|
||||
|
||||
/// <summary>
|
||||
/// What kept one retrieval from covering the whole data source.
|
||||
/// </summary>
|
||||
/// <param name="queryWrittenByUser">Whether the query is the user's own message, which decides who hears about its problems.</param>
|
||||
private sealed class RetrievalRun(bool queryWrittenByUser)
|
||||
{
|
||||
// Both channels search at the same time:
|
||||
private readonly Lock gapLock = new();
|
||||
private readonly HashSet<RetrievalGap> gaps = [];
|
||||
|
||||
public bool QueryWrittenByUser => queryWrittenByUser;
|
||||
|
||||
public void Add(RetrievalGap gap)
|
||||
{
|
||||
lock (this.gapLock)
|
||||
this.gaps.Add(gap);
|
||||
}
|
||||
|
||||
public IReadOnlyList<RetrievalGap> GetGaps()
|
||||
{
|
||||
lock (this.gapLock)
|
||||
return this.gaps.Order().ToList();
|
||||
}
|
||||
}
|
||||
|
||||
public Task<IReadOnlyList<IRetrievalContext>> RetrieveDataAsync(DataSourceLocalFile dataSource, IContent lastUserPrompt, ChatThread thread, CancellationToken token = default) =>
|
||||
this.RetrieveDataAsync(dataSource, lastUserPrompt, token);
|
||||
|
||||
public Task<IReadOnlyList<IRetrievalContext>> RetrieveDataAsync(DataSourceLocalDirectory dataSource, IContent lastUserPrompt, ChatThread thread, CancellationToken token = default) =>
|
||||
this.RetrieveDataAsync(dataSource, lastUserPrompt, token);
|
||||
|
||||
public Task<RetrievalPage> RetrieveDataAsync(DataSourceLocalFile dataSource, string query, int page, ChatThread thread, CancellationToken token = default) =>
|
||||
this.RetrievePageAsync(dataSource, query, page, new RetrievalRun(queryWrittenByUser: false), token);
|
||||
|
||||
public Task<RetrievalPage> RetrieveDataAsync(DataSourceLocalDirectory dataSource, string query, int page, ChatThread thread, CancellationToken token = default) =>
|
||||
this.RetrievePageAsync(dataSource, query, page, new RetrievalRun(queryWrittenByUser: false), token);
|
||||
|
||||
private async Task<IReadOnlyList<IRetrievalContext>> RetrieveDataAsync(IInternalDataSource dataSource, IContent lastUserPrompt, CancellationToken token)
|
||||
{
|
||||
var query = GetQueryText(lastUserPrompt);
|
||||
// The first page is what this retrieval has always returned:
|
||||
var firstPage = await this.RetrievePageAsync(dataSource, GetQueryText(lastUserPrompt), 1, new RetrievalRun(queryWrittenByUser: true), token);
|
||||
return firstPage.Contexts;
|
||||
}
|
||||
|
||||
private async Task<RetrievalPage> RetrievePageAsync(IInternalDataSource dataSource, string query, int page, RetrievalRun run, CancellationToken token)
|
||||
{
|
||||
var pageSize = (int)dataSource.MaxMatches;
|
||||
var window = RetrievalPaging.GetWindowSize(page, pageSize);
|
||||
if (string.IsNullOrWhiteSpace(query))
|
||||
{
|
||||
logger.LogDebug("Skipping local retrieval for data source '{DataSourceName}' ({DataSourceId}) because the latest prompt does not contain text.", dataSource.Name, dataSource.Id);
|
||||
return [];
|
||||
logger.LogDebug("Skipping local retrieval for data source '{DataSourceName}' ({DataSourceId}) because there is no text to search for.", dataSource.Name, dataSource.Id);
|
||||
return RetrievalPage.EMPTY;
|
||||
}
|
||||
|
||||
var maxMatches = (int)dataSource.MaxMatches;
|
||||
if (maxMatches == 0)
|
||||
return [];
|
||||
if (pageSize == 0)
|
||||
return RetrievalPage.EMPTY;
|
||||
|
||||
//
|
||||
// A data source waiting for its index is kept out of the selection before the RAG process
|
||||
@@ -86,31 +124,43 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
if (await embeddingService.IsAwaitingReindexAsync(dataSource, token))
|
||||
{
|
||||
logger.LogWarning("Skipping local retrieval for data source '{DataSourceName}' ({DataSourceId}) because its index has to be built anew.", dataSource.Name, dataSource.Id);
|
||||
await this.ReportRetrievalGapAsync(dataSource, "index-rebuilding", string.Format(TB("The data source '{0}' was left out of the answer: it is being indexed again and cannot be searched until that is finished."), dataSource.Name));
|
||||
return [];
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.NOT_SEARCHED, "index-rebuilding", string.Format(TB("The data source '{0}' was left out of the answer: it is being indexed again and cannot be searched until that is finished."), dataSource.Name));
|
||||
return RetrievalPage.EMPTY with { Gaps = run.GetGaps() };
|
||||
}
|
||||
|
||||
var collectionName = DataSourceEmbeddingNames.GetCollectionName(dataSource.Id);
|
||||
var vectorTask = this.SearchVectorAsync(dataSource, query, maxMatches, collectionName, token);
|
||||
var bm25Task = this.SearchBm25Async(dataSource, query, maxMatches, token);
|
||||
var vectorTask = this.SearchVectorAsync(dataSource, query, window, collectionName, run, token);
|
||||
var bm25Task = this.SearchBm25Async(dataSource, query, window, run, token);
|
||||
|
||||
await Task.WhenAll(vectorTask, bm25Task);
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
var hits = MergeResults(vectorTask.Result, bm25Task.Result, maxMatches);
|
||||
var (hits, hasMore) = RetrievalPaging.Merge(
|
||||
vectorTask.Result.Select((result, index) => FromVectorResult(result, index + 1)).ToList(),
|
||||
bm25Task.Result.Select((result, index) => FromBm25Result(result, index + 1)).ToList(),
|
||||
hit => hit.ChunkId,
|
||||
page,
|
||||
pageSize);
|
||||
|
||||
var gaps = run.GetGaps();
|
||||
logger.LogInformation(
|
||||
"Retrieved {MergedHits} local RAG hits for data source '{DataSourceName}' ({DataSourceId}). VectorCandidates={VectorHits}, BM25Candidates={BM25Hits}, RequestedPerChannel={RequestedPerChannel}.",
|
||||
"Retrieved {MergedHits} local RAG hits on page {Page} for data source '{DataSourceName}' ({DataSourceId}). VectorCandidates={VectorHits}, BM25Candidates={BM25Hits}, RequestedPerChannel={RequestedPerChannel}, HasMore={HasMore}, Gaps=[{Gaps}].",
|
||||
hits.Count,
|
||||
page,
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
vectorTask.Result.Count,
|
||||
bm25Task.Result.Count,
|
||||
maxMatches);
|
||||
window,
|
||||
hasMore,
|
||||
string.Join(", ", gaps));
|
||||
|
||||
return hits
|
||||
var contexts = hits
|
||||
.Where(hit => !string.IsNullOrWhiteSpace(hit.Text))
|
||||
.Select(hit => ToRetrievalContext(hit, dataSource))
|
||||
.ToList();
|
||||
|
||||
return new RetrievalPage(contexts, hasMore) { Gaps = gaps };
|
||||
}
|
||||
|
||||
private async Task<IReadOnlyList<VectorSearchResult>> SearchVectorAsync(
|
||||
@@ -118,6 +168,7 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
string query,
|
||||
int maxMatches,
|
||||
string collectionName,
|
||||
RetrievalRun run,
|
||||
CancellationToken token)
|
||||
{
|
||||
try
|
||||
@@ -130,18 +181,18 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
vectorStore.Name);
|
||||
await this.ReportRetrievalGapAsync(dataSource, "no-vector-store", string.Format(TB("The data source '{0}' was left out of the answer: its local index is not available."), dataSource.Name));
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.PARTLY_SEARCHED, "no-vector-store", string.Format(TB("The data source '{0}' was left out of the answer: its local index is not available."), dataSource.Name));
|
||||
return [];
|
||||
}
|
||||
|
||||
if (!DataSourceEmbeddingProviders.TryResolve(settingsManager, dataSource, out var embeddingProvider))
|
||||
{
|
||||
logger.LogWarning("Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the selected embedding provider is not available.", dataSource.Name, dataSource.Id);
|
||||
await this.ReportRetrievalGapAsync(dataSource, "no-embedding-provider", string.Format(TB("The data source '{0}' was left out of the answer: its embedding provider is not available. Please check it in the settings."), dataSource.Name));
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.PARTLY_SEARCHED, "no-embedding-provider", string.Format(TB("The data source '{0}' was left out of the answer: its embedding provider is not available. Please check it in the settings."), dataSource.Name));
|
||||
return [];
|
||||
}
|
||||
|
||||
if (!await this.QueryFitsEmbeddingProviderAsync(dataSource, embeddingProvider, query, token))
|
||||
if (!await this.QueryFitsEmbeddingProviderAsync(dataSource, embeddingProvider, query, run, token))
|
||||
return [];
|
||||
|
||||
var provider = embeddingProvider.CreateProvider();
|
||||
@@ -151,7 +202,7 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
if (vector is null || vector.Count == 0)
|
||||
{
|
||||
logger.LogWarning("Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because query embedding returned no vector.", dataSource.Name, dataSource.Id);
|
||||
await this.ReportRetrievalGapAsync(dataSource, "no-query-vector", string.Format(TB("The data source '{0}' was left out of the answer: its embedding provider '{1}' did not return a vector for your message."), dataSource.Name, embeddingProvider.Name));
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.PARTLY_SEARCHED, "no-query-vector", string.Format(TB("The data source '{0}' was left out of the answer: its embedding provider '{1}' did not return a vector to search with."), dataSource.Name, embeddingProvider.Name));
|
||||
return [];
|
||||
}
|
||||
|
||||
@@ -177,7 +228,7 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
exception,
|
||||
"Vector retrieval failed for data source '{DataSourceName}' ({DataSourceId}) because the embedding provider failed. FailureReason={FailureReason}, StatusCode={StatusCode}.",
|
||||
dataSource.Name, dataSource.Id, exception.FailureReason, exception.StatusCode);
|
||||
await this.ReportRetrievalGapAsync(dataSource, $"provider-{exception.FailureReason}", string.Format(TB("The data source '{0}' was left out of the answer. {1}"), dataSource.Name, exception.UserMessage));
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.PARTLY_SEARCHED, $"provider-{exception.FailureReason}", string.Format(TB("The data source '{0}' was left out of the answer. {1}"), dataSource.Name, exception.UserMessage));
|
||||
return [];
|
||||
}
|
||||
catch (VectorStoreUnreadableException exception)
|
||||
@@ -188,31 +239,40 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
// answer into one the user can do something about.
|
||||
//
|
||||
logger.LogWarning(exception, "Vector retrieval failed for data source '{DataSourceName}' ({DataSourceId}) because its vector store cannot be read.", dataSource.Name, dataSource.Id);
|
||||
await this.ReportRetrievalGapAsync(dataSource, "vector-store-unreadable", string.Format(TB("The data source '{0}' was left out of the answer: its index cannot be read anymore. You can repair it in your data source settings."), dataSource.Name));
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.PARTLY_SEARCHED, "vector-store-unreadable", string.Format(TB("The data source '{0}' was left out of the answer: its index cannot be read anymore. You can repair it in your data source settings."), dataSource.Name));
|
||||
return [];
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Vector retrieval failed for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
await this.ReportRetrievalGapAsync(dataSource, "vector-search-failed", string.Format(TB("The data source '{0}' was left out of the answer because searching it failed."), dataSource.Name));
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.PARTLY_SEARCHED, "vector-search-failed", string.Format(TB("The data source '{0}' was left out of the answer because searching it failed."), dataSource.Name));
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Tells the user once that a data source cannot take part in answering.
|
||||
/// Records that a data source cannot fully take part in answering, and tells the user once.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A failed search is not an error of the chat: the model still answers, only without what
|
||||
/// this data source knows. Saying so once is what keeps somebody from trusting an answer
|
||||
/// which was put together without half of its sources. Saying it with every prompt would be
|
||||
/// worse than saying nothing, which is why every gap is reported once per session.
|
||||
///
|
||||
/// The retrieval records every gap regardless, cf. RetrievalPage.Gaps: whoever asked for the
|
||||
/// page has to know each time, not once per session.
|
||||
/// </remarks>
|
||||
/// <param name="dataSource">The data source which could not be searched.</param>
|
||||
/// <param name="run">The retrieval this gap belongs to.</param>
|
||||
/// <param name="gap">What the gap means for the search.</param>
|
||||
/// <param name="gapKey">What kind of gap this is, so a different problem is reported again.</param>
|
||||
/// <param name="userMessage">What to tell the user.</param>
|
||||
private async Task ReportRetrievalGapAsync(IInternalDataSource dataSource, string gapKey, string userMessage)
|
||||
private async Task ReportRetrievalGapAsync(IInternalDataSource dataSource, RetrievalRun run, RetrievalGap gap, string gapKey, string userMessage)
|
||||
{
|
||||
run.Add(gap);
|
||||
if (!IsForTheUser(gap, run.QueryWrittenByUser))
|
||||
return;
|
||||
|
||||
lock (this.retrievalGapLock)
|
||||
{
|
||||
if (!this.reportedRetrievalGaps.Add($"{dataSource.Id}::{gapKey}"))
|
||||
@@ -222,23 +282,38 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
await MessageBus.INSTANCE.SendWarning(new(Icons.Material.Filled.SearchOff, userMessage));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Whether the user has to hear about a gap.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Problems of the data source are for the user, since only the user can fix them. Problems of
|
||||
/// the query are for whoever wrote it. When the model worked the query out, telling the user
|
||||
/// their message was too long would be wrong, and the model learns about it from the page and
|
||||
/// can search with a shorter one.
|
||||
/// </remarks>
|
||||
/// <param name="gap">What the gap means for the search.</param>
|
||||
/// <param name="queryWrittenByUser">Whether the query is the user's own message.</param>
|
||||
/// <returns>True when the user has to be told.</returns>
|
||||
internal static bool IsForTheUser(RetrievalGap gap, bool queryWrittenByUser) => gap is not RetrievalGap.QUERY_NOT_SEARCHABLE || queryWrittenByUser;
|
||||
|
||||
private async Task<bool> QueryFitsEmbeddingProviderAsync(
|
||||
IInternalDataSource dataSource,
|
||||
EmbeddingProvider embeddingProvider,
|
||||
string query,
|
||||
RetrievalRun run,
|
||||
CancellationToken token)
|
||||
{
|
||||
var providerTokenLimit = Math.Max(1, embeddingProvider.EffectiveTokenLimit);
|
||||
if (query.Length > RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH)
|
||||
{
|
||||
logger.LogWarning(
|
||||
"Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the latest prompt has {CharacterCount} characters and exceeds the safe tokenizer request length of {MaxCharacterCount}. ProviderTokenLimit={ProviderTokenLimit}.",
|
||||
"Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the query has {CharacterCount} characters and exceeds the safe tokenizer request length of {MaxCharacterCount}. ProviderTokenLimit={ProviderTokenLimit}.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
query.Length,
|
||||
RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH,
|
||||
providerTokenLimit);
|
||||
await this.ReportRetrievalGapAsync(dataSource, "query-too-long", string.Format(TB("The data source '{0}' was left out of the answer because your message is too long to search with."), dataSource.Name));
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.QUERY_NOT_SEARCHABLE, "query-too-long", string.Format(TB("The data source '{0}' was left out of the answer because your message is too long to search with."), dataSource.Name));
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -251,7 +326,7 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
dataSource.Id,
|
||||
embeddingProvider.Name,
|
||||
tokenCountResponse?.Message ?? "No response was returned by the tokenizer service.");
|
||||
await this.ReportRetrievalGapAsync(dataSource, "no-token-count", string.Format(TB("The data source '{0}' was left out of the answer: the tokenizer of its embedding provider '{1}' is not available."), dataSource.Name, embeddingProvider.Name));
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.PARTLY_SEARCHED, "no-token-count", string.Format(TB("The data source '{0}' was left out of the answer: the tokenizer of its embedding provider '{1}' is not available."), dataSource.Name, embeddingProvider.Name));
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -259,20 +334,20 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
if (queryTokenCount > providerTokenLimit)
|
||||
{
|
||||
logger.LogWarning(
|
||||
"Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the latest prompt has {QueryTokenCount} tokens, exceeding embedding provider '{EmbeddingProviderName}' limit of {ProviderTokenLimit} tokens.",
|
||||
"Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the query has {QueryTokenCount} tokens, exceeding embedding provider '{EmbeddingProviderName}' limit of {ProviderTokenLimit} tokens.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
queryTokenCount,
|
||||
embeddingProvider.Name,
|
||||
providerTokenLimit);
|
||||
await this.ReportRetrievalGapAsync(dataSource, "query-over-token-limit", string.Format(TB("The data source '{0}' was left out of the answer because your message is longer than its embedding provider '{1}' accepts."), dataSource.Name, embeddingProvider.Name));
|
||||
await this.ReportRetrievalGapAsync(dataSource, run, RetrievalGap.QUERY_NOT_SEARCHABLE, "query-over-token-limit", string.Format(TB("The data source '{0}' was left out of the answer because your message is longer than its embedding provider '{1}' accepts."), dataSource.Name, embeddingProvider.Name));
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
private async Task<IReadOnlyList<IndexStoreSearchResult>> SearchBm25Async(IInternalDataSource dataSource, string query, int maxMatches, CancellationToken token)
|
||||
private async Task<IReadOnlyList<IndexStoreSearchResult>> SearchBm25Async(IInternalDataSource dataSource, string query, int maxMatches, RetrievalRun run, CancellationToken token)
|
||||
{
|
||||
try
|
||||
{
|
||||
@@ -284,6 +359,7 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
indexStore.Name);
|
||||
run.Add(RetrievalGap.PARTLY_SEARCHED);
|
||||
return [];
|
||||
}
|
||||
|
||||
@@ -302,6 +378,7 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "BM25 retrieval failed for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
|
||||
run.Add(RetrievalGap.PARTLY_SEARCHED);
|
||||
return [];
|
||||
}
|
||||
}
|
||||
@@ -312,7 +389,7 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
return results;
|
||||
|
||||
logger.LogWarning(
|
||||
"Local RAG {SearchName} search returned {ReturnedHits} chunks for data source '{DataSourceName}' ({DataSourceId}), which exceeds the configured maximum {MaxMatches}. Truncating to the datasource limit.",
|
||||
"Local RAG {SearchName} search returned {ReturnedHits} chunks for data source '{DataSourceName}' ({DataSourceId}), which exceeds the requested maximum {MaxMatches}. Truncating to it.",
|
||||
searchName,
|
||||
results.Count,
|
||||
dataSource.Name,
|
||||
@@ -322,47 +399,6 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
return results.Take(maxMatches).ToList();
|
||||
}
|
||||
|
||||
private static IReadOnlyList<LocalRetrievalHit> MergeResults(
|
||||
IReadOnlyList<VectorSearchResult> vectorResults,
|
||||
IReadOnlyList<IndexStoreSearchResult> bm25Results,
|
||||
int maxMatches)
|
||||
{
|
||||
// Future reranking should replace this deterministic channel merge.
|
||||
var merged = new List<LocalRetrievalHit>(maxMatches * 2);
|
||||
var seenChunkIds = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
AppendHits(
|
||||
merged,
|
||||
seenChunkIds,
|
||||
vectorResults
|
||||
.Select((result, index) => FromVectorResult(result, index + 1)),
|
||||
maxMatches);
|
||||
|
||||
AppendHits(
|
||||
merged,
|
||||
seenChunkIds,
|
||||
bm25Results
|
||||
.Select((result, index) => FromBm25Result(result, index + 1)),
|
||||
maxMatches);
|
||||
|
||||
return merged;
|
||||
}
|
||||
|
||||
private static void AppendHits(List<LocalRetrievalHit> merged, HashSet<string> seenChunkIds, IEnumerable<LocalRetrievalHit> hits, int maxNewHits)
|
||||
{
|
||||
var added = 0;
|
||||
foreach (var hit in hits)
|
||||
{
|
||||
if (!string.IsNullOrWhiteSpace(hit.ChunkId) && !seenChunkIds.Add(hit.ChunkId))
|
||||
continue;
|
||||
|
||||
merged.Add(hit);
|
||||
added++;
|
||||
if (added >= maxNewHits)
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
private static LocalRetrievalHit FromVectorResult(VectorSearchResult result, int rank) =>
|
||||
new(
|
||||
RetrievalChannel.VECTOR,
|
||||
|
||||
@@ -39,9 +39,10 @@ public sealed class DataSourceService
|
||||
/// </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="retrievalMode">How the data sources are searched in effect, which decides whether any agent participates at all.</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, DataSourceOptions dataSourceOptions, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
|
||||
public async Task<AllowedSelectedDataSources> GetDataSources(AIStudio.Settings.Provider selectedLLMProvider, DataSourceOptions dataSourceOptions, DataSourceRetrievalMode retrievalMode, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
|
||||
{
|
||||
//
|
||||
// Case: Somehow the selected LLM provider was not set. The default provider
|
||||
@@ -55,7 +56,7 @@ public sealed class DataSourceService
|
||||
}
|
||||
|
||||
var usingTrustedProvider = selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager);
|
||||
var participatingProviders = this.GetParticipatingProviders(selectedLLMProvider.Id, dataSourceOptions,
|
||||
var participatingProviders = this.GetParticipatingProviders(selectedLLMProvider.Id, dataSourceOptions, retrievalMode,
|
||||
new("chat provider", usingTrustedProvider, selectedLLMProvider.GetConfidenceLevel(this.settingsManager)));
|
||||
return await this.GetDataSources(usingTrustedProvider, participatingProviders, previousSelectedDataSources);
|
||||
}
|
||||
@@ -67,9 +68,10 @@ public sealed class DataSourceService
|
||||
/// </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="retrievalMode">How the data sources are searched in effect, which decides whether any agent participates at all.</param>
|
||||
/// <param name="requestedDataSources">The data sources to check.</param>
|
||||
/// <returns>The requested data sources that are allowed for the provider.</returns>
|
||||
public async Task<IReadOnlyList<IDataSource>> GetAllowedDataSources(AIStudio.Settings.Provider selectedLLMProvider, DataSourceOptions dataSourceOptions, IReadOnlyCollection<IDataSource> requestedDataSources)
|
||||
public async Task<IReadOnlyList<IDataSource>> GetAllowedDataSources(AIStudio.Settings.Provider selectedLLMProvider, DataSourceOptions dataSourceOptions, DataSourceRetrievalMode retrievalMode, IReadOnlyCollection<IDataSource> requestedDataSources)
|
||||
{
|
||||
if (selectedLLMProvider == Settings.Provider.NONE)
|
||||
{
|
||||
@@ -78,7 +80,7 @@ public sealed class DataSourceService
|
||||
}
|
||||
|
||||
var usingTrustedProvider = selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager);
|
||||
var participatingProviders = this.GetParticipatingProviders(selectedLLMProvider.Id, dataSourceOptions,
|
||||
var participatingProviders = this.GetParticipatingProviders(selectedLLMProvider.Id, dataSourceOptions, retrievalMode,
|
||||
new("chat provider", usingTrustedProvider, selectedLLMProvider.GetConfidenceLevel(this.settingsManager)));
|
||||
var allowedDataSources = await this.GetAllowedDataSources(usingTrustedProvider, participatingProviders, requestedDataSources);
|
||||
|
||||
@@ -98,9 +100,10 @@ public sealed class DataSourceService
|
||||
/// </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="retrievalMode">How the data sources are searched in effect, which decides whether any agent participates at all.</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, DataSourceOptions dataSourceOptions, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
|
||||
public async Task<AllowedSelectedDataSources> GetDataSources(IProvider selectedLLMProvider, DataSourceOptions dataSourceOptions, DataSourceRetrievalMode retrievalMode, IReadOnlyCollection<IDataSource>? previousSelectedDataSources = null)
|
||||
{
|
||||
//
|
||||
// Case: Somehow the selected LLM provider was not set. The default provider
|
||||
@@ -114,24 +117,49 @@ public sealed class DataSourceService
|
||||
}
|
||||
|
||||
var usingTrustedProvider = selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager);
|
||||
var participatingProviders = this.GetParticipatingProviders(selectedLLMProvider.ConfiguredProviderId, dataSourceOptions,
|
||||
var participatingProviders = this.GetParticipatingProviders(selectedLLMProvider.ConfiguredProviderId, dataSourceOptions, retrievalMode,
|
||||
new("chat provider", usingTrustedProvider, selectedLLMProvider.GetConfidenceLevel(this.settingsManager)));
|
||||
return await this.GetDataSources(usingTrustedProvider, participatingProviders, previousSelectedDataSources);
|
||||
}
|
||||
|
||||
private IReadOnlyList<ParticipatingProvider> GetParticipatingProviders(string currentProviderId, DataSourceOptions dataSourceOptions, ParticipatingProvider currentProvider)
|
||||
private IReadOnlyList<ParticipatingProvider> GetParticipatingProviders(string currentProviderId, DataSourceOptions dataSourceOptions, DataSourceRetrievalMode retrievalMode, 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");
|
||||
var retrievalContextValidationEnabled = this.settingsManager.ConfigurationData.AgentRetrievalContextValidation.EnableRetrievalContextValidation;
|
||||
foreach (var (component, role) in GetParticipatingAgents(dataSourceOptions, retrievalMode, retrievalContextValidationEnabled))
|
||||
this.AddAgentProvider(providers, component, currentProviderId, role);
|
||||
|
||||
return providers;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Which agents get to see the data of the data sources, besides the chat provider.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Only the classic RAG process runs agents. With Semantic Search, the chat model picks the
|
||||
/// data sources and judges what it found itself, so the data reaches no other provider.
|
||||
/// Counting the providers of the agents there would hold back data sources which the chat
|
||||
/// provider alone may use.
|
||||
/// </remarks>
|
||||
/// <param name="dataSourceOptions">The active data source options.</param>
|
||||
/// <param name="retrievalMode">How the data sources are searched in effect.</param>
|
||||
/// <param name="retrievalContextValidationEnabled">Whether the validation of retrieval contexts is enabled in the settings.</param>
|
||||
/// <returns>The component of each participating agent, together with its role for the log.</returns>
|
||||
internal static IReadOnlyList<(Components Component, string Role)> GetParticipatingAgents(DataSourceOptions dataSourceOptions, DataSourceRetrievalMode retrievalMode, bool retrievalContextValidationEnabled)
|
||||
{
|
||||
if (retrievalMode is DataSourceRetrievalMode.SEMANTIC_SEARCH)
|
||||
return [];
|
||||
|
||||
var agents = new List<(Components Component, string Role)>(2);
|
||||
if (dataSourceOptions.AutomaticDataSourceSelection)
|
||||
agents.Add((Components.AGENT_DATA_SOURCE_SELECTION, "data source selection agent"));
|
||||
|
||||
if (dataSourceOptions.AutomaticValidation && retrievalContextValidationEnabled)
|
||||
agents.Add((Components.AGENT_RETRIEVAL_CONTEXT_VALIDATION, "retrieval context validation agent"));
|
||||
|
||||
return agents;
|
||||
}
|
||||
|
||||
private void AddAgentProvider(List<ParticipatingProvider> providers, Components component, string currentProviderId, string role)
|
||||
{
|
||||
var provider = this.settingsManager.GetPreselectedProvider(component, currentProviderId, true);
|
||||
|
||||
@@ -194,6 +194,7 @@ public sealed class DirectChatService(SettingsManager settingsManager, DataSourc
|
||||
DisableDataSources = false,
|
||||
AutomaticDataSourceSelection = false,
|
||||
AutomaticValidation = standardOptions.AutomaticValidation,
|
||||
RetrievalMode = standardOptions.RetrievalMode,
|
||||
PreselectedDataSourceIds = launcherDataSourceIds.Select(dataSourceId => dataSourceId.ToString()).ToList(),
|
||||
};
|
||||
}
|
||||
@@ -272,9 +273,11 @@ public sealed class DirectChatService(SettingsManager settingsManager, DataSourc
|
||||
//
|
||||
// The options the launched chat will run under are what this check runs against: they
|
||||
// decide which agent providers take part, and an agent with too little confidence makes
|
||||
// a data source unavailable.
|
||||
// a data source unavailable. Agents take part only when the chat searches with every
|
||||
// message, which the chat decides the same way:
|
||||
//
|
||||
availableDataSources = await dataSourceService.GetAllowedDataSources(provider, chosenOptions, requestedDataSources);
|
||||
var retrievalMode = await toolRegistry.GetEffectiveRetrievalModeAsync(chosenOptions, provider, Components.CHAT);
|
||||
availableDataSources = await dataSourceService.GetAllowedDataSources(provider, chosenOptions, retrievalMode.Mode, requestedDataSources);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
|
||||
Reference in new issue
Block a user