MindWorkAI/AI-Studio
1
0 Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2026-10-06 18:09:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added mailboxes as local data sources (#1022)
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-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 / Prepare & create release (push) Blocked by required conditions
Build and Release / Publish release (push) Blocked by required conditions
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 }}) (-x86_64-unknown-linux-gnu, linux-x64, ubuntu-22.04, x86_64-unknown-linux-gnu, appimage,updater, appimage) (push) Blocked by required conditions

Browse Source
This commit is contained in:
Thorsten Sommer authored and GitHub committed 2026-10-04 12:26:26 +02:00
1 parent 93be539e50
commit c4400c0ff6
343 files changed
+26051 -3122
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs
+7 -281
View File
@@ -1,8 +1,6 @@
using AIStudio.Chat;
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.Databases;
using AIStudio.Tools.Databases.IndexStore;
using AIStudio.Tools.Databases.VectorStore;
using AIStudio.Tools.PluginSystem;
@@ -11,19 +9,10 @@ using AIStudio.Tools.Rust;
namespace AIStudio.Tools.Services;
public sealed class DataSourceLocalRetrievalService(
SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider,
DataSourceEmbeddingService embeddingService, ILogger<DataSourceLocalRetrievalService> logger)
public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService indexSearch, ILogger<DataSourceLocalRetrievalService> logger)
{
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DataSourceLocalRetrievalService).Namespace, nameof(DataSourceLocalRetrievalService));
//
// Which gaps the user was already told about in this session. Retrieval runs for every single
// message, so without this one broken embedding provider would put a warning on every prompt.
//
private readonly HashSet<string> reportedRetrievalGaps = new(StringComparer.Ordinal);
private readonly Lock retrievalGapLock = new();
private enum RetrievalChannel
{
VECTOR,
@@ -54,31 +43,6 @@ 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);
@@ -121,20 +85,18 @@ public sealed class DataSourceLocalRetrievalService(
// chunks are still in place and the keyword search would happily answer from them while
// the vector search finds nothing.
//
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, 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));
if (await indexSearch.IsAwaitingReindexAsync(dataSource, run, token))
return RetrievalPage.EMPTY with { Gaps = run.GetGaps() };
}
var collectionName = DataSourceEmbeddingNames.GetCollectionName(dataSource.Id);
var vectorTask = this.SearchVectorAsync(dataSource, query, window, collectionName, run, token);
var bm25Task = this.SearchBm25Async(dataSource, query, window, run, token);
var vectorTask = indexSearch.SearchVectorsAsync(dataSource, query, window, filter: null, run, token);
var bm25Task = indexSearch.SearchKeywordsAsync(dataSource, window, indexStore => indexStore.SearchChunksAsync(dataSource.Id, query, window, token), run, token);
await Task.WhenAll(vectorTask, bm25Task);
token.ThrowIfCancellationRequested();
this.LogVectorResults(dataSource, vectorTask.Result);
this.LogBm25Results(dataSource, bm25Task.Result);
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(),
@@ -163,242 +125,6 @@ public sealed class DataSourceLocalRetrievalService(
return new RetrievalPage(contexts, hasMore) { Gaps = gaps };
}
private async Task<IReadOnlyList<VectorSearchResult>> SearchVectorAsync(
IInternalDataSource dataSource,
string query,
int maxMatches,
string collectionName,
RetrievalRun run,
CancellationToken token)
{
try
{
var vectorStore = await databaseClientProvider.GetVectorStoreAsync(token);
if (!vectorStore.IsAvailable)
{
logger.LogWarning(
"Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because vector store '{VectorStoreName}' is unavailable.",
dataSource.Name,
dataSource.Id,
vectorStore.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, 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, run, token))
return [];
var provider = embeddingProvider.CreateProvider();
var vectors = await provider.EmbedTextAsync(embeddingProvider.Model, settingsManager, token, [query]);
token.ThrowIfCancellationRequested();
var vector = vectors.FirstOrDefault();
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, 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 [];
}
var results = this.LimitSearchResults(
dataSource,
"vector",
await vectorStore.SearchEmbeddingAsync(collectionName, vector, maxMatches, token),
maxMatches);
this.LogVectorResults(dataSource, results);
return results;
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
throw;
}
catch (ProviderRequestException exception)
{
//
// The embedding provider named the cause and what to do about it. That sentence is
// worth far more to the user than the fact that a search came back empty:
//
logger.LogWarning(
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, 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)
{
//
// Its own gap key, because this is not a search which went wrong but an index which has
// to be built anew. Saying that once per session is what turns a silently shortened
// 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, 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, 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>
/// 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, 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}"))
return;
}
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 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, 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;
}
var tokenCountResponse = await rustService.GetTokenCount(embeddingProvider, query, token);
if (tokenCountResponse is not { Success: true })
{
logger.LogWarning(
"Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the token count for embedding provider '{EmbeddingProviderName}' could not be determined. Reason='{Reason}'.",
dataSource.Name,
dataSource.Id,
embeddingProvider.Name,
tokenCountResponse?.Message ?? "No response was returned by the tokenizer service.");
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;
}
var queryTokenCount = tokenCountResponse.Value.TokenCount;
if (queryTokenCount > providerTokenLimit)
{
logger.LogWarning(
"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, 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, RetrievalRun run, CancellationToken token)
{
try
{
var indexStore = await databaseClientProvider.GetIndexStoreAsync(token);
if (!indexStore.IsAvailable)
{
logger.LogWarning(
"Skipping BM25 retrieval for data source '{DataSourceName}' ({DataSourceId}) because local RAG index '{DatabaseName}' is unavailable.",
dataSource.Name,
dataSource.Id,
indexStore.Name);
run.Add(RetrievalGap.PARTLY_SEARCHED);
return [];
}
var results = this.LimitSearchResults(
dataSource,
"BM25",
await indexStore.SearchChunksAsync(dataSource.Id, query, maxMatches, token),
maxMatches);
this.LogBm25Results(dataSource, results);
return results;
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
throw;
}
catch (Exception exception)
{
logger.LogWarning(exception, "BM25 retrieval failed for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id);
run.Add(RetrievalGap.PARTLY_SEARCHED);
return [];
}
}
private IReadOnlyList<T> LimitSearchResults<T>(IInternalDataSource dataSource, string searchName, IReadOnlyList<T> results, int maxMatches)
{
if (results.Count <= maxMatches)
return results;
logger.LogWarning(
"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,
dataSource.Id,
maxMatches);
return results.Take(maxMatches).ToList();
}
private static LocalRetrievalHit FromVectorResult(VectorSearchResult result, int rank) =>
new(
RetrievalChannel.VECTOR,