mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-17 17:38:00 +00:00
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-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
93be539e50
commit
c4400c0ff6
343 files changed
+26051
-3122
No files matched your search
@@ -296,7 +296,7 @@ public sealed class AssistantPluginGenerationService(ToolRegistry toolRegistry,
|
||||
|
||||
var builder = new StringBuilder();
|
||||
foreach (var tool in activeTools)
|
||||
builder.AppendLine($"- {tool.Definition.Id}: {tool.Definition.Function.DescriptionForLLM}");
|
||||
builder.AppendLine($"- {tool.Id}: {tool.DescriptionForLLM}");
|
||||
|
||||
return builder.ToString().TrimEnd();
|
||||
}
|
||||
@@ -868,7 +868,7 @@ public sealed class AssistantPluginGenerationService(ToolRegistry toolRegistry,
|
||||
|
||||
var alreadyRequested = RequestedToolIds(previousVersion).ToHashSet(StringComparer.Ordinal);
|
||||
return toolIds
|
||||
.Where(toolId => !alreadyRequested.Contains(toolId) && toolRegistry.GetDefinition(toolId) is null)
|
||||
.Where(toolId => !alreadyRequested.Contains(toolId) && !toolRegistry.IsKnown(toolId))
|
||||
.ToList();
|
||||
}
|
||||
|
||||
|
||||
@@ -18,12 +18,13 @@ public sealed record ConfigurationPluginDeleteSummary(
|
||||
int DocumentAnalysisPolicies,
|
||||
int LockedSettings,
|
||||
int MandatoryInfos,
|
||||
int Introductions)
|
||||
int Introductions,
|
||||
int MailboxProviders)
|
||||
{
|
||||
/// <summary>
|
||||
/// An empty summary, used when the configuration plugin is not running and we cannot tell what it configured.
|
||||
/// </summary>
|
||||
public static readonly ConfigurationPluginDeleteSummary EMPTY = new(0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
|
||||
public static readonly ConfigurationPluginDeleteSummary EMPTY = new(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
|
||||
|
||||
/// <summary>
|
||||
/// True when the deletion affects anything beyond the plugin directory.
|
||||
@@ -38,5 +39,6 @@ public sealed record ConfigurationPluginDeleteSummary(
|
||||
this.DocumentAnalysisPolicies > 0 ||
|
||||
this.LockedSettings > 0 ||
|
||||
this.MandatoryInfos > 0 ||
|
||||
this.Introductions > 0;
|
||||
this.Introductions > 0 ||
|
||||
this.MailboxProviders > 0;
|
||||
}
|
||||
@@ -15,6 +15,7 @@ namespace AIStudio.Tools.Services;
|
||||
/// <param name="DeclaredSettings">How many settings the configuration takes over.</param>
|
||||
/// <param name="MandatoryInfos">How many mandatory information texts users must accept.</param>
|
||||
/// <param name="Introductions">How many introductions the configuration adds to the welcome page.</param>
|
||||
/// <param name="MailboxProviders">How many mail servers the configuration offers for new mailboxes.</param>
|
||||
public sealed record ConfigurationPluginImportSummary(
|
||||
IReadOnlyList<ConfigurationPluginDestination> Destinations,
|
||||
int ChatTemplates,
|
||||
@@ -22,7 +23,8 @@ public sealed record ConfigurationPluginImportSummary(
|
||||
int DocumentAnalysisPolicies,
|
||||
int DeclaredSettings,
|
||||
int MandatoryInfos,
|
||||
int Introductions)
|
||||
int Introductions,
|
||||
int MailboxProviders)
|
||||
{
|
||||
/// <summary>
|
||||
/// True when the configuration sets up anything at all.
|
||||
@@ -34,5 +36,6 @@ public sealed record ConfigurationPluginImportSummary(
|
||||
this.DocumentAnalysisPolicies > 0 ||
|
||||
this.DeclaredSettings > 0 ||
|
||||
this.MandatoryInfos > 0 ||
|
||||
this.Introductions > 0;
|
||||
this.Introductions > 0 ||
|
||||
this.MailboxProviders > 0;
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// What a data source waits for the user to decide, beyond the failures its list explains.
|
||||
/// </summary>
|
||||
public enum DataSourceAttention
|
||||
{
|
||||
/// <summary>
|
||||
/// Nothing to decide.
|
||||
/// </summary>
|
||||
NONE,
|
||||
|
||||
/// <summary>
|
||||
/// The server refused the sign-in, and AI Studio does not sign in again on its own.
|
||||
/// </summary>
|
||||
AUTH_FAILED,
|
||||
|
||||
/// <summary>
|
||||
/// A run would remove many documents from the index at once, and waits for the user to agree.
|
||||
/// </summary>
|
||||
MASS_REMOVAL_PENDING,
|
||||
|
||||
/// <summary>
|
||||
/// The organization allows only its own mail servers, and the mailbox is on another one. AI Studio does not connect to it until the organization allows the server.
|
||||
/// </summary>
|
||||
SERVER_NOT_ALLOWED,
|
||||
}
|
||||
@@ -12,13 +12,14 @@ namespace AIStudio.Tools.Services;
|
||||
/// rejected API key, an unreachable provider, and a file nobody may read were one and the same
|
||||
/// list entry. The surrounding fields are what lets the UI offer the matching way out.
|
||||
/// </remarks>
|
||||
/// <param name="FilePath">The file that failed or the name of the data source when the failure was not about one file.</param>
|
||||
/// <param name="DocumentKey">The key of the document that failed, for a file its full path, or the name of the data source when the failure was not about one document.</param>
|
||||
/// <param name="Reason">What to tell the user about it, ready to show.</param>
|
||||
/// <param name="OccurredAtUtc">When it happened, so the list still makes sense when the user looks at it later.</param>
|
||||
/// <param name="FailureReason">What kind of failure it was. Everything that did not come from a provider stays at NONE.</param>
|
||||
/// <param name="StatusCode">What the provider answered, where it answered at all.</param>
|
||||
/// <param name="EmbeddingProviderName">The embedding provider that was asked.</param>
|
||||
/// <param name="ExtractionCode">Why reading the file failed, where the failure was about reading it at all.</param>
|
||||
/// <param name="IsPermanent">Whether the file stays out of the index until it changes.</param>
|
||||
public sealed record DataSourceEmbeddingFailure(string FilePath, string Reason, DateTimeOffset OccurredAtUtc, ProviderRequestFailureReason FailureReason = ProviderRequestFailureReason.NONE,
|
||||
HttpStatusCode? StatusCode = null, string EmbeddingProviderName = "", FileExtractionErrorCode ExtractionCode = FileExtractionErrorCode.NONE, bool IsPermanent = false);
|
||||
/// <param name="ExtractionCode">Why reading the document failed, where the failure was about reading it at all.</param>
|
||||
/// <param name="IsPermanent">Whether the document stays out of the index until it changes.</param>
|
||||
/// <param name="DisplayName">How the document is called for the user, e.g. the name of a file or the subject of a mail. Empty where the run did not know it, e.g. for a failure stored by an earlier run.</param>
|
||||
public sealed record DataSourceEmbeddingFailure(string DocumentKey, string Reason, DateTimeOffset OccurredAtUtc, ProviderRequestFailureReason FailureReason = ProviderRequestFailureReason.NONE,
|
||||
HttpStatusCode? StatusCode = null, string EmbeddingProviderName = "", FileExtractionErrorCode ExtractionCode = FileExtractionErrorCode.NONE, bool IsPermanent = false, string DisplayName = "");
|
||||
@@ -1,3 +1,3 @@
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
public sealed record DataSourceEmbeddingOverview(DataSourceEmbeddingState State, int IndexedFiles, int TotalFiles, int FailedFiles);
|
||||
public sealed record DataSourceEmbeddingOverview(DataSourceEmbeddingState State, int IndexedDocuments, int TotalDocuments, int FailedDocuments);
|
||||
@@ -7,11 +7,11 @@ namespace AIStudio.Tools.Services;
|
||||
|
||||
internal static class DataSourceEmbeddingProviders
|
||||
{
|
||||
public static bool TryResolve(SettingsManager settingsManager, IDataSource dataSource, [NotNullWhen(true)] out EmbeddingProvider? embeddingProvider)
|
||||
public static bool TryResolve(SettingsManager settingsManager, IDataSourceBase dataSource, [NotNullWhen(true)] out EmbeddingProvider? embeddingProvider)
|
||||
{
|
||||
embeddingProvider = settingsManager.ConfigurationData.EmbeddingProviders.FirstOrDefault(provider =>
|
||||
dataSource is IInternalDataSource internalDataSource &&
|
||||
provider.Id.Equals(internalDataSource.EmbeddingId, StringComparison.OrdinalIgnoreCase));
|
||||
dataSource is IIndexedDataSource indexedDataSource &&
|
||||
provider.Id.Equals(indexedDataSource.EmbeddingId, StringComparison.OrdinalIgnoreCase));
|
||||
|
||||
return embeddingProvider != default && embeddingProvider.UsedLLMProvider is not LLMProviders.NONE;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// Why an indexing run of a data source was started.
|
||||
/// </summary>
|
||||
internal enum DataSourceEmbeddingRefreshMode
|
||||
{
|
||||
STARTUP_HASH_CHECK,
|
||||
HASH_CHECK,
|
||||
WATCHER_HASH_CHECK,
|
||||
MANUAL_RETRY,
|
||||
|
||||
/// <summary>
|
||||
/// Carries on where the last run of the same data source stopped after its share of the work.
|
||||
/// </summary>
|
||||
CONTINUATION,
|
||||
|
||||
/// <summary>
|
||||
/// The interval came round at which a data source is looked at again, since nothing reports its
|
||||
/// changes, e.g. a mailbox on a server.
|
||||
/// </summary>
|
||||
INTERVAL_CHECK,
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
using System.Globalization;
|
||||
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Mail;
|
||||
using AIStudio.Tools.Services.Indexing;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
public sealed partial class DataSourceEmbeddingService
|
||||
{
|
||||
internal const int DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH = 300;
|
||||
|
||||
/// <summary>
|
||||
/// What this build writes next to a chunk besides its text. Raise it whenever that changes.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A stored chunk keeps the metadata of the run which wrote it, and nothing recomputes it: the
|
||||
/// fingerprint of a file says whether the file changed, not whether we got better at reading
|
||||
/// it. Raising this number makes the embedding signature differ, which drops the index and
|
||||
/// builds it again — the only way corrected page numbers reach a data source somebody indexed
|
||||
/// earlier.
|
||||
///
|
||||
/// Version 2: the page of a chunk is taken from the runtime metadata instead of being read back
|
||||
/// out of the chunk text, which is what left Word and OpenDocument files, and passages
|
||||
/// continuing across a page break, without a page.
|
||||
/// </remarks>
|
||||
private const string CHUNK_METADATA_VERSION = "2";
|
||||
|
||||
/// <summary>
|
||||
/// What this build makes of a mail before it is cut into chunks. Raise it whenever that changes.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The counterpart of CHUNK_METADATA_VERSION for mailboxes alone. A mail on the server never
|
||||
/// changes, so nothing reads it again once it is indexed, however differently MailTextBuilder
|
||||
/// would write it today. Raising this number rebuilds the index of every mailbox and leaves
|
||||
/// every other data source alone.
|
||||
///
|
||||
/// Version 1: the header block and the text from the HTML part, without attachments.
|
||||
///
|
||||
/// Version 2: the text of attached documents follows, each attachment under a line naming it.
|
||||
/// The signature of a signed mail no longer counts as an attachment.
|
||||
/// </remarks>
|
||||
private const string MAIL_TEXT_VERSION = "2";
|
||||
|
||||
/// <summary>
|
||||
/// Works out how the text of a data source is cut for a given embedding provider.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Static, because the answer follows from its two arguments alone. That lets the embedding
|
||||
/// signature be built for a configuration which is not stored yet, which is what the dialogs ask
|
||||
/// before they save a change.
|
||||
/// </remarks>
|
||||
/// <param name="dataSource">The data source whose own chunk settings apply.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider whose token limit caps them.</param>
|
||||
/// <returns>The chunk size and overlap which are actually used.</returns>
|
||||
internal static ChunkingOptions GetChunkingOptions(IDataSourceBase dataSource, EmbeddingProvider embeddingProvider)
|
||||
{
|
||||
var providerMaxChunkTokenLength = Math.Max(1, embeddingProvider.EffectiveTokenLimit);
|
||||
var dataSourceMaxChunkTokenLength = dataSource is IIndexedDataSource { MaxChunkTokenLength: > 0 } indexedDataSource
|
||||
? indexedDataSource.MaxChunkTokenLength
|
||||
: 0;
|
||||
var maxChunkTokenLength = dataSourceMaxChunkTokenLength > 0
|
||||
? Math.Min(dataSourceMaxChunkTokenLength, providerMaxChunkTokenLength)
|
||||
: providerMaxChunkTokenLength;
|
||||
|
||||
var configuredOverlapTokenLength = dataSource is IIndexedDataSource overlapDataSource
|
||||
? overlapDataSource.ChunkOverlapTokenLength
|
||||
: DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH;
|
||||
var overlapTokenLength = Math.Clamp(configuredOverlapTokenLength, 0, Math.Max(0, maxChunkTokenLength - 1));
|
||||
|
||||
return new(maxChunkTokenLength, overlapTokenLength);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Describes how the vectors of a data source were made.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// What appears here decides when stored embeddings are thrown away: a signature differing from
|
||||
/// the persisted one drops the whole index and builds it again. So it names the embedding model,
|
||||
/// where it runs, how the text was cut for it, and the chunk metadata version — the things a
|
||||
/// vector actually depends on.
|
||||
///
|
||||
/// Two of them are less obvious than they look. The Hugging Face inference provider belongs to
|
||||
/// where the model runs: the same model name served by another backend is another vector source.
|
||||
/// And a custom tokenizer enters through its content, not through its path, because a tokenizer
|
||||
/// is stored under the name it came with — almost always tokenizer.json — so swapping one for
|
||||
/// another lands on the identical path, while moving the data directory changes every path
|
||||
/// without changing a single tokenizer.
|
||||
///
|
||||
/// The chunk settings enter only as what they amount to, never as what somebody typed. A data
|
||||
/// source storing 0 means "follow the embedding provider", and writing that provider's own limit
|
||||
/// into the field changes nothing about how the text is cut. Carrying the typed numbers as well
|
||||
/// made that a different signature, so opening the expert settings of a data source — which
|
||||
/// fills an empty limit with the provider's — threw the whole index away for nothing.
|
||||
///
|
||||
/// The confidence level a data source asks of a provider is deliberately not among them. It
|
||||
/// changes no vector, and it is enforced live on every request anyway: DataSourceService checks
|
||||
/// it against the participating chat providers and against the embedding provider, and this
|
||||
/// service checks it again before each indexing run. It was part of this signature once, which
|
||||
/// re-embedded every file of a data source whenever somebody raised or lowered it — real money
|
||||
/// at a cloud embedding provider, for nothing.
|
||||
///
|
||||
/// A mailbox appends the version of its mail text, cf. MAIL_TEXT_VERSION, and up to which size
|
||||
/// it reads attachments. A mail is never read again once it is indexed, so reading attachments
|
||||
/// from now on, or larger ones, reaches the mails indexed before only by indexing them anew. The
|
||||
/// dialog asks before it does that, as for every change of this signature. Nothing else is
|
||||
/// appended for the other kinds of data source, so their stored signatures stay valid.
|
||||
/// </remarks>
|
||||
internal static string BuildEmbeddingSignature(IDataSourceBase dataSource, EmbeddingProvider embeddingProvider, ChunkingOptions chunkingOptions)
|
||||
{
|
||||
var signature = string.Join('|',
|
||||
CHUNK_METADATA_VERSION,
|
||||
embeddingProvider.Id,
|
||||
embeddingProvider.UsedLLMProvider,
|
||||
embeddingProvider.Model.Id,
|
||||
embeddingProvider.Host,
|
||||
embeddingProvider.Hostname,
|
||||
embeddingProvider.HFInferenceProvider,
|
||||
embeddingProvider.TokenizerFingerprint,
|
||||
embeddingProvider.EffectiveTokenLimit,
|
||||
chunkingOptions.MaxChunkTokenLength,
|
||||
chunkingOptions.OverlapTokenLength);
|
||||
|
||||
if (dataSource is not DataSourceMailbox mailbox)
|
||||
return signature;
|
||||
|
||||
var attachments = MailAttachmentRules.GetMaxSizeMegabytes(mailbox) is { } maxSizeMegabytes
|
||||
? maxSizeMegabytes.ToString(CultureInfo.InvariantCulture)
|
||||
: "none";
|
||||
|
||||
return $"{signature}|mail:{MAIL_TEXT_VERSION}|attachments:{attachments}";
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Describes how the vectors of a data source were made, working the chunking out along the way.
|
||||
/// </summary>
|
||||
/// <param name="dataSource">The data source the vectors belong to.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider which makes them.</param>
|
||||
/// <returns>The signature of this pairing.</returns>
|
||||
internal static string BuildEmbeddingSignature(IDataSourceBase dataSource, EmbeddingProvider embeddingProvider) =>
|
||||
BuildEmbeddingSignature(dataSource, embeddingProvider, GetChunkingOptions(dataSource, embeddingProvider));
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
|
||||
@@ -22,7 +23,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
/// <param name="dataSourceId">The data source to build anew.</param>
|
||||
public async Task RepairDataSourceAsync(string dataSourceId)
|
||||
{
|
||||
if (!this.TryGetConfiguredDataSource(dataSourceId, out var dataSource) || !this.IsSupportedInternalDataSource(dataSource))
|
||||
if (!this.TryGetConfiguredIndexedSource(dataSourceId, out var dataSource))
|
||||
return;
|
||||
|
||||
logger.LogWarning(
|
||||
@@ -37,11 +38,99 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
await this.ResetPersistedStateAsync(dataSourceId, null, null, CancellationToken.None);
|
||||
this.statuses.TryRemove(dataSourceId, out _);
|
||||
this.PublishStatusChanged();
|
||||
PublishStatusChanged();
|
||||
|
||||
await this.QueueDataSourceAsync(dataSource, true, DataSourceEmbeddingRefreshMode.MANUAL_RETRY);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Records that the user agreed to the removal a mailbox held back, and syncs the mailbox to carry it out.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The count is the one the status showed the user, cf. DataSourceEmbeddingStatus.PendingRemovalCount.
|
||||
/// The agreement holds for exactly that number: should a sync have come up with another one in the
|
||||
/// meantime, nothing is agreed to, and the user is asked anew. The sync which carries the removal
|
||||
/// out signs in like any other, so a recorded sign-in failure still keeps it from the server.
|
||||
/// </remarks>
|
||||
/// <param name="dataSourceId">The mailbox.</param>
|
||||
/// <param name="removalCount">How many mails the user agreed to remove from the index.</param>
|
||||
/// <returns>True when the agreement was recorded. False when the number changed in the meantime, or the index cannot be reached.</returns>
|
||||
public async Task<bool> ApprovePendingMailRemovalAsync(string dataSourceId, int removalCount)
|
||||
{
|
||||
if (!this.TryGetConfiguredIndexedSource(dataSourceId, out var dataSource) || dataSource is not DataSourceMailbox)
|
||||
return false;
|
||||
|
||||
var indexStore = await databaseClientProvider.GetIndexStoreAsync(CancellationToken.None);
|
||||
if (!await indexStore.ApprovePendingMailRemovalAsync(dataSourceId, removalCount, CancellationToken.None))
|
||||
{
|
||||
logger.LogInformation("The removal held back for mailbox '{DataSourceId}' was not approved, since another number is held back by now or the index store is unavailable.", dataSourceId);
|
||||
return false;
|
||||
}
|
||||
|
||||
logger.LogInformation("The user approved removing {RemovalCount} mails from the index of mailbox '{DataSourceId}'.", removalCount, dataSourceId);
|
||||
await this.QueueDataSourceAsync(dataSource, true, DataSourceEmbeddingRefreshMode.HASH_CHECK);
|
||||
return true;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Deletes what the index keeps of mailboxes which are no longer configured.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Such leftovers come about in two ways. An older version of AI Studio does not know mailboxes
|
||||
/// and saves the settings without them. Or the stores could not be reached while the user
|
||||
/// deleted a mailbox. Either way, nothing would ever delete the index then, nor a refused
|
||||
/// sign-in kept next to it.
|
||||
///
|
||||
/// Only with settings which were loaded and may be written. Settings which could not be read
|
||||
/// look like settings without a single mailbox, and every index would be deleted for that. Both
|
||||
/// stores have to be there as well: the row in the index store is all that leads to the
|
||||
/// collection in the vector store, so deleting the one without the other would leave the other
|
||||
/// behind for good.
|
||||
///
|
||||
/// Runs before the first run, so none of the mailboxes is in use. The password stays in the OS
|
||||
/// keyring: its entry carries the name of the mailbox, which nothing here knows anymore.
|
||||
/// </remarks>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
private async Task DeleteOrphanedMailboxIndexesAsync(CancellationToken token)
|
||||
{
|
||||
if (!settingsManager.HasCompletedInitialSettingsLoad || settingsManager.SettingsWriteBlocked)
|
||||
{
|
||||
logger.LogInformation("Leaving the indexes of mailboxes alone, since the settings could not be loaded completely. SettingsWriteBlockReason={SettingsWriteBlockReason}.", settingsManager.SettingsWriteBlockReason);
|
||||
return;
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
var vectorStore = await databaseClientProvider.GetVectorStoreAsync(token);
|
||||
var indexStore = await databaseClientProvider.GetIndexStoreAsync(token);
|
||||
if (!vectorStore.IsAvailable || !indexStore.IsAvailable)
|
||||
return;
|
||||
|
||||
var configuredIds = settingsManager.ConfigurationData.Mailboxes
|
||||
.Select(mailbox => mailbox.Id)
|
||||
.ToHashSet(StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
foreach (var mailboxId in await indexStore.GetStoredMailboxIdsAsync(token))
|
||||
{
|
||||
if (configuredIds.Contains(mailboxId))
|
||||
continue;
|
||||
|
||||
await this.ResetPersistedStateAsync(mailboxId, vectorStore, indexStore, token);
|
||||
await indexStore.ClearMailboxAuthFailureAsync(mailboxId, token);
|
||||
logger.LogInformation("Deleted the index of mailbox '{DataSourceId}', which is no longer configured.", mailboxId);
|
||||
}
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
// Nothing depends on it: whatever is left over is tried again on the next start.
|
||||
logger.LogWarning(exception, "Could not delete the indexes of mailboxes which are no longer configured.");
|
||||
}
|
||||
}
|
||||
|
||||
private async Task ResetPersistedStateAsync(
|
||||
string dataSourceId,
|
||||
VectorStoreClient? vectorStore,
|
||||
|
||||
@@ -1,17 +1,19 @@
|
||||
using System.Collections.Concurrent;
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Settings.DataModel;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <remarks>
|
||||
/// A watcher here is whatever an indexer uses to notice that one of its data sources changed, a file
|
||||
/// system watcher for local files, an interval for mailboxes. The indexers own them; this service
|
||||
/// only decides when they run.
|
||||
/// </remarks>
|
||||
public sealed partial class DataSourceEmbeddingService
|
||||
{
|
||||
private const int WATCHER_DEBOUNCE_SECONDS = 2;
|
||||
|
||||
private readonly ConcurrentDictionary<string, DataSourceWatcherRegistration> watchers = new(StringComparer.OrdinalIgnoreCase);
|
||||
private readonly Dictionary<string, CancellationTokenSource> watcherDebounceTokens = new(StringComparer.OrdinalIgnoreCase);
|
||||
private readonly object watcherDebounceLock = new();
|
||||
|
||||
/// <summary>
|
||||
/// Lets every indexer track changes to its data sources, or stops all tracking.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Tracking runs only while local data sources refresh on their own, and not before the startup
|
||||
/// hash check is done.
|
||||
/// </remarks>
|
||||
private void RefreshWatchers()
|
||||
{
|
||||
if (!settingsManager.ConfigurationData.App.DataSourceIndexing.AutomaticRefresh)
|
||||
@@ -27,263 +29,56 @@ public sealed partial class DataSourceEmbeddingService
|
||||
return;
|
||||
}
|
||||
|
||||
var supportedSources = settingsManager.ConfigurationData.DataSources
|
||||
.Where(this.IsSupportedInternalDataSource)
|
||||
.ToDictionary(source => source.Id, StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
foreach (var existingWatcherId in this.watchers.Keys.Except(supportedSources.Keys, StringComparer.OrdinalIgnoreCase).ToList())
|
||||
this.RemoveWatcher(existingWatcherId);
|
||||
|
||||
foreach (var dataSource in supportedSources.Values)
|
||||
this.EnsureWatcher(dataSource);
|
||||
}
|
||||
|
||||
private void EnsureWatcher(IDataSource dataSource)
|
||||
{
|
||||
if (!settingsManager.ConfigurationData.App.DataSourceIndexing.AutomaticRefresh)
|
||||
return;
|
||||
|
||||
var configuration = GetWatchConfiguration(dataSource);
|
||||
if (configuration is null)
|
||||
return;
|
||||
|
||||
if (this.watchers.TryGetValue(dataSource.Id, out var existingRegistration))
|
||||
{
|
||||
if (IsSameWatchConfiguration(existingRegistration.Configuration, configuration))
|
||||
return;
|
||||
|
||||
this.RemoveWatcher(dataSource.Id);
|
||||
}
|
||||
|
||||
var watcher = this.CreateWatcher(dataSource.Id, configuration);
|
||||
if (watcher is null)
|
||||
return;
|
||||
|
||||
if (!this.watchers.TryAdd(dataSource.Id, new DataSourceWatcherRegistration(watcher, configuration)))
|
||||
watcher.Dispose();
|
||||
}
|
||||
|
||||
private FileSystemWatcher? CreateWatcher(string dataSourceId, DataSourceWatcherConfiguration configuration)
|
||||
{
|
||||
try
|
||||
{
|
||||
var watcher = new FileSystemWatcher(configuration.RootPath)
|
||||
{
|
||||
Filter = configuration.Filter,
|
||||
IncludeSubdirectories = configuration.IncludeSubdirectories,
|
||||
NotifyFilter = NotifyFilters.FileName | NotifyFilters.DirectoryName | NotifyFilters.LastWrite | NotifyFilters.CreationTime | NotifyFilters.Size,
|
||||
};
|
||||
|
||||
watcher.Changed += (_, args) => this.OnWatchedDataSourceChanged(dataSourceId, configuration, args);
|
||||
watcher.Deleted += (_, args) => this.OnWatchedDataSourceChanged(dataSourceId, configuration, args);
|
||||
watcher.Created += (_, args) => this.OnWatchedDataSourceChanged(dataSourceId, configuration, args);
|
||||
watcher.Renamed += (_, args) => this.OnWatchedDataSourceChanged(dataSourceId, configuration, args);
|
||||
watcher.Error += (_, args) =>
|
||||
{
|
||||
logger.LogWarning(args.GetException(), "The file watcher for data source '{DataSourceId}' failed. Recreating it.", dataSourceId);
|
||||
this.RemoveWatcher(dataSourceId);
|
||||
this.EnsureWatcher(dataSourceId);
|
||||
this.ScheduleWatchedDataSourceRefresh(dataSourceId);
|
||||
};
|
||||
watcher.EnableRaisingEvents = true;
|
||||
return watcher;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Failed to create file watcher for data source '{DataSourceId}' at '{RootPath}'.", dataSourceId, configuration.RootPath);
|
||||
return null;
|
||||
}
|
||||
var supportedSources = this.GetConfiguredIndexedSources();
|
||||
foreach (var indexer in this.indexers)
|
||||
indexer.TrackChanges(supportedSources.Where(indexer.Supports).ToList(), this.RequestRunAsync);
|
||||
}
|
||||
|
||||
private void RemoveWatcher(string dataSourceId)
|
||||
{
|
||||
this.CancelPendingWatcherRefresh(dataSourceId);
|
||||
|
||||
if (this.watchers.TryRemove(dataSourceId, out var registration))
|
||||
registration.Watcher.Dispose();
|
||||
foreach (var indexer in this.indexers)
|
||||
indexer.StopTracking(dataSourceId);
|
||||
}
|
||||
|
||||
|
||||
private void RemoveAllWatchers()
|
||||
{
|
||||
foreach (var watcherId in this.watchers.Keys.ToList())
|
||||
this.RemoveWatcher(watcherId);
|
||||
foreach (var indexer in this.indexers)
|
||||
indexer.StopTrackingAll();
|
||||
}
|
||||
|
||||
private void DisposeWatchers()
|
||||
{
|
||||
this.CancelAllPendingWatcherRefreshes();
|
||||
|
||||
foreach (var registration in this.watchers.Values)
|
||||
registration.Watcher.Dispose();
|
||||
|
||||
this.watchers.Clear();
|
||||
foreach (var indexer in this.indexers)
|
||||
indexer.Dispose();
|
||||
}
|
||||
|
||||
private void OnWatchedDataSourceChanged(string dataSourceId, DataSourceWatcherConfiguration configuration, FileSystemEventArgs args)
|
||||
/// <summary>
|
||||
/// Queues the run an indexer asked for after it noticed a change, or because its interval came round.
|
||||
/// </summary>
|
||||
/// <param name="dataSourceId">The id of the data source which changed.</param>
|
||||
/// <param name="refreshMode">Why the indexer asks.</param>
|
||||
private async Task RequestRunAsync(string dataSourceId, DataSourceEmbeddingRefreshMode refreshMode)
|
||||
{
|
||||
if (!this.IsRelevantWatcherEvent(configuration, args))
|
||||
if (!this.TryGetConfiguredIndexedSource(dataSourceId, out var dataSource))
|
||||
return;
|
||||
|
||||
if (this.statuses.TryGetValue(dataSourceId, out var status) && IsWaitingForSignIn(status))
|
||||
{
|
||||
logger.LogDebug(
|
||||
"Ignoring file system change for data source '{DataSourceId}' at '{Path}' (event={ChangeType}) because the path is not part of the RAG index.",
|
||||
dataSourceId,
|
||||
args.FullPath,
|
||||
args.ChangeType);
|
||||
logger.LogDebug("Skipped the requested run of data source '{DataSourceId}' because its sign-in failed and waits for the user. RefreshMode={RefreshMode}.", dataSourceId, refreshMode);
|
||||
return;
|
||||
}
|
||||
|
||||
logger.LogDebug(
|
||||
"Detected relevant file system change for data source '{DataSourceId}' at '{Path}' (event={ChangeType}). Scheduling a debounced embedding run.",
|
||||
dataSourceId,
|
||||
args.FullPath,
|
||||
args.ChangeType);
|
||||
|
||||
this.ScheduleWatchedDataSourceRefresh(dataSourceId);
|
||||
await this.QueueDataSourceAsync(dataSource, true, refreshMode);
|
||||
}
|
||||
|
||||
private void ScheduleWatchedDataSourceRefresh(string dataSourceId)
|
||||
{
|
||||
if (!settingsManager.ConfigurationData.App.DataSourceIndexing.AutomaticRefresh)
|
||||
return;
|
||||
|
||||
var debounceToken = new CancellationTokenSource();
|
||||
|
||||
lock (this.watcherDebounceLock)
|
||||
{
|
||||
if (this.watcherDebounceTokens.Remove(dataSourceId, out var existingToken))
|
||||
existingToken.Cancel();
|
||||
|
||||
this.watcherDebounceTokens[dataSourceId] = debounceToken;
|
||||
}
|
||||
|
||||
_ = Task.Run(async () =>
|
||||
{
|
||||
try
|
||||
{
|
||||
await Task.Delay(TimeSpan.FromSeconds(WATCHER_DEBOUNCE_SECONDS), debounceToken.Token);
|
||||
if (!this.TryCompletePendingWatcherRefresh(dataSourceId, debounceToken))
|
||||
return;
|
||||
|
||||
var dataSource = settingsManager.ConfigurationData.DataSources
|
||||
.FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase));
|
||||
|
||||
if (dataSource is not null)
|
||||
{
|
||||
logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) after file system changes settled. The hash pipeline will reindex only changed files.", dataSource.Name, dataSource.Id);
|
||||
await this.QueueDataSourceAsync(dataSource, true, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK);
|
||||
}
|
||||
}
|
||||
catch (OperationCanceledException)
|
||||
{
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Failed to queue watched data source '{DataSourceId}' after a file system change.", dataSourceId);
|
||||
}
|
||||
finally
|
||||
{
|
||||
debounceToken.Dispose();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
private void EnsureWatcher(string dataSourceId)
|
||||
{
|
||||
var dataSource = settingsManager.ConfigurationData.DataSources
|
||||
.FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase));
|
||||
|
||||
if (dataSource is not null)
|
||||
this.EnsureWatcher(dataSource);
|
||||
}
|
||||
|
||||
private void CancelPendingWatcherRefresh(string dataSourceId)
|
||||
{
|
||||
lock (this.watcherDebounceLock)
|
||||
{
|
||||
if (this.watcherDebounceTokens.Remove(dataSourceId, out var token))
|
||||
token.Cancel();
|
||||
}
|
||||
}
|
||||
|
||||
private void CancelAllPendingWatcherRefreshes()
|
||||
{
|
||||
lock (this.watcherDebounceLock)
|
||||
{
|
||||
foreach (var token in this.watcherDebounceTokens.Values)
|
||||
token.Cancel();
|
||||
|
||||
this.watcherDebounceTokens.Clear();
|
||||
}
|
||||
}
|
||||
|
||||
private bool TryCompletePendingWatcherRefresh(string dataSourceId, CancellationTokenSource debounceToken)
|
||||
{
|
||||
lock (this.watcherDebounceLock)
|
||||
{
|
||||
if (!this.watcherDebounceTokens.TryGetValue(dataSourceId, out var currentToken) || !ReferenceEquals(currentToken, debounceToken))
|
||||
return false;
|
||||
|
||||
this.watcherDebounceTokens.Remove(dataSourceId);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
private bool IsRelevantWatcherEvent(DataSourceWatcherConfiguration configuration, FileSystemEventArgs args)
|
||||
{
|
||||
if (args is RenamedEventArgs renamedArgs)
|
||||
{
|
||||
return this.IsRelevantWatcherPath(configuration, renamedArgs.FullPath, args.ChangeType)
|
||||
|| this.IsRelevantWatcherPath(configuration, renamedArgs.OldFullPath, args.ChangeType);
|
||||
}
|
||||
|
||||
return this.IsRelevantWatcherPath(configuration, args.FullPath, args.ChangeType);
|
||||
}
|
||||
|
||||
private bool IsRelevantWatcherPath(DataSourceWatcherConfiguration configuration, string path, WatcherChangeTypes changeType)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(path))
|
||||
return false;
|
||||
|
||||
var fileName = Path.GetFileName(path);
|
||||
if (string.IsNullOrWhiteSpace(fileName))
|
||||
return true;
|
||||
|
||||
if (!configuration.IncludeSubdirectories && !string.Equals(fileName, configuration.Filter, StringComparison.OrdinalIgnoreCase))
|
||||
return false;
|
||||
|
||||
if (Directory.Exists(path))
|
||||
return true;
|
||||
|
||||
if (IsSkippedRagFileName(fileName))
|
||||
return false;
|
||||
|
||||
if (this.IsSupportedRagFilePath(path))
|
||||
return true;
|
||||
|
||||
return changeType is WatcherChangeTypes.Deleted or WatcherChangeTypes.Renamed
|
||||
&& string.IsNullOrWhiteSpace(Path.GetExtension(path));
|
||||
}
|
||||
|
||||
private static DataSourceWatcherConfiguration? GetWatchConfiguration(IDataSource dataSource) => dataSource switch
|
||||
{
|
||||
DataSourceLocalDirectory localDirectory when Directory.Exists(localDirectory.Path) => new DataSourceWatcherConfiguration(
|
||||
localDirectory.Path,
|
||||
"*.*",
|
||||
true),
|
||||
DataSourceLocalFile localFile when File.Exists(localFile.FilePath) && !string.IsNullOrWhiteSpace(Path.GetDirectoryName(localFile.FilePath)) => new DataSourceWatcherConfiguration(
|
||||
Path.GetDirectoryName(localFile.FilePath)!,
|
||||
Path.GetFileName(localFile.FilePath),
|
||||
false),
|
||||
_ => null,
|
||||
};
|
||||
|
||||
private static bool IsSameWatchConfiguration(DataSourceWatcherConfiguration left, DataSourceWatcherConfiguration right)
|
||||
{
|
||||
return left.IncludeSubdirectories == right.IncludeSubdirectories
|
||||
&& string.Equals(left.RootPath, right.RootPath, StringComparison.OrdinalIgnoreCase)
|
||||
&& string.Equals(left.Filter, right.Filter, StringComparison.OrdinalIgnoreCase);
|
||||
}
|
||||
|
||||
private sealed record DataSourceWatcherConfiguration(string RootPath, string Filter, bool IncludeSubdirectories);
|
||||
|
||||
private sealed record DataSourceWatcherRegistration(FileSystemWatcher Watcher, DataSourceWatcherConfiguration Configuration);
|
||||
/// <summary>
|
||||
/// Whether a data source waits for the user since its sign-in failed.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Then only the user starts its next run, by saving a new password or with a retry. The indexer
|
||||
/// would not sign in on any other run either, but every one of them would make the row of the
|
||||
/// data source change from failed to queued and back, every interval anew.
|
||||
/// </remarks>
|
||||
/// <param name="status">The current status of the data source.</param>
|
||||
internal static bool IsWaitingForSignIn(DataSourceEmbeddingStatus status) => status.Attention is DataSourceAttention.AUTH_FAILED;
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -4,38 +4,49 @@ using AIStudio.Tools.PluginSystem;
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <remarks>
|
||||
/// CurrentFileBlock and CurrentFilePage are null rather than zero while nothing is known about
|
||||
/// them: a file which is only about to start has no first block, and not every kind of document
|
||||
/// has pages to count. Block numbers start at one, the way the page states them.
|
||||
/// Counts documents, whatever a document is for the kind of data source: a file for a folder, a
|
||||
/// mail for a mailbox. Which word the user reads is the business of the page showing it.
|
||||
///
|
||||
/// CurrentDocumentBlock and CurrentDocumentPage are null rather than zero while nothing is known
|
||||
/// about them: a document which is only about to start has no first block, and not every kind of
|
||||
/// document has pages to count. Block numbers start at one, the way the page states them.
|
||||
///
|
||||
/// VectorStoreUnreadable says why a data source failed, not only that it did. The UI needs that
|
||||
/// difference to offer the repair for this one case, and it is carried as its own flag so nothing
|
||||
/// has to read it back out of the message in LastError.
|
||||
/// has to read it back out of the message in LastError. Attention does the same for what the user
|
||||
/// has to decide, and PendingRemovalCount is the number a held back removal asks about: agreeing to
|
||||
/// it hands that very number back, cf. ApprovePendingMailRemovalAsync.
|
||||
///
|
||||
/// LastSyncUtc is when the data source was last worked through as a whole, for those which are
|
||||
/// synced rather than watched, such as mailboxes. It stays null for the others.
|
||||
/// </remarks>
|
||||
public sealed record DataSourceEmbeddingStatus(
|
||||
string DataSourceId,
|
||||
string DataSourceName,
|
||||
DataSourceType DataSourceType,
|
||||
DataSourceEmbeddingState State,
|
||||
int TotalFiles,
|
||||
int IndexedFiles,
|
||||
int FailedFiles,
|
||||
string CurrentFile,
|
||||
int TotalDocuments,
|
||||
int IndexedDocuments,
|
||||
int FailedDocuments,
|
||||
string CurrentDocument,
|
||||
string LastError,
|
||||
IReadOnlyList<DataSourceEmbeddingFailure> Failures,
|
||||
int PermanentlySkippedFiles = 0,
|
||||
int? CurrentFileBlock = null,
|
||||
int? CurrentFilePage = null,
|
||||
bool VectorStoreUnreadable = false)
|
||||
int PermanentlySkippedDocuments = 0,
|
||||
int? CurrentDocumentBlock = null,
|
||||
int? CurrentDocumentPage = null,
|
||||
bool VectorStoreUnreadable = false,
|
||||
DataSourceAttention Attention = DataSourceAttention.NONE,
|
||||
int? PendingRemovalCount = null,
|
||||
DateTimeOffset? LastSyncUtc = null)
|
||||
{
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DataSourceEmbeddingStatus).Namespace, nameof(DataSourceEmbeddingStatus));
|
||||
|
||||
/// <remarks>
|
||||
/// Files which were skipped for good are done, even though nothing was indexed of them.
|
||||
/// Documents which were skipped for good are done, even though nothing was indexed of them.
|
||||
/// Leaving them out would keep the bar short of the end for a data source which has nothing
|
||||
/// left to do.
|
||||
/// </remarks>
|
||||
public int ProgressPercent => this.TotalFiles <= 0 ? 0 : Math.Clamp((int)Math.Round((this.IndexedFiles + this.PermanentlySkippedFiles) * 100d / this.TotalFiles), 0, 100);
|
||||
public int ProgressPercent => this.TotalDocuments <= 0 ? 0 : Math.Clamp((int)Math.Round((this.IndexedDocuments + this.PermanentlySkippedDocuments) * 100d / this.TotalDocuments), 0, 100);
|
||||
|
||||
public string StateLabel => this.State switch
|
||||
{
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -19,7 +19,7 @@ internal static class EmbeddingChangeImpact
|
||||
/// <param name="before">The embedding provider as it is stored.</param>
|
||||
/// <param name="after">The embedding provider as it would be stored.</param>
|
||||
/// <returns>True when the stored index would be discarded.</returns>
|
||||
public static bool AffectsStoredIndex(IDataSource dataSource, EmbeddingProvider before, EmbeddingProvider after) =>
|
||||
public static bool AffectsStoredIndex(IDataSourceBase dataSource, EmbeddingProvider before, EmbeddingProvider after) =>
|
||||
!string.Equals(
|
||||
DataSourceEmbeddingService.BuildEmbeddingSignature(dataSource, before),
|
||||
DataSourceEmbeddingService.BuildEmbeddingSignature(dataSource, after),
|
||||
@@ -39,7 +39,7 @@ internal static class EmbeddingChangeImpact
|
||||
/// <param name="after">The data source as it would be stored.</param>
|
||||
/// <param name="afterProvider">The embedding provider it would point at.</param>
|
||||
/// <returns>True when the stored index would be discarded.</returns>
|
||||
public static bool AffectsStoredIndex(IDataSource before, EmbeddingProvider beforeProvider, IDataSource after, EmbeddingProvider afterProvider) =>
|
||||
public static bool AffectsStoredIndex(IDataSourceBase before, EmbeddingProvider beforeProvider, IDataSourceBase after, EmbeddingProvider afterProvider) =>
|
||||
!string.Equals(
|
||||
DataSourceEmbeddingService.BuildEmbeddingSignature(before, beforeProvider),
|
||||
DataSourceEmbeddingService.BuildEmbeddingSignature(after, afterProvider),
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// How large the chunks of a data source may become, and how much of each one the next repeats.
|
||||
/// </summary>
|
||||
/// <param name="MaxChunkTokenLength">The largest chunk in tokens of the embedding provider.</param>
|
||||
/// <param name="OverlapTokenLength">How many tokens of a chunk the next one starts with.</param>
|
||||
internal sealed record ChunkingOptions(int MaxChunkTokenLength, int OverlapTokenLength);
|
||||
@@ -0,0 +1,9 @@
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// One way to cut a text into smaller units, tried when the text is too large for a single chunk.
|
||||
/// </summary>
|
||||
/// <param name="Name">The name of the rule, for the log.</param>
|
||||
/// <param name="Split">Cuts a text into units, given the text and the segments it was made of. Null means a hard cut.</param>
|
||||
/// <param name="UsesSourceSegmentCounts">Whether the units are the segments themselves, so their token counts and pages can be reused.</param>
|
||||
internal sealed record ChunkingRule(string Name, Func<string, IReadOnlyList<string>, IReadOnlyList<string>>? Split, bool UsesSourceSegmentCounts = false);
|
||||
@@ -0,0 +1,8 @@
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// The rules a kind of text is cut by, from the coarsest to the finest.
|
||||
/// </summary>
|
||||
/// <param name="Name">The name of the strategy, for the log.</param>
|
||||
/// <param name="Rules">The rules, tried one after the other until every chunk fits.</param>
|
||||
internal sealed record ChunkingStrategy(string Name, IReadOnlyList<ChunkingRule> Rules);
|
||||
@@ -0,0 +1,401 @@
|
||||
using AIStudio.Provider;
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
|
||||
using static AIStudio.Tools.Services.Indexing.IndexingLogFormat;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// How far an indexing run got through the documents of its data source, and what became of each.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Counts the documents, keeps the list of failures and turns both into the status the embedding
|
||||
/// page shows. What happens to a document after it was indexed, or after it failed, is the same for
|
||||
/// every kind of data source, so that is decided here as well: the stores learn about it, the
|
||||
/// manifest follows, and the user hears about it the way the failure deserves.
|
||||
/// </remarks>
|
||||
/// <param name="context">The run the documents are indexed in.</param>
|
||||
/// <param name="totalDocuments">How many documents the data source has, readable or not.</param>
|
||||
/// <param name="failedInputs">How many of them could not even be looked at.</param>
|
||||
/// <param name="lastInputError">Why the last of those could not be looked at, or an empty string.</param>
|
||||
/// <param name="inputFailures">The failures of those which could not be looked at.</param>
|
||||
/// <param name="logger">The logger of the embedding service, so the log reads the same whoever writes it.</param>
|
||||
internal sealed class DocumentRunProgress(IndexedRunContext context, int totalDocuments, int failedInputs, string lastInputError, IEnumerable<DataSourceEmbeddingFailure> inputFailures, ILogger logger)
|
||||
{
|
||||
/// <summary>
|
||||
/// How often the block progress within one document is reported to the user interface at most.
|
||||
/// </summary>
|
||||
private static readonly TimeSpan BLOCK_PROGRESS_INTERVAL = TimeSpan.FromSeconds(3);
|
||||
|
||||
private readonly List<DataSourceEmbeddingFailure> failures = inputFailures.ToList();
|
||||
|
||||
//
|
||||
// Which kinds of provider failure the user was already told about in this run. A rejected
|
||||
// API key is the same problem for every one of a few thousand documents, and one message
|
||||
// is what it takes to send the user to the settings.
|
||||
//
|
||||
private readonly HashSet<ProviderRequestFailureReason> reportedFailureReasons = [];
|
||||
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DocumentRunProgress).Namespace, nameof(DocumentRunProgress));
|
||||
|
||||
public int TotalDocuments => totalDocuments;
|
||||
|
||||
/// <summary>
|
||||
/// Documents which were left alone because nothing changed since they were indexed.
|
||||
/// </summary>
|
||||
public int UnchangedDocuments { get; private set; }
|
||||
|
||||
/// <summary>
|
||||
/// Documents which were not read because they failed for a reason of their own before.
|
||||
/// </summary>
|
||||
public int PermanentlySkippedDocuments { get; private set; }
|
||||
|
||||
/// <summary>
|
||||
/// Documents which were indexed in this run.
|
||||
/// </summary>
|
||||
public int IndexedDocuments { get; private set; }
|
||||
|
||||
public int NewDocuments { get; private set; }
|
||||
|
||||
public int ChangedDocuments { get; private set; }
|
||||
|
||||
public int FailedDocuments { get; private set; } = failedInputs;
|
||||
|
||||
public string LastError { get; private set; } = lastInputError;
|
||||
|
||||
/// <summary>
|
||||
/// When the data source was last worked through as a whole, for those which are synced rather
|
||||
/// than watched. Every status of the run carries it.
|
||||
/// </summary>
|
||||
public DateTimeOffset? LastSyncUtc { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Documents which need nothing more in this run, whether they were indexed now or before.
|
||||
/// </summary>
|
||||
public int DoneDocuments => this.UnchangedDocuments + this.IndexedDocuments;
|
||||
|
||||
/// <summary>
|
||||
/// Counts documents which nothing changed about since they were indexed.
|
||||
/// </summary>
|
||||
/// <param name="count">How many of them.</param>
|
||||
public void RecordUnchanged(int count = 1) => this.UnchangedDocuments += count;
|
||||
|
||||
/// <summary>
|
||||
/// Counts a document which is not read again, because it failed for a reason of its own before
|
||||
/// and has not changed since.
|
||||
/// </summary>
|
||||
/// <param name="documentKey">The key of the document.</param>
|
||||
/// <param name="failure">Why it failed back then.</param>
|
||||
public void RecordStillUnreadable(string documentKey, PermanentIndexingFailureRecord failure)
|
||||
{
|
||||
this.PermanentlySkippedDocuments++;
|
||||
|
||||
// The stored reason keeps its place in the list, so the user still sees why:
|
||||
this.failures.Add(new DataSourceEmbeddingFailure(documentKey, failure.Message, failure.OccurredAtUtc, ExtractionCode: failure.Code, IsPermanent: true));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Tells the user interface where the run stands.
|
||||
/// </summary>
|
||||
/// <param name="currentDocument">The name of the document being worked on, or an empty string.</param>
|
||||
/// <param name="currentBlock">The block of that document being worked on, when known.</param>
|
||||
/// <param name="currentPage">The page that block is on, when known.</param>
|
||||
public void Publish(string currentDocument = "", int? currentBlock = null, int? currentPage = null) =>
|
||||
context.PublishStatus(this.CreateStatus(DataSourceEmbeddingState.RUNNING, currentDocument, this.LastError, currentBlock, currentPage));
|
||||
|
||||
/// <summary>
|
||||
/// Announces that work on a document starts, and hands out what reports its progress.
|
||||
/// </summary>
|
||||
/// <param name="document">The document.</param>
|
||||
/// <returns>Told about every block of the document, with its number and its page.</returns>
|
||||
public Action<int, int?> BeginDocument(EmbeddingDocument document)
|
||||
{
|
||||
this.Publish(document.DisplayName);
|
||||
|
||||
//
|
||||
// What the page says while one document is being worked on. Without it, a document of
|
||||
// several thousand pages leaves the same sentence standing for hours, and a progress
|
||||
// which never moves cannot be told apart from one which is stuck.
|
||||
//
|
||||
var lastBlockReportUtc = DateTimeOffset.MinValue;
|
||||
return (blockNumber, pageNumber) =>
|
||||
{
|
||||
//
|
||||
// The first block goes out at once, so the line is there instead of blank. After
|
||||
// that, at most one message every BLOCK_PROGRESS_INTERVAL: each one re-renders the
|
||||
// embedding page, the navigation bar and the table in the settings, and the blocks
|
||||
// of a large document arrive far faster than anybody can read them.
|
||||
//
|
||||
var nowUtc = DateTimeOffset.UtcNow;
|
||||
if (blockNumber > 1 && nowUtc - lastBlockReportUtc < BLOCK_PROGRESS_INTERVAL)
|
||||
return;
|
||||
|
||||
lastBlockReportUtc = nowUtc;
|
||||
this.Publish(document.DisplayName, blockNumber, pageNumber);
|
||||
};
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Records a document as indexed, in the index store as well as in the manifest.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Called once the kind of data source is sure the document did not change while it was read.
|
||||
/// Until then, its index row says it has no chunks.
|
||||
/// </remarks>
|
||||
/// <param name="document">The document.</param>
|
||||
/// <param name="chunkCount">How many chunks were stored for it.</param>
|
||||
/// <param name="isNew">Whether the index knew nothing about it before.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
public async Task RecordDocumentIndexedAsync(EmbeddingDocument document, int chunkCount, bool isNew, CancellationToken token)
|
||||
{
|
||||
var embeddedAtUtc = DateTimeOffset.UtcNow;
|
||||
var state = document.State with { ChunkCount = chunkCount, EmbeddedAtUtc = embeddedAtUtc };
|
||||
await context.IndexStore.UpsertFileAsync(context.DataSource.Id, state, token);
|
||||
context.Manifest.Files[document.Key] = new EmbeddedFileRecord(state.Fingerprint, state.FileSize, state.LastWriteUtc, embeddedAtUtc, chunkCount);
|
||||
await context.ForgetPermanentFailureAsync(document.Key, token);
|
||||
|
||||
this.IndexedDocuments++;
|
||||
if (isNew)
|
||||
this.NewDocuments++;
|
||||
else
|
||||
this.ChangedDocuments++;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Records why a document could not be indexed, and removes what the attempt left behind.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Never called for a cancelled run, nor for a vector store which cannot be read at all: those
|
||||
/// are not about one document, and whoever indexes has to let them through.
|
||||
/// </remarks>
|
||||
/// <param name="document">The document.</param>
|
||||
/// <param name="exception">What went wrong.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
public async Task RecordDocumentFailureAsync(EmbeddingDocument document, Exception exception, CancellationToken token)
|
||||
{
|
||||
var dataSource = context.DataSource;
|
||||
switch (exception)
|
||||
{
|
||||
case ProviderRequestException providerFailure:
|
||||
{
|
||||
//
|
||||
// The provider said what went wrong and what the user can do about it. That
|
||||
// sentence is what goes into the status, together with the classification the UI
|
||||
// needs to offer the matching way out.
|
||||
//
|
||||
this.FailedDocuments++;
|
||||
this.LastError = providerFailure.UserMessage;
|
||||
this.failures.Add(new DataSourceEmbeddingFailure(document.Key, providerFailure.UserMessage, DateTimeOffset.UtcNow, providerFailure.FailureReason, providerFailure.StatusCode, context.EmbeddingProvider.Name, DisplayName: document.DisplayName));
|
||||
context.Manifest.Files.Remove(document.Key);
|
||||
await context.ForgetPermanentFailureAsync(document.Key, token);
|
||||
await context.CleanupFailedDocumentAsync(document.Key, token);
|
||||
|
||||
logger.LogWarning(
|
||||
providerFailure,
|
||||
"Failed to embed file '{FilePath}' for data source '{DataSourceName}' because the embedding provider '{EmbeddingProviderName}' failed. FailureReason={FailureReason}, StatusCode={StatusCode}.",
|
||||
document.Key,
|
||||
dataSource.Name,
|
||||
context.EmbeddingProvider.Name,
|
||||
providerFailure.FailureReason,
|
||||
providerFailure.StatusCode);
|
||||
this.Publish(document.DisplayName);
|
||||
|
||||
// Once per kind of failure, not once per document:
|
||||
if (this.reportedFailureReasons.Add(providerFailure.FailureReason))
|
||||
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.CloudOff, providerFailure.UserMessage));
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
case FileExtractionException extractionFailure when extractionFailure.Code.IsPermanentIndexingFailure():
|
||||
{
|
||||
//
|
||||
// The document itself is why this failed, so trying it again changes nothing until
|
||||
// the document does. The reason is written into the index, and the fingerprint next
|
||||
// to it decides when to come back: an OCR run over a scanned PDF changes both size
|
||||
// and write time, which is exactly the moment the file deserves another attempt.
|
||||
//
|
||||
this.PermanentlySkippedDocuments++;
|
||||
var occurredAtUtc = DateTimeOffset.UtcNow;
|
||||
var indexingMessage = this.GetFailureMessage(extractionFailure.Code, document);
|
||||
this.failures.Add(new DataSourceEmbeddingFailure(document.Key, indexingMessage, occurredAtUtc, ExtractionCode: extractionFailure.Code, IsPermanent: true, DisplayName: document.DisplayName));
|
||||
context.Manifest.Files.Remove(document.Key);
|
||||
await context.CleanupFailedDocumentAsync(document.Key, token);
|
||||
|
||||
var state = document.State;
|
||||
context.Manifest.PermanentFailures[state.AbsolutePath] = new PermanentIndexingFailureRecord(state.Fingerprint, extractionFailure.Code, indexingMessage, occurredAtUtc);
|
||||
await context.IndexStore.UpsertPermanentFailureAsync(
|
||||
dataSource.Id,
|
||||
new PermanentIndexingFailure(state.ParentFileId, state.AbsolutePath, state.Fingerprint, extractionFailure.Code, indexingMessage, occurredAtUtc),
|
||||
token);
|
||||
|
||||
logger.LogInformation(
|
||||
extractionFailure,
|
||||
"Skipping file '{FilePath}' of data source '{DataSourceName}' ({DataSourceId}) from now on because reading it failed for a reason which lies in the file. FailureCode={FailureCode}, MetadataHashPrefix={MetadataHashPrefix}.",
|
||||
document.Key,
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
extractionFailure.Code,
|
||||
ShortHash(state.Fingerprint));
|
||||
this.Publish(document.DisplayName);
|
||||
break;
|
||||
}
|
||||
|
||||
default:
|
||||
{
|
||||
//
|
||||
// Everything which is not the provider's doing: a document which changed while it
|
||||
// was read, one which yielded no text, a vector store which refused to store. These
|
||||
// are about this one document, so they go into the list and not into a message
|
||||
// which would interrupt whatever the user is doing right now.
|
||||
//
|
||||
this.FailedDocuments++;
|
||||
var extractionCode = exception is FileExtractionException extractionFailure ? extractionFailure.Code : FileExtractionErrorCode.NONE;
|
||||
|
||||
//
|
||||
// Deliberately not the message of the exception: that one is written for the log
|
||||
// file, in English, and repeats the path which the list shows anyway.
|
||||
//
|
||||
var failureMessage = this.GetFailureMessage(extractionCode, document);
|
||||
this.LastError = failureMessage;
|
||||
this.failures.Add(new DataSourceEmbeddingFailure(document.Key, failureMessage, DateTimeOffset.UtcNow, EmbeddingProviderName: context.EmbeddingProvider.Name, ExtractionCode: extractionCode, DisplayName: document.DisplayName));
|
||||
context.Manifest.Files.Remove(document.Key);
|
||||
await context.ForgetPermanentFailureAsync(document.Key, token);
|
||||
await context.CleanupFailedDocumentAsync(document.Key, token);
|
||||
|
||||
logger.LogWarning(exception, "Failed to embed file '{FilePath}' for data source '{DataSourceName}'.", document.Key, dataSource.Name);
|
||||
this.Publish(document.DisplayName);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Finishes the run: the collection is tidied up, the data source is marked as worked through,
|
||||
/// and the user interface learns how it went.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The hash is written last on purpose. It is what says that a run got through the whole data
|
||||
/// source, so a run which stops before this point leaves the data source marked as unfinished.
|
||||
/// </remarks>
|
||||
/// <param name="sourceHash">The hash of the data source as this run found it.</param>
|
||||
/// <param name="reason">Why the collection is optimized now, for the log.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
public async Task CompleteRunAsync(string sourceHash, string reason, CancellationToken token)
|
||||
{
|
||||
await this.StoreSourceHashAsync(sourceHash, reason, token);
|
||||
|
||||
//
|
||||
// Documents which were skipped for good do not make a run unsuccessful: nothing is left to
|
||||
// try, and a data source made of nothing but scanned images would otherwise ask for
|
||||
// attention forever.
|
||||
//
|
||||
var hasFailures = this.FailedDocuments > 0;
|
||||
var lastError = hasFailures
|
||||
? string.IsNullOrWhiteSpace(this.LastError)
|
||||
? context.DataSource is DataSourceMailbox
|
||||
? TB("Some mails could not be indexed. The list below says which ones and why.")
|
||||
: TB("Some files could not be indexed. The list below says which ones and why.")
|
||||
: this.LastError
|
||||
: string.Empty;
|
||||
|
||||
context.PublishStatus(this.CreateStatus(hasFailures ? DataSourceEmbeddingState.FAILED : DataSourceEmbeddingState.COMPLETED, string.Empty, lastError, null, null));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Ends a run before it got through the whole data source, keeping what it did for the next one.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Stores everything CompleteRunAsync stores, the hash included. For a data source which is
|
||||
/// worked through in several runs, the hash therefore says that its index can be searched, not
|
||||
/// that the data source was worked through as a whole: that is what LastSyncUtc tells. Without
|
||||
/// it, a mailbox would stay out of every search for the hours its first sync takes.
|
||||
///
|
||||
/// The user interface learns nothing here. Whatever comes next tells it: the run being queued
|
||||
/// again, or the question the user has to answer first.
|
||||
/// </remarks>
|
||||
/// <param name="sourceHash">The hash of the data source as this run found it.</param>
|
||||
/// <param name="reason">Why the collection is optimized now, for the log.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
public Task PauseRunAsync(string sourceHash, string reason, CancellationToken token) => this.StoreSourceHashAsync(sourceHash, reason, token);
|
||||
|
||||
/// <summary>
|
||||
/// Tells the user interface how the data source stands as the index holds it, without a run.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// For a data source which is not looked at right now, e.g. a mailbox while AI Studio starts:
|
||||
/// its server is asked nothing then. Nothing is written.
|
||||
/// </remarks>
|
||||
/// <param name="workedThrough">Whether a run got through the whole data source at some point, which decides between completed and idle.</param>
|
||||
public void PublishStoredState(bool workedThrough) =>
|
||||
context.PublishStatus(this.CreateStatus(workedThrough ? DataSourceEmbeddingState.COMPLETED : DataSourceEmbeddingState.IDLE, string.Empty, string.Empty, null, null));
|
||||
|
||||
/// <summary>
|
||||
/// Ends a run which cannot go on, for a reason which is not about any one document.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// E.g. the server of a mailbox which cannot be reached or refuses the sign-in, or a removal
|
||||
/// the user has to agree to first. Nothing is written: whatever the run indexed so far stays,
|
||||
/// and the data source keeps counting as not worked through, so the next run picks up where
|
||||
/// this one stopped.
|
||||
/// </remarks>
|
||||
/// <param name="message">Why the run stopped, and what the user can do about it.</param>
|
||||
/// <param name="attention">What the user has to decide, if anything.</param>
|
||||
/// <param name="pendingRemovalCount">The number of documents a held back removal asks about.</param>
|
||||
public void PublishRunFailure(string message, DataSourceAttention attention = DataSourceAttention.NONE, int? pendingRemovalCount = null)
|
||||
{
|
||||
this.LastError = message;
|
||||
this.failures.Add(new DataSourceEmbeddingFailure(context.DataSource.Name, message, DateTimeOffset.UtcNow));
|
||||
context.PublishStatus(this.CreateStatus(DataSourceEmbeddingState.FAILED, string.Empty, message, null, null, attention, pendingRemovalCount));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Stores what a run found about the data source as a whole, and tidies up the collection.
|
||||
/// </summary>
|
||||
private async Task StoreSourceHashAsync(string sourceHash, string reason, CancellationToken token)
|
||||
{
|
||||
context.Manifest.SourceHash = sourceHash;
|
||||
token.ThrowIfCancellationRequested();
|
||||
await context.OptimizeCollectionIfNeededAsync(reason, token);
|
||||
|
||||
token.ThrowIfCancellationRequested();
|
||||
await context.IndexStore.UpdateDataSourceHashAsync(context.DataSource.Id, sourceHash, token);
|
||||
token.ThrowIfCancellationRequested();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// What the user reads about a document which could not be indexed.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A mail only ever fails as a whole: its text is not read from a file, and an attachment which
|
||||
/// cannot be read costs the mail nothing but that attachment. So one sentence serves every
|
||||
/// reason, and the log holds the details. The sentences about files would speak of a file, and
|
||||
/// of a change which never comes to a mail.
|
||||
/// </remarks>
|
||||
/// <param name="code">Why reading the document failed, NONE when it was not about reading it.</param>
|
||||
/// <param name="document">The document.</param>
|
||||
/// <returns>The message, ready to show.</returns>
|
||||
private string GetFailureMessage(FileExtractionErrorCode code, EmbeddingDocument document) => context.DataSource is DataSourceMailbox
|
||||
? string.Format(TB("The mail '{0}' could not be indexed. AI Studio tries again during the next sync."), document.DisplayName)
|
||||
: code.ToIndexingUserMessage(document.DisplayName);
|
||||
|
||||
private DataSourceEmbeddingStatus CreateStatus(DataSourceEmbeddingState state, string currentDocument, string lastError, int? currentBlock, int? currentPage, DataSourceAttention attention = DataSourceAttention.NONE, int? pendingRemovalCount = null) => new(
|
||||
context.DataSource.Id,
|
||||
context.DataSource.Name,
|
||||
context.DataSource.Type,
|
||||
state,
|
||||
totalDocuments,
|
||||
this.DoneDocuments,
|
||||
this.FailedDocuments,
|
||||
currentDocument,
|
||||
lastError,
|
||||
this.failures.ToList(),
|
||||
this.PermanentlySkippedDocuments,
|
||||
currentBlock,
|
||||
currentPage,
|
||||
Attention: attention,
|
||||
PendingRemovalCount: pendingRemovalCount,
|
||||
LastSyncUtc: this.LastSyncUtc);
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// One chunk as the chunking produced it, together with the page it starts on.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The page is carried rather than read back out of the chunk text. The runtime states it, and
|
||||
/// the chunking knows which source segment a chunk begins in, so nothing has to be derived from
|
||||
/// a marker in the text — which is what used to leave Word files and continued passages without
|
||||
/// a page.
|
||||
/// </remarks>
|
||||
/// <param name="Text">The chunk itself, overlap prefix included.</param>
|
||||
/// <param name="PageNumber">The page the chunk's own content starts on, or null when it has none.</param>
|
||||
internal sealed record EmbeddingChunk(string Text, int? PageNumber);
|
||||
@@ -0,0 +1,17 @@
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// One document of a data source, ready to be embedded.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Everything the shared part of a run needs to know about a document, whatever kind of data source
|
||||
/// it comes from. Where the text comes from, and how it is cut into chunks, stays with the kind of
|
||||
/// data source: the chunks are only read once the old vectors of the document are gone.
|
||||
/// </remarks>
|
||||
/// <param name="Key">What the vectors and the index row of the document are filed under. For a file, its full path.</param>
|
||||
/// <param name="State">The index row of the document, written before its chunks with a chunk count of zero.</param>
|
||||
/// <param name="DisplayName">How the document is called in messages for the user. For a file, its name.</param>
|
||||
/// <param name="StreamChunks">Reads the document and yields its chunks.</param>
|
||||
internal sealed record EmbeddingDocument(string Key, EmbeddingStateFile State, string DisplayName, Func<CancellationToken, IAsyncEnumerable<EmbeddingChunk>> StreamChunks);
|
||||
@@ -0,0 +1,36 @@
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Tools.Security;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// Reads the text of a file through the runtime, in the pieces the chunking cuts along.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The runtime filters prompt injections while it reads, so the text arrives filtered.
|
||||
/// </remarks>
|
||||
internal static class ExtractedFileText
|
||||
{
|
||||
/// <summary>
|
||||
/// Reads the text of a file.
|
||||
/// </summary>
|
||||
/// <param name="rustService">The runtime.</param>
|
||||
/// <param name="filePath">The path of the file.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider whose tokenizer counts the pieces.</param>
|
||||
/// <param name="reportAs">The source the user is told about when passages were filtered out, or null for the file itself.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The text, empty when the file has none.</returns>
|
||||
public static async Task<SegmentedText> ReadAsync(RustService rustService, string filePath, EmbeddingProvider embeddingProvider, PromptInjectionSource? reportAs, CancellationToken token)
|
||||
{
|
||||
var segments = new List<TextSegment>();
|
||||
|
||||
await foreach (var segment in rustService.StreamArbitraryFileDataWithTokenCounts(filePath, embeddingProvider, reportAs, token))
|
||||
{
|
||||
var normalized = TextChunker.NormalizeSegment(segment.Content);
|
||||
if (!string.IsNullOrWhiteSpace(normalized))
|
||||
segments.Add(new(normalized, segment.TokenCount, segment.PageNumber));
|
||||
}
|
||||
|
||||
return new(string.Join("\n", segments.Select(segment => segment.Text)).Trim(), segments);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
using System.Collections.Concurrent;
|
||||
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Settings.DataModel;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
internal sealed partial class FileSourceIndexer
|
||||
{
|
||||
private const int WATCHER_DEBOUNCE_SECONDS = 2;
|
||||
|
||||
private readonly ConcurrentDictionary<string, DataSourceWatcherRegistration> watchers = new(StringComparer.OrdinalIgnoreCase);
|
||||
private readonly Dictionary<string, CancellationTokenSource> watcherDebounceTokens = new(StringComparer.OrdinalIgnoreCase);
|
||||
private readonly object watcherDebounceLock = new();
|
||||
|
||||
/// <summary>
|
||||
/// The data sources handed in last, by their id. A watcher which failed is created anew from here.
|
||||
/// </summary>
|
||||
private IReadOnlyDictionary<string, IIndexedDataSource> trackedSources = new Dictionary<string, IIndexedDataSource>(StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
/// <summary>
|
||||
/// How to ask the embedding service for a run, as handed in last.
|
||||
/// </summary>
|
||||
private Func<string, DataSourceEmbeddingRefreshMode, Task> requestRunCallback = (_, _) => Task.CompletedTask;
|
||||
|
||||
/// <inheritdoc />
|
||||
public void TrackChanges(IReadOnlyCollection<IIndexedDataSource> dataSources, Func<string, DataSourceEmbeddingRefreshMode, Task> requestRun)
|
||||
{
|
||||
this.requestRunCallback = requestRun;
|
||||
|
||||
var supportedSources = dataSources
|
||||
.Where(this.Supports)
|
||||
.ToDictionary(source => source.Id, StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
this.trackedSources = supportedSources;
|
||||
|
||||
foreach (var existingWatcherId in this.watchers.Keys.Except(supportedSources.Keys, StringComparer.OrdinalIgnoreCase).ToList())
|
||||
this.StopTracking(existingWatcherId);
|
||||
|
||||
foreach (var dataSource in supportedSources.Values)
|
||||
this.EnsureWatcher(dataSource);
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public void StopTracking(string dataSourceId)
|
||||
{
|
||||
this.CancelPendingWatcherRefresh(dataSourceId);
|
||||
|
||||
if (this.watchers.TryRemove(dataSourceId, out var registration))
|
||||
registration.Watcher.Dispose();
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public void StopTrackingAll()
|
||||
{
|
||||
foreach (var watcherId in this.watchers.Keys.ToList())
|
||||
this.StopTracking(watcherId);
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public void Dispose()
|
||||
{
|
||||
this.CancelAllPendingWatcherRefreshes();
|
||||
|
||||
foreach (var registration in this.watchers.Values)
|
||||
registration.Watcher.Dispose();
|
||||
|
||||
this.watchers.Clear();
|
||||
}
|
||||
|
||||
private void EnsureWatcher(IDataSourceBase dataSource)
|
||||
{
|
||||
if (!settingsManager.ConfigurationData.App.DataSourceIndexing.AutomaticRefresh)
|
||||
return;
|
||||
|
||||
var configuration = GetWatchConfiguration(dataSource);
|
||||
if (configuration is null)
|
||||
return;
|
||||
|
||||
if (this.watchers.TryGetValue(dataSource.Id, out var existingRegistration))
|
||||
{
|
||||
if (IsSameWatchConfiguration(existingRegistration.Configuration, configuration))
|
||||
return;
|
||||
|
||||
this.StopTracking(dataSource.Id);
|
||||
}
|
||||
|
||||
var watcher = this.CreateWatcher(dataSource.Id, configuration);
|
||||
if (watcher is null)
|
||||
return;
|
||||
|
||||
if (!this.watchers.TryAdd(dataSource.Id, new DataSourceWatcherRegistration(watcher, configuration)))
|
||||
watcher.Dispose();
|
||||
}
|
||||
|
||||
private FileSystemWatcher? CreateWatcher(string dataSourceId, DataSourceWatcherConfiguration configuration)
|
||||
{
|
||||
try
|
||||
{
|
||||
var watcher = new FileSystemWatcher(configuration.RootPath)
|
||||
{
|
||||
Filter = configuration.Filter,
|
||||
IncludeSubdirectories = configuration.IncludeSubdirectories,
|
||||
NotifyFilter = NotifyFilters.FileName | NotifyFilters.DirectoryName | NotifyFilters.LastWrite | NotifyFilters.CreationTime | NotifyFilters.Size,
|
||||
};
|
||||
|
||||
watcher.Changed += (_, args) => this.OnWatchedDataSourceChanged(dataSourceId, configuration, args);
|
||||
watcher.Deleted += (_, args) => this.OnWatchedDataSourceChanged(dataSourceId, configuration, args);
|
||||
watcher.Created += (_, args) => this.OnWatchedDataSourceChanged(dataSourceId, configuration, args);
|
||||
watcher.Renamed += (_, args) => this.OnWatchedDataSourceChanged(dataSourceId, configuration, args);
|
||||
watcher.Error += (_, args) =>
|
||||
{
|
||||
logger.LogWarning(args.GetException(), "The file watcher for data source '{DataSourceId}' failed. Recreating it.", dataSourceId);
|
||||
this.StopTracking(dataSourceId);
|
||||
this.EnsureWatcher(dataSourceId);
|
||||
this.ScheduleWatchedDataSourceRefresh(dataSourceId);
|
||||
};
|
||||
watcher.EnableRaisingEvents = true;
|
||||
return watcher;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Failed to create file watcher for data source '{DataSourceId}' at '{RootPath}'.", dataSourceId, configuration.RootPath);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private void OnWatchedDataSourceChanged(string dataSourceId, DataSourceWatcherConfiguration configuration, FileSystemEventArgs args)
|
||||
{
|
||||
if (!IsRelevantWatcherEvent(configuration, args))
|
||||
{
|
||||
logger.LogDebug(
|
||||
"Ignoring file system change for data source '{DataSourceId}' at '{Path}' (event={ChangeType}) because the path is not part of the RAG index.",
|
||||
dataSourceId,
|
||||
args.FullPath,
|
||||
args.ChangeType);
|
||||
return;
|
||||
}
|
||||
|
||||
logger.LogDebug(
|
||||
"Detected relevant file system change for data source '{DataSourceId}' at '{Path}' (event={ChangeType}). Scheduling a debounced embedding run.",
|
||||
dataSourceId,
|
||||
args.FullPath,
|
||||
args.ChangeType);
|
||||
|
||||
this.ScheduleWatchedDataSourceRefresh(dataSourceId);
|
||||
}
|
||||
|
||||
private void ScheduleWatchedDataSourceRefresh(string dataSourceId)
|
||||
{
|
||||
if (!settingsManager.ConfigurationData.App.DataSourceIndexing.AutomaticRefresh)
|
||||
return;
|
||||
|
||||
var debounceToken = new CancellationTokenSource();
|
||||
|
||||
lock (this.watcherDebounceLock)
|
||||
{
|
||||
if (this.watcherDebounceTokens.Remove(dataSourceId, out var existingToken))
|
||||
existingToken.Cancel();
|
||||
|
||||
this.watcherDebounceTokens[dataSourceId] = debounceToken;
|
||||
}
|
||||
|
||||
_ = Task.Run(async () =>
|
||||
{
|
||||
try
|
||||
{
|
||||
await Task.Delay(TimeSpan.FromSeconds(WATCHER_DEBOUNCE_SECONDS), debounceToken.Token);
|
||||
if (!this.TryCompletePendingWatcherRefresh(dataSourceId, debounceToken))
|
||||
return;
|
||||
|
||||
if (this.trackedSources.TryGetValue(dataSourceId, out var dataSource))
|
||||
{
|
||||
logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) after file system changes settled. The hash pipeline will reindex only changed files.", dataSource.Name, dataSource.Id);
|
||||
await this.requestRunCallback(dataSource.Id, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK);
|
||||
}
|
||||
}
|
||||
catch (OperationCanceledException)
|
||||
{
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Failed to queue watched data source '{DataSourceId}' after a file system change.", dataSourceId);
|
||||
}
|
||||
finally
|
||||
{
|
||||
debounceToken.Dispose();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
private void EnsureWatcher(string dataSourceId)
|
||||
{
|
||||
if (this.trackedSources.TryGetValue(dataSourceId, out var dataSource))
|
||||
this.EnsureWatcher(dataSource);
|
||||
}
|
||||
|
||||
private void CancelPendingWatcherRefresh(string dataSourceId)
|
||||
{
|
||||
lock (this.watcherDebounceLock)
|
||||
{
|
||||
if (this.watcherDebounceTokens.Remove(dataSourceId, out var token))
|
||||
token.Cancel();
|
||||
}
|
||||
}
|
||||
|
||||
private void CancelAllPendingWatcherRefreshes()
|
||||
{
|
||||
lock (this.watcherDebounceLock)
|
||||
{
|
||||
foreach (var token in this.watcherDebounceTokens.Values)
|
||||
token.Cancel();
|
||||
|
||||
this.watcherDebounceTokens.Clear();
|
||||
}
|
||||
}
|
||||
|
||||
private bool TryCompletePendingWatcherRefresh(string dataSourceId, CancellationTokenSource debounceToken)
|
||||
{
|
||||
lock (this.watcherDebounceLock)
|
||||
{
|
||||
if (!this.watcherDebounceTokens.TryGetValue(dataSourceId, out var currentToken) || !ReferenceEquals(currentToken, debounceToken))
|
||||
return false;
|
||||
|
||||
this.watcherDebounceTokens.Remove(dataSourceId);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
private static bool IsRelevantWatcherEvent(DataSourceWatcherConfiguration configuration, FileSystemEventArgs args)
|
||||
{
|
||||
if (args is RenamedEventArgs renamedArgs)
|
||||
{
|
||||
return IsRelevantWatcherPath(configuration, renamedArgs.FullPath, args.ChangeType)
|
||||
|| IsRelevantWatcherPath(configuration, renamedArgs.OldFullPath, args.ChangeType);
|
||||
}
|
||||
|
||||
return IsRelevantWatcherPath(configuration, args.FullPath, args.ChangeType);
|
||||
}
|
||||
|
||||
private static bool IsRelevantWatcherPath(DataSourceWatcherConfiguration configuration, string path, WatcherChangeTypes changeType)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(path))
|
||||
return false;
|
||||
|
||||
var fileName = Path.GetFileName(path);
|
||||
if (string.IsNullOrWhiteSpace(fileName))
|
||||
return true;
|
||||
|
||||
if (!configuration.IncludeSubdirectories && !string.Equals(fileName, configuration.Filter, StringComparison.OrdinalIgnoreCase))
|
||||
return false;
|
||||
|
||||
if (Directory.Exists(path))
|
||||
return true;
|
||||
|
||||
if (IsSkippedRagFileName(fileName))
|
||||
return false;
|
||||
|
||||
if (IsSupportedRagFilePath(path))
|
||||
return true;
|
||||
|
||||
return changeType is WatcherChangeTypes.Deleted or WatcherChangeTypes.Renamed
|
||||
&& string.IsNullOrWhiteSpace(Path.GetExtension(path));
|
||||
}
|
||||
|
||||
private static DataSourceWatcherConfiguration? GetWatchConfiguration(IDataSourceBase dataSource) => dataSource switch
|
||||
{
|
||||
DataSourceLocalDirectory localDirectory when Directory.Exists(localDirectory.Path) => new DataSourceWatcherConfiguration(
|
||||
localDirectory.Path,
|
||||
"*.*",
|
||||
true),
|
||||
DataSourceLocalFile localFile when File.Exists(localFile.FilePath) && !string.IsNullOrWhiteSpace(Path.GetDirectoryName(localFile.FilePath)) => new DataSourceWatcherConfiguration(
|
||||
Path.GetDirectoryName(localFile.FilePath)!,
|
||||
Path.GetFileName(localFile.FilePath),
|
||||
false),
|
||||
_ => null,
|
||||
};
|
||||
|
||||
private static bool IsSameWatchConfiguration(DataSourceWatcherConfiguration left, DataSourceWatcherConfiguration right)
|
||||
{
|
||||
return left.IncludeSubdirectories == right.IncludeSubdirectories
|
||||
&& string.Equals(left.RootPath, right.RootPath, StringComparison.OrdinalIgnoreCase)
|
||||
&& string.Equals(left.Filter, right.Filter, StringComparison.OrdinalIgnoreCase);
|
||||
}
|
||||
|
||||
private sealed record DataSourceWatcherConfiguration(string RootPath, string Filter, bool IncludeSubdirectories);
|
||||
|
||||
private sealed record DataSourceWatcherRegistration(FileSystemWatcher Watcher, DataSourceWatcherConfiguration Configuration);
|
||||
}
|
||||
@@ -0,0 +1,650 @@
|
||||
using System.Security.Cryptography;
|
||||
using System.Text;
|
||||
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using AIStudio.Tools.Rust;
|
||||
using AIStudio.Tools.Security;
|
||||
|
||||
using static AIStudio.Tools.Services.Indexing.IndexingLogFormat;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// Indexes local files and folders.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A file is a document under its full path. Whether it changed is told by a fingerprint over its
|
||||
/// path, name, size and write time, so an unchanged folder is worked through without reading a
|
||||
/// single file. Reading one goes through the runtime, which extracts the text and filters it.
|
||||
///
|
||||
/// Changes are noticed by a file system watcher per data source. A burst of changes is waited out
|
||||
/// before a run is asked for, so saving a document twice in a row costs one run, not two.
|
||||
/// </remarks>
|
||||
/// <param name="settingsManager">The settings, read for whether local data sources refresh on their own.</param>
|
||||
/// <param name="rustService">The runtime, which extracts the text of the files.</param>
|
||||
/// <param name="guardService">The prompt injection filter, whose findings are reported once per run.</param>
|
||||
/// <param name="textChunker">Cuts the text of a file into chunks.</param>
|
||||
/// <param name="logger">The logger of the embedding service, so the log reads the same whoever writes it.</param>
|
||||
internal sealed partial class FileSourceIndexer(SettingsManager settingsManager, RustService rustService, PromptInjectionGuardService guardService, TextChunker textChunker, ILogger logger) : IIndexedSourceIndexer
|
||||
{
|
||||
private const string OFFICE_LOCK_FILE_PREFIX = "~$";
|
||||
private const bool IMAGE_EMBEDDING_ENABLED = false;
|
||||
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(FileSourceIndexer).Namespace, nameof(FileSourceIndexer));
|
||||
|
||||
private enum RagFileIndexingDecision
|
||||
{
|
||||
INDEXABLE,
|
||||
EXCLUDED,
|
||||
UNSUPPORTED,
|
||||
}
|
||||
|
||||
private sealed record DataSourceMetadataSnapshot(string SourceHash, IReadOnlyDictionary<string, string> FileHashes);
|
||||
|
||||
/// <inheritdoc />
|
||||
public bool Supports(IDataSourceBase dataSource) => dataSource is DataSourceLocalDirectory or DataSourceLocalFile;
|
||||
|
||||
/// <inheritdoc />
|
||||
public async Task<IndexedRunOutcome> ProcessAsync(IndexedRunContext context, DataSourceEmbeddingRefreshMode refreshMode, CancellationToken token)
|
||||
{
|
||||
if (context.DataSource is not IDataSource dataSource || !this.Supports(dataSource))
|
||||
throw new ArgumentException("The file indexer reads local files and folders only.", nameof(context));
|
||||
|
||||
var manifest = context.Manifest;
|
||||
|
||||
var inputFiles = this.GetInputFiles(dataSource);
|
||||
var indexedFiles = inputFiles.Files;
|
||||
var totalFiles = indexedFiles.Count + inputFiles.FailedFiles;
|
||||
|
||||
foreach (var failure in inputFiles.Failures)
|
||||
{
|
||||
logger.LogWarning(
|
||||
"Cannot index data source input '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}). Reason='{Reason}'.",
|
||||
failure.DocumentKey,
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
failure.Reason);
|
||||
}
|
||||
|
||||
logger.LogInformation(
|
||||
"Prepared data source '{DataSourceName}' ({DataSourceId}) for embedding. AccessibleFiles={AccessibleFiles}, FailedFiles={FailedFiles}, Collection='{CollectionName}'.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
indexedFiles.Count,
|
||||
inputFiles.FailedFiles,
|
||||
context.CollectionName);
|
||||
|
||||
var metadataSnapshot = this.BuildDataSourceMetadataSnapshot(dataSource, indexedFiles);
|
||||
var removedMissingFiles = await this.RemoveMissingFileEmbeddingsAsync(context, indexedFiles, token);
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
logger.LogInformation(
|
||||
"Compared data source hash for '{DataSourceName}' ({DataSourceId}). StoredSourceHashPrefix={StoredSourceHashPrefix}, CurrentSourceHashPrefix={CurrentSourceHashPrefix}, StoredFileRecords={StoredFileRecords}, CurrentFiles={CurrentFiles}, RemovedMissingFiles={RemovedMissingFiles}.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
ShortHash(manifest.SourceHash),
|
||||
ShortHash(metadataSnapshot.SourceHash),
|
||||
manifest.Files.Count,
|
||||
indexedFiles.Count,
|
||||
removedMissingFiles);
|
||||
|
||||
var progress = new DocumentRunProgress(context, totalFiles, inputFiles.FailedFiles, inputFiles.LastError, inputFiles.Failures, logger);
|
||||
if (CanSkipDataSourceByHash(manifest, metadataSnapshot, indexedFiles))
|
||||
{
|
||||
logger.LogInformation(
|
||||
"Skipping data source '{DataSourceName}' ({DataSourceId}) because the persisted data source hash and all persisted file hashes match. RefreshMode={RefreshMode}, PermanentlySkippedFiles={PermanentlySkippedFiles}.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
refreshMode,
|
||||
manifest.PermanentFailures.Count);
|
||||
|
||||
//
|
||||
// The files which were skipped for good are none of the indexed ones, and their stored
|
||||
// reasons belong into the list even on a run which read nothing at all:
|
||||
//
|
||||
progress.RecordUnchanged(indexedFiles.Count - manifest.PermanentFailures.Count);
|
||||
foreach (var (filePath, permanentFailure) in manifest.PermanentFailures)
|
||||
progress.RecordStillUnreadable(filePath, permanentFailure);
|
||||
|
||||
await progress.CompleteRunAsync(metadataSnapshot.SourceHash, "data source finished after removing missing files", token);
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
token.ThrowIfCancellationRequested();
|
||||
progress.Publish();
|
||||
|
||||
//
|
||||
// Everything the runtime filters out of these files is reported once for the whole data
|
||||
// source. A run over a few thousand documents which removes something in forty of them
|
||||
// is one thing that happened to the user, not forty. The scope ends with this method, so
|
||||
// the report arrives when the run is finished rather than in the middle of it.
|
||||
//
|
||||
await using var promptInjectionReportingScope = guardService.BeginAction();
|
||||
|
||||
foreach (var file in indexedFiles)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
var fingerprint = metadataSnapshot.FileHashes[file.FullName];
|
||||
if (manifest.Files.TryGetValue(file.FullName, out var existingRecord) &&
|
||||
string.Equals(existingRecord.Fingerprint, fingerprint, StringComparison.Ordinal))
|
||||
{
|
||||
logger.LogDebug(
|
||||
"Skipping unchanged file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) because the persisted metadata hash matches. MetadataHashPrefix={MetadataHashPrefix}, LastWriteUtc={LastWriteUtc:O}, FileSize={FileSize}.",
|
||||
file.FullName,
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
ShortHash(fingerprint),
|
||||
file.LastWriteTimeUtc,
|
||||
file.Length);
|
||||
progress.RecordUnchanged();
|
||||
progress.Publish();
|
||||
continue;
|
||||
}
|
||||
|
||||
//
|
||||
// A file which failed for a reason of its own is not read again until it changes.
|
||||
// Without this, a folder holding hundreds of scanned documents without a text layer
|
||||
// would spend half an hour on every start to arrive at the result we already have:
|
||||
//
|
||||
if (manifest.PermanentFailures.TryGetValue(file.FullName, out var permanentFailure) &&
|
||||
string.Equals(permanentFailure.Fingerprint, fingerprint, StringComparison.Ordinal))
|
||||
{
|
||||
logger.LogDebug(
|
||||
"Skipping file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) because reading it failed permanently before. FailureCode={FailureCode}, MetadataHashPrefix={MetadataHashPrefix}, OccurredAtUtc={OccurredAtUtc:O}.",
|
||||
file.FullName,
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
permanentFailure.Code,
|
||||
ShortHash(fingerprint),
|
||||
permanentFailure.OccurredAtUtc);
|
||||
progress.RecordStillUnreadable(file.FullName, permanentFailure);
|
||||
progress.Publish();
|
||||
continue;
|
||||
}
|
||||
|
||||
var document = this.CreateFileDocument(context, dataSource, file, fingerprint);
|
||||
var reportBlockProgress = progress.BeginDocument(document);
|
||||
|
||||
try
|
||||
{
|
||||
logger.LogInformation(
|
||||
"Embedding file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) because {EmbeddingReason}. CurrentMetadataHashPrefix={CurrentMetadataHashPrefix}. Progress={CompletedFiles}/{TotalFiles}.",
|
||||
file.FullName,
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
GetFileEmbeddingReason(file, fingerprint, existingRecord),
|
||||
ShortHash(fingerprint),
|
||||
progress.DoneDocuments + 1,
|
||||
totalFiles);
|
||||
var startedAtUtc = DateTimeOffset.UtcNow;
|
||||
var chunkCount = await context.IndexDocumentAsync(document, reportBlockProgress, token);
|
||||
token.ThrowIfCancellationRequested();
|
||||
var fingerprintAfterEmbedding = BuildFileMetadataHash(file);
|
||||
if (!string.Equals(fingerprint, fingerprintAfterEmbedding, StringComparison.Ordinal))
|
||||
throw new IOException(string.Format(TB("The file '{0}' changed while it was being indexed. What was indexed of it is discarded, and the file is tried again during the next run."), file.FullName));
|
||||
|
||||
await progress.RecordDocumentIndexedAsync(document, chunkCount, existingRecord is null, token);
|
||||
logger.LogInformation(
|
||||
"Embedded file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) successfully. Chunks={ChunkCount}, DurationMs={DurationMs}.",
|
||||
file.FullName,
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
chunkCount,
|
||||
(DateTimeOffset.UtcNow - startedAtUtc).TotalMilliseconds);
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (VectorStoreUnreadableException)
|
||||
{
|
||||
//
|
||||
// Not about this one file: the store of the whole data source cannot be opened, so
|
||||
// every remaining file would fail the same way. Carrying on would fill the list
|
||||
// with one entry per file and hide the single cause behind them.
|
||||
//
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
await progress.RecordDocumentFailureAsync(document, exception, token);
|
||||
}
|
||||
}
|
||||
|
||||
await progress.CompleteRunAsync(metadataSnapshot.SourceHash, "data source embedding run finished", token);
|
||||
logger.LogInformation(
|
||||
"Finished background embeddings for data source '{DataSourceName}' ({DataSourceId}). RefreshMode={RefreshMode}, Embedded={EmbeddedFiles}, New={NewFiles}, Changed={ChangedFiles}, Skipped={SkippedFiles}, PermanentlySkipped={PermanentlySkippedFiles}, RemovedMissing={RemovedMissingFiles}, Failed={FailedFiles}, Total={TotalFiles}, SourceHashPrefix={SourceHashPrefix}.",
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
refreshMode,
|
||||
progress.IndexedDocuments,
|
||||
progress.NewDocuments,
|
||||
progress.ChangedDocuments,
|
||||
progress.UnchangedDocuments,
|
||||
progress.PermanentlySkippedDocuments,
|
||||
removedMissingFiles,
|
||||
progress.FailedDocuments,
|
||||
totalFiles,
|
||||
ShortHash(metadataSnapshot.SourceHash));
|
||||
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Whether a file name marks a file which is never indexed, whatever its type.
|
||||
/// </summary>
|
||||
/// <param name="fileName">The name of the file.</param>
|
||||
/// <returns>True for shortcuts and the lock files of office suites.</returns>
|
||||
private static bool IsSkippedRagFileName(string fileName)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(fileName, FileTypes.SHORTCUT)
|
||||
|| fileName.StartsWith(OFFICE_LOCK_FILE_PREFIX, StringComparison.Ordinal);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Whether a file has a type whose text can be indexed.
|
||||
/// </summary>
|
||||
/// <param name="filePath">The path of the file.</param>
|
||||
/// <returns>True for every document type the runtime can read.</returns>
|
||||
private static bool IsSupportedRagFilePath(string filePath)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(filePath, FileTypes.DOCUMENT);
|
||||
}
|
||||
|
||||
private async Task<int> RemoveMissingFileEmbeddingsAsync(IndexedRunContext context, IReadOnlyCollection<FileInfo> indexedFiles, CancellationToken token)
|
||||
{
|
||||
var manifest = context.Manifest;
|
||||
var existingPaths = indexedFiles
|
||||
.Select(file => file.FullName)
|
||||
.ToHashSet(StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
var removedFiles = 0;
|
||||
foreach (var removedFilePath in manifest.Files.Keys.Except(existingPaths, StringComparer.OrdinalIgnoreCase).ToList())
|
||||
{
|
||||
await context.DeleteDocumentPointsAsync(removedFilePath, token);
|
||||
await context.IndexStore.DeleteFileAsync(context.DataSource.Id, removedFilePath, token);
|
||||
manifest.Files.Remove(removedFilePath);
|
||||
removedFiles++;
|
||||
logger.LogInformation(
|
||||
"Removed stale embeddings for deleted file '{FilePath}' from data source '{DataSourceName}' ({DataSourceId}).",
|
||||
removedFilePath,
|
||||
context.DataSource.Name,
|
||||
context.DataSource.Id);
|
||||
}
|
||||
|
||||
//
|
||||
// A file which is gone needs no mark keeping it out of the index. Without this, the table
|
||||
// would grow with every document the user ever deleted:
|
||||
//
|
||||
foreach (var removedFilePath in manifest.PermanentFailures.Keys.Except(existingPaths, StringComparer.OrdinalIgnoreCase).ToList())
|
||||
await context.ForgetPermanentFailureAsync(removedFilePath, token);
|
||||
|
||||
return removedFiles;
|
||||
}
|
||||
|
||||
/// <remarks>
|
||||
/// A file counts as settled when it was indexed or when it was skipped for good, both with a
|
||||
/// matching fingerprint. Counting only the indexed ones would let a single unreadable document
|
||||
/// send the whole folder through the slow path on every run.
|
||||
/// </remarks>
|
||||
private static bool CanSkipDataSourceByHash(DataSourceEmbeddingManifest manifest, DataSourceMetadataSnapshot metadataSnapshot, IReadOnlyCollection<FileInfo> indexedFiles)
|
||||
{
|
||||
if (!string.Equals(manifest.SourceHash, metadataSnapshot.SourceHash, StringComparison.Ordinal))
|
||||
return false;
|
||||
|
||||
if (manifest.Files.Count + manifest.PermanentFailures.Count != indexedFiles.Count)
|
||||
return false;
|
||||
|
||||
foreach (var file in indexedFiles)
|
||||
{
|
||||
if (!metadataSnapshot.FileHashes.TryGetValue(file.FullName, out var currentHash))
|
||||
return false;
|
||||
|
||||
if (manifest.Files.TryGetValue(file.FullName, out var existingRecord))
|
||||
{
|
||||
if (!string.Equals(existingRecord.Fingerprint, currentHash, StringComparison.Ordinal))
|
||||
return false;
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!manifest.PermanentFailures.TryGetValue(file.FullName, out var permanentFailure))
|
||||
return false;
|
||||
|
||||
if (!string.Equals(permanentFailure.Fingerprint, currentHash, StringComparison.Ordinal))
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
private static string GetFileEmbeddingReason(FileInfo file, string currentHash, EmbeddedFileRecord? existingRecord)
|
||||
{
|
||||
if (existingRecord is null)
|
||||
return "no stored file hash exists";
|
||||
|
||||
var reasons = new List<string>();
|
||||
if (!string.Equals(existingRecord.Fingerprint, currentHash, StringComparison.Ordinal))
|
||||
reasons.Add($"stored hash {ShortHash(existingRecord.Fingerprint)} differs from current hash {ShortHash(currentHash)}");
|
||||
|
||||
if (existingRecord.FileSize != file.Length)
|
||||
reasons.Add($"file size changed from {existingRecord.FileSize} to {file.Length} bytes");
|
||||
|
||||
if (existingRecord.LastWriteUtc != new DateTimeOffset(file.LastWriteTimeUtc))
|
||||
reasons.Add($"last modified time changed from {existingRecord.LastWriteUtc:O} to {file.LastWriteTimeUtc:O}");
|
||||
|
||||
return reasons.Count == 0
|
||||
? "the file hash changed"
|
||||
: string.Join("; ", reasons);
|
||||
}
|
||||
|
||||
private async IAsyncEnumerable<EmbeddingChunk> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
||||
{
|
||||
var options = DataSourceEmbeddingService.GetChunkingOptions(dataSource, embeddingProvider);
|
||||
var strategy = TextChunker.GetStrategyForFile(filePath);
|
||||
var content = await ExtractedFileText.ReadAsync(rustService, filePath, embeddingProvider, null, token);
|
||||
|
||||
await foreach (var chunk in textChunker.SplitAsync(content, strategy, options, embeddingProvider, token))
|
||||
yield return chunk;
|
||||
}
|
||||
|
||||
private FileEnumerationResult GetInputFiles(IDataSource dataSource)
|
||||
{
|
||||
var result = new FileEnumerationResult();
|
||||
|
||||
switch (dataSource)
|
||||
{
|
||||
case DataSourceLocalFile localFile when File.Exists(localFile.FilePath):
|
||||
var file = new FileInfo(localFile.FilePath);
|
||||
switch (this.GetRagFileIndexingDecision(file))
|
||||
{
|
||||
case RagFileIndexingDecision.INDEXABLE:
|
||||
result.Files.Add(file);
|
||||
break;
|
||||
|
||||
case RagFileIndexingDecision.EXCLUDED:
|
||||
logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", file.FullName);
|
||||
break;
|
||||
|
||||
default:
|
||||
result.AddFailure(localFile.FilePath, string.Format(TB("The file '{0}' has a type AI Studio cannot index."), localFile.FilePath));
|
||||
break;
|
||||
}
|
||||
|
||||
return result;
|
||||
|
||||
case DataSourceLocalDirectory localDirectory when Directory.Exists(localDirectory.Path):
|
||||
this.EnumerateAccessibleFiles(localDirectory.Path, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
switch (dataSource)
|
||||
{
|
||||
case DataSourceLocalFile localFile:
|
||||
result.AddFailure(localFile.FilePath, string.Format(TB("The file '{0}' does not exist."), localFile.FilePath));
|
||||
break;
|
||||
|
||||
case DataSourceLocalDirectory localDirectory:
|
||||
result.AddFailure(localDirectory.Path, string.Format(TB("The folder '{0}' does not exist."), localDirectory.Path));
|
||||
break;
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
private void EnumerateAccessibleFiles(string rootPath, FileEnumerationResult result)
|
||||
{
|
||||
var pendingDirectories = new Stack<string>();
|
||||
pendingDirectories.Push(rootPath);
|
||||
|
||||
while (pendingDirectories.Count > 0)
|
||||
{
|
||||
var currentPath = pendingDirectories.Pop();
|
||||
IEnumerable<string> subDirectories;
|
||||
IEnumerable<string> files;
|
||||
|
||||
try
|
||||
{
|
||||
subDirectories = Directory.EnumerateDirectories(currentPath);
|
||||
files = Directory.EnumerateFiles(currentPath);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Cannot access directory '{DirectoryPath}' while indexing.", currentPath);
|
||||
result.AddFailure(currentPath, string.Format(TB("The folder '{0}' could not be opened. Please check whether you are allowed to read it."), currentPath));
|
||||
continue;
|
||||
}
|
||||
|
||||
foreach (var filePath in files)
|
||||
{
|
||||
FileInfo fileInfo;
|
||||
try
|
||||
{
|
||||
fileInfo = new FileInfo(filePath);
|
||||
if (!fileInfo.Exists)
|
||||
continue;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", filePath);
|
||||
result.AddFailure(filePath, string.Format(TB("The file '{0}' could not be read. Please check whether you are allowed to read it."), filePath));
|
||||
continue;
|
||||
}
|
||||
|
||||
switch (this.GetRagFileIndexingDecision(fileInfo))
|
||||
{
|
||||
case RagFileIndexingDecision.INDEXABLE:
|
||||
result.Files.Add(fileInfo);
|
||||
break;
|
||||
|
||||
case RagFileIndexingDecision.EXCLUDED:
|
||||
logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", fileInfo.FullName);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
foreach (var subDirectory in subDirectories)
|
||||
{
|
||||
if (this.IsSkippedRagDirectory(subDirectory))
|
||||
continue;
|
||||
|
||||
pendingDirectories.Push(subDirectory);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private string TryGetRelativePath(IDataSource dataSource, FileInfo file) => dataSource switch
|
||||
{
|
||||
DataSourceLocalDirectory localDirectory => Path.GetRelativePath(localDirectory.Path, file.FullName),
|
||||
_ => file.Name
|
||||
};
|
||||
|
||||
private bool IsImageFilePath(string filePath)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(filePath, FileTypes.IMAGE);
|
||||
}
|
||||
|
||||
private RagFileIndexingDecision GetRagFileIndexingDecision(FileInfo file)
|
||||
{
|
||||
if (this.IsSkippedRagFile(file))
|
||||
return RagFileIndexingDecision.EXCLUDED;
|
||||
|
||||
if (!IMAGE_EMBEDDING_ENABLED && this.IsImageFilePath(file.FullName))
|
||||
return RagFileIndexingDecision.EXCLUDED;
|
||||
|
||||
return IsSupportedRagFilePath(file.FullName)
|
||||
? RagFileIndexingDecision.INDEXABLE
|
||||
: RagFileIndexingDecision.UNSUPPORTED;
|
||||
}
|
||||
|
||||
private bool IsSkippedRagFile(FileInfo file)
|
||||
{
|
||||
if (IsSkippedRagFileName(file.Name))
|
||||
return true;
|
||||
|
||||
try
|
||||
{
|
||||
return file.Attributes.HasFlag(FileAttributes.ReparsePoint)
|
||||
|| file.Attributes.HasFlag(FileAttributes.Offline)
|
||||
|| file.Attributes.HasFlag(FileAttributes.Temporary)
|
||||
|| file.Attributes.HasFlag(FileAttributes.System);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", file.FullName);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
private bool IsSkippedRagDirectory(string path)
|
||||
{
|
||||
try
|
||||
{
|
||||
var directory = new DirectoryInfo(path);
|
||||
return directory.Attributes.HasFlag(FileAttributes.ReparsePoint)
|
||||
|| directory.Attributes.HasFlag(FileAttributes.Offline)
|
||||
|| directory.Attributes.HasFlag(FileAttributes.System);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Cannot inspect directory '{DirectoryPath}' while indexing.", path);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
private DataSourceMetadataSnapshot BuildDataSourceMetadataSnapshot(IDataSource dataSource, IReadOnlyList<FileInfo> indexedFiles)
|
||||
{
|
||||
var fileHashes = indexedFiles
|
||||
.OrderBy(file => file.FullName, StringComparer.OrdinalIgnoreCase)
|
||||
.ToDictionary(file => file.FullName, BuildFileMetadataHash, StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
var sourceHash = dataSource switch
|
||||
{
|
||||
DataSourceLocalFile localFile => indexedFiles.Count > 0
|
||||
? fileHashes[indexedFiles[0].FullName]
|
||||
: BuildMetadataHash("file", localFile.FilePath, Path.GetFileName(localFile.FilePath), "missing", "0"),
|
||||
|
||||
DataSourceLocalDirectory localDirectory => this.BuildDirectoryMetadataHash(localDirectory, indexedFiles, fileHashes),
|
||||
|
||||
_ => BuildMetadataHash(dataSource.Type.ToString(), dataSource.Id, dataSource.Name)
|
||||
};
|
||||
|
||||
return new(sourceHash, fileHashes);
|
||||
}
|
||||
|
||||
private string BuildDirectoryMetadataHash(DataSourceLocalDirectory dataSource, IReadOnlyList<FileInfo> indexedFiles, IReadOnlyDictionary<string, string> fileHashes)
|
||||
{
|
||||
var directory = new DirectoryInfo(dataSource.Path);
|
||||
directory.Refresh();
|
||||
|
||||
var totalSize = 0L;
|
||||
var latestFileWriteTicks = 0L;
|
||||
foreach (var file in indexedFiles)
|
||||
{
|
||||
file.Refresh();
|
||||
if (!file.Exists)
|
||||
continue;
|
||||
|
||||
totalSize += file.Length;
|
||||
latestFileWriteTicks = Math.Max(latestFileWriteTicks, file.LastWriteTimeUtc.Ticks);
|
||||
}
|
||||
|
||||
var latestWriteTicks = Math.Max(directory.LastWriteTimeUtc.Ticks, latestFileWriteTicks);
|
||||
var parts = new List<string>
|
||||
{
|
||||
"directory",
|
||||
directory.FullName,
|
||||
directory.Name,
|
||||
latestWriteTicks.ToString(),
|
||||
totalSize.ToString(),
|
||||
indexedFiles.Count.ToString()
|
||||
};
|
||||
|
||||
foreach (var file in indexedFiles.OrderBy(file => file.FullName, StringComparer.OrdinalIgnoreCase))
|
||||
{
|
||||
parts.Add(this.TryGetRelativePath(dataSource, file));
|
||||
parts.Add(fileHashes[file.FullName]);
|
||||
}
|
||||
|
||||
return BuildMetadataHash(parts);
|
||||
}
|
||||
|
||||
private static string BuildFileMetadataHash(FileInfo file)
|
||||
{
|
||||
file.Refresh();
|
||||
if (!file.Exists)
|
||||
{
|
||||
return BuildMetadataHash(
|
||||
"file",
|
||||
file.FullName,
|
||||
file.Name,
|
||||
"missing",
|
||||
"0");
|
||||
}
|
||||
|
||||
return BuildMetadataHash(
|
||||
"file",
|
||||
file.FullName,
|
||||
file.Name,
|
||||
file.LastWriteTimeUtc.Ticks.ToString(),
|
||||
file.Length.ToString());
|
||||
}
|
||||
|
||||
private static string BuildMetadataHash(params string[] parts)
|
||||
{
|
||||
return BuildMetadataHash((IEnumerable<string>)parts);
|
||||
}
|
||||
|
||||
private static string BuildMetadataHash(IEnumerable<string> parts)
|
||||
{
|
||||
var fingerprintSource = new StringBuilder();
|
||||
foreach (var part in parts)
|
||||
fingerprintSource.Append(part.Length).Append(':').Append(part).Append('|');
|
||||
|
||||
var bytes = SHA256.HashData(Encoding.UTF8.GetBytes(fingerprintSource.ToString()));
|
||||
return Convert.ToHexString(bytes);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Describes one file as a document for the shared part of an indexing run.
|
||||
/// </summary>
|
||||
/// <param name="context">The run the file is indexed in.</param>
|
||||
/// <param name="dataSource">The data source the file belongs to.</param>
|
||||
/// <param name="file">The file.</param>
|
||||
/// <param name="fingerprint">The fingerprint of the file as it is about to be read.</param>
|
||||
/// <returns>The document.</returns>
|
||||
private EmbeddingDocument CreateFileDocument(IndexedRunContext context, IDataSource dataSource, FileInfo file, string fingerprint) => new(
|
||||
file.FullName,
|
||||
this.CreateEmbeddingStateFile(dataSource, file, fingerprint, 0, DateTimeOffset.UtcNow),
|
||||
file.Name,
|
||||
token => this.StreamEmbeddingChunksAsync(file.FullName, dataSource, context.EmbeddingProvider, token));
|
||||
|
||||
private EmbeddingStateFile CreateEmbeddingStateFile(IDataSource dataSource, FileInfo file, string fingerprint, int chunkCount, DateTimeOffset embeddedAtUtc)
|
||||
{
|
||||
file.Refresh();
|
||||
var absolutePath = Path.GetFullPath(file.FullName);
|
||||
return new(
|
||||
IndexedDocumentIds.CreateParentId(dataSource.Id, absolutePath),
|
||||
absolutePath,
|
||||
file.Name,
|
||||
this.TryGetRelativePath(dataSource, file),
|
||||
GetFileType(file),
|
||||
fingerprint,
|
||||
file.Exists ? file.Length : 0,
|
||||
file.Exists ? new DateTimeOffset(file.CreationTimeUtc) : DateTimeOffset.UnixEpoch,
|
||||
file.Exists ? new DateTimeOffset(file.LastWriteTimeUtc) : DateTimeOffset.UnixEpoch,
|
||||
embeddedAtUtc,
|
||||
chunkCount);
|
||||
}
|
||||
|
||||
private static string GetFileType(FileInfo file)
|
||||
{
|
||||
var extension = file.Extension.TrimStart('.').ToLowerInvariant();
|
||||
return string.IsNullOrWhiteSpace(extension) ? "unknown" : extension;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
using AIStudio.Settings;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// Indexes one kind of data source: knows how to find its documents and how to read them.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Everything else is the same for every kind and stays out of here: the embedding service queues
|
||||
/// the runs, prepares each one in an IndexedRunContext and reports the progress, and the context
|
||||
/// embeds and stores the documents. A new kind of data source needs nothing but an indexer of its
|
||||
/// own, and tables of its own for whatever it has to remember beyond its documents.
|
||||
///
|
||||
/// What an indexer does bring is how it notices that a data source changed. Files are watched by
|
||||
/// the file system; another kind may have to look again at an interval. Whether anything is tracked
|
||||
/// at all is not the indexer's to decide: the embedding service starts and stops the tracking.
|
||||
/// </remarks>
|
||||
internal interface IIndexedSourceIndexer : IDisposable
|
||||
{
|
||||
/// <summary>
|
||||
/// Whether this indexer reads the documents of a data source.
|
||||
/// </summary>
|
||||
/// <param name="dataSource">The data source.</param>
|
||||
/// <returns>True when this indexer is the one for it.</returns>
|
||||
public bool Supports(IDataSourceBase dataSource);
|
||||
|
||||
/// <summary>
|
||||
/// Works through the documents of the data source of a run, from finding them to completing the run.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A run may stop after its share of the work and leave the rest to another one, so a data source
|
||||
/// which takes hours does not hold up all others for that long. Whether and when that run comes
|
||||
/// is the embedding service's to decide.
|
||||
/// </remarks>
|
||||
/// <param name="context">The prepared run, whose data source this indexer supports.</param>
|
||||
/// <param name="refreshMode">Why the run was started.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>Whether another run has to carry on.</returns>
|
||||
public Task<IndexedRunOutcome> ProcessAsync(IndexedRunContext context, DataSourceEmbeddingRefreshMode refreshMode, CancellationToken token);
|
||||
|
||||
/// <summary>
|
||||
/// Keeps track of changes to the given data sources, and of nothing else.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Called whenever the configured data sources may have changed, always with all of those this
|
||||
/// indexer supports. A data source which was tracked before and is missing now is no longer
|
||||
/// tracked.
|
||||
/// </remarks>
|
||||
/// <param name="dataSources">The data sources to track.</param>
|
||||
/// <param name="requestRun">Asks the embedding service for a run of the data source with the given id.</param>
|
||||
public void TrackChanges(IReadOnlyCollection<IIndexedDataSource> dataSources, Func<string, DataSourceEmbeddingRefreshMode, Task> requestRun);
|
||||
|
||||
/// <summary>
|
||||
/// Stops tracking changes to one data source, and drops what was about to be reported for it.
|
||||
/// </summary>
|
||||
/// <param name="dataSourceId">The id of the data source.</param>
|
||||
public void StopTracking(string dataSourceId);
|
||||
|
||||
/// <summary>
|
||||
/// Stops tracking changes to any data source.
|
||||
/// </summary>
|
||||
public void StopTrackingAll();
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
using System.Security.Cryptography;
|
||||
using System.Text;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// The ids under which documents and their chunks are stored.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Derived rather than random, so the same document of the same data source always ends up under
|
||||
/// the same ids, in the vector store as well as in the index store. The formats must stay as they
|
||||
/// are: every stored id was made with them.
|
||||
/// </remarks>
|
||||
internal static class IndexedDocumentIds
|
||||
{
|
||||
/// <summary>
|
||||
/// The id of one chunk, which is also the id of its point in the vector store.
|
||||
/// </summary>
|
||||
/// <param name="dataSourceId">The data source the document belongs to.</param>
|
||||
/// <param name="fingerprint">The fingerprint of the document as it was read.</param>
|
||||
/// <param name="chunkIndex">The position of the chunk within the document.</param>
|
||||
/// <returns>The chunk id.</returns>
|
||||
public static string CreateChunkId(string dataSourceId, string fingerprint, int chunkIndex) =>
|
||||
CreateStableGuid($"{dataSourceId}:chunk:{fingerprint}:{chunkIndex}");
|
||||
|
||||
/// <summary>
|
||||
/// The id of a document's row in the index store, which its chunks point at.
|
||||
/// </summary>
|
||||
/// <param name="dataSourceId">The data source the document belongs to.</param>
|
||||
/// <param name="absolutePath">What the row is stored under. For a file, its absolute path.</param>
|
||||
/// <returns>The document id.</returns>
|
||||
public static string CreateParentId(string dataSourceId, string absolutePath) =>
|
||||
CreateStableGuid($"{dataSourceId}:parent-file:{absolutePath}");
|
||||
|
||||
private static string CreateStableGuid(string source)
|
||||
{
|
||||
var hash = SHA256.HashData(Encoding.UTF8.GetBytes(source));
|
||||
var guidBytes = hash[..16].ToArray();
|
||||
|
||||
guidBytes[6] = (byte)((guidBytes[6] & 0x0F) | 0x40);
|
||||
guidBytes[8] = (byte)((guidBytes[8] & 0x3F) | 0x80);
|
||||
|
||||
return new Guid(guidBytes).ToString();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,370 @@
|
||||
using AIStudio.Provider;
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// One indexing run of a data source: what it works with, and what it does to the stores.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Worked out once before the run starts, after the stores, the embedding provider and the stored
|
||||
/// manifest were all found to be usable. Whatever kind of data source is indexed, these are the same
|
||||
/// things, and so is cleaning up after a document. Keeping them here lets each kind of data source
|
||||
/// bring only what is its own: how it finds its documents and how it reads them.
|
||||
///
|
||||
/// A document is addressed by its key, which is what its vectors and its row in the index store are
|
||||
/// filed under. For a file, that is its full path.
|
||||
/// </remarks>
|
||||
/// <param name="dataSource">The data source being indexed.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider the data source points at.</param>
|
||||
/// <param name="provider">The provider instance which embeds the text.</param>
|
||||
/// <param name="vectorStore">The vector store, known to be available.</param>
|
||||
/// <param name="indexStore">The index store, known to be available.</param>
|
||||
/// <param name="manifest">What the index stores about the data source, made to match the current embedding configuration.</param>
|
||||
/// <param name="settingsManager">The settings, which the provider reads when it embeds.</param>
|
||||
/// <param name="publishStatus">Hands a new status of the data source to the embedding service.</param>
|
||||
/// <param name="logger">The logger of the embedding service, so the log reads the same whoever writes it.</param>
|
||||
internal sealed class IndexedRunContext(IIndexedDataSource dataSource, EmbeddingProvider embeddingProvider, IProvider provider, VectorStoreClient vectorStore, IndexStoreClient indexStore, DataSourceEmbeddingManifest manifest, SettingsManager settingsManager, Action<DataSourceEmbeddingStatus> publishStatus, ILogger logger)
|
||||
{
|
||||
/// <summary>
|
||||
/// After how many stored chunks the collection is optimized while a run is still going.
|
||||
/// </summary>
|
||||
private const int VECTOR_STORE_OPTIMIZATION_CHUNK_THRESHOLD = 100_000;
|
||||
|
||||
private long storedChunksSinceLastOptimization;
|
||||
private bool hasPendingChanges;
|
||||
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(IndexedRunContext).Namespace, nameof(IndexedRunContext));
|
||||
|
||||
/// <summary>
|
||||
/// One chunk on its way to the stores, with the id it is stored under.
|
||||
/// </summary>
|
||||
private sealed record EmbeddingChunkDraft(string ChunkId, string Text, int ChunkIndex, int? PageNumber);
|
||||
|
||||
public IIndexedDataSource DataSource => dataSource;
|
||||
|
||||
public EmbeddingProvider EmbeddingProvider => embeddingProvider;
|
||||
|
||||
public IProvider Provider => provider;
|
||||
|
||||
public VectorStoreClient VectorStore => vectorStore;
|
||||
|
||||
public IndexStoreClient IndexStore => indexStore;
|
||||
|
||||
public DataSourceEmbeddingManifest Manifest => manifest;
|
||||
|
||||
public string CollectionName { get; } = DataSourceEmbeddingNames.GetCollectionName(dataSource.Id);
|
||||
|
||||
/// <summary>
|
||||
/// Tells the user interface about a new status of the data source.
|
||||
/// </summary>
|
||||
/// <param name="status">The status.</param>
|
||||
public void PublishStatus(DataSourceEmbeddingStatus status) => publishStatus(status);
|
||||
|
||||
/// <summary>
|
||||
/// Embeds one document and stores it, in place of whatever was stored for it before.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The old vectors and the old index row go first, then the row is written anew with a chunk
|
||||
/// count of zero, so the chunks have something to point at while they arrive batch by batch.
|
||||
/// The final row, with the real chunk count, is written by whoever decides that the document
|
||||
/// was indexed: only the kind of data source knows whether it changed in the meantime.
|
||||
/// </remarks>
|
||||
/// <param name="document">The document to index.</param>
|
||||
/// <param name="reportBlockProgress">Told about every chunk, with its number and its page.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The number of chunks stored for the document.</returns>
|
||||
public async Task<int> IndexDocumentAsync(EmbeddingDocument document, Action<int, int?> reportBlockProgress, CancellationToken token)
|
||||
{
|
||||
logger.LogDebug(
|
||||
"Resetting stored embeddings for file '{FilePath}' in collection '{CollectionName}' before re-indexing.",
|
||||
document.Key,
|
||||
this.CollectionName);
|
||||
await this.DeleteDocumentPointsAsync(document.Key, token);
|
||||
await indexStore.DeleteFileAsync(dataSource.Id, document.Key, token);
|
||||
await indexStore.UpsertFileAsync(dataSource.Id, document.State, token);
|
||||
|
||||
var embeddingBatchSize = Math.Max(1, embeddingProvider.EffectiveEmbeddingBatchSize);
|
||||
var batch = new List<EmbeddingChunkDraft>(embeddingBatchSize);
|
||||
var totalChunkCount = 0;
|
||||
|
||||
await foreach (var chunk in document.StreamChunks(token))
|
||||
{
|
||||
batch.Add(new(IndexedDocumentIds.CreateChunkId(dataSource.Id, document.State.Fingerprint, totalChunkCount), chunk.Text, totalChunkCount, chunk.PageNumber));
|
||||
totalChunkCount++;
|
||||
reportBlockProgress(totalChunkCount, chunk.PageNumber);
|
||||
|
||||
if (batch.Count >= embeddingBatchSize)
|
||||
await this.FlushBatchAsync(document, batch, token);
|
||||
}
|
||||
|
||||
if (batch.Count > 0)
|
||||
await this.FlushBatchAsync(document, batch, token);
|
||||
|
||||
//
|
||||
// The extraction itself did not report a failure, but nothing usable came out of it. For
|
||||
// the index this is the same case as a scanned page without a text layer, which is why it
|
||||
// carries a code of its own instead of an unclassified exception:
|
||||
//
|
||||
if (totalChunkCount == 0)
|
||||
throw new FileExtractionException(FileExtractionErrorCode.NO_CONTENT, string.Format(TB("No text could be read from '{0}'."), document.Key));
|
||||
|
||||
logger.LogDebug(
|
||||
"Generated {ChunkCount} chunks for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
|
||||
totalChunkCount,
|
||||
document.Key,
|
||||
dataSource.Name,
|
||||
dataSource.Id);
|
||||
|
||||
return totalChunkCount;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Removes the vectors of one document from the collection.
|
||||
/// </summary>
|
||||
/// <param name="documentKey">The key of the document.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
public async Task DeleteDocumentPointsAsync(string documentKey, CancellationToken token)
|
||||
{
|
||||
await vectorStore.DeleteEmbeddingByFile(this.CollectionName, documentKey, token);
|
||||
this.hasPendingChanges = true;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Removes whatever a failed attempt left behind of one document.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Never throws for a store which refuses: the document already failed, and that failure is the
|
||||
/// one the user has to hear about.
|
||||
/// </remarks>
|
||||
/// <param name="documentKey">The key of the document.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
public async Task CleanupFailedDocumentAsync(string documentKey, CancellationToken token)
|
||||
{
|
||||
try
|
||||
{
|
||||
await this.DeleteDocumentPointsAsync(documentKey, token);
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(
|
||||
exception,
|
||||
"Could not remove vector points while cleaning up failed embedding for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
|
||||
documentKey,
|
||||
dataSource.Name,
|
||||
dataSource.Id);
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
await indexStore.DeleteFileAsync(dataSource.Id, documentKey, token);
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(
|
||||
exception,
|
||||
"Could not remove embedding state while cleaning up failed embedding for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
|
||||
documentKey,
|
||||
dataSource.Name,
|
||||
dataSource.Id);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Drops the mark which keeps a document out of the index, in the store as well as in the manifest.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Called whenever a document was read, and whenever it failed for a reason outside of itself. The
|
||||
/// state heals on its own that way: a document which becomes readable, or a drive which comes
|
||||
/// back, leaves nothing behind.
|
||||
/// </remarks>
|
||||
/// <param name="documentKey">The key of the document.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
public async Task ForgetPermanentFailureAsync(string documentKey, CancellationToken token)
|
||||
{
|
||||
if (!manifest.PermanentFailures.Remove(documentKey))
|
||||
return;
|
||||
|
||||
await indexStore.DeletePermanentFailureAsync(dataSource.Id, documentKey, token);
|
||||
logger.LogDebug(
|
||||
"Removed the permanent indexing failure of file '{FilePath}' from data source '{DataSourceName}' ({DataSourceId}).",
|
||||
documentKey,
|
||||
dataSource.Name,
|
||||
dataSource.Id);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Counts chunks which reached the collection, and optimizes it once enough of them came together.
|
||||
/// </summary>
|
||||
/// <param name="chunkCount">How many chunks were just stored.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
public async Task RecordStoredChunksAsync(int chunkCount, CancellationToken token)
|
||||
{
|
||||
if (chunkCount > 0)
|
||||
{
|
||||
this.hasPendingChanges = true;
|
||||
this.storedChunksSinceLastOptimization += chunkCount;
|
||||
}
|
||||
|
||||
if (this.storedChunksSinceLastOptimization >= VECTOR_STORE_OPTIMIZATION_CHUNK_THRESHOLD)
|
||||
await this.OptimizeCollectionIfNeededAsync("stored chunk threshold reached", token);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Optimizes the collection when anything in it changed since the last time.
|
||||
/// </summary>
|
||||
/// <param name="reason">Why it is asked for now, for the log.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
public async Task OptimizeCollectionIfNeededAsync(string reason, CancellationToken token)
|
||||
{
|
||||
if (!this.hasPendingChanges)
|
||||
return;
|
||||
|
||||
logger.LogInformation(
|
||||
"Optimizing embedding collection '{CollectionName}' for data source '{DataSourceName}' ({DataSourceId}). Reason='{Reason}', StoredChunksSinceLastOptimization={StoredChunksSinceLastOptimization}, ChunkThreshold={ChunkThreshold}.",
|
||||
this.CollectionName,
|
||||
dataSource.Name,
|
||||
dataSource.Id,
|
||||
reason,
|
||||
this.storedChunksSinceLastOptimization,
|
||||
VECTOR_STORE_OPTIMIZATION_CHUNK_THRESHOLD);
|
||||
|
||||
await vectorStore.OptimizeVectorStore(this.CollectionName, token);
|
||||
this.storedChunksSinceLastOptimization = 0;
|
||||
this.hasPendingChanges = false;
|
||||
}
|
||||
|
||||
private async Task FlushBatchAsync(EmbeddingDocument document, List<EmbeddingChunkDraft> batch, CancellationToken token)
|
||||
{
|
||||
logger.LogDebug(
|
||||
"Requesting embeddings for batch of {ChunkCount} chunks from file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
|
||||
batch.Count,
|
||||
document.Key,
|
||||
dataSource.Name,
|
||||
dataSource.Id);
|
||||
|
||||
var texts = batch.Select(item => item.Text).ToList();
|
||||
IReadOnlyList<IReadOnlyList<float>> vectors;
|
||||
try
|
||||
{
|
||||
vectors = await provider.EmbedTextAsync(embeddingProvider.Model, settingsManager, token, texts);
|
||||
token.ThrowIfCancellationRequested();
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (ProviderRequestException)
|
||||
{
|
||||
//
|
||||
// The provider already named the cause and what to do about it. Wrapping that in a
|
||||
// sentence about a batch of chunks would replace the one thing the user can act on
|
||||
// with the fact that something failed:
|
||||
//
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
//
|
||||
// The messages of these exceptions end up in the log, never in front of the user, so
|
||||
// they name the document by its key: the name a user reads may be the subject of a mail.
|
||||
//
|
||||
throw new InvalidOperationException(string.Format(TB("The embedding provider was not able to embed {0} part(s) of '{1}'. The provider reported: {2}"), batch.Count, document.Key, exception.Message), exception);
|
||||
}
|
||||
|
||||
if (vectors.Count != batch.Count)
|
||||
throw new InvalidOperationException(string.Format(TB("The embedding provider answered with {0} vectors for {1} parts of '{2}'. Please select another embedding model or provider."), vectors.Count, batch.Count, document.Key));
|
||||
|
||||
var vectorSize = vectors.FirstOrDefault()?.Count ?? 0;
|
||||
if (vectorSize <= 0)
|
||||
throw new InvalidOperationException(TB("The embedding provider answered with an empty vector. Please select another embedding model or provider."));
|
||||
|
||||
if (vectors.Any(vector => vector.Count != vectorSize))
|
||||
throw new InvalidOperationException(TB("The embedding provider answered with vectors of different sizes. Please select another embedding model or provider."));
|
||||
|
||||
if (vectors.Any(vector => vector.Any(value => !float.IsFinite(value))))
|
||||
throw new InvalidOperationException(TB("The embedding provider answered with a vector containing an invalid number. Please select another embedding model or provider."));
|
||||
|
||||
if (manifest.VectorSize > 0 && manifest.VectorSize != vectorSize)
|
||||
throw new InvalidOperationException(string.Format(TB("The size of the embedding vectors changed from {0} to {1}. Please save the data source again to index it from scratch."), manifest.VectorSize, vectorSize));
|
||||
|
||||
if (manifest.VectorSize == 0)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
var ensureResult = await vectorStore.EnsureVectorStoreExists(this.CollectionName, dataSource.Name, vectorSize, token);
|
||||
if (!ensureResult.Created)
|
||||
{
|
||||
logger.LogWarning(
|
||||
"Vector store '{CollectionName}' exists for data source '{DataSourceName}' ({DataSourceId}) although no persisted embedding state exists. Replacing the orphaned store before indexing.",
|
||||
this.CollectionName,
|
||||
dataSource.Name,
|
||||
dataSource.Id);
|
||||
await vectorStore.DeleteVectorStore(this.CollectionName, token);
|
||||
ensureResult = await vectorStore.EnsureVectorStoreExists(this.CollectionName, dataSource.Name, vectorSize, token);
|
||||
if (!ensureResult.Created)
|
||||
throw new InvalidOperationException(string.Format(TB("The local index '{0}' could not be created again. Please restart AI Studio and try once more."), this.CollectionName));
|
||||
}
|
||||
|
||||
await indexStore.UpdateVectorSizeAsync(dataSource.Id, vectorSize, token);
|
||||
manifest.VectorSize = vectorSize;
|
||||
logger.LogInformation(
|
||||
"Created embedding collection '{CollectionName}' with vector size {VectorSize} for data source '{DataSourceName}' ({DataSourceId}).",
|
||||
this.CollectionName,
|
||||
vectorSize,
|
||||
dataSource.Name,
|
||||
dataSource.Id);
|
||||
}
|
||||
|
||||
token.ThrowIfCancellationRequested();
|
||||
var embeddedAtUtc = DateTimeOffset.UtcNow;
|
||||
var state = document.State;
|
||||
var points = batch.Select((item, index) => new VectorStoragePoint(
|
||||
item.ChunkId,
|
||||
vectors[index],
|
||||
dataSource.Id,
|
||||
dataSource.Type.ToString(),
|
||||
item.ChunkId,
|
||||
state.ParentFileId,
|
||||
document.Key,
|
||||
state.AbsolutePath,
|
||||
state.FileName,
|
||||
state.RelativePath,
|
||||
state.FileType,
|
||||
item.PageNumber,
|
||||
item.ChunkIndex,
|
||||
item.Text,
|
||||
state.Fingerprint,
|
||||
state.CreationUtc,
|
||||
state.LastWriteUtc,
|
||||
embeddedAtUtc)).ToList();
|
||||
|
||||
await vectorStore.InsertEmbedding(this.CollectionName, points, token);
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
var chunks = batch
|
||||
.Select(chunk => new EmbeddingStateChunk(chunk.ChunkId, state.ParentFileId, chunk.PageNumber, chunk.ChunkIndex, chunk.Text, embeddedAtUtc))
|
||||
.ToList();
|
||||
|
||||
await indexStore.UpsertChunksAsync(dataSource.Id, chunks, token);
|
||||
await this.RecordStoredChunksAsync(batch.Count, token);
|
||||
|
||||
logger.LogDebug(
|
||||
"Stored {ChunkCount} embedded chunks for file '{FilePath}' in collection '{CollectionName}'.",
|
||||
batch.Count,
|
||||
document.Key,
|
||||
this.CollectionName);
|
||||
|
||||
batch.Clear();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// How an indexing run ended, as far as the queue is concerned.
|
||||
/// </summary>
|
||||
internal enum IndexedRunOutcome
|
||||
{
|
||||
/// <summary>
|
||||
/// Nothing is left to do until the data source changes, or somebody asks for a run.
|
||||
/// </summary>
|
||||
DONE,
|
||||
|
||||
/// <summary>
|
||||
/// The run stopped after its share of the work, so other data sources get their turn. The
|
||||
/// embedding service queues another run, which carries on where this one stopped.
|
||||
/// </summary>
|
||||
MORE_TO_DO,
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// How the indexing writes values into the log.
|
||||
/// </summary>
|
||||
internal static class IndexingLogFormat
|
||||
{
|
||||
/// <summary>
|
||||
/// Shortens a hash or a signature to what is needed to tell two of them apart in the log.
|
||||
/// </summary>
|
||||
/// <param name="value">The hash or signature.</param>
|
||||
/// <returns>Its first twelve characters, or a marker when it is empty.</returns>
|
||||
public static string ShortHash(string value)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(value))
|
||||
return "<empty>";
|
||||
|
||||
return value.Length <= 12 ? value : value[..12];
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,198 @@
|
||||
using AIStudio.Settings;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// Asks for a run of each tracked data source at a fixed interval, and soon again for one whose
|
||||
/// server was out of reach.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// For a kind of data source whose changes nothing reports, e.g. a mailbox on a server, so an indexer
|
||||
/// has to look again from time to time. The indexer hands its tracking over to this, and the
|
||||
/// embedding service decides as for every indexer whether anything is tracked at all.
|
||||
///
|
||||
/// One timer serves all data sources of an indexer, and they come round together. The queue works
|
||||
/// through them one after the other anyway. Which data sources are tracked never moves the rounds,
|
||||
/// so editing the data sources cannot keep putting them off. One added later joins the next round;
|
||||
/// whoever adds it queues it right away regardless.
|
||||
///
|
||||
/// The first round comes shortly after tracking starts rather than a whole interval later, so the
|
||||
/// index is not that much behind after AI Studio started. Whether a round or a retry is due goes by
|
||||
/// the clock, cf. IntervalRunSchedule; the timer only asks once per check period. A timer alone
|
||||
/// would count only the time the computer is awake.
|
||||
/// </remarks>
|
||||
/// <param name="interval">The time between two rounds.</param>
|
||||
/// <param name="firstRoundDelay">The time from the start of the tracking to the first round.</param>
|
||||
/// <param name="refreshMode">Why the runs are asked for, as the embedding service gets to know.</param>
|
||||
/// <param name="logger">The logger of the embedding service.</param>
|
||||
internal sealed class IntervalRunRequester(TimeSpan interval, TimeSpan firstRoundDelay, DataSourceEmbeddingRefreshMode refreshMode, ILogger logger) : IDisposable
|
||||
{
|
||||
/// <summary>
|
||||
/// How often the timer asks what is due. After the computer woke up, a round comes no later than this.
|
||||
/// </summary>
|
||||
private static readonly TimeSpan CHECK_PERIOD = TimeSpan.FromMinutes(1);
|
||||
|
||||
private readonly Lock stateLock = new();
|
||||
private readonly IntervalRunSchedule schedule = new(interval, firstRoundDelay, CHECK_PERIOD);
|
||||
private Func<string, DataSourceEmbeddingRefreshMode, Task>? requestRun;
|
||||
private Timer? timer;
|
||||
|
||||
/// <summary>
|
||||
/// Tracks exactly the given data sources from now on.
|
||||
/// </summary>
|
||||
/// <param name="dataSources">The data sources to track.</param>
|
||||
/// <param name="requestRunCallback">Asks the embedding service for a run of the data source with the given id.</param>
|
||||
public void Track(IReadOnlyCollection<IIndexedDataSource> dataSources, Func<string, DataSourceEmbeddingRefreshMode, Task> requestRunCallback)
|
||||
{
|
||||
lock (this.stateLock)
|
||||
{
|
||||
this.schedule.Track(dataSources.Select(dataSource => dataSource.Id), DateTimeOffset.UtcNow);
|
||||
this.requestRun = requestRunCallback;
|
||||
|
||||
if (!this.schedule.IsTrackingAny)
|
||||
{
|
||||
this.StopTimer();
|
||||
return;
|
||||
}
|
||||
|
||||
//
|
||||
// The first check comes when the first round is due, and every further one a check
|
||||
// period later.
|
||||
//
|
||||
this.timer ??= new Timer(this.OnCheckDue, null, firstRoundDelay, CHECK_PERIOD);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Stops tracking one data source.
|
||||
/// </summary>
|
||||
/// <param name="dataSourceId">The id of the data source.</param>
|
||||
public void Stop(string dataSourceId)
|
||||
{
|
||||
lock (this.stateLock)
|
||||
{
|
||||
this.schedule.Stop(dataSourceId);
|
||||
if (!this.schedule.IsTrackingAny)
|
||||
this.StopTimer();
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Stops tracking any data source.
|
||||
/// </summary>
|
||||
public void StopAll()
|
||||
{
|
||||
lock (this.stateLock)
|
||||
{
|
||||
this.schedule.StopAll();
|
||||
this.StopTimer();
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Notes that the server of a data source was out of reach, so it is tried again soon.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Only for a network which is down or a host out of reach. A server which answered, even with a
|
||||
/// refusal, is reached.
|
||||
/// </remarks>
|
||||
/// <param name="dataSourceId">The id of the data source.</param>
|
||||
public void RecordServerOutOfReach(string dataSourceId)
|
||||
{
|
||||
TimeSpan? delay;
|
||||
lock (this.stateLock)
|
||||
delay = this.schedule.RecordServerOutOfReach(dataSourceId, DateTimeOffset.UtcNow);
|
||||
|
||||
if (delay is { } pause)
|
||||
logger.LogInformation("The server of data source '{DataSourceId}' was out of reach. It is tried again in {Minutes} minute(s).", dataSourceId, pause.TotalMinutes);
|
||||
else
|
||||
logger.LogDebug("The server of data source '{DataSourceId}' was out of reach. No earlier retry is planned for it.", dataSourceId);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Notes that the server of a data source could be reached.
|
||||
/// </summary>
|
||||
/// <param name="dataSourceId">The id of the data source.</param>
|
||||
public void RecordServerReached(string dataSourceId)
|
||||
{
|
||||
lock (this.stateLock)
|
||||
this.schedule.RecordServerReached(dataSourceId);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Asks for a run of every tracked data source, which is what one round does.
|
||||
/// </summary>
|
||||
internal async Task RequestRunsAsync()
|
||||
{
|
||||
string[] ids;
|
||||
Func<string, DataSourceEmbeddingRefreshMode, Task>? callback;
|
||||
lock (this.stateLock)
|
||||
{
|
||||
ids = [..this.schedule.DataSourceIds];
|
||||
callback = this.requestRun;
|
||||
}
|
||||
|
||||
await this.RequestAsync(ids, callback);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Asks for the runs which are due now: a round, or the retries of servers which were out of reach.
|
||||
/// </summary>
|
||||
private async Task RequestDueRunsAsync()
|
||||
{
|
||||
(bool IsRound, IReadOnlyList<string> DataSourceIds) due;
|
||||
Func<string, DataSourceEmbeddingRefreshMode, Task>? callback;
|
||||
lock (this.stateLock)
|
||||
{
|
||||
due = this.schedule.TakeDue(DateTimeOffset.UtcNow);
|
||||
callback = this.requestRun;
|
||||
}
|
||||
|
||||
if (!due.IsRound)
|
||||
foreach (var id in due.DataSourceIds)
|
||||
logger.LogInformation("Asking again for a run of data source '{DataSourceId}', since its server was out of reach.", id);
|
||||
|
||||
await this.RequestAsync(due.DataSourceIds, callback);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Asks for a run of each of the given data sources.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A request which fails costs only its own data source the run.
|
||||
/// </remarks>
|
||||
private async Task RequestAsync(IReadOnlyList<string> ids, Func<string, DataSourceEmbeddingRefreshMode, Task>? callback)
|
||||
{
|
||||
if (callback is null)
|
||||
return;
|
||||
|
||||
foreach (var id in ids)
|
||||
{
|
||||
try
|
||||
{
|
||||
await callback(id, refreshMode);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Could not queue the data source '{DataSourceId}' for its regular run.", id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Asks for what is due when the timer elapses. RequestAsync catches every failure itself.
|
||||
/// </summary>
|
||||
private void OnCheckDue(object? state) => _ = this.RequestDueRunsAsync();
|
||||
|
||||
private void StopTimer()
|
||||
{
|
||||
this.timer?.Dispose();
|
||||
this.timer = null;
|
||||
}
|
||||
|
||||
#region Implementation of IDisposable
|
||||
|
||||
public void Dispose() => this.StopAll();
|
||||
|
||||
#endregion
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// Decides which data sources are due for a run: all of them once an interval came round, and one
|
||||
/// whose server was out of reach a little earlier than that.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Everything here goes by the clock of the system. A timer counts only the time the computer is
|
||||
/// awake, so after a night asleep, the rest of an interval would still have to pass before the next
|
||||
/// round. By the clock, that round is due as soon as the computer wakes up.
|
||||
///
|
||||
/// A server out of reach is tried again a few times with growing pauses, instead of only with the
|
||||
/// next round. After waking up or logging in, a VPN tunnel is often up a minute or two later than
|
||||
/// the network, and its mailbox would otherwise wait for a whole interval. The retries start anew
|
||||
/// once the server could be reached, when the tracking starts, and after the computer was asleep,
|
||||
/// so a server which was out of reach all evening is tried again soon the next morning. Once they
|
||||
/// are used up, only the rounds try again: a computer which is offline for hours does not keep
|
||||
/// trying every few minutes.
|
||||
///
|
||||
/// The owner asks what is due once per check period and passes the current time to every call,
|
||||
/// which keeps this free of timers. It is not thread-safe; the owner locks.
|
||||
/// </remarks>
|
||||
/// <param name="interval">The time between two rounds.</param>
|
||||
/// <param name="firstRoundDelay">The time from the start of the tracking to the first round.</param>
|
||||
/// <param name="checkPeriod">How often the owner asks what is due. A gap of more than twice as long means the computer was asleep.</param>
|
||||
internal sealed class IntervalRunSchedule(TimeSpan interval, TimeSpan firstRoundDelay, TimeSpan checkPeriod)
|
||||
{
|
||||
/// <summary>
|
||||
/// The pauses before the retries after a server was out of reach, in this order.
|
||||
/// </summary>
|
||||
internal static readonly TimeSpan[] RETRY_DELAYS = [TimeSpan.FromMinutes(1), TimeSpan.FromMinutes(2), TimeSpan.FromMinutes(4), TimeSpan.FromMinutes(8)];
|
||||
|
||||
/// <summary>
|
||||
/// How much earlier than planned a run counts as due.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The timer of the owner and the clock of the system are two different clocks. A check planned
|
||||
/// for the very moment a run is due may come a little early by the clock, and would otherwise put
|
||||
/// the run off by a whole check period.
|
||||
/// </remarks>
|
||||
private static readonly TimeSpan EARLINESS_TOLERANCE = TimeSpan.FromSeconds(1);
|
||||
|
||||
private readonly HashSet<string> dataSourceIds = new(StringComparer.OrdinalIgnoreCase);
|
||||
private readonly Dictionary<string, RetryState> retries = new(StringComparer.OrdinalIgnoreCase);
|
||||
private DateTimeOffset nextRoundUtc;
|
||||
private DateTimeOffset lastCheckUtc;
|
||||
|
||||
/// <summary>
|
||||
/// Whether any data source is tracked.
|
||||
/// </summary>
|
||||
public bool IsTrackingAny => this.dataSourceIds.Count > 0;
|
||||
|
||||
/// <summary>
|
||||
/// The ids of the tracked data sources.
|
||||
/// </summary>
|
||||
public IReadOnlyCollection<string> DataSourceIds => this.dataSourceIds;
|
||||
|
||||
/// <summary>
|
||||
/// Tracks exactly the given data sources from now on.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Which data sources are tracked never moves the rounds, so editing them cannot keep putting the
|
||||
/// rounds off. Only when nothing was tracked before does the first round come after the first
|
||||
/// round delay.
|
||||
/// </remarks>
|
||||
/// <param name="ids">The ids of the data sources.</param>
|
||||
/// <param name="now">The current time.</param>
|
||||
public void Track(IEnumerable<string> ids, DateTimeOffset now)
|
||||
{
|
||||
var wasTrackingAny = this.IsTrackingAny;
|
||||
this.dataSourceIds.Clear();
|
||||
this.dataSourceIds.UnionWith(ids);
|
||||
|
||||
if (!wasTrackingAny)
|
||||
{
|
||||
this.nextRoundUtc = now + firstRoundDelay;
|
||||
this.lastCheckUtc = now;
|
||||
this.retries.Clear();
|
||||
return;
|
||||
}
|
||||
|
||||
foreach (var id in this.retries.Keys.Where(id => !this.dataSourceIds.Contains(id)).ToList())
|
||||
this.retries.Remove(id);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Stops tracking one data source.
|
||||
/// </summary>
|
||||
/// <param name="id">The id of the data source.</param>
|
||||
public void Stop(string id)
|
||||
{
|
||||
this.dataSourceIds.Remove(id);
|
||||
this.retries.Remove(id);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Stops tracking any data source.
|
||||
/// </summary>
|
||||
public void StopAll()
|
||||
{
|
||||
this.dataSourceIds.Clear();
|
||||
this.retries.Clear();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Takes what is due now: a round of every tracked data source, or the retries whose pause is over.
|
||||
/// </summary>
|
||||
/// <param name="now">The current time.</param>
|
||||
/// <returns>Whether a round is due, and the ids of the data sources to run.</returns>
|
||||
public (bool IsRound, IReadOnlyList<string> DataSourceIds) TakeDue(DateTimeOffset now)
|
||||
{
|
||||
//
|
||||
// The owner asks once per check period. A much longer gap means the computer was asleep,
|
||||
// and right after waking up, a server may be out of reach for a little while.
|
||||
//
|
||||
if (now - this.lastCheckUtc > checkPeriod * 2)
|
||||
this.ChangeEveryRetry(retry => retry with { Attempts = 0 });
|
||||
|
||||
this.lastCheckUtc = now;
|
||||
|
||||
//
|
||||
// A clock which was set back must not put the next round off by more than one interval.
|
||||
//
|
||||
if (this.nextRoundUtc - now > interval)
|
||||
this.nextRoundUtc = now + interval;
|
||||
|
||||
if (IsDue(this.nextRoundUtc, now))
|
||||
{
|
||||
this.nextRoundUtc = now + interval;
|
||||
|
||||
//
|
||||
// The round runs every data source anyway, a pending retry included.
|
||||
//
|
||||
this.ChangeEveryRetry(retry => retry with { DueUtc = null });
|
||||
return (true, [..this.dataSourceIds]);
|
||||
}
|
||||
|
||||
var dueRetries = this.retries.Where(entry => entry.Value.DueUtc is { } dueUtc && IsDue(dueUtc, now)).Select(entry => entry.Key).ToList();
|
||||
foreach (var id in dueRetries)
|
||||
this.retries[id] = this.retries[id] with { DueUtc = null };
|
||||
|
||||
return (false, dueRetries);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Notes that the server of a data source was out of reach, and plans the next retry.
|
||||
/// </summary>
|
||||
/// <param name="id">The id of the data source.</param>
|
||||
/// <param name="now">The current time.</param>
|
||||
/// <returns>The pause before the retry, or null when the retries are used up or the data source is not tracked.</returns>
|
||||
public TimeSpan? RecordServerOutOfReach(string id, DateTimeOffset now)
|
||||
{
|
||||
if (!this.dataSourceIds.Contains(id))
|
||||
return null;
|
||||
|
||||
//
|
||||
// A data source without an entry has no retries behind it, which is just what the default
|
||||
// value says.
|
||||
//
|
||||
var attempts = this.retries.GetValueOrDefault(id).Attempts;
|
||||
if (attempts >= RETRY_DELAYS.Length)
|
||||
return null;
|
||||
|
||||
var delay = RETRY_DELAYS[attempts];
|
||||
this.retries[id] = new RetryState(attempts + 1, now + delay);
|
||||
return delay;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Notes that the server of a data source could be reached, so its next outage starts the retries anew.
|
||||
/// </summary>
|
||||
/// <param name="id">The id of the data source.</param>
|
||||
public void RecordServerReached(string id) => this.retries.Remove(id);
|
||||
|
||||
private static bool IsDue(DateTimeOffset dueUtc, DateTimeOffset now) => now >= dueUtc - EARLINESS_TOLERANCE;
|
||||
|
||||
/// <summary>
|
||||
/// Changes the retries of every data source.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Goes over a copy of the entries, since writing a changed value back would otherwise break the
|
||||
/// enumeration. There are as many entries as data sources whose server was out of reach.
|
||||
/// </remarks>
|
||||
/// <param name="change">How a retry changes.</param>
|
||||
private void ChangeEveryRetry(Func<RetryState, RetryState> change)
|
||||
{
|
||||
foreach (var (id, retry) in this.retries.ToList())
|
||||
this.retries[id] = change(retry);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// The retries of one data source whose server was out of reach. The default value stands for none.
|
||||
/// </summary>
|
||||
/// <param name="Attempts">How many retries were planned since the series started.</param>
|
||||
/// <param name="DueUtc">When the next retry is due, or null when none is planned.</param>
|
||||
private readonly record struct RetryState(int Attempts, DateTimeOffset? DueUtc);
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
using AIStudio.Settings;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// The files an attachment is written to while the runtime reads its text.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The runtime reads documents from disk only. An attachment therefore lands in a file of its own
|
||||
/// for as long as its text is read, and is deleted right after. Their names are random: the
|
||||
/// runtime logs the path it reads, and the name of an attachment may name a person.
|
||||
///
|
||||
/// Whatever a crash left behind is deleted when AI Studio starts again, before the first mailbox
|
||||
/// is synced. That is why the files live in a directory of their own within the data directory of
|
||||
/// AI Studio, rather than in the temporary directory of the OS: deleting everything in it touches
|
||||
/// no file of another program.
|
||||
/// </remarks>
|
||||
internal static class MailAttachmentFiles
|
||||
{
|
||||
private const string DIRECTORY_NAME = "mailAttachments";
|
||||
|
||||
/// <summary>
|
||||
/// Makes up the path of a file for one attachment, and creates the directory it lies in.
|
||||
/// </summary>
|
||||
/// <param name="extension">The extension of the attachment, checked against the known document types beforehand, without a dot.</param>
|
||||
/// <returns>The path, where no file exists yet.</returns>
|
||||
/// <exception cref="InvalidOperationException">The data directory is not known yet.</exception>
|
||||
public static string CreatePath(string extension)
|
||||
{
|
||||
var directory = GetDirectory() ?? throw new InvalidOperationException("The data directory is not known yet.");
|
||||
Directory.CreateDirectory(directory);
|
||||
return Path.Join(directory, $"{Guid.NewGuid():N}.{extension}");
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Deletes the file of one attachment, should it exist.
|
||||
/// </summary>
|
||||
/// <param name="path">The path from CreatePath.</param>
|
||||
/// <param name="logger">The logger, for a file which cannot be deleted. The next start deletes it.</param>
|
||||
public static void Delete(string path, ILogger logger)
|
||||
{
|
||||
try
|
||||
{
|
||||
File.Delete(path);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "The temporary file '{Path}' of a mail attachment could not be deleted. It is deleted on the next start.", path);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Deletes every file an earlier session of AI Studio left behind.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Only to be called before any mailbox is synced, since it deletes the files of a running sync as well.
|
||||
/// </remarks>
|
||||
/// <param name="logger">The logger.</param>
|
||||
public static void DeleteLeftovers(ILogger logger)
|
||||
{
|
||||
if (GetDirectory() is not { } directory || !Directory.Exists(directory))
|
||||
return;
|
||||
|
||||
try
|
||||
{
|
||||
Directory.Delete(directory, true);
|
||||
logger.LogInformation("Deleted the temporary files of mail attachments an earlier session left behind.");
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "The temporary files of mail attachments in '{Directory}' could not be deleted.", directory);
|
||||
}
|
||||
}
|
||||
|
||||
private static string? GetDirectory() => string.IsNullOrWhiteSpace(SettingsManager.DataDirectory)
|
||||
? null
|
||||
: Path.Join(SettingsManager.DataDirectory, DIRECTORY_NAME);
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
using AIStudio.Tools.Mail;
|
||||
|
||||
using MailKit;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// One attachment of a mail as a sync read it.
|
||||
/// </summary>
|
||||
/// <param name="Part">The attachment, from the structure of the mail.</param>
|
||||
/// <param name="Name">Its file name, filtered for prompt injections, or empty when it has none.</param>
|
||||
/// <param name="TextState">Whether its text was read and, when not, why.</param>
|
||||
/// <param name="Content">Its text, which the runtime filtered while reading it. Empty unless the text was read.</param>
|
||||
/// <param name="Strategy">How its text is cut into chunks, by its file type.</param>
|
||||
internal sealed record MailAttachmentText(BodyPartBasic Part, string Name, MailPartTextState TextState, SegmentedText Content, ChunkingStrategy Strategy)
|
||||
{
|
||||
private static readonly SegmentedText NO_TEXT = new(string.Empty, []);
|
||||
|
||||
/// <summary>
|
||||
/// An attachment whose text was not read.
|
||||
/// </summary>
|
||||
/// <param name="part">The attachment, from the structure of the mail.</param>
|
||||
/// <param name="name">Its file name, filtered for prompt injections.</param>
|
||||
/// <param name="textState">Why its text was not read.</param>
|
||||
public static MailAttachmentText WithoutText(BodyPartBasic part, string name, MailPartTextState textState) => new(part, name, textState, NO_TEXT, TextChunker.DOCUMENT_STRATEGY);
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
using System.Runtime.CompilerServices;
|
||||
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Mail;
|
||||
using AIStudio.Tools.Security;
|
||||
|
||||
using MailKit;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <remarks>
|
||||
/// The attachments of a mail: which of them are read, how, and how their text joins the chunks of
|
||||
/// the mail.
|
||||
/// </remarks>
|
||||
internal sealed partial class MailboxIndexer
|
||||
{
|
||||
/// <summary>
|
||||
/// Reads the text of one attachment, should the rules of the mailbox allow it.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// An attachment which cannot be read costs the mail nothing but that attachment: the mail is
|
||||
/// indexed with the reason stored in place of the text. Only a lost connection fails the mail,
|
||||
/// since every other mail would fail the same way.
|
||||
/// </remarks>
|
||||
/// <param name="context">The run the mail is indexed in.</param>
|
||||
/// <param name="mailbox">The mailbox the mail belongs to.</param>
|
||||
/// <param name="connector">The connection, with the folder of the mail open.</param>
|
||||
/// <param name="uid">The UID of the mail.</param>
|
||||
/// <param name="part">The attachment, from the structure of the mail.</param>
|
||||
/// <param name="name">The file name of the attachment, filtered for prompt injections.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The attachment with its text, or with the reason why there is none.</returns>
|
||||
/// <exception cref="MailboxConnectionException">The server could not deliver the attachment.</exception>
|
||||
private async Task<MailAttachmentText> ReadAttachmentAsync(IndexedRunContext context, DataSourceMailbox mailbox, ImapMailboxConnector connector, UniqueId uid, BodyPartBasic part, string name, CancellationToken token)
|
||||
{
|
||||
if (MailAttachmentRules.GetReasonToSkip(part, MailAttachmentRules.GetMaxSizeMegabytes(mailbox)) is { } reasonToSkip)
|
||||
return MailAttachmentText.WithoutText(part, name, reasonToSkip);
|
||||
|
||||
string? path = null;
|
||||
try
|
||||
{
|
||||
path = MailAttachmentFiles.CreatePath(MailAttachmentRules.GetExtension(part));
|
||||
await using (var file = new FileStream(path, FileMode.CreateNew, FileAccess.Write, FileShare.None))
|
||||
await connector.FetchAttachmentAsync(uid, part, file, token);
|
||||
|
||||
// Filtered by the runtime while it reads, and reported as content of the mailbox, not of a file the user never saw:
|
||||
var content = await ExtractedFileText.ReadAsync(rustService, path, context.EmbeddingProvider, PromptInjectionSource.MailContent(mailbox.Name), token);
|
||||
return new(part, name, MailPartTextState.EXTRACTED, content, TextChunker.GetStrategyForFile(path));
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (MailboxConnectionException)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "The text of an attachment of a mail in mailbox '{MailboxId}' could not be read. The mail is indexed without it.", mailbox.Id);
|
||||
return MailAttachmentText.WithoutText(part, name, MailPartTextState.EXTRACTION_FAILED);
|
||||
}
|
||||
finally
|
||||
{
|
||||
if (path is not null)
|
||||
MailAttachmentFiles.Delete(path, logger);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Cuts a mail into chunks: first its own text, then the text of each attachment.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Each attachment is cut on its own, by the strategy of its file type, and its first chunk
|
||||
/// names it. A chunk of an attachment carries no page: the chunks of a mail are told apart by
|
||||
/// the mail alone, so a page would not say which attachment it is a page of.
|
||||
/// </remarks>
|
||||
/// <param name="mailContent">The header block and the text of the mail.</param>
|
||||
/// <param name="attachments">The attachments, those whose text was not read included.</param>
|
||||
/// <param name="options">How large a chunk may become.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider whose tokenizer measures the chunks.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The chunks of the mail, in order.</returns>
|
||||
private async IAsyncEnumerable<EmbeddingChunk> StreamMailChunksAsync(SegmentedText mailContent, IReadOnlyList<MailAttachmentText> attachments, ChunkingOptions options, EmbeddingProvider embeddingProvider, [EnumeratorCancellation] CancellationToken token)
|
||||
{
|
||||
await foreach (var chunk in textChunker.SplitAsync(mailContent, TextChunker.DOCUMENT_STRATEGY, options, embeddingProvider, token))
|
||||
yield return chunk;
|
||||
|
||||
foreach (var attachment in attachments)
|
||||
{
|
||||
if (attachment.TextState is not MailPartTextState.EXTRACTED || attachment.Content.Text.Length is 0)
|
||||
continue;
|
||||
|
||||
// In English like the header block, cf. MailTextBuilder:
|
||||
var heading = attachment.Name.Length > 0 ? $"Attachment: {attachment.Name}" : "Attachment without a name";
|
||||
await foreach (var chunk in textChunker.SplitAsync(attachment.Content, attachment.Strategy, options, embeddingProvider, token, heading))
|
||||
yield return chunk with { PageNumber = null };
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,272 @@
|
||||
using System.Text;
|
||||
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
using AIStudio.Tools.Mail;
|
||||
using AIStudio.Tools.Security;
|
||||
|
||||
using MailKit;
|
||||
|
||||
using MimeKit;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <remarks>
|
||||
/// One mail at a time: what it is called in the index, how its text is read, and what the index
|
||||
/// keeps of it beyond its chunks.
|
||||
/// </remarks>
|
||||
internal sealed partial class MailboxIndexer
|
||||
{
|
||||
/// <summary>
|
||||
/// What the index stores as the type of a mail, where a file has its extension.
|
||||
/// </summary>
|
||||
private const string MAIL_FILE_TYPE = "mail";
|
||||
|
||||
/// <summary>
|
||||
/// The content type of a header block on its own, cf. RFC 6522.
|
||||
/// </summary>
|
||||
private const string HEADER_BLOCK_CONTENT_TYPE = "text/rfc822-headers";
|
||||
|
||||
/// <summary>
|
||||
/// How many fields of a mail pass the filter ahead of the names of its attachments: the header
|
||||
/// block, the body and the subject.
|
||||
/// </summary>
|
||||
private const int FILTERED_MAIL_FIELDS = 3;
|
||||
|
||||
/// <summary>
|
||||
/// Links a mail the index holds already to its place in this folder, or reads and indexes it.
|
||||
/// </summary>
|
||||
private async Task SyncNewMailAsync(IndexedRunContext context, DataSourceMailbox mailbox, ImapMailboxConnector connector, string folderPath, IMessageSummary summary, DocumentRunProgress progress, ISet<string> encounteredKeys, CancellationToken token)
|
||||
{
|
||||
if (summary.Headers is not { } headers)
|
||||
{
|
||||
logger.LogWarning("The server delivered a mail of mailbox '{MailboxId}' without its header block. The mail is skipped.", mailbox.Id);
|
||||
return;
|
||||
}
|
||||
|
||||
var key = MailContentKey.Create(MailSummaryReader.ReadIdentity(summary));
|
||||
var mailId = IndexedDocumentIds.CreateParentId(mailbox.Id, key);
|
||||
var location = new MailLocationRecord(folderPath, summary.UniqueId.Id, MailSummaryReader.ReadFlags(summary.Flags));
|
||||
encounteredKeys.Add(key);
|
||||
|
||||
//
|
||||
// A mail the index holds already, which was moved or copied here, or numbered anew by the
|
||||
// server. It keeps its chunks and only gains a location.
|
||||
//
|
||||
if (await context.IndexStore.AddMailLocationAsync(mailbox.Id, mailId, location, token))
|
||||
{
|
||||
progress.RecordUnchanged();
|
||||
progress.Publish();
|
||||
return;
|
||||
}
|
||||
|
||||
var mailHash = MailSummaryReader.ComputeMailHash(summary);
|
||||
if (context.Manifest.PermanentFailures.TryGetValue(key, out var permanentFailure) && string.Equals(permanentFailure.Fingerprint, mailHash, StringComparison.Ordinal))
|
||||
{
|
||||
progress.RecordStillUnreadable(key, permanentFailure);
|
||||
progress.Publish();
|
||||
return;
|
||||
}
|
||||
|
||||
//
|
||||
// Named after the subject as the server delivered it. The name is for the user alone, who
|
||||
// reads it on the embeddings page as in any mail program. A model only ever gets to read
|
||||
// the filtered subject, which the index stores.
|
||||
//
|
||||
var subject = MailTextNormalization.NormalizeHeaderValue(headers[HeaderId.Subject]);
|
||||
var displayName = subject.Length > 0 ? subject : TB("(no subject)");
|
||||
|
||||
var foundAtUtc = DateTimeOffset.UtcNow;
|
||||
var isNew = !context.Manifest.Files.ContainsKey(key);
|
||||
var document = this.CreateMailDocument(context, mailbox, summary, key, mailId, mailHash, displayName, foundAtUtc, null, []);
|
||||
|
||||
try
|
||||
{
|
||||
var (text, attachments) = await this.ReadMailAsync(context, connector, mailbox, summary, token);
|
||||
document = this.CreateMailDocument(context, mailbox, summary, key, mailId, mailHash, displayName, foundAtUtc, text, attachments);
|
||||
|
||||
var reportBlockProgress = progress.BeginDocument(document);
|
||||
var chunkCount = await context.IndexDocumentAsync(document, reportBlockProgress, token);
|
||||
await progress.RecordDocumentIndexedAsync(document, chunkCount, isNew, token);
|
||||
|
||||
// Only after the last chunk: indexing the document deleted whatever the index kept of the mail.
|
||||
await context.IndexStore.UpsertMailAsync(mailbox.Id, CreateMailRecord(mailId, summary, text, attachments, mailHash, location, foundAtUtc), token);
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (MailboxConnectionException)
|
||||
{
|
||||
// Not about this one mail: the connection is gone, and every other mail would fail the same way.
|
||||
throw;
|
||||
}
|
||||
catch (VectorStoreUnreadableException)
|
||||
{
|
||||
// Not about this one mail either: the store of the whole mailbox cannot be opened.
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
await progress.RecordDocumentFailureAsync(document, exception, token);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Reads the text of a mail and of its attachments, filtered for prompt injections.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Of an encrypted mail, nothing but the header block is read: its text parts and its
|
||||
/// attachments stay on the server. The subject and the names of the attachments are filtered
|
||||
/// on their own as well, since the index stores them apart from the text, and whatever reads
|
||||
/// them there may hand them on to a model. A passage the filter removes from them is therefore
|
||||
/// reported twice.
|
||||
/// </remarks>
|
||||
private async Task<(MailText Text, IReadOnlyList<MailAttachmentText> Attachments)> ReadMailAsync(IndexedRunContext context, ImapMailboxConnector connector, DataSourceMailbox mailbox, IMessageSummary summary, CancellationToken token)
|
||||
{
|
||||
var textParts = MailEncryptionDetection.Detect(summary.Body) is MailEncryptionKind.NONE ? await connector.FetchTextPartsAsync(summary, token) : null;
|
||||
var text = MailTextBuilder.Build(MailSummaryReader.ReadTextSource(summary, textParts));
|
||||
|
||||
// The text itself may reveal an encryption the structure did not, cf. MailTextBuilder:
|
||||
var attachmentParts = text.EncryptionKind is MailEncryptionKind.NONE ? MailSummaryReader.ReadAttachments(summary) : [];
|
||||
|
||||
var source = PromptInjectionSource.MailContent(mailbox.Name);
|
||||
var filtered = await guardService.SanitizeAsync([
|
||||
new(text.HeaderBlock, source),
|
||||
new(text.Body, source),
|
||||
new(text.Subject, source),
|
||||
..attachmentParts.Select(part => new PromptInjectionText(MailTextNormalization.NormalizeHeaderValue(part.FileName), source)),
|
||||
]);
|
||||
|
||||
var attachments = new List<MailAttachmentText>(attachmentParts.Count);
|
||||
for (var index = 0; index < attachmentParts.Count; index++)
|
||||
attachments.Add(await this.ReadAttachmentAsync(context, mailbox, connector, summary.UniqueId, attachmentParts[index], filtered[FILTERED_MAIL_FIELDS + index], token));
|
||||
|
||||
return (text with { HeaderBlock = filtered[0], Body = filtered[1], Subject = filtered[2] }, attachments);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Describes a mail as a document for the shared part of an indexing run.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A mail has no path of its own. Where it lies is kept as its locations, which change without
|
||||
/// the mail being embedded again, so a folder stored here would soon name the wrong one.
|
||||
/// </remarks>
|
||||
/// <param name="context">The run the mail is indexed in.</param>
|
||||
/// <param name="mailbox">The mailbox the mail belongs to.</param>
|
||||
/// <param name="summary">The summary of the mail.</param>
|
||||
/// <param name="key">The content key of the mail.</param>
|
||||
/// <param name="mailId">The id of the mail, which is the id of its document.</param>
|
||||
/// <param name="mailHash">The hash of the mail, which is the fingerprint of its document.</param>
|
||||
/// <param name="displayName">How the mail is called in messages for the user.</param>
|
||||
/// <param name="foundAtUtc">When AI Studio found the mail.</param>
|
||||
/// <param name="text">The filtered text of the mail, or null while it is not read yet. Such a document has no chunks and only serves to record a failure.</param>
|
||||
/// <param name="attachments">The attachments of the mail, empty while it is not read yet.</param>
|
||||
/// <returns>The document.</returns>
|
||||
private EmbeddingDocument CreateMailDocument(IndexedRunContext context, DataSourceMailbox mailbox, IMessageSummary summary, string key, string mailId, string mailHash, string displayName, DateTimeOffset foundAtUtc, MailText? text, IReadOnlyList<MailAttachmentText> attachments)
|
||||
{
|
||||
var (sentAtUtc, receivedAtUtc) = ReadDates(summary, foundAtUtc);
|
||||
var state = new EmbeddingStateFile(
|
||||
mailId,
|
||||
key,
|
||||
text?.Subject ?? string.Empty,
|
||||
string.Empty,
|
||||
MAIL_FILE_TYPE,
|
||||
mailHash,
|
||||
summary.Size ?? 0,
|
||||
sentAtUtc ?? receivedAtUtc,
|
||||
receivedAtUtc,
|
||||
DateTimeOffset.UtcNow,
|
||||
0);
|
||||
|
||||
//
|
||||
// The mail itself is one piece of text: the first chunk starts with the header block, so a
|
||||
// search for a sender or a subject finds the mail. The attachments follow it.
|
||||
//
|
||||
var fullText = text?.FullText ?? string.Empty;
|
||||
var content = new SegmentedText(fullText, fullText.Length is 0 ? [] : [new TextSegment(fullText, null, null)]);
|
||||
var chunkingOptions = DataSourceEmbeddingService.GetChunkingOptions(mailbox, context.EmbeddingProvider);
|
||||
|
||||
return new(key, state, displayName, chunkToken => this.StreamMailChunksAsync(content, attachments, chunkingOptions, context.EmbeddingProvider, chunkToken));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Puts together what the index keeps about a mail beyond its chunks.
|
||||
/// </summary>
|
||||
/// <param name="mailId">The id of the mail, which is the id of its document.</param>
|
||||
/// <param name="summary">The summary of the mail, with its header block and its structure.</param>
|
||||
/// <param name="text">The filtered text of the mail.</param>
|
||||
/// <param name="attachments">The attachments of the mail. Each one is kept, with its text or with the reason why there is none.</param>
|
||||
/// <param name="mailHash">The hash of the mail, cf. MailSummaryReader.ComputeMailHash.</param>
|
||||
/// <param name="location">Where the mail was found.</param>
|
||||
/// <param name="foundAtUtc">When AI Studio found the mail. The index keeps the earliest time it found the mail.</param>
|
||||
/// <returns>The mail as the index keeps it.</returns>
|
||||
internal static MailRecord CreateMailRecord(string mailId, IMessageSummary summary, MailText text, IReadOnlyList<MailAttachmentText> attachments, string mailHash, MailLocationRecord location, DateTimeOffset foundAtUtc)
|
||||
{
|
||||
var headers = summary.Headers ?? throw new ArgumentException("The mail was fetched without its header block.", nameof(summary));
|
||||
var (sentAtUtc, receivedAtUtc) = ReadDates(summary, foundAtUtc);
|
||||
|
||||
//
|
||||
// The header block as the server delivered it, encoded words and all, for whoever has to
|
||||
// judge the mail later. It is never cut into chunks, and it is filtered when it is read.
|
||||
//
|
||||
var headerBlock = MailSummaryReader.ReadHeaderBlock(headers);
|
||||
List<MailPartRecord> parts = [new(MailPartKind.HEADERS, string.Empty, HEADER_BLOCK_CONTENT_TYPE, Encoding.UTF8.GetByteCount(headerBlock), headerBlock, MailPartTextState.EXTRACTED)];
|
||||
|
||||
var bodyPart = text.BodySource switch
|
||||
{
|
||||
MailBodySource.HTML => summary.HtmlBody,
|
||||
MailBodySource.PLAIN_TEXT => summary.TextBody,
|
||||
_ => null,
|
||||
};
|
||||
|
||||
if (bodyPart is not null && text.Body.Length > 0)
|
||||
parts.Add(new(MailPartKind.BODY, string.Empty, bodyPart.ContentType.MimeType, bodyPart.Octets, text.Body, MailPartTextState.EXTRACTED));
|
||||
|
||||
//
|
||||
// Every attachment, whether its text was read or not: that a mail has attachments, and
|
||||
// what they are called, is worth searching for either way.
|
||||
//
|
||||
parts.AddRange(attachments.Select(attachment => new MailPartRecord(
|
||||
MailPartKind.ATTACHMENT,
|
||||
attachment.Name,
|
||||
attachment.Part.ContentType.MimeType,
|
||||
attachment.Part.Octets,
|
||||
attachment.TextState is MailPartTextState.EXTRACTED ? attachment.Content.Text : null,
|
||||
attachment.TextState)));
|
||||
|
||||
return new(
|
||||
mailId,
|
||||
MailHeaders.ReadMessageIds(headers, HeaderId.MessageId).FirstOrDefault(IsStorableMessageId) ?? string.Empty,
|
||||
MailHeaders.ReadMessageIds(headers, HeaderId.InReplyTo).FirstOrDefault(IsStorableMessageId) ?? string.Empty,
|
||||
MailHeaders.ReadMessageIds(headers, HeaderId.References).Where(IsStorableMessageId).ToList(),
|
||||
sentAtUtc,
|
||||
receivedAtUtc,
|
||||
text.Importance,
|
||||
text.EncryptionKind,
|
||||
mailHash,
|
||||
foundAtUtc,
|
||||
MailSummaryReader.ReadAddresses(headers),
|
||||
parts,
|
||||
[location]);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// When the sender says the mail was written, and when it arrived at the server.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A server has to report when a mail arrived. Should one not do so, the date the sender gives
|
||||
/// comes closest, and after that the moment AI Studio found the mail.
|
||||
/// </remarks>
|
||||
private static (DateTimeOffset? SentAtUtc, DateTimeOffset ReceivedAtUtc) ReadDates(IMessageSummary summary, DateTimeOffset foundAtUtc)
|
||||
{
|
||||
var sentAtUtc = summary.Headers is { } headers ? MailHeaders.ReadDate(headers)?.ToUniversalTime() : null;
|
||||
return (sentAtUtc, summary.InternalDate?.ToUniversalTime() ?? sentAtUtc ?? foundAtUtc);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Whether the index can keep a Message-ID, which it stores separated by spaces, cf. UpsertMailAsync.
|
||||
/// </summary>
|
||||
private static bool IsStorableMessageId(string messageId) => messageId.Length > 0 && !messageId.Any(char.IsWhiteSpace);
|
||||
}
|
||||
@@ -0,0 +1,646 @@
|
||||
using System.Globalization;
|
||||
using System.Security.Cryptography;
|
||||
using System.Text;
|
||||
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.Mail;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using AIStudio.Tools.Security;
|
||||
|
||||
using static AIStudio.Tools.Services.Indexing.IndexingLogFormat;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// Indexes e-mail mailboxes on IMAP servers.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A mail is a document under its content key, cf. MailContentKey, the same in every folder it lies
|
||||
/// in. Where it lies is kept apart from it, as its locations: a mail which moves, or which the server
|
||||
/// numbers anew, only changes its locations and is never embedded again.
|
||||
///
|
||||
/// A run signs in once and goes through the folders twice. The first pass asks every folder which
|
||||
/// mails belong into the index and settles all which needs no text: the flags, and what is gone.
|
||||
/// The second pass fetches the new mails, the newest first, links those the index holds already,
|
||||
/// and reads and embeds the others. What is gone is removed only at the very end, after every folder
|
||||
/// had its chance to link a mail which merely moved there. Even then, a mail is only orphaned, and
|
||||
/// deleted one run later, should it not turn up again by then.
|
||||
///
|
||||
/// A run reads a limited number of mails, for a limited time, and leaves the rest to the next run.
|
||||
/// The first sync of a large mailbox takes hours, and the other data sources wait in the same
|
||||
/// queue. Nothing is removed before a run got through the whole mailbox, and a removal of a large
|
||||
/// part of it waits for the user to agree, cf. MailRemovalGuard.
|
||||
///
|
||||
/// Signing in is guarded, because every refused attempt brings an account closer to being locked.
|
||||
/// A refused sign-in is stored, and no run signs in again until the user saves a new password,
|
||||
/// tests the connection, or explicitly asks for another try. That this record can be read at all
|
||||
/// rests on the index store being available, which the embedding service makes sure of before
|
||||
/// every run.
|
||||
///
|
||||
/// Starting AI Studio asks no server anything: the run at startup shows what the index holds. Logs
|
||||
/// name the mailbox by its id, never by a subject, an address, a folder, an attachment or what a
|
||||
/// server answered.
|
||||
///
|
||||
/// An organization may allow only its own mail servers, cf. MailServerPolicy. A mailbox on another
|
||||
/// server is then not even signed in to, and its index stays as it is.
|
||||
/// </remarks>
|
||||
/// <param name="settingsManager">The settings, which say which mail servers the organization allows.</param>
|
||||
/// <param name="rustService">The runtime, which holds the password in the OS keyring and reads the text of attachments.</param>
|
||||
/// <param name="guardService">The prompt injection filter, which every mail passes before it is embedded.</param>
|
||||
/// <param name="textChunker">Cuts the text of a mail into chunks.</param>
|
||||
/// <param name="logger">The logger of the embedding service, so the log reads the same whoever writes it.</param>
|
||||
internal sealed partial class MailboxIndexer(SettingsManager settingsManager, RustService rustService, PromptInjectionGuardService guardService, TextChunker textChunker, ILogger logger) : IIndexedSourceIndexer
|
||||
{
|
||||
/// <summary>
|
||||
/// How many mails are fetched at once, before their text is read one after the other.
|
||||
/// </summary>
|
||||
private const int SUMMARY_BATCH_SIZE = 100;
|
||||
|
||||
/// <summary>
|
||||
/// How many mails one run reads at most, whether embedding them works or not.
|
||||
/// </summary>
|
||||
private const int MAX_READ_MAILS_PER_RUN = 500;
|
||||
|
||||
/// <summary>
|
||||
/// How long one run takes at most, before it leaves the rest to the next one.
|
||||
/// </summary>
|
||||
private static readonly TimeSpan MAX_RUN_DURATION = TimeSpan.FromMinutes(5);
|
||||
|
||||
/// <summary>
|
||||
/// How often every mailbox is synced while AI Studio runs. It goes by the clock, so a computer
|
||||
/// which slept longer than this syncs right after waking up.
|
||||
/// </summary>
|
||||
private static readonly TimeSpan SYNC_INTERVAL = TimeSpan.FromMinutes(16);
|
||||
|
||||
/// <summary>
|
||||
/// How long after the tracking started the mailboxes are synced for the first time. The run at
|
||||
/// startup asks no server anything, so this is the first sync after AI Studio started.
|
||||
/// </summary>
|
||||
private static readonly TimeSpan FIRST_SYNC_DELAY = TimeSpan.FromMinutes(1);
|
||||
|
||||
private readonly IntervalRunRequester syncRequester = new(SYNC_INTERVAL, FIRST_SYNC_DELAY, DataSourceEmbeddingRefreshMode.INTERVAL_CHECK, logger);
|
||||
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(MailboxIndexer).Namespace, nameof(MailboxIndexer));
|
||||
|
||||
/// <summary>
|
||||
/// One folder after the first pass of a run.
|
||||
/// </summary>
|
||||
/// <param name="Folder">The folder as the server lists it.</param>
|
||||
/// <param name="ServerState">How the folder stood on the server during the first pass.</param>
|
||||
/// <param name="StoredAtStart">The folder as the first pass stored it.</param>
|
||||
/// <param name="Plan">What is left to do in the folder.</param>
|
||||
/// <param name="IndexedUidCount">How many of its mails belong into the index.</param>
|
||||
/// <param name="StoredLocationCount">How many locations the index held in the folder before the run.</param>
|
||||
private sealed record PlannedFolder(MailServerFolder Folder, MailFolderState ServerState, MailFolderRecord StoredAtStart, MailFolderSyncPlan Plan, int IndexedUidCount, int StoredLocationCount);
|
||||
|
||||
/// <summary>
|
||||
/// What a run found gone, once it got through the whole mailbox.
|
||||
/// </summary>
|
||||
/// <param name="GoneFolderPaths">The stored folders which no longer belong to the mailbox.</param>
|
||||
/// <param name="Count">How many locations would go, those of the gone folders included.</param>
|
||||
/// <param name="IndexedCount">How many locations the index held before the run.</param>
|
||||
private sealed record PlannedRemovals(IReadOnlyList<string> GoneFolderPaths, int Count, int IndexedCount);
|
||||
|
||||
/// <inheritdoc />
|
||||
public bool Supports(IDataSourceBase dataSource) => dataSource is DataSourceMailbox;
|
||||
|
||||
/// <inheritdoc />
|
||||
public async Task<IndexedRunOutcome> ProcessAsync(IndexedRunContext context, DataSourceEmbeddingRefreshMode refreshMode, CancellationToken token)
|
||||
{
|
||||
if (context.DataSource is not DataSourceMailbox mailbox)
|
||||
throw new ArgumentException("The mailbox indexer reads mailboxes only.", nameof(context));
|
||||
|
||||
//
|
||||
// A server the organization does not allow comes before everything else the mailbox might
|
||||
// wait for: no sign-in will be tried, so neither a refused one nor a held-back removal can
|
||||
// be dealt with. Checked at startup as well, which needs no server:
|
||||
//
|
||||
var policy = MailServerPolicy.Read(settingsManager);
|
||||
if (!policy.IsAllowed(mailbox.Host))
|
||||
{
|
||||
logger.LogInformation("Not connecting to mailbox '{MailboxId}' because the organization allows only its own mail servers, and this one is none of them.", mailbox.Id);
|
||||
(await this.CreateStoredStateProgressAsync(context, token)).PublishRunFailure(MailboxConnectionFailure.SERVER_NOT_ALLOWED.GetDescription(), DataSourceAttention.SERVER_NOT_ALLOWED);
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
if (refreshMode is DataSourceEmbeddingRefreshMode.STARTUP_HASH_CHECK)
|
||||
{
|
||||
logger.LogInformation("Showing the stored index of mailbox '{MailboxId}' without asking its server, since AI Studio is starting.", mailbox.Id);
|
||||
await this.PublishStoredStateAsync(context, token);
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
//
|
||||
// Only an explicit request of the user signs in despite a refused sign-in, and only once:
|
||||
// a retry on the embeddings page, or the repair of the index.
|
||||
//
|
||||
var authFailure = await context.IndexStore.GetMailboxAuthFailureAsync(mailbox.Id, token);
|
||||
if (authFailure is not null && refreshMode is not DataSourceEmbeddingRefreshMode.MANUAL_RETRY)
|
||||
{
|
||||
logger.LogInformation("Not signing in to mailbox '{MailboxId}' because the server refused a sign-in on {FailedAtUtc:O}, and nobody has dealt with that yet.", mailbox.Id, authFailure.FailedAtUtc);
|
||||
(await this.CreateStoredStateProgressAsync(context, token)).PublishRunFailure(GetAuthFailureMessage(authFailure), DataSourceAttention.AUTH_FAILED);
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
var password = await this.ReadPasswordAsync(mailbox);
|
||||
if (password is null)
|
||||
{
|
||||
(await this.CreateStoredStateProgressAsync(context, token)).PublishRunFailure(TB("The password of the mailbox could not be read from the operating system. Please enter it again in the settings of the mailbox."));
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
await using var connector = new ImapMailboxConnector();
|
||||
try
|
||||
{
|
||||
await connector.ConnectAsync(mailbox, password, policy, token);
|
||||
}
|
||||
catch (MailboxConnectionException e)
|
||||
{
|
||||
logger.LogWarning("Signing in to mailbox '{MailboxId}' failed: {Failure} ({ExceptionType}).", mailbox.Id, e.Failure, e.InnerException?.GetType().Name ?? "no inner exception");
|
||||
this.RecordConnectionFailure(mailbox.Id, e.Failure);
|
||||
var progress = await this.CreateStoredStateProgressAsync(context, token);
|
||||
|
||||
//
|
||||
// A network which is down says nothing about the password, so only a refusal is
|
||||
// recorded. From now on, no run signs in on its own anymore, not even after a restart.
|
||||
//
|
||||
if (e.Failure is MailboxConnectionFailure.AUTHENTICATION_FAILED)
|
||||
{
|
||||
var failure = MailboxAuthFailure.FromServerAnswer(e.InnerException?.Message ?? string.Empty);
|
||||
await context.IndexStore.UpsertMailboxAuthFailureAsync(mailbox.Id, failure, token);
|
||||
progress.PublishRunFailure(GetAuthFailureMessage(failure), DataSourceAttention.AUTH_FAILED);
|
||||
}
|
||||
else
|
||||
progress.PublishRunFailure(e.Failure.GetDescription(), e.Failure is MailboxConnectionFailure.SERVER_NOT_ALLOWED ? DataSourceAttention.SERVER_NOT_ALLOWED : DataSourceAttention.NONE);
|
||||
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
this.syncRequester.RecordServerReached(mailbox.Id);
|
||||
if (authFailure is not null)
|
||||
{
|
||||
logger.LogInformation("Signing in to mailbox '{MailboxId}' worked again, so the refused sign-in on record is cleared.", mailbox.Id);
|
||||
await context.IndexStore.ClearMailboxAuthFailureAsync(mailbox.Id, token);
|
||||
}
|
||||
|
||||
return await this.SyncAsync(context, mailbox, connector, refreshMode, token);
|
||||
}
|
||||
|
||||
#region Implementation of IIndexedSourceIndexer's tracking
|
||||
|
||||
/// <inheritdoc />
|
||||
/// <remarks>
|
||||
/// A server reports no changes to a mailbox, so every mailbox is synced at an interval,
|
||||
/// cf. SYNC_INTERVAL, and a mailbox whose server was out of reach is tried again soon. A mailbox
|
||||
/// whose sign-in failed is left out by the embedding service, and it would not sign in anyway.
|
||||
/// </remarks>
|
||||
public void TrackChanges(IReadOnlyCollection<IIndexedDataSource> dataSources, Func<string, DataSourceEmbeddingRefreshMode, Task> requestRun) => this.syncRequester.Track(dataSources, requestRun);
|
||||
|
||||
/// <inheritdoc />
|
||||
public void StopTracking(string dataSourceId) => this.syncRequester.Stop(dataSourceId);
|
||||
|
||||
/// <inheritdoc />
|
||||
public void StopTrackingAll() => this.syncRequester.StopAll();
|
||||
|
||||
/// <inheritdoc />
|
||||
public void Dispose() => this.syncRequester.Dispose();
|
||||
|
||||
#endregion
|
||||
|
||||
/// <summary>
|
||||
/// Works through the folders of a mailbox, once the connector is signed in.
|
||||
/// </summary>
|
||||
private async Task<IndexedRunOutcome> SyncAsync(IndexedRunContext context, DataSourceMailbox mailbox, ImapMailboxConnector connector, DataSourceEmbeddingRefreshMode refreshMode, CancellationToken token)
|
||||
{
|
||||
var runStartedUtc = DateTimeOffset.UtcNow;
|
||||
DocumentRunProgress? progress = null;
|
||||
|
||||
try
|
||||
{
|
||||
var selection = MailFolderSelection.Select(await connector.GetFoldersAsync(token), mailbox.RootFolder);
|
||||
if (!selection.RootFolderFound)
|
||||
{
|
||||
//
|
||||
// Never the same as an empty mailbox. Read that way, renaming a folder on the server
|
||||
// would remove every mail from the index, and embedding them again costs hours.
|
||||
//
|
||||
logger.LogWarning("The server no longer lists the root folder of mailbox '{MailboxId}'. The sync stops, and nothing is removed from the index.", mailbox.Id);
|
||||
(await this.CreateStoredStateProgressAsync(context, token)).PublishRunFailure(string.Format(TB("The server no longer lists the folder '{0}' to which this mailbox is limited. Perhaps it was renamed or deleted. Nothing was removed from the index. Rename the folder back on the server, or add the mailbox anew."), mailbox.RootFolder));
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
var syncState = await context.IndexStore.GetMailboxSyncStateAsync(mailbox.Id, token);
|
||||
var storedFolders = (await context.IndexStore.GetMailFoldersAsync(mailbox.Id, token)).ToDictionary(folder => folder.Path, StringComparer.Ordinal);
|
||||
var plannedFolders = await this.PlanFoldersAsync(context, mailbox, connector, selection, storedFolders, token);
|
||||
|
||||
//
|
||||
// Counted by location: a mail which lies in two folders counts twice, as it does on the
|
||||
// server. Telling the mails apart would mean fetching every one of them.
|
||||
//
|
||||
progress = new DocumentRunProgress(context, plannedFolders.Sum(folder => folder.IndexedUidCount), 0, string.Empty, [], logger)
|
||||
{
|
||||
LastSyncUtc = syncState.LastSyncCompletedUtc,
|
||||
};
|
||||
|
||||
progress.RecordUnchanged(plannedFolders.Sum(folder => folder.Plan.KeptUids.Count));
|
||||
progress.Publish();
|
||||
|
||||
//
|
||||
// Whatever the prompt injection filter removes from the mails of this run is reported
|
||||
// once, when the second pass is done, rather than once per mail.
|
||||
//
|
||||
var encounteredKeys = new HashSet<string>(StringComparer.Ordinal);
|
||||
var workedThrough = true;
|
||||
await using (guardService.BeginAction())
|
||||
{
|
||||
foreach (var folder in plannedFolders)
|
||||
{
|
||||
if (await this.SyncNewMailsAsync(context, mailbox, connector, folder, progress, encounteredKeys, runStartedUtc, token))
|
||||
continue;
|
||||
|
||||
workedThrough = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
var sourceHash = BuildSourceHash(plannedFolders);
|
||||
if (!workedThrough)
|
||||
return await this.PauseAsync(mailbox, progress, sourceHash, refreshMode, token);
|
||||
|
||||
//
|
||||
// Every folder is through, so what is gone can be told now: a mail missing from one
|
||||
// folder had its chance to turn up in another.
|
||||
//
|
||||
var removals = await CollectRemovalsAsync(context, mailbox, plannedFolders, storedFolders.Keys, token);
|
||||
if (MailRemovalGuard.Decide(removals.IndexedCount, removals.Count, syncState) is MailRemovalDecision.HOLD_BACK)
|
||||
{
|
||||
logger.LogWarning("Syncing mailbox '{MailboxId}' would remove {RemovalCount} of {IndexedCount} location(s) from the index at once. The removal waits for the user to agree.", mailbox.Id, removals.Count, removals.IndexedCount);
|
||||
await context.IndexStore.HoldBackMailRemovalAsync(mailbox.Id, removals.Count, token);
|
||||
await progress.PauseRunAsync(sourceHash, "mailbox sync waits for a removal to be approved", token);
|
||||
progress.PublishRunFailure(GetRemovalMessage(removals.Count), DataSourceAttention.MASS_REMOVAL_PENDING, removals.Count);
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
var deletedMails = await this.ApplyRemovalsAsync(context, mailbox, plannedFolders, removals, runStartedUtc, token);
|
||||
await ForgetVanishedFailuresAsync(context, encounteredKeys, token);
|
||||
|
||||
var completedUtc = DateTimeOffset.UtcNow;
|
||||
await context.IndexStore.CompleteMailboxSyncAsync(mailbox.Id, completedUtc, token);
|
||||
progress.LastSyncUtc = completedUtc;
|
||||
|
||||
await progress.CompleteRunAsync(sourceHash, "mailbox sync finished", token);
|
||||
logger.LogInformation(
|
||||
"Finished syncing mailbox '{MailboxId}'. RefreshMode={RefreshMode}, Folders={FolderCount}, Embedded={EmbeddedMails}, LinkedOrUnchanged={UnchangedMails}, PermanentlySkipped={PermanentlySkippedMails}, Failed={FailedMails}, RemovedLocations={RemovedLocations}, Deleted={DeletedMails}, SourceHashPrefix={SourceHashPrefix}.",
|
||||
mailbox.Id,
|
||||
refreshMode,
|
||||
plannedFolders.Count,
|
||||
progress.IndexedDocuments,
|
||||
progress.UnchangedDocuments,
|
||||
progress.PermanentlySkippedDocuments,
|
||||
progress.FailedDocuments,
|
||||
removals.Count,
|
||||
deletedMails,
|
||||
ShortHash(sourceHash));
|
||||
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
catch (MailboxConnectionException e)
|
||||
{
|
||||
//
|
||||
// Nothing was removed yet, since that happens at the very end. What this run indexed
|
||||
// stays, and the next run picks up from there.
|
||||
//
|
||||
logger.LogWarning("Syncing mailbox '{MailboxId}' stopped: {Failure} ({ExceptionType}). What was indexed so far stays.", mailbox.Id, e.Failure, e.InnerException?.GetType().Name ?? "no inner exception");
|
||||
this.RecordConnectionFailure(mailbox.Id, e.Failure);
|
||||
await context.OptimizeCollectionIfNeededAsync("mailbox sync stopped", token);
|
||||
progress ??= await this.CreateStoredStateProgressAsync(context, token);
|
||||
progress.PublishRunFailure(e.Failure.GetDescription());
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Tells the sync requester whether a failed connection reached the server at all.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A server out of reach is tried again soon, since a VPN tunnel is often up only a little
|
||||
/// later. A server which answered, even with a refusal, waits for the next round; settings
|
||||
/// which never let a connection start, and a server the organization does not allow, say
|
||||
/// nothing about the server.
|
||||
/// </remarks>
|
||||
/// <param name="mailboxId">The id of the mailbox.</param>
|
||||
/// <param name="failure">Why the connection failed.</param>
|
||||
private void RecordConnectionFailure(string mailboxId, MailboxConnectionFailure failure)
|
||||
{
|
||||
if (failure is MailboxConnectionFailure.NETWORK_UNAVAILABLE)
|
||||
this.syncRequester.RecordServerOutOfReach(mailboxId);
|
||||
else if (failure is not (MailboxConnectionFailure.INVALID_SETTINGS or MailboxConnectionFailure.SERVER_NOT_ALLOWED))
|
||||
this.syncRequester.RecordServerReached(mailboxId);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Ends a run which did its share of the work before it got through the mailbox.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Only a run which got something done hands on to another one. When every mail it read
|
||||
/// failed, e.g. because the embedding provider is down, the next run would fail all the same,
|
||||
/// and the one after it, without end. Such a run ends with its failures instead, and waits for
|
||||
/// the next time the mailbox is queued.
|
||||
/// </remarks>
|
||||
private async Task<IndexedRunOutcome> PauseAsync(DataSourceMailbox mailbox, DocumentRunProgress progress, string sourceHash, DataSourceEmbeddingRefreshMode refreshMode, CancellationToken token)
|
||||
{
|
||||
if (progress.IndexedDocuments is 0 && progress.FailedDocuments > 0)
|
||||
{
|
||||
logger.LogWarning("Syncing mailbox '{MailboxId}' stops without carrying on, since none of the {FailedMails} mail(s) it read could be indexed.", mailbox.Id, progress.FailedDocuments);
|
||||
await progress.CompleteRunAsync(sourceHash, "mailbox sync stopped without progress", token);
|
||||
return IndexedRunOutcome.DONE;
|
||||
}
|
||||
|
||||
await progress.PauseRunAsync(sourceHash, "mailbox sync paused", token);
|
||||
logger.LogInformation(
|
||||
"Syncing mailbox '{MailboxId}' carries on in another run. RefreshMode={RefreshMode}, Embedded={EmbeddedMails}, Failed={FailedMails}, Done={DoneMails}/{TotalMails}.",
|
||||
mailbox.Id,
|
||||
refreshMode,
|
||||
progress.IndexedDocuments,
|
||||
progress.FailedDocuments,
|
||||
progress.DoneDocuments,
|
||||
progress.TotalDocuments);
|
||||
|
||||
return IndexedRunOutcome.MORE_TO_DO;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// The first pass: finds out what there is to do in every folder, and settles all of it which needs no text.
|
||||
/// </summary>
|
||||
private async Task<IReadOnlyList<PlannedFolder>> PlanFoldersAsync(IndexedRunContext context, DataSourceMailbox mailbox, ImapMailboxConnector connector, MailFolderSelection selection, IReadOnlyDictionary<string, MailFolderRecord> storedFolders, CancellationToken token)
|
||||
{
|
||||
var receivedSince = mailbox.MaxAge.GetReceivedSince(DateTimeOffset.UtcNow);
|
||||
var plannedFolders = new List<PlannedFolder>(selection.Folders.Count);
|
||||
foreach (var folder in selection.Folders)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
var serverState = await connector.OpenFolderAsync(folder.FullName, token);
|
||||
var storedFolder = storedFolders.GetValueOrDefault(folder.FullName);
|
||||
var storedLocations = await context.IndexStore.GetMailLocationsAsync(mailbox.Id, folder.FullName, token);
|
||||
var indexedUids = await connector.SearchIndexedMailsAsync(receivedSince, token);
|
||||
var plan = MailFolderSyncPlan.Create(storedFolder, storedLocations, serverState, indexedUids);
|
||||
|
||||
//
|
||||
// Stored before anything else happens in the folder: a new UIDVALIDITY drops the stored
|
||||
// locations, and it has to do so before the first location under it is linked. Where
|
||||
// the last complete pass got to stays until this pass is complete as well.
|
||||
//
|
||||
var storedAtStart = new MailFolderRecord(
|
||||
folder.FullName,
|
||||
folder.SpecialUse,
|
||||
serverState.UidValidity,
|
||||
plan.UidValidityChanged ? null : storedFolder?.UidNext,
|
||||
plan.UidValidityChanged ? null : storedFolder?.HighestModSeq,
|
||||
serverState.MessageCount,
|
||||
serverState.UnseenCount,
|
||||
storedFolder?.InitialSyncCompletedUtc);
|
||||
|
||||
await context.IndexStore.UpsertMailFolderAsync(mailbox.Id, storedAtStart, token);
|
||||
if (plan.UidValidityChanged)
|
||||
logger.LogInformation("The server numbered folder {FolderNumber} of mailbox '{MailboxId}' anew. Its mails are linked again by their keys, without embedding them again.", plannedFolders.Count + 1, mailbox.Id);
|
||||
|
||||
if (plan.ChecksFlags)
|
||||
{
|
||||
var fetchedFlags = await connector.FetchFlagsAsync(plan.KeptUids, plan.FlagsChangedSinceModSeq, token);
|
||||
await context.IndexStore.UpdateMailFlagsAsync(mailbox.Id, folder.FullName, MailFolderSyncPlan.GetChangedFlags(storedLocations, fetchedFlags), token);
|
||||
}
|
||||
|
||||
logger.LogDebug(
|
||||
"Planned folder {FolderNumber}/{FolderCount} of mailbox '{MailboxId}'. New={NewMails}, Gone={GoneMails}, Kept={KeptMails}, ChecksFlags={ChecksFlags}.",
|
||||
plannedFolders.Count + 1,
|
||||
selection.Folders.Count,
|
||||
mailbox.Id,
|
||||
plan.NewUids.Count,
|
||||
plan.GoneUids.Count,
|
||||
plan.KeptUids.Count,
|
||||
plan.ChecksFlags);
|
||||
|
||||
plannedFolders.Add(new(folder, serverState, storedAtStart, plan, indexedUids.Count, storedLocations.Count));
|
||||
}
|
||||
|
||||
return plannedFolders;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// The second pass through one folder: links or indexes the mails the index does not hold there yet.
|
||||
/// </summary>
|
||||
/// <returns>True when the folder is through, false when the run did its share of the work before.</returns>
|
||||
private async Task<bool> SyncNewMailsAsync(IndexedRunContext context, DataSourceMailbox mailbox, ImapMailboxConnector connector, PlannedFolder folder, DocumentRunProgress progress, ISet<string> encounteredKeys, DateTimeOffset runStartedUtc, CancellationToken token)
|
||||
{
|
||||
if (folder.Plan.NewUids.Count > 0)
|
||||
{
|
||||
//
|
||||
// The server may have numbered the folder anew since the first pass. Its UIDs would
|
||||
// name other mails then, so the folder waits for the next run, which notices it.
|
||||
//
|
||||
var serverState = await connector.OpenFolderAsync(folder.Folder.FullName, token);
|
||||
if (serverState.UidValidity != folder.ServerState.UidValidity)
|
||||
{
|
||||
logger.LogInformation("The server numbered a folder of mailbox '{MailboxId}' anew during the sync. The folder is synced during the next run.", mailbox.Id);
|
||||
return true;
|
||||
}
|
||||
|
||||
foreach (var uidBatch in folder.Plan.NewUids.Chunk(SUMMARY_BATCH_SIZE))
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
if (HasDoneItsShare(progress, runStartedUtc))
|
||||
return false;
|
||||
|
||||
var summaries = await connector.FetchSummariesAsync(uidBatch, token);
|
||||
foreach (var summary in summaries.OrderByDescending(summary => summary.UniqueId.Id))
|
||||
{
|
||||
if (HasDoneItsShare(progress, runStartedUtc))
|
||||
return false;
|
||||
|
||||
await this.SyncNewMailAsync(context, mailbox, connector, folder.Folder.FullName, summary, progress, encounteredKeys, token);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// The folder is through. From now on, only what changed after the first pass of this run
|
||||
// needs to be asked about.
|
||||
//
|
||||
await context.IndexStore.UpsertMailFolderAsync(mailbox.Id, folder.StoredAtStart with
|
||||
{
|
||||
UidNext = folder.ServerState.UidNext,
|
||||
HighestModSeq = folder.ServerState.HighestModSeq,
|
||||
InitialSyncCompletedUtc = folder.StoredAtStart.InitialSyncCompletedUtc ?? DateTimeOffset.UtcNow,
|
||||
}, token);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Whether a run did its share of the work, so other data sources get their turn.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Counted are the mails which were read, whether embedding them worked or not. Linking a mail
|
||||
/// the index holds already costs next to nothing, so only the time limits that.
|
||||
/// </remarks>
|
||||
private static bool HasDoneItsShare(DocumentRunProgress progress, DateTimeOffset runStartedUtc) =>
|
||||
progress.IndexedDocuments + progress.FailedDocuments >= MAX_READ_MAILS_PER_RUN || DateTimeOffset.UtcNow - runStartedUtc >= MAX_RUN_DURATION;
|
||||
|
||||
/// <summary>
|
||||
/// Works out what a run which got through the whole mailbox would remove.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Counted by location, as during the run. A mail which lies in two folders and leaves both
|
||||
/// counts twice, so the user is asked rather too often than too rarely.
|
||||
/// </remarks>
|
||||
private static async Task<PlannedRemovals> CollectRemovalsAsync(IndexedRunContext context, DataSourceMailbox mailbox, IReadOnlyList<PlannedFolder> plannedFolders, IEnumerable<string> storedFolderPaths, CancellationToken token)
|
||||
{
|
||||
// Folders which the server no longer lists, or which no longer lie below the root folder:
|
||||
var plannedPaths = plannedFolders.Select(folder => folder.Folder.FullName).ToHashSet(StringComparer.Ordinal);
|
||||
var goneFolderPaths = storedFolderPaths.Where(path => !plannedPaths.Contains(path)).ToList();
|
||||
|
||||
var goneFolderLocationCount = 0;
|
||||
foreach (var folderPath in goneFolderPaths)
|
||||
goneFolderLocationCount += (await context.IndexStore.GetMailLocationsAsync(mailbox.Id, folderPath, token)).Count;
|
||||
|
||||
return new(
|
||||
goneFolderPaths,
|
||||
plannedFolders.Sum(folder => folder.Plan.GoneUids.Count) + goneFolderLocationCount,
|
||||
plannedFolders.Sum(folder => folder.StoredLocationCount) + goneFolderLocationCount);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Forgets what is gone from the server, now that every folder had its chance to link a mail which only moved.
|
||||
/// </summary>
|
||||
/// <returns>How many mails were deleted from the index, which lost their last location before this run.</returns>
|
||||
private async Task<int> ApplyRemovalsAsync(IndexedRunContext context, DataSourceMailbox mailbox, IReadOnlyList<PlannedFolder> plannedFolders, PlannedRemovals removals, DateTimeOffset runStartedUtc, CancellationToken token)
|
||||
{
|
||||
foreach (var folder in plannedFolders.Where(folder => folder.Plan.GoneUids.Count > 0))
|
||||
await context.IndexStore.RemoveMailLocationsAsync(mailbox.Id, folder.Folder.FullName, folder.Plan.GoneUids, token);
|
||||
|
||||
foreach (var folderPath in removals.GoneFolderPaths)
|
||||
await context.IndexStore.DeleteMailFolderAsync(mailbox.Id, folderPath, token);
|
||||
|
||||
//
|
||||
// A mail which lost its last location during this run gets one more run to turn up again,
|
||||
// e.g. when it was moved into a folder this run had already been through:
|
||||
//
|
||||
var orphanedKeys = await context.IndexStore.GetOrphanedMailsAsync(mailbox.Id, runStartedUtc, token);
|
||||
foreach (var key in orphanedKeys)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
await context.DeleteDocumentPointsAsync(key, token);
|
||||
await context.IndexStore.DeleteFileAsync(mailbox.Id, key, token);
|
||||
context.Manifest.Files.Remove(key);
|
||||
}
|
||||
|
||||
return orphanedKeys.Count;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Forgets the marks of mails skipped for good which did not turn up during this run.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Such a mail never gets a location, so every run comes across it among the new mails, for as
|
||||
/// long as it belongs into the index. One which did not turn up is gone, and its mark would stay
|
||||
/// forever.
|
||||
/// </remarks>
|
||||
private static async Task ForgetVanishedFailuresAsync(IndexedRunContext context, IReadOnlySet<string> encounteredKeys, CancellationToken token)
|
||||
{
|
||||
foreach (var key in context.Manifest.PermanentFailures.Keys.Where(key => !encounteredKeys.Contains(key)).ToList())
|
||||
await context.ForgetPermanentFailureAsync(key, token);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Tells the user interface how the mailbox stands, from the index alone.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A refused sign-in and a removal held back are stored, so they are shown again after a
|
||||
/// restart, before any server is asked.
|
||||
/// </remarks>
|
||||
private async Task PublishStoredStateAsync(IndexedRunContext context, CancellationToken token)
|
||||
{
|
||||
var dataSourceId = context.DataSource.Id;
|
||||
var progress = await this.CreateStoredStateProgressAsync(context, token);
|
||||
if (await context.IndexStore.GetMailboxAuthFailureAsync(dataSourceId, token) is { } authFailure)
|
||||
{
|
||||
progress.PublishRunFailure(GetAuthFailureMessage(authFailure), DataSourceAttention.AUTH_FAILED);
|
||||
return;
|
||||
}
|
||||
|
||||
var syncState = await context.IndexStore.GetMailboxSyncStateAsync(dataSourceId, token);
|
||||
if (syncState is { PendingRemovalCount: { } pendingRemovalCount, PendingRemovalApprovedUtc: null })
|
||||
{
|
||||
progress.PublishRunFailure(GetRemovalMessage(pendingRemovalCount), DataSourceAttention.MASS_REMOVAL_PENDING, pendingRemovalCount);
|
||||
return;
|
||||
}
|
||||
|
||||
progress.PublishStoredState(syncState.LastSyncCompletedUtc is not null);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Counts what the index holds of a mailbox, for a status which comes without a run.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Counted by location, as during a run, cf. SyncAsync.
|
||||
/// </remarks>
|
||||
private async Task<DocumentRunProgress> CreateStoredStateProgressAsync(IndexedRunContext context, CancellationToken token)
|
||||
{
|
||||
var dataSourceId = context.DataSource.Id;
|
||||
var locationCount = 0;
|
||||
foreach (var folder in await context.IndexStore.GetMailFoldersAsync(dataSourceId, token))
|
||||
locationCount += (await context.IndexStore.GetMailLocationsAsync(dataSourceId, folder.Path, token)).Count;
|
||||
|
||||
var permanentFailures = context.Manifest.PermanentFailures;
|
||||
var progress = new DocumentRunProgress(context, locationCount + permanentFailures.Count, 0, string.Empty, [], logger)
|
||||
{
|
||||
LastSyncUtc = (await context.IndexStore.GetMailboxSyncStateAsync(dataSourceId, token)).LastSyncCompletedUtc,
|
||||
};
|
||||
|
||||
progress.RecordUnchanged(locationCount);
|
||||
foreach (var (key, failure) in permanentFailures)
|
||||
progress.RecordStillUnreadable(key, failure);
|
||||
|
||||
return progress;
|
||||
}
|
||||
|
||||
private async Task<string?> ReadPasswordAsync(DataSourceMailbox mailbox)
|
||||
{
|
||||
var requestedSecret = await rustService.GetSecret(mailbox, SecretStoreType.DATA_SOURCE, isTrying: true);
|
||||
if (requestedSecret.Success)
|
||||
return await requestedSecret.Secret.Decrypt(Program.ENCRYPTION);
|
||||
|
||||
logger.LogWarning("Could not read the password of mailbox '{MailboxId}' from the OS keyring.", mailbox.Id);
|
||||
return null;
|
||||
}
|
||||
|
||||
private static string GetAuthFailureMessage(MailboxAuthFailure failure) => string.Format(
|
||||
TB("Signing in to the mailbox failed on {0}. Presumably your password changed. AI Studio does not try again on its own, so that your account is not locked. Please enter your current password in the settings of the mailbox."),
|
||||
failure.FailedAtUtc.ToLocalTime().ToString("g", I18N.I.Culture));
|
||||
|
||||
private static string GetRemovalMessage(int removalCount) => string.Format(
|
||||
TB("This sync would remove {0} mails from the index of AI Studio at once, so it waits for you to agree. On the server, the mails stay as they are. Should they come back later, e.g. because you choose a larger period again, they have to be embedded anew, which takes time and, with a cloud provider, money."),
|
||||
removalCount.CompactCount());
|
||||
|
||||
/// <summary>
|
||||
/// Hashes how the folders of a mailbox stood on the server when a run worked through them.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Nothing compares it to decide what to do, since every run asks the server anyway. It says
|
||||
/// that the index of the mailbox can be searched, cf. DocumentRunProgress.PauseRunAsync.
|
||||
/// </remarks>
|
||||
private static string BuildSourceHash(IReadOnlyList<PlannedFolder> plannedFolders)
|
||||
{
|
||||
var source = new StringBuilder("mailbox");
|
||||
foreach (var folder in plannedFolders)
|
||||
{
|
||||
var state = folder.ServerState;
|
||||
source.Append('\n').Append(string.Create(CultureInfo.InvariantCulture, $"{folder.Folder.FullName}|{state.UidValidity}|{state.UidNext}|{state.HighestModSeq}|{state.MessageCount}"));
|
||||
}
|
||||
|
||||
return Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(source.ToString())));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// A whole text, together with the pieces its source delivered it in.
|
||||
/// </summary>
|
||||
/// <param name="Text">The pieces joined into one text.</param>
|
||||
/// <param name="SourceSegments">The pieces, in order.</param>
|
||||
internal sealed record SegmentedText(string Text, IReadOnlyList<TextSegment> SourceSegments);
|
||||
+114
-553
@@ -1,108 +1,131 @@
|
||||
using System.Security.Cryptography;
|
||||
using System.Text;
|
||||
using System.Text.RegularExpressions;
|
||||
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using AIStudio.Tools.Rust;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
public sealed partial class DataSourceEmbeddingService
|
||||
/// <summary>
|
||||
/// Cuts a text into chunks which fit the embedding provider, with some overlap between them.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Knows nothing about where the text came from. Whoever reads a document hands in its text in the
|
||||
/// pieces the source delivered, together with the strategy which suits that kind of text, and gets
|
||||
/// the chunks back one by one. Every size is measured with the tokenizer of the embedding provider,
|
||||
/// since only its count decides whether a chunk fits.
|
||||
/// </remarks>
|
||||
/// <param name="rustService">The runtime, which counts the tokens.</param>
|
||||
/// <param name="logger">The logger of the embedding service, so the log reads the same whoever writes it.</param>
|
||||
internal sealed partial class TextChunker(RustService rustService, ILogger logger)
|
||||
{
|
||||
private const string OFFICE_LOCK_FILE_PREFIX = "~$";
|
||||
internal const int DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH = 300;
|
||||
private const bool IMAGE_EMBEDDING_ENABLED = false;
|
||||
/// <summary>
|
||||
/// For documents: pages, then headings, paragraphs, lines and words.
|
||||
/// </summary>
|
||||
public static readonly ChunkingStrategy DOCUMENT_STRATEGY = new("document", [
|
||||
new("Page or extracted section", SplitBySourceSegments, true),
|
||||
new("Heading", SplitByDocumentHeadings),
|
||||
new("Paragraph", SplitByParagraphs),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
]);
|
||||
|
||||
/// <summary>
|
||||
/// What this build writes next to a chunk besides its text. Raise it whenever that changes.
|
||||
/// For presentations: slides, then lines and words.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A stored chunk keeps the metadata of the run which wrote it, and nothing recomputes it: the
|
||||
/// fingerprint of a file says whether the file changed, not whether we got better at reading
|
||||
/// it. Raising this number makes the embedding signature differ, which drops the index and
|
||||
/// builds it again — the only way corrected page numbers reach a data source somebody indexed
|
||||
/// earlier.
|
||||
///
|
||||
/// Version 2: the page of a chunk is taken from the runtime metadata instead of being read back
|
||||
/// out of the chunk text, which is what left Word and OpenDocument files, and passages
|
||||
/// continuing across a page break, without a page.
|
||||
/// </remarks>
|
||||
private const string CHUNK_METADATA_VERSION = "2";
|
||||
|
||||
private enum RagFileIndexingDecision
|
||||
{
|
||||
INDEXABLE,
|
||||
EXCLUDED,
|
||||
UNSUPPORTED,
|
||||
}
|
||||
|
||||
private sealed record ExtractedFileSegment(string Text, int? TokenCount, int? PageNumber);
|
||||
|
||||
private sealed record ExtractedFileContent(string Text, IReadOnlyList<ExtractedFileSegment> SourceSegments);
|
||||
public static readonly ChunkingStrategy PRESENTATION_STRATEGY = new("presentation", [
|
||||
new("Slide", SplitBySourceSegments, true),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
]);
|
||||
|
||||
/// <summary>
|
||||
/// One chunk as the chunking produced it, together with the page it starts on.
|
||||
/// For tables and spreadsheets: rows or sheets, then lines and words.
|
||||
/// </summary>
|
||||
public static readonly ChunkingStrategy TABLE_STRATEGY = new("table", [
|
||||
new("Row or sheet", SplitBySourceSegments, true),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
]);
|
||||
|
||||
/// <summary>
|
||||
/// For source code: extracted sections, then lines and words.
|
||||
/// </summary>
|
||||
public static readonly ChunkingStrategy SOURCE_CODE_STRATEGY = new("source-code", [
|
||||
new("Extracted section", SplitBySourceSegments, true),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
]);
|
||||
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(TextChunker).Namespace, nameof(TextChunker));
|
||||
|
||||
/// <summary>
|
||||
/// Brings the line breaks of one piece of text into the form the chunking expects.
|
||||
/// </summary>
|
||||
/// <param name="input">The piece as the source delivered it.</param>
|
||||
/// <returns>The piece with Unix line breaks and without surrounding whitespace.</returns>
|
||||
public static string NormalizeSegment(string input)
|
||||
{
|
||||
return input
|
||||
.Replace("\r\n", "\n", StringComparison.Ordinal)
|
||||
.Replace('\r', '\n')
|
||||
.Trim();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Picks the strategy which suits a file, by its type.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The page is carried rather than read back out of the chunk text. The runtime states it, and
|
||||
/// the chunking knows which source segment a chunk begins in, so nothing has to be derived from
|
||||
/// a marker in the text — which is what used to leave Word files and continued passages without
|
||||
/// a page.
|
||||
/// For every reader of files alike, be the file on disk or attached to a mail. HTML counts as a
|
||||
/// document, not as source code: the runtime turns it into text before it gets here.
|
||||
/// </remarks>
|
||||
/// <param name="Text">The chunk itself, overlap prefix included.</param>
|
||||
/// <param name="PageNumber">The page the chunk's own content starts on, or null when it has none.</param>
|
||||
private sealed record EmbeddingChunk(string Text, int? PageNumber);
|
||||
|
||||
private sealed record EmbeddingChunkDraft(string ChunkId, string Text, int ChunkIndex, int? PageNumber);
|
||||
|
||||
internal sealed record ChunkingOptions(int MaxChunkTokenLength, int OverlapTokenLength);
|
||||
|
||||
private sealed record ChunkingStrategy(string Name, IReadOnlyList<ChunkingRule> Rules);
|
||||
|
||||
private sealed record ChunkingRule(string Name, Func<string, IReadOnlyList<string>, IReadOnlyList<string>>? Split, bool UsesSourceSegmentCounts = false);
|
||||
|
||||
private sealed record DataSourceMetadataSnapshot(string SourceHash, IReadOnlyDictionary<string, string> FileHashes);
|
||||
|
||||
private async IAsyncEnumerable<EmbeddingChunk> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
||||
/// <param name="filePath">The path or the name of the file.</param>
|
||||
/// <returns>The strategy to cut its text by.</returns>
|
||||
public static ChunkingStrategy GetStrategyForFile(string filePath)
|
||||
{
|
||||
var options = GetChunkingOptions(dataSource, embeddingProvider);
|
||||
var strategy = this.GetChunkingStrategy(filePath);
|
||||
var content = await this.ReadExtractedFileContentAsync(filePath, embeddingProvider, token);
|
||||
if (FileTypes.IsAllowedPath(filePath, FileTypes.POWER_POINT))
|
||||
return PRESENTATION_STRATEGY;
|
||||
|
||||
await foreach (var chunk in this.SplitByChunkingStrategyAsync(content, strategy, options, embeddingProvider, token))
|
||||
yield return chunk;
|
||||
if (FileTypes.IsAllowedPath(filePath, FileTypes.TABULAR, FileTypes.SPREADSHEET))
|
||||
return TABLE_STRATEGY;
|
||||
|
||||
if (!FileTypes.IsAllowedPath(filePath, FileTypes.HTML) && FileTypes.IsAllowedPath(filePath, FileTypes.SOURCE_CODE))
|
||||
return SOURCE_CODE_STRATEGY;
|
||||
|
||||
return DOCUMENT_STRATEGY;
|
||||
}
|
||||
|
||||
private async Task<ExtractedFileContent> ReadExtractedFileContentAsync(string filePath, EmbeddingProvider embeddingProvider, CancellationToken token)
|
||||
{
|
||||
var segments = new List<ExtractedFileSegment>();
|
||||
|
||||
await foreach (var segment in rustService.StreamArbitraryFileDataWithTokenCounts(filePath, embeddingProvider, token))
|
||||
{
|
||||
var normalized = NormalizeChunkSegment(segment.Content);
|
||||
if (!string.IsNullOrWhiteSpace(normalized))
|
||||
segments.Add(new(normalized, segment.TokenCount, segment.PageNumber));
|
||||
}
|
||||
|
||||
return new(string.Join("\n", segments.Select(segment => segment.Text)).Trim(), segments);
|
||||
}
|
||||
|
||||
private async IAsyncEnumerable<EmbeddingChunk> SplitByChunkingStrategyAsync(ExtractedFileContent content, ChunkingStrategy strategy, ChunkingOptions options, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
||||
/// <summary>
|
||||
/// Cuts a text into chunks.
|
||||
/// </summary>
|
||||
/// <param name="content">The text, in the pieces its source delivered it in.</param>
|
||||
/// <param name="strategy">The rules to cut it by.</param>
|
||||
/// <param name="options">How large a chunk may become, and how much the next one repeats.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider whose tokenizer measures the chunks.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <param name="heading">
|
||||
/// A line the first chunk starts with, e.g. which attachment of a mail the text comes from, or
|
||||
/// empty for none. It counts against the size of that chunk, like the overlap of every later one.
|
||||
/// </param>
|
||||
/// <returns>The chunks, in the order of the text.</returns>
|
||||
public async IAsyncEnumerable<EmbeddingChunk> SplitAsync(SegmentedText content, ChunkingStrategy strategy, ChunkingOptions options, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token, string heading = "")
|
||||
{
|
||||
var estimatedTokenCount = SumTokenCounts(content.SourceSegments);
|
||||
|
||||
// The whole text starts where the first segment starts, so that is the page it is on until
|
||||
// the splitting reaches a segment boundary:
|
||||
var firstPageNumber = content.SourceSegments.Count > 0 ? content.SourceSegments[0].PageNumber : null;
|
||||
await foreach (var chunk in this.SplitTextByRulesAsync(content.Text, content.SourceSegments, strategy, 0, options, embeddingProvider, firstPageNumber, token, estimatedTokenCount: estimatedTokenCount))
|
||||
await foreach (var chunk in this.SplitTextByRulesAsync(content.Text, content.SourceSegments, strategy, 0, options, embeddingProvider, firstPageNumber, token, heading, estimatedTokenCount))
|
||||
yield return chunk;
|
||||
}
|
||||
|
||||
private async IAsyncEnumerable<EmbeddingChunk> SplitTextByRulesAsync(
|
||||
string text,
|
||||
IReadOnlyList<ExtractedFileSegment> sourceSegments,
|
||||
IReadOnlyList<TextSegment> sourceSegments,
|
||||
ChunkingStrategy strategy,
|
||||
int ruleIndex,
|
||||
ChunkingOptions options,
|
||||
@@ -296,7 +319,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
return largestValidUnitCount;
|
||||
}
|
||||
|
||||
private static int? SumTokenCounts(IReadOnlyList<ExtractedFileSegment> segments)
|
||||
private static int? SumTokenCounts(IReadOnlyList<TextSegment> segments)
|
||||
{
|
||||
var result = 0L;
|
||||
foreach (var segment in segments)
|
||||
@@ -312,7 +335,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private static IReadOnlyList<int>? EstimateSplitUnitTokenCounts(
|
||||
IReadOnlyList<string> units,
|
||||
IReadOnlyList<ExtractedFileSegment> sourceSegments,
|
||||
IReadOnlyList<TextSegment> sourceSegments,
|
||||
bool usesSourceSegmentCounts,
|
||||
int? sourceTokenCount)
|
||||
{
|
||||
@@ -582,74 +605,6 @@ public sealed partial class DataSourceEmbeddingService
|
||||
throw new InvalidOperationException(string.Format(TB("The tokens of the text could not be counted for the embedding provider '{0}'. {1}"), embeddingProvider.Name, message));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Works out how the text of a data source is cut for a given embedding provider.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Static, because the answer follows from its two arguments alone. That lets the embedding
|
||||
/// signature be built for a configuration which is not stored yet, which is what the dialogs ask
|
||||
/// before they save a change.
|
||||
/// </remarks>
|
||||
/// <param name="dataSource">The data source whose own chunk settings apply.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider whose token limit caps them.</param>
|
||||
/// <returns>The chunk size and overlap which are actually used.</returns>
|
||||
internal static ChunkingOptions GetChunkingOptions(IDataSource dataSource, EmbeddingProvider embeddingProvider)
|
||||
{
|
||||
var providerMaxChunkTokenLength = Math.Max(1, embeddingProvider.EffectiveTokenLimit);
|
||||
var dataSourceMaxChunkTokenLength = dataSource is IInternalDataSource { MaxChunkTokenLength: > 0 } internalDataSource
|
||||
? internalDataSource.MaxChunkTokenLength
|
||||
: 0;
|
||||
var maxChunkTokenLength = dataSourceMaxChunkTokenLength > 0
|
||||
? Math.Min(dataSourceMaxChunkTokenLength, providerMaxChunkTokenLength)
|
||||
: providerMaxChunkTokenLength;
|
||||
|
||||
var configuredOverlapTokenLength = dataSource is IInternalDataSource overlapDataSource
|
||||
? overlapDataSource.ChunkOverlapTokenLength
|
||||
: DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH;
|
||||
var overlapTokenLength = Math.Clamp(configuredOverlapTokenLength, 0, Math.Max(0, maxChunkTokenLength - 1));
|
||||
|
||||
return new(maxChunkTokenLength, overlapTokenLength);
|
||||
}
|
||||
|
||||
private ChunkingStrategy GetChunkingStrategy(string filePath)
|
||||
{
|
||||
if (this.IsPresentationFilePath(filePath))
|
||||
return new("presentation", [
|
||||
new("Slide", SplitBySourceSegments, true),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
]);
|
||||
|
||||
if (this.IsDelimitedTableFilePath(filePath) || this.IsSpreadsheetFilePath(filePath))
|
||||
return new("table", [
|
||||
new("Row or sheet", SplitBySourceSegments, true),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
]);
|
||||
|
||||
if (this.IsSourceCodeFilePath(filePath))
|
||||
return GetSourceCodeChunkingStrategy();
|
||||
|
||||
return new("document", [
|
||||
new("Page or extracted section", SplitBySourceSegments, true),
|
||||
new("Heading", SplitByDocumentHeadings),
|
||||
new("Paragraph", SplitByParagraphs),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
]);
|
||||
}
|
||||
|
||||
private static ChunkingStrategy GetSourceCodeChunkingStrategy() =>
|
||||
new("source-code", [
|
||||
new("Extracted section", SplitBySourceSegments, true),
|
||||
new("Line break", SplitByLineBreaks),
|
||||
new("Whitespace", SplitByWhitespace),
|
||||
new("Hard cut", null),
|
||||
]);
|
||||
|
||||
private static List<string> NormalizeSplitUnits(IReadOnlyList<string> units, string fallbackText)
|
||||
{
|
||||
var result = units
|
||||
@@ -699,7 +654,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private static IReadOnlyList<string> SplitByParagraphs(string text, IReadOnlyList<string> sourceSegments)
|
||||
{
|
||||
var matches = Regex.Matches(text, @"\n[ \t]*\n", RegexOptions.CultureInvariant);
|
||||
var matches = ParagraphBreakRegex().Matches(text);
|
||||
if (matches.Count == 0)
|
||||
return [text];
|
||||
|
||||
@@ -740,7 +695,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
|
||||
private static IReadOnlyList<string> SplitByWhitespace(string text, IReadOnlyList<string> sourceSegments)
|
||||
{
|
||||
var matches = Regex.Matches(text, @"\S+\s*", RegexOptions.CultureInvariant);
|
||||
var matches = WordWithTrailingWhitespaceRegex().Matches(text);
|
||||
if (matches.Count == 0)
|
||||
return [text];
|
||||
|
||||
@@ -773,7 +728,7 @@ public sealed partial class DataSourceEmbeddingService
|
||||
if (string.IsNullOrWhiteSpace(trimmed))
|
||||
return false;
|
||||
|
||||
if (Regex.IsMatch(trimmed, @"^#{1,6}\s+\S", RegexOptions.CultureInvariant))
|
||||
if (MarkdownHeadingRegex().IsMatch(trimmed))
|
||||
return true;
|
||||
|
||||
if (!string.IsNullOrWhiteSpace(previousLine) || !string.IsNullOrWhiteSpace(nextLine))
|
||||
@@ -785,415 +740,21 @@ public sealed partial class DataSourceEmbeddingService
|
||||
if (trimmed.Contains("|", StringComparison.Ordinal) || trimmed.EndsWith(".", StringComparison.Ordinal))
|
||||
return false;
|
||||
|
||||
return Regex.IsMatch(trimmed, @"^(\d+(\.\d+)*\.?\s+\S|(?i:chapter|section)\s+\S|[A-Z0-9][A-Z0-9 ,:;'/&()_-]{2,})$", RegexOptions.CultureInvariant);
|
||||
return PlainHeadingRegex().IsMatch(trimmed);
|
||||
}
|
||||
|
||||
private FileEnumerationResult GetInputFiles(IDataSource dataSource)
|
||||
{
|
||||
var result = new FileEnumerationResult();
|
||||
[GeneratedRegex(@"\n[ \t]*\n", RegexOptions.CultureInvariant)]
|
||||
private static partial Regex ParagraphBreakRegex();
|
||||
|
||||
switch (dataSource)
|
||||
{
|
||||
case DataSourceLocalFile localFile when File.Exists(localFile.FilePath):
|
||||
var file = new FileInfo(localFile.FilePath);
|
||||
switch (this.GetRagFileIndexingDecision(file))
|
||||
{
|
||||
case RagFileIndexingDecision.INDEXABLE:
|
||||
result.Files.Add(file);
|
||||
break;
|
||||
[GeneratedRegex(@"\S+\s*", RegexOptions.CultureInvariant)]
|
||||
private static partial Regex WordWithTrailingWhitespaceRegex();
|
||||
|
||||
case RagFileIndexingDecision.EXCLUDED:
|
||||
logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", file.FullName);
|
||||
break;
|
||||
|
||||
default:
|
||||
result.AddFailure(localFile.FilePath, string.Format(TB("The file '{0}' has a type AI Studio cannot index."), localFile.FilePath));
|
||||
break;
|
||||
}
|
||||
|
||||
return result;
|
||||
|
||||
case DataSourceLocalDirectory localDirectory when Directory.Exists(localDirectory.Path):
|
||||
this.EnumerateAccessibleFiles(localDirectory.Path, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
switch (dataSource)
|
||||
{
|
||||
case DataSourceLocalFile localFile:
|
||||
result.AddFailure(localFile.FilePath, string.Format(TB("The file '{0}' does not exist."), localFile.FilePath));
|
||||
break;
|
||||
|
||||
case DataSourceLocalDirectory localDirectory:
|
||||
result.AddFailure(localDirectory.Path, string.Format(TB("The folder '{0}' does not exist."), localDirectory.Path));
|
||||
break;
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
private void EnumerateAccessibleFiles(string rootPath, FileEnumerationResult result)
|
||||
{
|
||||
var pendingDirectories = new Stack<string>();
|
||||
pendingDirectories.Push(rootPath);
|
||||
|
||||
while (pendingDirectories.Count > 0)
|
||||
{
|
||||
var currentPath = pendingDirectories.Pop();
|
||||
IEnumerable<string> subDirectories;
|
||||
IEnumerable<string> files;
|
||||
|
||||
try
|
||||
{
|
||||
subDirectories = Directory.EnumerateDirectories(currentPath);
|
||||
files = Directory.EnumerateFiles(currentPath);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Cannot access directory '{DirectoryPath}' while indexing.", currentPath);
|
||||
result.AddFailure(currentPath, string.Format(TB("The folder '{0}' could not be opened. Please check whether you are allowed to read it."), currentPath));
|
||||
continue;
|
||||
}
|
||||
|
||||
foreach (var filePath in files)
|
||||
{
|
||||
FileInfo fileInfo;
|
||||
try
|
||||
{
|
||||
fileInfo = new FileInfo(filePath);
|
||||
if (!fileInfo.Exists)
|
||||
continue;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", filePath);
|
||||
result.AddFailure(filePath, string.Format(TB("The file '{0}' could not be read. Please check whether you are allowed to read it."), filePath));
|
||||
continue;
|
||||
}
|
||||
|
||||
switch (this.GetRagFileIndexingDecision(fileInfo))
|
||||
{
|
||||
case RagFileIndexingDecision.INDEXABLE:
|
||||
result.Files.Add(fileInfo);
|
||||
break;
|
||||
|
||||
case RagFileIndexingDecision.EXCLUDED:
|
||||
logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", fileInfo.FullName);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
foreach (var subDirectory in subDirectories)
|
||||
{
|
||||
if (this.IsSkippedRagDirectory(subDirectory))
|
||||
continue;
|
||||
|
||||
pendingDirectories.Push(subDirectory);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private string TryGetRelativePath(IDataSource dataSource, FileInfo file) => dataSource switch
|
||||
{
|
||||
DataSourceLocalDirectory localDirectory => Path.GetRelativePath(localDirectory.Path, file.FullName),
|
||||
_ => file.Name
|
||||
};
|
||||
|
||||
private static string NormalizeChunkSegment(string input)
|
||||
{
|
||||
return input
|
||||
.Replace("\r\n", "\n", StringComparison.Ordinal)
|
||||
.Replace('\r', '\n')
|
||||
.Trim();
|
||||
}
|
||||
|
||||
private bool IsImageFilePath(string filePath)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(filePath, FileTypes.IMAGE);
|
||||
}
|
||||
|
||||
private bool IsPresentationFilePath(string filePath)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(filePath, FileTypes.POWER_POINT);
|
||||
}
|
||||
|
||||
private bool IsDelimitedTableFilePath(string filePath)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(filePath, FileTypes.TABULAR);
|
||||
}
|
||||
|
||||
private bool IsSpreadsheetFilePath(string filePath)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(filePath, FileTypes.SPREADSHEET);
|
||||
}
|
||||
|
||||
private bool IsSourceCodeFilePath(string filePath)
|
||||
{
|
||||
return !this.IsHtmlFilePath(filePath) && FileTypes.IsAllowedPath(filePath, FileTypes.SOURCE_CODE);
|
||||
}
|
||||
|
||||
private bool IsHtmlFilePath(string filePath)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(filePath, FileTypes.HTML);
|
||||
}
|
||||
|
||||
private bool IsSupportedRagFilePath(string filePath)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(filePath, FileTypes.DOCUMENT);
|
||||
}
|
||||
|
||||
private RagFileIndexingDecision GetRagFileIndexingDecision(FileInfo file)
|
||||
{
|
||||
if (this.IsSkippedRagFile(file))
|
||||
return RagFileIndexingDecision.EXCLUDED;
|
||||
|
||||
if (!IMAGE_EMBEDDING_ENABLED && this.IsImageFilePath(file.FullName))
|
||||
return RagFileIndexingDecision.EXCLUDED;
|
||||
|
||||
return this.IsSupportedRagFilePath(file.FullName)
|
||||
? RagFileIndexingDecision.INDEXABLE
|
||||
: RagFileIndexingDecision.UNSUPPORTED;
|
||||
}
|
||||
|
||||
private bool IsSkippedRagFile(FileInfo file)
|
||||
{
|
||||
if (IsSkippedRagFileName(file.Name))
|
||||
return true;
|
||||
|
||||
try
|
||||
{
|
||||
return file.Attributes.HasFlag(FileAttributes.ReparsePoint)
|
||||
|| file.Attributes.HasFlag(FileAttributes.Offline)
|
||||
|| file.Attributes.HasFlag(FileAttributes.Temporary)
|
||||
|| file.Attributes.HasFlag(FileAttributes.System);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", file.FullName);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
private static bool IsSkippedRagFileName(string fileName)
|
||||
{
|
||||
return FileTypes.IsAllowedPath(fileName, FileTypes.SHORTCUT)
|
||||
|| fileName.StartsWith(OFFICE_LOCK_FILE_PREFIX, StringComparison.Ordinal);
|
||||
}
|
||||
|
||||
private bool IsSkippedRagDirectory(string path)
|
||||
{
|
||||
try
|
||||
{
|
||||
var directory = new DirectoryInfo(path);
|
||||
return directory.Attributes.HasFlag(FileAttributes.ReparsePoint)
|
||||
|| directory.Attributes.HasFlag(FileAttributes.Offline)
|
||||
|| directory.Attributes.HasFlag(FileAttributes.System);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Cannot inspect directory '{DirectoryPath}' while indexing.", path);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
[GeneratedRegex(@"^#{1,6}\s+\S", RegexOptions.CultureInvariant)]
|
||||
private static partial Regex MarkdownHeadingRegex();
|
||||
|
||||
/// <summary>
|
||||
/// Describes how the vectors of a data source were made.
|
||||
/// A heading without Markdown: a numbered one, a chapter or section, or a line in capitals.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// What appears here decides when stored embeddings are thrown away: a signature differing from
|
||||
/// the persisted one drops the whole index and builds it again. So it names the embedding model,
|
||||
/// where it runs, how the text was cut for it, and the chunk metadata version — the things a
|
||||
/// vector actually depends on.
|
||||
///
|
||||
/// Two of them are less obvious than they look. The Hugging Face inference provider belongs to
|
||||
/// where the model runs: the same model name served by another backend is another vector source.
|
||||
/// And a custom tokenizer enters through its content, not through its path, because a tokenizer
|
||||
/// is stored under the name it came with — almost always tokenizer.json — so swapping one for
|
||||
/// another lands on the identical path, while moving the data directory changes every path
|
||||
/// without changing a single tokenizer.
|
||||
///
|
||||
/// The chunk settings enter only as what they amount to, never as what somebody typed. A data
|
||||
/// source storing 0 means "follow the embedding provider", and writing that provider's own limit
|
||||
/// into the field changes nothing about how the text is cut. Carrying the typed numbers as well
|
||||
/// made that a different signature, so opening the expert settings of a data source — which
|
||||
/// fills an empty limit with the provider's — threw the whole index away for nothing.
|
||||
///
|
||||
/// The confidence level a data source asks of a provider is deliberately not among them. It
|
||||
/// changes no vector, and it is enforced live on every request anyway: DataSourceService checks
|
||||
/// it against the participating chat providers and against the embedding provider, and this
|
||||
/// service checks it again before each indexing run. It was part of this signature once, which
|
||||
/// re-embedded every file of a data source whenever somebody raised or lowered it — real money
|
||||
/// at a cloud embedding provider, for nothing.
|
||||
/// </remarks>
|
||||
internal static string BuildEmbeddingSignature(IDataSource dataSource, EmbeddingProvider embeddingProvider, ChunkingOptions chunkingOptions)
|
||||
{
|
||||
return string.Join('|',
|
||||
CHUNK_METADATA_VERSION,
|
||||
embeddingProvider.Id,
|
||||
embeddingProvider.UsedLLMProvider,
|
||||
embeddingProvider.Model.Id,
|
||||
embeddingProvider.Host,
|
||||
embeddingProvider.Hostname,
|
||||
embeddingProvider.HFInferenceProvider,
|
||||
embeddingProvider.TokenizerFingerprint,
|
||||
embeddingProvider.EffectiveTokenLimit,
|
||||
chunkingOptions.MaxChunkTokenLength,
|
||||
chunkingOptions.OverlapTokenLength);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Describes how the vectors of a data source were made, working the chunking out along the way.
|
||||
/// </summary>
|
||||
/// <param name="dataSource">The data source the vectors belong to.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider which makes them.</param>
|
||||
/// <returns>The signature of this pairing.</returns>
|
||||
internal static string BuildEmbeddingSignature(IDataSource dataSource, EmbeddingProvider embeddingProvider) =>
|
||||
BuildEmbeddingSignature(dataSource, embeddingProvider, GetChunkingOptions(dataSource, embeddingProvider));
|
||||
|
||||
private DataSourceMetadataSnapshot BuildDataSourceMetadataSnapshot(IDataSource dataSource, IReadOnlyList<FileInfo> indexedFiles)
|
||||
{
|
||||
var fileHashes = indexedFiles
|
||||
.OrderBy(file => file.FullName, StringComparer.OrdinalIgnoreCase)
|
||||
.ToDictionary(file => file.FullName, BuildFileMetadataHash, StringComparer.OrdinalIgnoreCase);
|
||||
|
||||
var sourceHash = dataSource switch
|
||||
{
|
||||
DataSourceLocalFile localFile => indexedFiles.Count > 0
|
||||
? fileHashes[indexedFiles[0].FullName]
|
||||
: BuildMetadataHash("file", localFile.FilePath, Path.GetFileName(localFile.FilePath), "missing", "0"),
|
||||
|
||||
DataSourceLocalDirectory localDirectory => this.BuildDirectoryMetadataHash(localDirectory, indexedFiles, fileHashes),
|
||||
|
||||
_ => BuildMetadataHash(dataSource.Type.ToString(), dataSource.Id, dataSource.Name)
|
||||
};
|
||||
|
||||
return new(sourceHash, fileHashes);
|
||||
}
|
||||
|
||||
private string BuildDirectoryMetadataHash(DataSourceLocalDirectory dataSource, IReadOnlyList<FileInfo> indexedFiles, IReadOnlyDictionary<string, string> fileHashes)
|
||||
{
|
||||
var directory = new DirectoryInfo(dataSource.Path);
|
||||
directory.Refresh();
|
||||
|
||||
var totalSize = 0L;
|
||||
var latestFileWriteTicks = 0L;
|
||||
foreach (var file in indexedFiles)
|
||||
{
|
||||
file.Refresh();
|
||||
if (!file.Exists)
|
||||
continue;
|
||||
|
||||
totalSize += file.Length;
|
||||
latestFileWriteTicks = Math.Max(latestFileWriteTicks, file.LastWriteTimeUtc.Ticks);
|
||||
}
|
||||
|
||||
var latestWriteTicks = Math.Max(directory.LastWriteTimeUtc.Ticks, latestFileWriteTicks);
|
||||
var parts = new List<string>
|
||||
{
|
||||
"directory",
|
||||
directory.FullName,
|
||||
directory.Name,
|
||||
latestWriteTicks.ToString(),
|
||||
totalSize.ToString(),
|
||||
indexedFiles.Count.ToString()
|
||||
};
|
||||
|
||||
foreach (var file in indexedFiles.OrderBy(file => file.FullName, StringComparer.OrdinalIgnoreCase))
|
||||
{
|
||||
parts.Add(this.TryGetRelativePath(dataSource, file));
|
||||
parts.Add(fileHashes[file.FullName]);
|
||||
}
|
||||
|
||||
return BuildMetadataHash(parts);
|
||||
}
|
||||
|
||||
private static string BuildFileMetadataHash(FileInfo file)
|
||||
{
|
||||
file.Refresh();
|
||||
if (!file.Exists)
|
||||
{
|
||||
return BuildMetadataHash(
|
||||
"file",
|
||||
file.FullName,
|
||||
file.Name,
|
||||
"missing",
|
||||
"0");
|
||||
}
|
||||
|
||||
return BuildMetadataHash(
|
||||
"file",
|
||||
file.FullName,
|
||||
file.Name,
|
||||
file.LastWriteTimeUtc.Ticks.ToString(),
|
||||
file.Length.ToString());
|
||||
}
|
||||
|
||||
private static string BuildMetadataHash(params string[] parts)
|
||||
{
|
||||
return BuildMetadataHash((IEnumerable<string>)parts);
|
||||
}
|
||||
|
||||
private static string BuildMetadataHash(IEnumerable<string> parts)
|
||||
{
|
||||
var fingerprintSource = new StringBuilder();
|
||||
foreach (var part in parts)
|
||||
fingerprintSource.Append(part.Length).Append(':').Append(part).Append('|');
|
||||
|
||||
var bytes = SHA256.HashData(Encoding.UTF8.GetBytes(fingerprintSource.ToString()));
|
||||
return Convert.ToHexString(bytes);
|
||||
}
|
||||
|
||||
private EmbeddingStateFile CreateEmbeddingStateFile(IDataSource dataSource, FileInfo file, string fingerprint, int chunkCount, DateTimeOffset embeddedAtUtc)
|
||||
{
|
||||
file.Refresh();
|
||||
var absolutePath = Path.GetFullPath(file.FullName);
|
||||
return new(
|
||||
this.CreateParentFileId(dataSource.Id, absolutePath),
|
||||
absolutePath,
|
||||
file.Name,
|
||||
this.TryGetRelativePath(dataSource, file),
|
||||
GetFileType(file),
|
||||
fingerprint,
|
||||
file.Exists ? file.Length : 0,
|
||||
file.Exists ? new DateTimeOffset(file.CreationTimeUtc) : DateTimeOffset.UnixEpoch,
|
||||
file.Exists ? new DateTimeOffset(file.LastWriteTimeUtc) : DateTimeOffset.UnixEpoch,
|
||||
embeddedAtUtc,
|
||||
chunkCount);
|
||||
}
|
||||
|
||||
private IReadOnlyList<EmbeddingStateChunk> CreateEmbeddingStateChunks(EmbeddingStateFile parentFile, IReadOnlyList<EmbeddingChunkDraft> batch, DateTimeOffset embeddedAtUtc)
|
||||
{
|
||||
return batch
|
||||
.Select(chunk => new EmbeddingStateChunk(
|
||||
chunk.ChunkId,
|
||||
parentFile.ParentFileId,
|
||||
chunk.PageNumber,
|
||||
chunk.ChunkIndex,
|
||||
chunk.Text,
|
||||
embeddedAtUtc))
|
||||
.ToList();
|
||||
}
|
||||
|
||||
private static string GetFileType(FileInfo file)
|
||||
{
|
||||
var extension = file.Extension.TrimStart('.').ToLowerInvariant();
|
||||
return string.IsNullOrWhiteSpace(extension) ? "unknown" : extension;
|
||||
}
|
||||
|
||||
private string CreatePointId(string dataSourceId, string fingerprint, int chunkIndex) =>
|
||||
CreateStableGuid($"{dataSourceId}:chunk:{fingerprint}:{chunkIndex}");
|
||||
|
||||
private string CreateParentFileId(string dataSourceId, string absolutePath) =>
|
||||
CreateStableGuid($"{dataSourceId}:parent-file:{absolutePath}");
|
||||
|
||||
private static string CreateStableGuid(string source)
|
||||
{
|
||||
var hash = SHA256.HashData(Encoding.UTF8.GetBytes(source));
|
||||
var guidBytes = hash[..16].ToArray();
|
||||
|
||||
guidBytes[6] = (byte)((guidBytes[6] & 0x0F) | 0x40);
|
||||
guidBytes[8] = (byte)((guidBytes[8] & 0x3F) | 0x80);
|
||||
|
||||
return new Guid(guidBytes).ToString();
|
||||
}
|
||||
}
|
||||
[GeneratedRegex(@"^(\d+(\.\d+)*\.?\s+\S|(?i:chapter|section)\s+\S|[A-Z0-9][A-Z0-9 ,:;'/&()_-]{2,})$", RegexOptions.CultureInvariant)]
|
||||
private static partial Regex PlainHeadingRegex();
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
/// <summary>
|
||||
/// One piece of a text as its source delivered it: a page, a slide, a sheet or a section.
|
||||
/// </summary>
|
||||
/// <param name="Text">The text of the piece.</param>
|
||||
/// <param name="TokenCount">Its token count for the embedding provider, when the source already knows it.</param>
|
||||
/// <param name="PageNumber">The page it is on, or null when the source has no pages.</param>
|
||||
internal sealed record TextSegment(string Text, int? TokenCount, int? PageNumber);
|
||||
@@ -0,0 +1,299 @@
|
||||
using AIStudio.Provider;
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Tools.Databases;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using AIStudio.Tools.RAG;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// Searches the local index of one data source: its vectors by meaning, its keyword index by words.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Shared by everything which searches what AI Studio indexed itself, the local files and folders
|
||||
/// as well as the mailboxes. Each of them decides what to search for and what to make of the
|
||||
/// matches. What they share is how a search which cannot cover the whole data source turns into a
|
||||
/// gap the user hears about, instead of into an answer quietly put together without it.<br/><br/>
|
||||
/// Nothing here logs what a match holds. The name of a file may go into the log, the subject of a
|
||||
/// mail must not, so whoever searches logs the matches itself.
|
||||
/// </remarks>
|
||||
public sealed class LocalIndexSearchService(SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider, DataSourceEmbeddingService embeddingService, ILogger<LocalIndexSearchService> logger)
|
||||
{
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(LocalIndexSearchService).Namespace, nameof(LocalIndexSearchService));
|
||||
|
||||
//
|
||||
// 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();
|
||||
|
||||
/// <summary>
|
||||
/// Whether the index of a data source has to be built anew before it can be searched.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Whoever searches asks this before either channel does, since both of them read what the
|
||||
/// rebuild is about to discard.
|
||||
/// </remarks>
|
||||
/// <param name="dataSource">The data source to search.</param>
|
||||
/// <param name="run">The retrieval, which records the gap when the data source cannot be searched.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>True when the data source cannot be searched until its index is rebuilt.</returns>
|
||||
public async Task<bool> IsAwaitingReindexAsync(IIndexedDataSource dataSource, RetrievalRun run, CancellationToken token)
|
||||
{
|
||||
if (!await embeddingService.IsAwaitingReindexAsync(dataSource, token))
|
||||
return false;
|
||||
|
||||
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));
|
||||
return true;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Searches the vectors of a data source for the chunks closest in meaning to the query.
|
||||
/// </summary>
|
||||
/// <param name="dataSource">The data source to search.</param>
|
||||
/// <param name="query">What to search for.</param>
|
||||
/// <param name="maxMatches">How many chunks to return at most.</param>
|
||||
/// <param name="filter">The chunks the search may return, or null to search all of them.</param>
|
||||
/// <param name="run">The retrieval, which records what kept the search from covering the data source.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The chunks, best first. Empty when nothing matched, or when the search could not run, which the run records then.</returns>
|
||||
public async Task<IReadOnlyList<VectorSearchResult>> SearchVectorsAsync(IIndexedDataSource dataSource, string query, int maxMatches, VectorSearchFilter? filter, RetrievalRun run, CancellationToken token)
|
||||
{
|
||||
// Nothing can match, so there is no need to send the query to the embedding provider:
|
||||
if (filter is { MatchesNothing: true })
|
||||
return [];
|
||||
|
||||
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 collectionName = DataSourceEmbeddingNames.GetCollectionName(dataSource.Id);
|
||||
var results = filter is null
|
||||
? await vectorStore.SearchEmbeddingAsync(collectionName, vector, maxMatches, token)
|
||||
: await vectorStore.SearchEmbeddingAsync(collectionName, vector, maxMatches, filter, token);
|
||||
|
||||
return this.LimitSearchResults(dataSource, "vector", results, maxMatches);
|
||||
}
|
||||
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>
|
||||
/// Searches the keyword index of a data source with BM25.
|
||||
/// </summary>
|
||||
/// <param name="dataSource">The data source to search.</param>
|
||||
/// <param name="maxMatches">How many chunks to return at most.</param>
|
||||
/// <param name="search">The search to run against the keyword index, asking for at most maxMatches chunks of this data source.</param>
|
||||
/// <param name="run">The retrieval, which records what kept the search from covering the data source.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The chunks, best first. Empty when nothing matched, or when the search could not run, which the run records then.</returns>
|
||||
public async Task<IReadOnlyList<IndexStoreSearchResult>> SearchKeywordsAsync(IIndexedDataSource dataSource, int maxMatches, Func<IndexStoreClient, Task<IReadOnlyList<IndexStoreSearchResult>>> search, 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 [];
|
||||
}
|
||||
|
||||
return this.LimitSearchResults(dataSource, "BM25", await search(indexStore), maxMatches);
|
||||
}
|
||||
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 [];
|
||||
}
|
||||
}
|
||||
|
||||
/// <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(IIndexedDataSource 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(
|
||||
IIndexedDataSource 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 IReadOnlyList<T> LimitSearchResults<T>(IIndexedDataSource 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();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.RAG;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// How many mails of a mailbox meet some conditions, as far as the index could tell.
|
||||
/// </summary>
|
||||
/// <param name="Count">The number of mails, or null when the mailbox could not be counted.</param>
|
||||
/// <param name="Gaps">What kept the count from covering the whole mailbox, empty when nothing did.</param>
|
||||
public sealed record MailCountOutcome(MailCountResult? Count, IReadOnlyList<RetrievalGap> Gaps);
|
||||
@@ -0,0 +1,8 @@
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// One chunk a search of a mailbox found, by the mail it belongs to.
|
||||
/// </summary>
|
||||
/// <param name="MailId">The id of the mail.</param>
|
||||
/// <param name="Text">The text of the chunk.</param>
|
||||
internal sealed record MailPassage(string MailId, string Text);
|
||||
@@ -0,0 +1,13 @@
|
||||
using AIStudio.Settings.DataModel;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// Everything the index keeps about one mail, together with the mailbox it belongs to.
|
||||
/// </summary>
|
||||
/// <param name="Mailbox">The mailbox, which the provider of the chat may read.</param>
|
||||
/// <param name="Summary">What a list of mails shows about the mail, e.g., its subject and its folders.</param>
|
||||
/// <param name="Mail">Its addresses, its parts as AI Studio read them, and its places on the server.</param>
|
||||
/// <param name="InReplyToMailId">The id of the mail it replies to, or null when the index does not hold that one.</param>
|
||||
public sealed record MailReading(DataSourceMailbox Mailbox, MailSummary Summary, MailRecord Mail, string? InReplyToMailId);
|
||||
@@ -0,0 +1,10 @@
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// One mail a search of a mailbox found.
|
||||
/// </summary>
|
||||
/// <param name="Summary">What a list of mails shows about the mail.</param>
|
||||
/// <param name="Passage">The passage of the mail which matched the query best, or null when the mails were listed without a query.</param>
|
||||
public sealed record MailSearchHit(MailSummary Summary, string? Passage);
|
||||
@@ -0,0 +1,14 @@
|
||||
using AIStudio.Tools.RAG;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// One page of the mails a search found in a mailbox.
|
||||
/// </summary>
|
||||
/// <param name="Hits">The mails: the best matches first, or without a query the most recently received first.</param>
|
||||
/// <param name="HasMore">Whether the next page is worth asking for.</param>
|
||||
/// <param name="Gaps">What kept the search from covering the whole mailbox, empty when nothing did.</param>
|
||||
public sealed record MailSearchPage(IReadOnlyList<MailSearchHit> Hits, bool HasMore, IReadOnlyList<RetrievalGap> Gaps)
|
||||
{
|
||||
public static readonly MailSearchPage EMPTY = new([], false, []);
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// How far the index of a mailbox reaches, so that an answer from it can say what it does not cover.
|
||||
/// </summary>
|
||||
/// <param name="ReceivedSinceUtc">Since when mails are indexed, or null when all of them are. Flagged mails are indexed regardless of their age.</param>
|
||||
/// <param name="LastCompleteSyncUtc">When the last complete sync ended, or null while the first one is still running.</param>
|
||||
/// <param name="SignInRefusedAtUtc">When the server refused to let AI Studio sign in, or null when it did not. Until the user deals with that, no new mail arrives in the index.</param>
|
||||
/// <param name="PendingRemovalCount">How many mails the index still holds although a sync would have removed them, waiting for the user to agree, or null when there are none.</param>
|
||||
/// <param name="Folders">The folders of the mailbox and how far their sync got, ordered by path.</param>
|
||||
public sealed record MailboxCoverage(DateTimeOffset? ReceivedSinceUtc, DateTimeOffset? LastCompleteSyncUtc, DateTimeOffset? SignInRefusedAtUtc, int? PendingRemovalCount, IReadOnlyList<MailFolderRecord> Folders);
|
||||
@@ -0,0 +1,402 @@
|
||||
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.Mail;
|
||||
using AIStudio.Tools.RAG;
|
||||
|
||||
namespace AIStudio.Tools.Services;
|
||||
|
||||
/// <summary>
|
||||
/// Searches, counts and reads the mails of the mailboxes, for the mail tools.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Every way into a mailbox goes through here, and each one checks on its own whether the provider
|
||||
/// of the chat may read the mailbox. The tools offer only the mailboxes it may read, but the settings
|
||||
/// may change between preparing a request and running one of its calls.<br/><br/>
|
||||
/// Everything comes from the index, nothing from the server, so nothing here signs in anywhere. The
|
||||
/// log names a mailbox, but never a subject, an address, or a folder.
|
||||
/// </remarks>
|
||||
public sealed class MailboxRetrievalService(SettingsManager settingsManager, DatabaseClientProvider databaseClientProvider, LocalIndexSearchService indexSearch, ILogger<MailboxRetrievalService> logger)
|
||||
{
|
||||
/// <summary>
|
||||
/// How many chunks each channel delivers for every mail a page needs.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The channels find chunks, a page lists mails. A mail which matches tends to match with several
|
||||
/// of its chunks, its header block and its text, maybe an attachment as well, so a window of as
|
||||
/// many chunks as mails would come up short. Paging stops at RetrievalPaging.MAX_RESULT_WINDOW
|
||||
/// mails, which makes four hundred chunks per channel at most.
|
||||
/// </remarks>
|
||||
private const int CHUNKS_PER_MAIL = 4;
|
||||
|
||||
/// <summary>
|
||||
/// Whether mailboxes exist in this installation right now: both of their previews are switched on.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Mailboxes are a preview of their own, on top of local RAG. The mail tools are available
|
||||
/// exactly as long as this holds, see IToolImplementation.IsAvailable.
|
||||
/// </remarks>
|
||||
public bool AreMailboxesEnabled => PreviewFeatures.PRE_RAG_2024.IsEnabled(settingsManager) && PreviewFeatures.PRE_MAILBOXES_2026.IsEnabled(settingsManager);
|
||||
|
||||
/// <summary>
|
||||
/// The mailboxes which the provider of a chat may read, in the order the tools offer them.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The same mailboxes always come in the same order, by their number and then by their id: the
|
||||
/// providers cache a request from its beginning, and the tools describing the mailboxes are part
|
||||
/// of it, cf. SemanticSearchTool.InOfferOrder. A mailbox on a server the organization does not
|
||||
/// allow is left out, although its index is still there, cf. MailServerPolicy.
|
||||
/// </remarks>
|
||||
/// <param name="chatProviderConfidence">How much the provider of the chat is trusted.</param>
|
||||
/// <returns>The mailboxes, none while one of the previews is switched off.</returns>
|
||||
public IReadOnlyList<DataSourceMailbox> GetReadableMailboxes(ConfidenceLevel chatProviderConfidence)
|
||||
{
|
||||
if (!this.AreMailboxesEnabled)
|
||||
return [];
|
||||
|
||||
var policy = MailServerPolicy.Read(settingsManager);
|
||||
return settingsManager.ConfigurationData.Mailboxes
|
||||
.Where(mailbox => policy.IsAllowed(mailbox.Host))
|
||||
.Where(mailbox => IsReadable(mailbox.ConfidenceLevel, chatProviderConfidence, this.GetEmbeddingProviderConfidence(mailbox)))
|
||||
.OrderBy(mailbox => mailbox.Num)
|
||||
.ThenBy(mailbox => mailbox.Id, StringComparer.Ordinal)
|
||||
.ToList();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Whether the providers which see the content of a mailbox may see it.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Two providers see it while it is searched: the provider of the chat reads what is found, and
|
||||
/// the embedding provider gets the query, which the model may well have written from a mail. The
|
||||
/// embedding provider has to meet the level even though it embedded the mails long before: its
|
||||
/// confidence may have been lowered since. A mailbox whose embedding provider is gone cannot be
|
||||
/// read at all, because nobody can tell whether a replacement would meet the level.
|
||||
/// </remarks>
|
||||
/// <param name="mailboxConfidence">The confidence level the mailbox requires.</param>
|
||||
/// <param name="chatProviderConfidence">How much the provider of the chat is trusted.</param>
|
||||
/// <param name="embeddingProviderConfidence">How much the embedding provider of the mailbox is trusted, or null when it is not available.</param>
|
||||
/// <returns>True when the mailbox may be read.</returns>
|
||||
internal static bool IsReadable(ConfidenceLevel mailboxConfidence, ConfidenceLevel chatProviderConfidence, ConfidenceLevel? embeddingProviderConfidence) =>
|
||||
embeddingProviderConfidence is { } embeddingConfidence
|
||||
&& chatProviderConfidence.AllowsMailboxConfidenceLevel(mailboxConfidence)
|
||||
&& embeddingConfidence.AllowsMailboxConfidenceLevel(mailboxConfidence);
|
||||
|
||||
/// <summary>
|
||||
/// Reads how far the index of a mailbox reaches.
|
||||
/// </summary>
|
||||
/// <param name="chatProviderConfidence">How much the provider of the chat is trusted.</param>
|
||||
/// <param name="mailboxId">The mailbox.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The coverage, or null when the index cannot be read right now.</returns>
|
||||
/// <exception cref="MailboxNotReadableException">The mailbox is not configured, or the provider of the chat may not read it.</exception>
|
||||
public async Task<MailboxCoverage?> GetCoverageAsync(ConfidenceLevel chatProviderConfidence, string mailboxId, CancellationToken token)
|
||||
{
|
||||
var mailbox = this.RequireReadableMailbox(mailboxId, chatProviderConfidence);
|
||||
try
|
||||
{
|
||||
var indexStore = await databaseClientProvider.GetIndexStoreAsync(token);
|
||||
if (!indexStore.IsAvailable)
|
||||
{
|
||||
logger.LogWarning("Cannot tell how far the index of mailbox '{MailboxName}' ({MailboxId}) reaches, because local RAG index '{DatabaseName}' is unavailable.", mailbox.Name, mailbox.Id, indexStore.Name);
|
||||
return null;
|
||||
}
|
||||
|
||||
return await ReadCoverageAsync(indexStore, mailbox, DateTimeOffset.UtcNow, token);
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Cannot tell how far the index of mailbox '{MailboxName}' ({MailboxId}) reaches.", mailbox.Name, mailbox.Id);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/// <param name="indexStore">The index store, which has to be available.</param>
|
||||
/// <param name="mailbox">The mailbox.</param>
|
||||
/// <param name="now">The point in time the period of the mailbox is counted back from.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
internal static async Task<MailboxCoverage> ReadCoverageAsync(IndexStoreClient indexStore, DataSourceMailbox mailbox, DateTimeOffset now, CancellationToken token)
|
||||
{
|
||||
var syncState = await indexStore.GetMailboxSyncStateAsync(mailbox.Id, token);
|
||||
var refusedSignIn = await indexStore.GetMailboxAuthFailureAsync(mailbox.Id, token);
|
||||
var folders = await indexStore.GetMailFoldersAsync(mailbox.Id, token);
|
||||
|
||||
return new(mailbox.MaxAge.GetReceivedSince(now), syncState.LastSyncCompletedUtc, refusedSignIn?.FailedAtUtc, syncState.PendingRemovalCount, folders);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Searches one page of the mails of a mailbox which meet the conditions.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// With a query, the mails are searched by meaning and by words, and each mail comes with the
|
||||
/// passage which matched best. Without one, the mails which meet the conditions are listed, the
|
||||
/// most recently received first. Either way, a page holds the mails of the page of each channel
|
||||
/// and is cut the same way, see RetrievalPaging. The query comes from the model, so its problems
|
||||
/// are reported to the model and not to the user.
|
||||
/// </remarks>
|
||||
/// <param name="chatProviderConfidence">How much the provider of the chat is trusted.</param>
|
||||
/// <param name="mailboxId">The mailbox.</param>
|
||||
/// <param name="query">What to search for, or null to list the mails.</param>
|
||||
/// <param name="filter">The conditions the mails have to meet.</param>
|
||||
/// <param name="page">The page, starting at 1, up to RetrievalPaging.GetLastPage for the page size of the mailbox.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The page; empty with a gap when the mailbox could not be searched.</returns>
|
||||
/// <exception cref="MailboxNotReadableException">The mailbox is not configured, or the provider of the chat may not read it.</exception>
|
||||
/// <exception cref="ArgumentOutOfRangeException">The page is below 1 or beyond the last page.</exception>
|
||||
public async Task<MailSearchPage> SearchAsync(ConfidenceLevel chatProviderConfidence, string mailboxId, string? query, MailFilter filter, int page, CancellationToken token)
|
||||
{
|
||||
var mailbox = this.RequireReadableMailbox(mailboxId, chatProviderConfidence);
|
||||
var pageSize = (int)mailbox.MaxMatches;
|
||||
|
||||
// Checks the page before anything is searched:
|
||||
_ = RetrievalPaging.GetWindowSize(page, pageSize);
|
||||
if (pageSize == 0)
|
||||
return MailSearchPage.EMPTY;
|
||||
|
||||
var run = new RetrievalRun(queryWrittenByUser: false);
|
||||
if (await indexSearch.IsAwaitingReindexAsync(mailbox, run, token))
|
||||
return MailSearchPage.EMPTY with { Gaps = run.GetGaps() };
|
||||
|
||||
var byRelevance = !string.IsNullOrWhiteSpace(query);
|
||||
try
|
||||
{
|
||||
var indexStore = await databaseClientProvider.GetIndexStoreAsync(token);
|
||||
if (!indexStore.IsAvailable)
|
||||
{
|
||||
logger.LogWarning("Skipping the search of mailbox '{MailboxName}' ({MailboxId}) because local RAG index '{DatabaseName}' is unavailable.", mailbox.Name, mailbox.Id, indexStore.Name);
|
||||
run.Add(RetrievalGap.NOT_SEARCHED);
|
||||
return MailSearchPage.EMPTY with { Gaps = run.GetGaps() };
|
||||
}
|
||||
|
||||
var result = string.IsNullOrWhiteSpace(query)
|
||||
? await ListNewestFirstAsync(indexStore, mailbox.Id, filter, page, pageSize, token)
|
||||
: await this.SearchByRelevanceAsync(indexStore, mailbox, query, filter, page, pageSize, run, token);
|
||||
|
||||
var gaps = run.GetGaps();
|
||||
logger.LogInformation(
|
||||
"Searched mailbox '{MailboxName}' ({MailboxId}). ByRelevance={ByRelevance}, Page={Page}, Mails={MailCount}, HasMore={HasMore}, Gaps=[{Gaps}].",
|
||||
mailbox.Name,
|
||||
mailbox.Id,
|
||||
byRelevance,
|
||||
page,
|
||||
result.Hits.Count,
|
||||
result.HasMore,
|
||||
string.Join(", ", gaps));
|
||||
|
||||
return result with { Gaps = gaps };
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Searching mailbox '{MailboxName}' ({MailboxId}) failed. ByRelevance={ByRelevance}.", mailbox.Name, mailbox.Id, byRelevance);
|
||||
run.Add(RetrievalGap.NOT_SEARCHED);
|
||||
return MailSearchPage.EMPTY with { Gaps = run.GetGaps() };
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Lists one page of the mails which meet the conditions, the most recently received first.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The window of the page is fetched from the start, as with a search by relevance, although
|
||||
/// the index could skip to the page directly. That keeps one rule for how far a mailbox can be
|
||||
/// paged through and when a further page is worth asking for, whichever way it is searched.
|
||||
/// </remarks>
|
||||
/// <param name="indexStore">The index store, which has to be available.</param>
|
||||
/// <param name="mailboxId">The mailbox.</param>
|
||||
/// <param name="filter">The conditions.</param>
|
||||
/// <param name="page">The page, starting at 1.</param>
|
||||
/// <param name="pageSize">How many mails a page lists.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
internal static async Task<MailSearchPage> ListNewestFirstAsync(IndexStoreClient indexStore, string mailboxId, MailFilter filter, int page, int pageSize, CancellationToken token)
|
||||
{
|
||||
var window = await indexStore.QueryMailsAsync(mailboxId, filter, 0, RetrievalPaging.GetWindowSize(page, pageSize), token);
|
||||
var (mailIds, hasMore) = RetrievalPaging.Cut(window, page, pageSize);
|
||||
var summaries = await indexStore.GetMailSummariesAsync(mailboxId, mailIds, token);
|
||||
|
||||
return new(summaries.Select(summary => new MailSearchHit(summary, null)).ToList(), hasMore, []);
|
||||
}
|
||||
|
||||
private async Task<MailSearchPage> SearchByRelevanceAsync(IndexStoreClient indexStore, DataSourceMailbox mailbox, string query, MailFilter filter, int page, int pageSize, RetrievalRun run, CancellationToken token)
|
||||
{
|
||||
var chunkWindow = RetrievalPaging.GetWindowSize(page, pageSize) * CHUNKS_PER_MAIL;
|
||||
|
||||
//
|
||||
// Without a condition, every chunk of the collection may match, and the collection holds
|
||||
// nothing but this mailbox. Sending the ids of all its chunks along would only make the
|
||||
// request large, so the whole collection is searched, and the summaries leave out the few
|
||||
// mails which lost their last place on the server since the last sync. With a condition,
|
||||
// the vector search sees the chunks of the mails which meet it, and nothing else:
|
||||
//
|
||||
var vectorFilter = filter.HasConditions ? new VectorSearchFilter(await indexStore.GetMailChunkIdsAsync(mailbox.Id, filter, token)) : null;
|
||||
var vectorTask = indexSearch.SearchVectorsAsync(mailbox, query, chunkWindow, vectorFilter, run, token);
|
||||
var keywordTask = indexSearch.SearchKeywordsAsync(mailbox, chunkWindow, store => store.SearchMailChunksAsync(mailbox.Id, query, filter, chunkWindow, token), run, token);
|
||||
|
||||
await Task.WhenAll(vectorTask, keywordTask);
|
||||
token.ThrowIfCancellationRequested();
|
||||
|
||||
var vectorChunks = vectorTask.Result;
|
||||
var keywordChunks = keywordTask.Result;
|
||||
var (passages, hasMore) = RetrievalPaging.Merge(
|
||||
BestPassagePerMail(vectorChunks.Select(chunk => new MailPassage(chunk.ParentFileId, chunk.Text))),
|
||||
BestPassagePerMail(keywordChunks.Select(chunk => new MailPassage(chunk.ParentFileId, chunk.ChunkText))),
|
||||
passage => passage.MailId,
|
||||
page,
|
||||
pageSize);
|
||||
|
||||
hasMore = hasMore || MayHoldMoreMails(page, pageSize, chunkWindow, vectorChunks.Count, keywordChunks.Count);
|
||||
|
||||
var passageTexts = passages
|
||||
.DistinctBy(passage => passage.MailId, StringComparer.Ordinal)
|
||||
.ToDictionary(passage => passage.MailId, passage => passage.Text, StringComparer.Ordinal);
|
||||
|
||||
var summaries = await indexStore.GetMailSummariesAsync(mailbox.Id, passages.Select(passage => passage.MailId).ToList(), token);
|
||||
return new(summaries.Select(summary => new MailSearchHit(summary, passageTexts[summary.MailId])).ToList(), hasMore, []);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Keeps the best chunk of every mail, which turns a list of chunks into a list of mails.
|
||||
/// </summary>
|
||||
/// <param name="chunks">The chunks a channel found, the best first.</param>
|
||||
/// <returns>One passage per mail, in the order of their best chunks.</returns>
|
||||
internal static IReadOnlyList<MailPassage> BestPassagePerMail(IEnumerable<MailPassage> chunks)
|
||||
{
|
||||
var seenMails = new HashSet<string>(StringComparer.Ordinal);
|
||||
return chunks.Where(chunk => !string.IsNullOrWhiteSpace(chunk.MailId) && seenMails.Add(chunk.MailId)).ToList();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Whether a channel may hold further mails which its window of chunks did not reach.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A channel which filled its whole window may have more to show, even when its chunks belong
|
||||
/// to fewer mails than the page needed. Saying there is more when there is not costs one page
|
||||
/// which turns out empty; saying the opposite would hide mails.
|
||||
/// </remarks>
|
||||
/// <param name="page">The page, starting at 1.</param>
|
||||
/// <param name="pageSize">How many mails a page lists per channel.</param>
|
||||
/// <param name="chunkWindow">How many chunks each channel was asked for.</param>
|
||||
/// <param name="chunkCounts">How many chunks each channel delivered.</param>
|
||||
/// <returns>True when the next page is worth asking for.</returns>
|
||||
internal static bool MayHoldMoreMails(int page, int pageSize, int chunkWindow, params int[] chunkCounts) =>
|
||||
page < RetrievalPaging.GetLastPage(pageSize) && chunkCounts.Any(count => count >= chunkWindow);
|
||||
|
||||
/// <summary>
|
||||
/// Counts the mails of a mailbox which meet the conditions.
|
||||
/// </summary>
|
||||
/// <param name="chatProviderConfidence">How much the provider of the chat is trusted.</param>
|
||||
/// <param name="mailboxId">The mailbox.</param>
|
||||
/// <param name="filter">The conditions.</param>
|
||||
/// <param name="grouping">How to break the number down.</param>
|
||||
/// <param name="maxGroups">How many groups to return at most, the largest ones.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The count; without a number and with a gap when the mailbox could not be counted.</returns>
|
||||
/// <exception cref="MailboxNotReadableException">The mailbox is not configured, or the provider of the chat may not read it.</exception>
|
||||
public async Task<MailCountOutcome> CountAsync(ConfidenceLevel chatProviderConfidence, string mailboxId, MailFilter filter, MailCountGrouping grouping, int maxGroups, CancellationToken token)
|
||||
{
|
||||
var mailbox = this.RequireReadableMailbox(mailboxId, chatProviderConfidence);
|
||||
|
||||
// While the index is built anew, it holds only part of the mails, and any number would be too low:
|
||||
var run = new RetrievalRun(queryWrittenByUser: false);
|
||||
if (await indexSearch.IsAwaitingReindexAsync(mailbox, run, token))
|
||||
return new(null, run.GetGaps());
|
||||
|
||||
try
|
||||
{
|
||||
var indexStore = await databaseClientProvider.GetIndexStoreAsync(token);
|
||||
if (!indexStore.IsAvailable)
|
||||
{
|
||||
logger.LogWarning("Skipping the count of mailbox '{MailboxName}' ({MailboxId}) because local RAG index '{DatabaseName}' is unavailable.", mailbox.Name, mailbox.Id, indexStore.Name);
|
||||
return new(null, [RetrievalGap.NOT_SEARCHED]);
|
||||
}
|
||||
|
||||
var count = await indexStore.CountMailsAsync(mailbox.Id, filter, grouping, maxGroups, token);
|
||||
logger.LogInformation("Counted mailbox '{MailboxName}' ({MailboxId}). Grouping={Grouping}, Total={TotalCount}, Groups={GroupCount}.", mailbox.Name, mailbox.Id, grouping, count.TotalCount, count.Groups.Count);
|
||||
return new(count, []);
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
{
|
||||
throw;
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Counting mailbox '{MailboxName}' ({MailboxId}) failed. Grouping={Grouping}.", mailbox.Name, mailbox.Id, grouping);
|
||||
return new(null, [RetrievalGap.NOT_SEARCHED]);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Reads a mail from whichever mailbox holds it, among those the provider of the chat may read.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The id of a mail is derived from the id of its mailbox, so no two mailboxes share one. A mail
|
||||
/// of a mailbox the provider may not read is not found, exactly like a mail nobody ever indexed.
|
||||
/// </remarks>
|
||||
/// <param name="chatProviderConfidence">How much the provider of the chat is trusted.</param>
|
||||
/// <param name="mailId">The id of the mail.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The mail, or null when no mailbox the provider may read holds it.</returns>
|
||||
/// <exception cref="InvalidOperationException">The index cannot be read right now, so nobody can tell whether the mail exists.</exception>
|
||||
public async Task<MailReading?> ReadAsync(ConfidenceLevel chatProviderConfidence, string mailId, CancellationToken token)
|
||||
{
|
||||
var mailboxes = this.GetReadableMailboxes(chatProviderConfidence);
|
||||
if (mailboxes.Count == 0)
|
||||
return null;
|
||||
|
||||
var indexStore = await databaseClientProvider.GetIndexStoreAsync(token);
|
||||
if (!indexStore.IsAvailable)
|
||||
throw new InvalidOperationException($"No mail can be read, because local RAG index '{indexStore.Name}' is unavailable.");
|
||||
|
||||
return await ReadMailAsync(indexStore, mailboxes, mailId, token);
|
||||
}
|
||||
|
||||
/// <remarks>
|
||||
/// A mail which lost its last place on the server is not read: it is gone, or about to turn up
|
||||
/// under another id once the next sync found where it went.
|
||||
/// </remarks>
|
||||
/// <param name="indexStore">The index store, which has to be available.</param>
|
||||
/// <param name="mailboxes">The mailboxes to look in.</param>
|
||||
/// <param name="mailId">The id of the mail.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
internal static async Task<MailReading?> ReadMailAsync(IndexStoreClient indexStore, IReadOnlyList<DataSourceMailbox> mailboxes, string mailId, CancellationToken token)
|
||||
{
|
||||
foreach (var mailbox in mailboxes)
|
||||
{
|
||||
var summaries = await indexStore.GetMailSummariesAsync(mailbox.Id, [mailId], token);
|
||||
if (summaries.Count == 0)
|
||||
continue;
|
||||
|
||||
var mail = await indexStore.GetMailAsync(mailbox.Id, mailId, token);
|
||||
if (mail is null)
|
||||
continue;
|
||||
|
||||
var inReplyToMailId = await indexStore.FindMailByMessageIdAsync(mailbox.Id, mail.InReplyTo, token);
|
||||
return new(mailbox, summaries[0], mail, inReplyToMailId == mailId ? null : inReplyToMailId);
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
private DataSourceMailbox RequireReadableMailbox(string mailboxId, ConfidenceLevel chatProviderConfidence)
|
||||
{
|
||||
foreach (var mailbox in this.GetReadableMailboxes(chatProviderConfidence))
|
||||
if (mailbox.Id == mailboxId)
|
||||
return mailbox;
|
||||
|
||||
logger.LogWarning("The mailbox '{MailboxId}' is not configured, the provider of the chat may not read it, or the organization does not allow its server. Its mails stay closed.", mailboxId);
|
||||
throw new MailboxNotReadableException(mailboxId);
|
||||
}
|
||||
|
||||
private ConfidenceLevel? GetEmbeddingProviderConfidence(DataSourceMailbox mailbox) =>
|
||||
DataSourceEmbeddingProviders.TryResolve(settingsManager, mailbox, out var embeddingProvider) ? embeddingProvider.GetConfidenceLevel(settingsManager) : null;
|
||||
}
|
||||
@@ -55,7 +55,8 @@ public sealed partial class PluginInstallService
|
||||
DocumentAnalysisPolicies: CountObjects(PluginConfigurationObjectType.DOCUMENT_ANALYSIS_POLICY),
|
||||
LockedSettings: lockedSettings,
|
||||
MandatoryInfos: configurationPlugin.MandatoryInfos.Count,
|
||||
Introductions: configurationPlugin.Introductions.Count);
|
||||
Introductions: configurationPlugin.Introductions.Count,
|
||||
MailboxProviders: configurationPlugin.MailboxProviders.Count);
|
||||
|
||||
int CountObjects(PluginConfigurationObjectType type) => configObjects.Count(configObject => configObject.Type == type);
|
||||
}
|
||||
|
||||
@@ -192,7 +192,8 @@ public sealed partial class PluginInstallService
|
||||
DocumentAnalysisPolicies: CountObjects(PluginConfigurationObjectType.DOCUMENT_ANALYSIS_POLICY),
|
||||
DeclaredSettings: configurationPlugin.DeclaredSettingsCount,
|
||||
MandatoryInfos: configurationPlugin.MandatoryInfos.Count,
|
||||
Introductions: configurationPlugin.Introductions.Count);
|
||||
Introductions: configurationPlugin.Introductions.Count,
|
||||
MailboxProviders: configurationPlugin.MailboxProviders.Count);
|
||||
|
||||
int CountObjects(PluginConfigurationObjectType type) => configObjects.Count(configObject => configObject.Type == type);
|
||||
}
|
||||
|
||||
@@ -265,16 +265,29 @@ public sealed partial class RustService
|
||||
|
||||
public async IAsyncEnumerable<string> StreamArbitraryFileData(string path, bool extractImages = false, [EnumeratorCancellation] CancellationToken token = default)
|
||||
{
|
||||
await foreach (var segment in this.StreamArbitraryFileDataCore(path, extractImages, false, string.Empty, token))
|
||||
await foreach (var segment in this.StreamArbitraryFileDataCore(path, extractImages, false, string.Empty, null, token))
|
||||
yield return segment.Content;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Reads the content of a file for embedding, each piece together with its token count.
|
||||
/// </summary>
|
||||
/// <param name="path">The path of the file to read.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider whose tokenizer counts the pieces.</param>
|
||||
/// <param name="reportAs">
|
||||
/// The source the user is told about when passages were filtered out, or null for the file
|
||||
/// itself. A file which only exists while it is read, e.g. the attachment of a mail, names
|
||||
/// where its content came from instead.
|
||||
/// </param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The pieces of the content, in order.</returns>
|
||||
public async IAsyncEnumerable<ArbitraryFileDataSegment> StreamArbitraryFileDataWithTokenCounts(
|
||||
string path,
|
||||
EmbeddingProvider embeddingProvider,
|
||||
PromptInjectionSource? reportAs = null,
|
||||
[EnumeratorCancellation] CancellationToken token = default)
|
||||
{
|
||||
await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, embeddingProvider.TokenizerPath, token))
|
||||
await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, embeddingProvider.TokenizerPath, reportAs, token))
|
||||
{
|
||||
if (segment.TokenCount is { } tokenCount)
|
||||
{
|
||||
@@ -308,6 +321,7 @@ public sealed partial class RustService
|
||||
bool extractImages,
|
||||
bool includeTokenCount,
|
||||
string tokenizerPath,
|
||||
PromptInjectionSource? reportAs,
|
||||
[EnumeratorCancellation] CancellationToken token)
|
||||
{
|
||||
var streamId = Guid.NewGuid().ToString();
|
||||
@@ -447,7 +461,7 @@ public sealed partial class RustService
|
||||
// indexing run could filter documents without ever saying so.
|
||||
//
|
||||
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
|
||||
await guardService.ReportAsync(new(PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount));
|
||||
await guardService.ReportAsync(new(reportAs ?? PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount));
|
||||
}
|
||||
|
||||
private bool TryLogSseErrorMessage(string jsonContent, string path)
|
||||
|
||||
Reference in new issue
Block a user