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

This commit is contained in:
Thorsten Sommer authored and GitHub committed 2026-10-04 12:26:26 +02:00
1 parent 93be539e50
commit c4400c0ff6
343 files changed
+26051 -3122

No files matched your search

@@ -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);
@@ -0,0 +1,760 @@
using System.Text.RegularExpressions;
using AIStudio.Settings;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.Rust;
namespace AIStudio.Tools.Services.Indexing;
/// <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)
{
/// <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>
/// For presentations: slides, then lines and words.
/// </summary>
public static readonly ChunkingStrategy PRESENTATION_STRATEGY = new("presentation", [
new("Slide", SplitBySourceSegments, true),
new("Line break", SplitByLineBreaks),
new("Whitespace", SplitByWhitespace),
new("Hard cut", null),
]);
/// <summary>
/// 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>
/// 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="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)
{
if (FileTypes.IsAllowedPath(filePath, FileTypes.POWER_POINT))
return PRESENTATION_STRATEGY;
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;
}
/// <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, heading, estimatedTokenCount))
yield return chunk;
}
private async IAsyncEnumerable<EmbeddingChunk> SplitTextByRulesAsync(
string text,
IReadOnlyList<TextSegment> sourceSegments,
ChunkingStrategy strategy,
int ruleIndex,
ChunkingOptions options,
EmbeddingProvider embeddingProvider,
int? currentPageNumber,
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
string requiredOverlapPrefix = "",
int? estimatedTokenCount = null)
{
text = text.Trim();
if (string.IsNullOrWhiteSpace(text))
yield break;
var tokenCount = estimatedTokenCount;
var textWithOverlap = AddOverlapPrefix(text, requiredOverlapPrefix);
if (textWithOverlap.Length <= RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH &&
(estimatedTokenCount is null || estimatedTokenCount <= options.MaxChunkTokenLength))
{
tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, textWithOverlap, token);
if (tokenCount <= options.MaxChunkTokenLength)
{
yield return new(textWithOverlap, currentPageNumber);
yield break;
}
}
if (ruleIndex >= strategy.Rules.Count)
{
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, currentPageNumber, token, requiredOverlapPrefix, estimatedTokenCount))
yield return hardChunk;
yield break;
}
var rule = strategy.Rules[ruleIndex];
if (rule.Split is null)
{
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, currentPageNumber, token, requiredOverlapPrefix, estimatedTokenCount))
yield return hardChunk;
yield break;
}
var units = NormalizeSplitUnits(rule.Split(text, sourceSegments.Select(segment => segment.Text).ToList()), text);
if (units.Count <= 1)
{
await foreach (var chunk in this.SplitTextByRulesAsync(text, sourceSegments, strategy, ruleIndex + 1, options, embeddingProvider, currentPageNumber, token, requiredOverlapPrefix, estimatedTokenCount))
yield return chunk;
yield break;
}
logger.LogDebug(
"Splitting content for embedding provider '{EmbeddingProviderName}' with strategy '{ChunkingStrategy}' and rule '{ChunkingRule}'. EstimatedTokenCount={EstimatedTokenCount}, MaxChunkTokenLength={MaxChunkTokenLength}, OverlapTokenLength={OverlapTokenLength}.",
embeddingProvider.Name,
strategy.Name,
rule.Name,
tokenCount,
options.MaxChunkTokenLength,
options.OverlapTokenLength);
var index = 0;
var overlapPrefix = requiredOverlapPrefix;
var unitTokenCounts = EstimateSplitUnitTokenCounts(units, sourceSegments, rule.UsesSourceSegmentCounts, estimatedTokenCount);
//
// The first rule of every strategy cuts along the segments the runtime delivered, so there
// a unit is a segment and carries that segment's page. Every later rule cuts inside a
// single segment, where all units share the page they were handed. This is what ties a
// chunk to a page without anybody reading the text.
//
var unitsAreSourceSegments = rule.UsesSourceSegmentCounts && sourceSegments.Count == units.Count;
int? PageOfUnit(int unitIndex) => unitsAreSourceSegments ? sourceSegments[unitIndex].PageNumber ?? currentPageNumber : currentPageNumber;
while (index < units.Count)
{
token.ThrowIfCancellationRequested();
var unitCount = await this.FindLargestUnitCountWithinMaxChunkLengthAsync(units, unitTokenCounts, index, embeddingProvider, options.MaxChunkTokenLength, token, overlapPrefix);
if (unitCount > 0)
{
var rawChunk = string.Concat(units.Skip(index).Take(unitCount)).Trim();
var chunk = AddOverlapPrefix(rawChunk, overlapPrefix);
overlapPrefix = string.Empty;
//
// The page of the first unit this chunk covers, not of the overlap prefix in front
// of it: the prefix repeats what the chunk before already said, while the page has
// to name where this chunk's own content begins.
//
if (!string.IsNullOrWhiteSpace(chunk))
yield return new(chunk, PageOfUnit(index));
var nextIndex = index + unitCount;
if (nextIndex >= units.Count)
yield break;
var nextStartIndex = await this.CalculateNextStartIndexAsync(units, index, nextIndex, options, embeddingProvider, token);
if (nextStartIndex < nextIndex)
{
logger.LogDebug(
"Applied delimiter overlap while chunking. Strategy='{ChunkingStrategy}', Rule='{ChunkingRule}', PreviousStartUnitIndex={PreviousStartUnitIndex}, PreviousEndUnitIndex={PreviousEndUnitIndex}, NextStartUnitIndex={NextStartUnitIndex}, OverlapUnits={OverlapUnits}, OverlapTokenLength={OverlapTokenLength}.",
strategy.Name,
rule.Name,
index,
nextIndex,
nextStartIndex,
nextIndex - nextStartIndex,
options.OverlapTokenLength);
index = nextStartIndex;
}
else
{
overlapPrefix = await this.CreateOverlapPrefixAsync(chunk, strategy, rule, options, embeddingProvider, token);
index = nextIndex;
}
continue;
}
string? lastSplitUnit = null;
var unitTokenCount = unitTokenCounts?[index];
var unitPageNumber = PageOfUnit(index);
await foreach (var splitUnit in this.SplitTextByRulesAsync(units[index], [new(units[index], unitTokenCount, unitPageNumber)], strategy, ruleIndex + 1, options, embeddingProvider, unitPageNumber, token, overlapPrefix, unitTokenCount))
{
lastSplitUnit = splitUnit.Text;
yield return splitUnit;
}
overlapPrefix = lastSplitUnit is null
? string.Empty
: await this.CreateOverlapPrefixAsync(lastSplitUnit, strategy, rule, options, embeddingProvider, token);
index++;
}
}
private async Task<int> FindLargestUnitCountWithinMaxChunkLengthAsync(IReadOnlyList<string> units, IReadOnlyList<int>? estimatedUnitTokenCounts, int startUnitIndex, EmbeddingProvider embeddingProvider, int maxChunkTokenLength, CancellationToken token, string overlapPrefix = "")
{
var minimumCandidateUnitCount = 1;
var availableUnitCount = units.Count - startUnitIndex;
var maximumCandidateUnitCount = availableUnitCount;
var largestValidUnitCount = 0;
if (estimatedUnitTokenCounts is not null)
{
maximumCandidateUnitCount = 0;
var cumulativeEstimatedTokenCount = 0L;
for (var unitIndex = startUnitIndex; unitIndex < units.Count; unitIndex++)
{
cumulativeEstimatedTokenCount += estimatedUnitTokenCounts[unitIndex];
if (cumulativeEstimatedTokenCount > maxChunkTokenLength)
break;
maximumCandidateUnitCount++;
}
if (maximumCandidateUnitCount == 0)
maximumCandidateUnitCount = 1;
}
while (true)
{
var searchedMaximumCandidateUnitCount = maximumCandidateUnitCount;
while (minimumCandidateUnitCount <= maximumCandidateUnitCount)
{
token.ThrowIfCancellationRequested();
var candidateUnitCount = minimumCandidateUnitCount + (maximumCandidateUnitCount - minimumCandidateUnitCount) / 2;
var candidateText = AddOverlapPrefix(string.Concat(units.Skip(startUnitIndex).Take(candidateUnitCount)).Trim(), overlapPrefix);
var candidateFits = candidateText.Length <= RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH &&
await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidateText, token) <= maxChunkTokenLength;
if (candidateFits)
{
largestValidUnitCount = candidateUnitCount;
minimumCandidateUnitCount = candidateUnitCount + 1;
}
else
maximumCandidateUnitCount = candidateUnitCount - 1;
}
if (largestValidUnitCount < searchedMaximumCandidateUnitCount || largestValidUnitCount >= availableUnitCount)
break;
minimumCandidateUnitCount = searchedMaximumCandidateUnitCount + 1;
maximumCandidateUnitCount = (int)Math.Min(
availableUnitCount,
Math.Max(minimumCandidateUnitCount, (long)searchedMaximumCandidateUnitCount * 2));
}
return largestValidUnitCount;
}
private static int? SumTokenCounts(IReadOnlyList<TextSegment> segments)
{
var result = 0L;
foreach (var segment in segments)
{
if (segment.TokenCount is null)
return null;
result += segment.TokenCount.Value;
}
return (int)Math.Min(result, int.MaxValue);
}
private static IReadOnlyList<int>? EstimateSplitUnitTokenCounts(
IReadOnlyList<string> units,
IReadOnlyList<TextSegment> sourceSegments,
bool usesSourceSegmentCounts,
int? sourceTokenCount)
{
if (usesSourceSegmentCounts && sourceSegments.Count == units.Count && sourceSegments.All(segment => segment.TokenCount is not null))
return sourceSegments.Select(segment => segment.TokenCount.GetValueOrDefault()).ToList();
if (sourceTokenCount is null)
return null;
var totalLength = Math.Max(1, units.Sum(unit => unit.Length));
var result = new List<int>(units.Count);
var allocatedTokenCount = 0;
var consumedLength = 0L;
foreach (var unit in units)
{
consumedLength += unit.Length;
var tokenCountAtBoundary = (int)Math.Min(sourceTokenCount.Value, sourceTokenCount.Value * consumedLength / totalLength);
result.Add(Math.Max(0, tokenCountAtBoundary - allocatedTokenCount));
allocatedTokenCount = tokenCountAtBoundary;
}
if (result.Count > 0 && allocatedTokenCount < sourceTokenCount.Value)
result[^1] += sourceTokenCount.Value - allocatedTokenCount;
return result;
}
private async Task<string> CreateOverlapPrefixAsync(string chunk, ChunkingStrategy strategy, ChunkingRule rule, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
{
return await this.CreateOverlapPrefixAsync(chunk, strategy.Name, rule.Name, options, embeddingProvider, token);
}
private async Task<string> CreateOverlapPrefixAsync(string chunk, string strategyName, string ruleName, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
{
if (options.OverlapTokenLength <= 0)
return string.Empty;
chunk = chunk.Trim();
if (string.IsNullOrWhiteSpace(chunk))
return string.Empty;
var chunkTokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, chunk, token);
if (chunkTokenCount <= options.OverlapTokenLength)
{
logger.LogDebug(
"Applied whole-chunk overlap while chunking because the previous chunk is smaller than the requested overlap. Strategy='{ChunkingStrategy}', Rule='{ChunkingRule}', RequestedOverlapTokenLength={RequestedOverlapTokenLength}, ActualOverlapTokenCount={ActualOverlapTokenCount}.",
strategyName,
ruleName,
options.OverlapTokenLength,
chunkTokenCount);
return chunk;
}
var overlapStartIndex = await this.CalculateHardCutOverlapStartIndexAsync(chunk, 0, chunk.Length, options, embeddingProvider, token);
if (overlapStartIndex >= chunk.Length)
overlapStartIndex = FindLastNonWhitespaceStartIndex(chunk);
if (overlapStartIndex >= chunk.Length)
return string.Empty;
var overlapPrefix = chunk[overlapStartIndex..].Trim();
if (string.IsNullOrWhiteSpace(overlapPrefix))
return string.Empty;
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, overlapPrefix, token);
logger.LogDebug(
"Applied hard-cut overlap while chunking because delimiter overlap was not available. Strategy='{ChunkingStrategy}', Rule='{ChunkingRule}', RequestedOverlapTokenLength={RequestedOverlapTokenLength}, ActualOverlapTokenCount={ActualOverlapTokenCount}.",
strategyName,
ruleName,
options.OverlapTokenLength,
tokenCount);
return overlapPrefix;
}
private async Task<int> CalculateNextStartIndexAsync(IReadOnlyList<string> units, int chunkStartIndex, int chunkEndIndex, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
{
if (options.OverlapTokenLength <= 0)
return chunkEndIndex;
var bestStartIndex = chunkEndIndex;
var bestDistance = int.MaxValue;
for (var candidateStartIndex = chunkEndIndex - 1; candidateStartIndex > chunkStartIndex; candidateStartIndex--)
{
token.ThrowIfCancellationRequested();
var candidate = string.Concat(units.Skip(candidateStartIndex).Take(chunkEndIndex - candidateStartIndex)).Trim();
if (string.IsNullOrWhiteSpace(candidate))
continue;
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
var distance = Math.Abs(tokenCount - options.OverlapTokenLength);
if (distance < bestDistance)
{
bestStartIndex = candidateStartIndex;
bestDistance = distance;
}
if (tokenCount >= options.OverlapTokenLength && bestStartIndex < chunkEndIndex)
break;
}
return bestStartIndex <= chunkStartIndex ? chunkEndIndex : bestStartIndex;
}
/// <remarks>
/// The hard cut is only ever reached inside a single piece of text which no rule could split
/// any further, so every chunk it produces sits on the page that piece was handed.
/// </remarks>
private async IAsyncEnumerable<EmbeddingChunk> SplitTextByHardCutAsync(
string text,
ChunkingOptions options,
EmbeddingProvider embeddingProvider,
int? currentPageNumber,
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
string requiredOverlapPrefix = "",
int? estimatedTokenCount = null)
{
text = text.Trim();
var startIndex = 0;
var overlapPrefix = requiredOverlapPrefix;
while (startIndex < text.Length)
{
token.ThrowIfCancellationRequested();
while (startIndex < text.Length && char.IsWhiteSpace(text[startIndex]))
startIndex++;
if (startIndex >= text.Length)
yield break;
var bestEndIndex = startIndex;
var maximumCandidateEndIndex = Math.Min(text.Length, startIndex + RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH);
if (estimatedTokenCount > options.MaxChunkTokenLength)
{
var estimatedChunkLength = Math.Max(1L, (long)text.Length * options.MaxChunkTokenLength / estimatedTokenCount.Value);
maximumCandidateEndIndex = (int)Math.Min(text.Length, startIndex + estimatedChunkLength);
}
while (true)
{
var minimumCandidateEndIndex = bestEndIndex + 1;
var currentMaximumCandidateEndIndex = maximumCandidateEndIndex;
while (minimumCandidateEndIndex <= currentMaximumCandidateEndIndex)
{
var candidateEndIndex = minimumCandidateEndIndex + (currentMaximumCandidateEndIndex - minimumCandidateEndIndex) / 2;
var candidate = AddOverlapPrefix(text[startIndex..candidateEndIndex].Trim(), overlapPrefix);
var candidateFits = candidate.Length <= RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH &&
await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token) <= options.MaxChunkTokenLength;
if (candidateFits)
{
bestEndIndex = candidateEndIndex;
minimumCandidateEndIndex = candidateEndIndex + 1;
}
else
currentMaximumCandidateEndIndex = candidateEndIndex - 1;
}
if (bestEndIndex < maximumCandidateEndIndex || bestEndIndex >= text.Length ||
maximumCandidateEndIndex - startIndex >= RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH)
break;
var previousCandidateLength = maximumCandidateEndIndex - startIndex;
maximumCandidateEndIndex = (int)Math.Min(
Math.Min(text.Length, startIndex + (long)RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH),
startIndex + Math.Max(previousCandidateLength + 1L, previousCandidateLength * 2L));
}
if (bestEndIndex == startIndex)
{
if (!string.IsNullOrWhiteSpace(overlapPrefix))
{
var smallestOverlapPrefix = GetSmallestOverlapPrefix(overlapPrefix);
if (!string.IsNullOrWhiteSpace(smallestOverlapPrefix) && !string.Equals(smallestOverlapPrefix, overlapPrefix, StringComparison.Ordinal))
{
logger.LogDebug(
"Reduced hard-cut overlap because the configured overlap leaves no room for new content. RequestedOverlapTokenLength={RequestedOverlapTokenLength}, MaxChunkTokenLength={MaxChunkTokenLength}.",
options.OverlapTokenLength,
options.MaxChunkTokenLength);
overlapPrefix = smallestOverlapPrefix;
continue;
}
}
var smallestCandidate = AddOverlapPrefix(text[startIndex..Math.Min(startIndex + 1, text.Length)].Trim(), overlapPrefix);
var smallestCandidateTokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, smallestCandidate, token);
throw new InvalidOperationException(string.Format(TB("The chunk size configured for the embedding provider '{0}' is too small: the smallest piece the text can be cut into still has {1} tokens, while the limit is {2}."), embeddingProvider.Name, smallestCandidateTokenCount, options.MaxChunkTokenLength));
}
var chunk = AddOverlapPrefix(text[startIndex..bestEndIndex].Trim(), overlapPrefix);
if (!string.IsNullOrWhiteSpace(chunk))
yield return new(chunk, currentPageNumber);
if (bestEndIndex >= text.Length)
yield break;
overlapPrefix = await this.CreateOverlapPrefixAsync(chunk, "hard-cut", "Hard cut", options, embeddingProvider, token);
startIndex = bestEndIndex;
}
}
private async Task<int> CalculateHardCutOverlapStartIndexAsync(string text, int chunkStartIndex, int chunkEndIndex, ChunkingOptions options, EmbeddingProvider embeddingProvider, CancellationToken token)
{
if (options.OverlapTokenLength <= 0 || chunkEndIndex - chunkStartIndex <= 1)
return chunkEndIndex;
var low = chunkStartIndex + 1;
var high = chunkEndIndex - 1;
var bestStartIndex = chunkEndIndex;
while (low <= high)
{
token.ThrowIfCancellationRequested();
var mid = low + (high - low) / 2;
var candidate = text[mid..chunkEndIndex].Trim();
var tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, candidate, token);
if (tokenCount <= options.OverlapTokenLength)
{
bestStartIndex = mid;
high = mid - 1;
}
else
low = mid + 1;
}
return bestStartIndex <= chunkStartIndex ? chunkEndIndex : bestStartIndex;
}
private static int FindLastNonWhitespaceStartIndex(string text)
{
for (var index = text.Length - 1; index >= 0; index--)
{
if (!char.IsWhiteSpace(text[index]))
return index;
}
return text.Length;
}
private static string GetSmallestOverlapPrefix(string text)
{
var index = FindLastNonWhitespaceStartIndex(text);
return index >= text.Length ? string.Empty : text[index..].Trim();
}
private static string AddOverlapPrefix(string chunk, string overlapPrefix)
{
if (string.IsNullOrWhiteSpace(overlapPrefix))
return chunk.Trim();
return $"{overlapPrefix.TrimEnd()}\n{chunk.TrimStart()}".Trim();
}
private async Task<int> GetEmbeddingTokenCountAsync(EmbeddingProvider embeddingProvider, string text, CancellationToken token)
{
var response = await rustService.GetTokenCount(embeddingProvider, text, token);
if (response is { Success: true })
return response.Value.TokenCount;
var message = response?.Message ?? "No response was returned by the tokenizer service.";
throw new InvalidOperationException(string.Format(TB("The tokens of the text could not be counted for the embedding provider '{0}'. {1}"), embeddingProvider.Name, message));
}
private static List<string> NormalizeSplitUnits(IReadOnlyList<string> units, string fallbackText)
{
var result = units
.Where(unit => !string.IsNullOrWhiteSpace(unit))
.ToList();
return result.Count == 0 ? [fallbackText] : result;
}
private static IReadOnlyList<string> SplitBySourceSegments(string text, IReadOnlyList<string> sourceSegments)
{
return sourceSegments.Count > 1
? sourceSegments.Select(segment => segment + "\n").ToList()
: [text];
}
private static IReadOnlyList<string> SplitByDocumentHeadings(string text, IReadOnlyList<string> sourceSegments)
{
var lines = ReadLines(text);
if (lines.Count < 2)
return [text];
var result = new List<string>();
var segmentStart = 0;
for (var i = 0; i < lines.Count; i++)
{
var (lineStart, _, lineText) = lines[i];
if (lineStart == 0)
continue;
var previousLine = i > 0 ? lines[i - 1].Text : string.Empty;
var nextLine = i + 1 < lines.Count ? lines[i + 1].Text : string.Empty;
if (!IsDocumentHeadingLine(lineText, previousLine, nextLine))
continue;
result.Add(text[segmentStart..lineStart]);
segmentStart = lineStart;
}
if (segmentStart == 0)
return [text];
result.Add(text[segmentStart..]);
return result;
}
private static IReadOnlyList<string> SplitByParagraphs(string text, IReadOnlyList<string> sourceSegments)
{
var matches = ParagraphBreakRegex().Matches(text);
if (matches.Count == 0)
return [text];
var result = new List<string>();
var start = 0;
foreach (Match match in matches)
{
var end = match.Index + match.Length;
result.Add(text[start..end]);
start = end;
}
if (start < text.Length)
result.Add(text[start..]);
return result;
}
private static IReadOnlyList<string> SplitByLineBreaks(string text, IReadOnlyList<string> sourceSegments)
{
var result = new List<string>();
var start = 0;
for (var i = 0; i < text.Length; i++)
{
if (text[i] != '\n')
continue;
result.Add(text[start..(i + 1)]);
start = i + 1;
}
if (start < text.Length)
result.Add(text[start..]);
return result.Count == 0 ? [text] : result;
}
private static IReadOnlyList<string> SplitByWhitespace(string text, IReadOnlyList<string> sourceSegments)
{
var matches = WordWithTrailingWhitespaceRegex().Matches(text);
if (matches.Count == 0)
return [text];
return matches.Select(match => match.Value).ToList();
}
private static List<(int Start, int End, string Text)> ReadLines(string text)
{
var result = new List<(int Start, int End, string Text)>();
var start = 0;
for (var i = 0; i < text.Length; i++)
{
if (text[i] != '\n')
continue;
result.Add((start, i + 1, text[start..(i + 1)]));
start = i + 1;
}
if (start < text.Length)
result.Add((start, text.Length, text[start..]));
return result;
}
private static bool IsDocumentHeadingLine(string line, string previousLine, string nextLine)
{
var trimmed = line.Trim();
if (string.IsNullOrWhiteSpace(trimmed))
return false;
if (MarkdownHeadingRegex().IsMatch(trimmed))
return true;
if (!string.IsNullOrWhiteSpace(previousLine) || !string.IsNullOrWhiteSpace(nextLine))
return false;
if (trimmed.Length is < 3 or > 120)
return false;
if (trimmed.Contains("|", StringComparison.Ordinal) || trimmed.EndsWith(".", StringComparison.Ordinal))
return false;
return PlainHeadingRegex().IsMatch(trimmed);
}
[GeneratedRegex(@"\n[ \t]*\n", RegexOptions.CultureInvariant)]
private static partial Regex ParagraphBreakRegex();
[GeneratedRegex(@"\S+\s*", RegexOptions.CultureInvariant)]
private static partial Regex WordWithTrailingWhitespaceRegex();
[GeneratedRegex(@"^#{1,6}\s+\S", RegexOptions.CultureInvariant)]
private static partial Regex MarkdownHeadingRegex();
/// <summary>
/// A heading without Markdown: a numbered one, a chapter or section, or a line in capitals.
/// </summary>
[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);