Added a way to read what the AI received from your sources (#1100)

Co-authored-by: Thorsten Sommer <SommerEngineering@users.noreply.github.com>
This commit is contained in:
Paul KoudelkaandThorsten Sommer authored and GitHub committed 2026-10-26 10:00:01 +01:00
1 parent e0d8e18f34
commit cbc4ce5413
229 files changed
+12739 -1269

No files matched your search

@@ -5,5 +5,5 @@ namespace AIStudio.Tools.Services;
/// </summary>
/// <param name="Content">The extracted text.</param>
/// <param name="TokenCount">The number of tokens of that text.</param>
/// <param name="PageNumber">The page that text came from, or null when it has none. Presentations and spreadsheets have none.</param>
public sealed record ArbitraryFileDataSegment(string Content, int TokenCount, int? PageNumber);
/// <param name="Location">Where that text came from: its page, its sheet and row, or its line. Presentations have none.</param>
public sealed record ArbitraryFileDataSegment(string Content, int TokenCount, SourceLocation Location);
@@ -5,6 +5,7 @@ using System.Text;
using AIStudio.Chat;
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Tools.Databases;
namespace AIStudio.Tools.Services;
@@ -28,7 +29,7 @@ using Provider = AIStudio.Settings.Provider;
/// provider, the built-in one does the counting, and two tokenizers disagree by a few percent on
/// prose and by more than that on code.
/// </remarks>
public sealed class ConversationTokenCounter(RustService rustService, ILogger<ConversationTokenCounter> logger)
public sealed class ConversationTokenCounter(RustService rustService, DatabaseClientProvider databaseClientProvider, ILogger<ConversationTokenCounter> logger)
{
/// <summary>
/// How much text goes into one counting request.
@@ -267,6 +268,29 @@ public sealed class ConversationTokenCounter(RustService rustService, ILogger<Co
/// </remarks>
private async Task<int> CountDocumentAsync(Provider provider, FileAttachment document, CancellationToken token)
{
//
// A document sent before goes out as the text it was stored with, see FileAttachment.ProvenanceId,
// so that is what counts, whatever its file says by now. The entry never changes, so its id is
// key enough:
//
if (document.ProvenanceId is { } provenanceId)
{
var storedKey = $"{provider.TokenizerPath}{KEY_SEPARATOR}{provenanceId}";
if (this.counted.TryGetValue(storedKey, out var storedTokens))
return storedTokens;
var store = await databaseClientProvider.GetProvenanceStoreAsync(token);
var entries = await store.GetEntriesAsync([provenanceId], token);
if (entries.TryGetValue(provenanceId, out var entry))
{
storedTokens = await this.CountTextAsync(provider, entry.Content, token);
this.counted[storedKey] = storedTokens;
return storedTokens;
}
// Without its entry, the request reads the file instead, and so does the count.
}
var file = new FileInfo(document.FilePath);
if (!file.Exists)
return 0;
@@ -5,6 +5,7 @@ using AIStudio.Settings;
using AIStudio.Tools.AIJobs;
using AIStudio.Tools.Databases;
using AIStudio.Tools.Databases.IndexStore;
using AIStudio.Tools.Databases.ProvenanceStore;
using AIStudio.Tools.Databases.VectorStore;
using AIStudio.Tools.DataBackup;
using AIStudio.Tools.Media;
@@ -153,7 +154,7 @@ public sealed class DataBackupService(ILogger<DataBackupService> logger, Setting
}
/// <summary>
/// Copies the settings and the indexes, and writes them into the backup together with the rest of the data.
/// Copies the settings, the provenance store and the indexes, and writes them into the backup together with the rest of the data.
/// </summary>
/// <param name="archivePath">The file of the backup.</param>
/// <param name="progress">Receives how far writing got.</param>
@@ -162,6 +163,7 @@ public sealed class DataBackupService(ILogger<DataBackupService> logger, Setting
private async Task<DataBackupManifest> WriteBackupAsync(string archivePath, IProgress<DataBackupProgress> progress, CancellationToken token)
{
var indexCopiesDirectory = string.Empty;
var provenanceCopyDirectory = string.Empty;
try
{
var dataDirectory = SettingsManager.DataDirectory ?? throw new InvalidOperationException("The data directory is not known yet.");
@@ -170,9 +172,14 @@ public sealed class DataBackupService(ILogger<DataBackupService> logger, Setting
var settingsJson = JsonSerializer.Serialize(settings, SettingsManager.JSON_OPTIONS);
var organizationPluginFolders = DataBackupContents.FindOrganizationPluginsAmongUserPlugins(PluginFactory.AvailablePlugins.OfType<IAvailablePlugin>(), dataDirectory);
//
// The texts the chats received are copied once the chats themselves are written, cf.
// DataBackupSources.CopyProvenanceStoreAsync: an answer completed meanwhile then still
// finds its texts in the backup.
//
indexCopiesDirectory = await this.CopyIndexesAsync(dataDirectory, token);
var sources = new DataBackupSources(dataDirectory, configDirectory, settingsJson, indexCopiesDirectory, organizationPluginFolders);
var sources = new DataBackupSources(dataDirectory, configDirectory, settingsJson, indexCopiesDirectory, organizationPluginFolders, async copyToken => provenanceCopyDirectory = await this.CopyProvenanceStoreAsync(dataDirectory, copyToken));
var manifest = new DataBackupManifest(
DataBackupFormat.FORMAT_VERSION,
Assembly.GetExecutingAssembly().GetCustomAttribute<MetaDataAttribute>()?.Version ?? "unknown",
@@ -188,7 +195,8 @@ public sealed class DataBackupService(ILogger<DataBackupService> logger, Setting
}
finally
{
this.TryDeleteIndexCopies(indexCopiesDirectory);
this.TryDeleteCopies(indexCopiesDirectory);
this.TryDeleteCopies(provenanceCopyDirectory);
}
}
@@ -253,11 +261,56 @@ public sealed class DataBackupService(ILogger<DataBackupService> logger, Setting
}
catch
{
this.TryDeleteIndexCopies(preparation.StagingPath);
this.TryDeleteCopies(preparation.StagingPath);
throw;
}
}
/// <summary>
/// Copies the provenance store, for every backup.
/// </summary>
/// <remarks>
/// The store belongs to the chats, which every backup holds: without it, the chats of a backup
/// would refer to contents it lacks. So it is copied independently of the indexes, also when the
/// backup leaves those out.<br/><br/>
///
/// A store which AI Studio could not open is taken from the data folder as it lies there:
/// nothing writes to it then, and the backup keeps whatever there is. So is a store which could
/// not be copied, e.g., since it is damaged: that must not cost the backup of the chats.
/// </remarks>
/// <param name="dataDirectory">The data folder, which holds the directory of the copies.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The directory of the copy, in the layout of the data folder, or an empty string when the backup takes the store from the data folder.</returns>
private async Task<string> CopyProvenanceStoreAsync(string dataDirectory, CancellationToken token)
{
var provenanceStore = await databaseClientProvider.GetProvenanceStoreAsync(token);
if (provenanceStore is NoProvenanceStoreClient)
{
logger.LogWarning("The provenance store is not available, so the backup takes it from the data folder as it lies there.");
return string.Empty;
}
// In the same directory as the copies of the vector stores, which the next start clears in case AI Studio stops halfway:
var copyDirectory = Path.Join(DataRestoreFormat.DirectoryIn(dataDirectory, DataBackupFormat.STAGING_DIRECTORY), $"provenance-{Guid.NewGuid():N}");
try
{
var provenanceEntry = DataBackupContents.ENTRIES.First(entry => entry.Kind is DataBackupEntryKind.CHAT_PROVENANCE);
await provenanceStore.BackupToDirectoryAsync(Path.Join(copyDirectory, provenanceEntry.RelativePath.Replace(DataBackupContents.SEPARATOR, Path.DirectorySeparatorChar)), token);
return copyDirectory;
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
this.TryDeleteCopies(copyDirectory);
throw;
}
catch (Exception exception)
{
logger.LogWarning(exception, "Copying the provenance store failed, so the backup takes it from the data folder as it lies there.");
this.TryDeleteCopies(copyDirectory);
return string.Empty;
}
}
private DataBackupResult Refuse(string archivePath, string issue)
{
logger.LogWarning("Refused to start a backup into '{ArchivePath}': {Issue}", archivePath, issue);
@@ -291,21 +344,21 @@ public sealed class DataBackupService(ILogger<DataBackupService> logger, Setting
}
/// <summary>
/// Removes the copies of the indexes. Should that fail, the next start removes them.
/// Removes the copies a backup made of the indexes or of the provenance store. Should that fail, the next start removes them.
/// </summary>
private void TryDeleteIndexCopies(string indexCopiesDirectory)
private void TryDeleteCopies(string copiesDirectory)
{
if (string.IsNullOrWhiteSpace(indexCopiesDirectory))
if (string.IsNullOrWhiteSpace(copiesDirectory))
return;
try
{
if (Directory.Exists(indexCopiesDirectory))
Directory.Delete(indexCopiesDirectory, true);
if (Directory.Exists(copiesDirectory))
Directory.Delete(copiesDirectory, true);
}
catch (Exception exception)
{
logger.LogWarning(exception, "Could not remove the copies of the indexes at '{IndexCopiesDirectory}'.", indexCopiesDirectory);
logger.LogWarning(exception, "Could not remove the copies of the backup at '{CopiesDirectory}'.", copiesDirectory);
}
}
}
@@ -13,29 +13,42 @@ public sealed partial class DataSourceEmbeddingService
internal const int DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH = 300;
/// <summary>
/// What this build writes next to a chunk besides its text. Raise it whenever that changes.
/// The version of the chunk metadata, part of every embedding signature and frozen at 2. Raise it
/// only for a change which makes every stored vector wrong.
/// </summary>
/// <remarks>
/// A stored chunk keeps the metadata of the run which wrote it, and nothing recomputes it: the
/// fingerprint of a file says whether the file changed, not whether we got better at reading
/// it. Raising this number makes the embedding signature differ, which drops the index and
/// builds it again — the only way corrected page numbers reach a data source somebody indexed
/// earlier.
/// Raising this number makes the embedding signature of every data source differ, files and
/// mailboxes alike. That drops every index and sends every document to the embedding provider
/// once more, which is right only when no stored vector is valid anymore, e.g., because every
/// chunk would get another text. A better way of reading some files, or a new value stored next
/// to a chunk, does not justify it: a chunk which lacks the value reads as not knowing it, and
/// gains it once its document is indexed again.
///
/// Version 2: the page of a chunk is taken from the runtime metadata instead of being read back
/// out of the chunk text, which is what left Word and OpenDocument files, and passages
/// continuing across a page break, without a page.
/// continuing across a page break, without a page. It rebuilt every mailbox as well, although
/// mails have no pages, and fetched every mail from the server once more. That is why this
/// number is frozen.
///
/// Not raised for the page a chunk ends on, which chunks store since v26.10.1. It adds to what
/// a chunk says and corrects nothing: a chunk without it still names its start page, which is
/// right, and gains its end page once its file is indexed again. Rebuilding every index for it
/// would send every document to the embedding provider once more, for a page range.
/// </remarks>
private const string CHUNK_METADATA_VERSION = "2";
/// <summary>
/// What this build makes of a mail before it is cut into chunks. Raise it whenever that changes.
/// The version of what this build makes of a mail before it is cut into chunks, part of the
/// embedding signature of every mailbox and frozen at 2. Raise it only for a change which makes
/// every stored vector of a mailbox wrong.
/// </summary>
/// <remarks>
/// The counterpart of CHUNK_METADATA_VERSION for mailboxes alone. A mail on the server never
/// changes, so nothing reads it again once it is indexed, however differently MailTextBuilder
/// would write it today. Raising this number rebuilds the index of every mailbox and leaves
/// every other data source alone.
/// would write it today. Raising this number rebuilds the index of every mailbox, which fetches
/// every mail from the server and embeds it once more, and leaves every other data source alone.
/// A change which alters the text of some mails, such as those with attachments, does not
/// justify that.
///
/// Version 1: the header block and the text from the HTML part, without attachments.
///
@@ -1,9 +1,9 @@
using AIStudio.Chat;
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.Databases;
using AIStudio.Tools.Databases.IndexStore;
using AIStudio.Tools.Databases.VectorStore;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.RAG;
using AIStudio.Tools.Rust;
@@ -11,8 +11,6 @@ namespace AIStudio.Tools.Services;
public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService indexSearch, ILogger<DataSourceLocalRetrievalService> logger)
{
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DataSourceLocalRetrievalService).Namespace, nameof(DataSourceLocalRetrievalService));
private enum RetrievalChannel
{
VECTOR,
@@ -36,7 +34,7 @@ public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService inde
string FileName,
string RelativePath,
string FileType,
int? PageNumber,
ChunkLocation Location,
int ChunkIndex,
string Text,
double Score,
@@ -136,7 +134,7 @@ public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService inde
result.FileName,
result.RelativePath,
result.FileType,
result.PageNumber,
result.Location,
result.ChunkIndex,
result.Text,
result.Score,
@@ -153,7 +151,7 @@ public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService inde
result.FileName,
result.RelativePath,
result.FileType,
result.PageNumber,
result.Location,
result.ChunkIndex,
result.ChunkText,
result.Score,
@@ -164,36 +162,76 @@ public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService inde
var sourceName = FirstNonEmpty(hit.FileName, dataSource.Name);
var path = FirstNonEmpty(hit.AbsolutePath, hit.RelativePath);
var referenceLink = string.IsNullOrWhiteSpace(path) ? string.Empty : BuildReferenceLink(path, hit);
var type = GetRetrievalContentType(hit.FileType);
var located = LocatePassage(hit.Location, type, hit.ChunkIndex + 1);
return new RetrievalTextContext
{
DataSourceName = sourceName,
Category = RetrievalContentCategory.TEXT,
Type = GetRetrievalContentType(hit.FileType),
Type = type,
Path = path,
Links = [],
MatchedText = hit.Text,
SurroundingContent = [],
ReferenceTitle = BuildReferenceTitle(hit, dataSource),
ReferenceTitle = BuildReferenceTitle(sourceName, located),
ReferenceLink = referenceLink,
PageNumber = hit.PageNumber is > 0 ? hit.PageNumber : null,
PageNumber = located.PageNumber,
EndPageNumber = located.EndPageNumber,
SheetName = located.SheetName,
StartRow = located.StartRow,
EndRow = located.EndRow,
StartLine = located.StartLine,
EndLine = located.EndLine,
ChunkNumber = located.ChunkNumber,
};
}
private static string BuildReferenceTitle(LocalRetrievalHit hit, IInternalDataSource dataSource)
/// <summary>
/// Where a hit is in its document, as its passage names it.
/// </summary>
/// <remarks>
/// An end only names a range when the passage goes on beyond where it starts. A chunk indexed
/// before an end, a row or a line was kept has none, and names what it has. The lines of a
/// table, such as a CSV file, are its rows. A transcript names no lines: it is a file the user
/// never sees, so its lines would point nowhere.
/// </remarks>
/// <param name="location">Where the chunk of the hit is.</param>
/// <param name="type">What kind of content the hit is.</param>
/// <param name="chunkNumber">Which part of its document the hit is, counted from one.</param>
/// <returns>A passage holding nothing but where the hit is.</returns>
internal static SourcePassage LocatePassage(ChunkLocation location, RetrievalContentType type, int chunkNumber)
{
var sourceName = FirstNonEmpty(hit.FileName, dataSource.Name);
return BuildLocatedReferenceTitle(sourceName, hit.ChunkIndex, hit.PageNumber);
var pageNumber = location.PageNumber is > 0 ? location.PageNumber : null;
var rows = RangeOf(location.StartRow, location.EndRow);
var lines = type is RetrievalContentType.TEXT_TRANSCRIPT ? PassageRange.NONE : RangeOf(location.StartLine, location.EndLine);
if (type is RetrievalContentType.TEXT_SPREADSHEET && rows.Start is null)
{
rows = lines;
lines = PassageRange.NONE;
}
return new SourcePassage
{
Type = type,
PageNumber = pageNumber,
EndPageNumber = pageNumber is not null && location.EndPageNumber > pageNumber ? location.EndPageNumber : null,
SheetName = rows.Start is null ? null : location.SheetName,
StartRow = rows.Start,
EndRow = rows.End,
StartLine = lines.Start,
EndLine = lines.End,
ChunkNumber = chunkNumber,
};
}
private static string BuildLocatedReferenceTitle(string sourceName, int chunkIndex, int? pageNumber)
{
var location = pageNumber is > 0
? string.Format(TB("Page {0}"), pageNumber)
: string.Format(TB("Chunk {0}"), chunkIndex + 1);
private static PassageRange RangeOf(int? start, int? end) => start is > 0 ? new(start, end > start ? end : null) : PassageRange.NONE;
return $"{sourceName} ({location})";
}
//
// Worded like the location of every passage, so that the title reads the same before and after
// the passages of an answer rename it, cf. SourcePassageExtensions.NamePassagesInTitles:
//
private static string BuildReferenceTitle(string sourceName, SourcePassage located) => $"{sourceName} ({located.GetLocationLabel()})";
/// <remarks>
/// A known page is written as the fragment `#page=N`, which is what the PDF open parameters
@@ -205,8 +243,8 @@ public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService inde
{
var link = NormalizeLocalReferencePath(path);
var separator = link.Contains('#', StringComparison.Ordinal) ? "&" : "#";
return hit.PageNumber is > 0
? $"{link}{separator}page={hit.PageNumber}"
return hit.Location.PageNumber is > 0
? $"{link}{separator}page={hit.Location.PageNumber}"
: $"{link}{separator}chunk={hit.ChunkIndex}";
}
@@ -224,8 +262,12 @@ public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService inde
}
}
private static RetrievalContentType GetRetrievalContentType(string fileType)
internal static RetrievalContentType GetRetrievalContentType(string fileType)
{
// An audio or a video file is indexed by its transcript, so that is what its passages are:
if (FileTypes.IsAllowedExtension(fileType, FileTypes.AUDIO, FileTypes.VIDEO))
return RetrievalContentType.TEXT_TRANSCRIPT;
if (FileTypes.IsAllowedExtension(fileType, FileTypes.TABULAR, FileTypes.SPREADSHEET))
return RetrievalContentType.TEXT_SPREADSHEET;
@@ -266,7 +308,7 @@ public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService inde
result.Result.ParentFileId,
result.Result.FileName,
FirstNonEmpty(result.Result.AbsolutePath, result.Result.FilePath),
BuildLocatedReferenceTitle(FirstNonEmpty(result.Result.FileName, dataSource.Name), result.Result.ChunkIndex, result.Result.PageNumber));
BuildReferenceTitle(FirstNonEmpty(result.Result.FileName, dataSource.Name), LocatePassage(result.Result.Location, GetRetrievalContentType(result.Result.FileType), result.Result.ChunkIndex + 1)));
}
}
@@ -290,7 +332,7 @@ public sealed class DataSourceLocalRetrievalService(LocalIndexSearchService inde
result.Result.ParentFileId,
result.Result.FileName,
result.Result.AbsolutePath,
BuildLocatedReferenceTitle(FirstNonEmpty(result.Result.FileName, dataSource.Name), result.Result.ChunkIndex, result.Result.PageNumber));
BuildReferenceTitle(FirstNonEmpty(result.Result.FileName, dataSource.Name), LocatePassage(result.Result.Location, GetRetrievalContentType(result.Result.FileType), result.Result.ChunkIndex + 1)));
}
}
}
@@ -1,4 +1,4 @@
namespace AIStudio.Tools.Services;
public sealed record EmbeddedFileRecord(string Fingerprint, long FileSize, DateTimeOffset LastWriteUtc, DateTimeOffset EmbeddedAtUtc, int ChunkCount, string? TranscriptKey = null,
string? RelativePath = null, string? ParentFileId = null);
string? RelativePath = null, string? ParentFileId = null, string? ContentHash = null, int Generation = 0);
@@ -15,10 +15,24 @@ The parts:
- **`IIndexedSourceIndexer`** - one per kind of data source. `Supports` claims the data sources, `ProcessAsync`
finds and reads the documents of one run, and `TrackChanges` / `StopTracking` notice changes on their own.
`FileSourceIndexer` is the reference: a file system watcher per data source, fingerprints over path, size and
write time.
write time, and once a fingerprint changed, the hash of the content (`FileContentHash`), so a file saved again
without edits keeps its chunks without being read.
- **`IndexedRunContext`** - one prepared run: both stores, the embedding provider, the manifest, the
collection, and the scope for reporting prompt injections. `IndexDocumentAsync` embeds and stores one
document; the cleanup methods remove what a failed attempt left behind.
collection, and the scope for reporting prompt injections. `IndexDocumentAsync` embeds one document into a
new generation, next to the version it had stored; `DocumentRunProgress.RecordDocumentIndexedAsync` completes
it (`CompleteDocumentAsync`), and `CleanupFailedDocumentAsync` sorts out a failed attempt. A document indexed
before keeps its stored version whatever the failure, so never remove its manifest entry on a failure. Once
the embedding provider fails, the attempt asks it nothing more and reads the rest of the document for the
chunks with a stored vector; a document with nothing stored stops at the failure, since none of its chunks
could take over a vector. Either way, the attempt stays next to the stored version, so the next one reuses
its vectors; one which stopped for any other reason goes. Between attempts, a document keeps two generations at
most: the one its row names, and an attempt which stays for its vectors. A run of the file indexer removes the
others of every file which did not change since (`IndexedRunContext.RemoveLeftoverGenerationsAsync`), since no
attempt would ever come for it. A chunk whose text the document had stored before keeps its vector (`ReusableVectors`),
found by the hash of exactly the text that goes to the embedding provider (`ChunkTextHash`), so hash whatever
is sent, and send whatever is stored. The new chunks belong to the next generation, which is part of their
ids (`IndexedDocumentIds.CreateChunkIdSource`), so they never overwrite an old chunk, not even for a mail,
whose fingerprint never changes.
- **`EmbeddingDocument`** - one document: its key, its index row, its display name, and how to read its chunks.
- **`DocumentRunProgress`** - counts the documents, records indexed and failed ones in the stores, publishes the
status, and completes or pauses the run.
@@ -30,7 +44,9 @@ The parts:
which documents failed during an earlier step, so the next steps do not read them again, and which failures
were reported to the user already.
- **`TextChunker`** - cuts text into chunks the embedding provider accepts. Pick one of its strategies; do not
write a chunker of your own.
write a chunker of your own. It also says where each chunk starts and ends (`SourceLocation`: page, sheet and
row, or line), from the segments the runtime delivered; never read a location back out of the chunk text.
Only the pages go into the vector store, everything else about a location into the index store alone.
To add a kind of data source:
1. Write its indexer in `Tools/Services/Indexing/` and create it in `DataSourceEmbeddingService.CreateIndexers`,
@@ -59,8 +75,17 @@ Rules which are easy to break:
`Path` APIs: on Windows, `Path.GetFullPath` reads a key like `mail:…` as a file with an alternate data stream.
- **Ids and signature are pinned.** The formats in `IndexedDocumentIds` and the embedding signature
(`DataSourceEmbeddingService.BuildEmbeddingSignature`) are fixed by tests, because every stored chunk and every
index depends on them. When the metadata stored next to a chunk changes, raise `CHUNK_METADATA_VERSION`
deliberately: that rebuilds every index.
index depends on them. `CHUNK_METADATA_VERSION` and `MAIL_TEXT_VERSION` are frozen: raising one rebuilds every
index it is part of, so raise it only when every stored vector is wrong. A change to how documents are read, or
to what is stored next to a chunk, raises neither. A new value which old chunks may lack without being wrong,
such as the end page of a chunk, reads as unknown, and a chunk gains it once its document is indexed again.
- **A search finds the generation the row of a document names, and nothing else.** A row whose chunk count is
zero names none yet. The keyword search joins on it (`SqliteIndexStoreClientImplementation.MatchChunks`), and
the vector search leaves out `IndexStoreClient.GetHiddenChunksAsync`, which the runtime matches against the
`generation` of every point. A new way to search the index, or to list chunks for a search, follows the same
rule, or it finds failed attempts and documents which are not complete. Write a chunk into the index store
before its point into the vector store, and remove the point before the chunk: a point the index store does
not know cannot be hidden.
- **The service decides when, the indexer decides how.** Whether changes are tracked at all depends on the
automatic refresh setting and the startup hash check, and only the embedding service decides that.
- **A pause stores the source hash only when a partial index may be searched.** A stored source hash tells the
@@ -0,0 +1,8 @@
namespace AIStudio.Tools.Services.Indexing;
/// <summary>
/// Where a chunk starts and where it ends in its document, cf. TextChunker.Locate.
/// </summary>
/// <param name="Start">Where the chunk starts.</param>
/// <param name="End">Where the chunk ends.</param>
internal readonly record struct ChunkSpan(SourceLocation Start, SourceLocation End);
@@ -215,22 +215,21 @@ internal sealed class DocumentRunProgress(IndexedRunContext context, int totalDo
/// </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.
/// Until then, the document keeps the version it had, or its index row says it has no chunks
/// when it is new.<br/><br/>
/// The row names the new generation first, which every search goes by from then on, and only
/// then do the other generations go. A crash in between leaves chunks which no search finds,
/// and which the next indexing of the document removes.
/// </remarks>
/// <param name="document">The document.</param>
/// <param name="chunkCount">How many chunks were stored for it.</param>
/// <param name="indexed">What indexing the document stored.</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)
public async Task RecordDocumentIndexedAsync(EmbeddingDocument document, IndexedDocument indexed, 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, state.TranscriptKey, state.RelativePath, state.ParentFileId);
await context.ForgetPermanentFailureAsync(document.Key, token);
context.Pass.Forget(document.Key);
await this.StoreDocumentStateAsync(document.Key, document.State with { ChunkCount = indexed.ChunkCount, Generation = indexed.Generation, EmbeddedAtUtc = DateTimeOffset.UtcNow }, token);
await context.CompleteDocumentAsync(document.Key, indexed, token);
this.readDocuments++;
this.IndexedDocuments++;
if (isNew)
this.NewDocuments++;
@@ -238,10 +237,54 @@ internal sealed class DocumentRunProgress(IndexedRunContext context, int totalDo
this.ChangedDocuments++;
}
/// <summary>
/// Records a document whose fingerprint changed while its content did not, e.g., a file saved again without edits.
/// </summary>
/// <remarks>
/// Its chunks and vectors stay as they are, under the ids they have. Only the new fingerprint
/// is written, into the vector store as well as into the index store, the way a relocation
/// writes it, so the next run finds the document unchanged. Nothing is read through the
/// runtime, and nothing goes to the embedding provider. An attempt which failed in between and
/// stays for its vectors is of no use any longer, and goes.
/// </remarks>
/// <param name="document">The document, as it would be indexed now.</param>
/// <param name="storedRecord">What the manifest knows about the document from its last indexing.</param>
/// <param name="token">The cancellation token.</param>
public async Task RecordContentUnchangedAsync(EmbeddingDocument document, EmbeddedFileRecord storedRecord, CancellationToken token)
{
var state = document.State with { ChunkCount = storedRecord.ChunkCount, EmbeddedAtUtc = storedRecord.EmbeddedAtUtc, TranscriptKey = storedRecord.TranscriptKey, Generation = storedRecord.Generation };
await context.VectorStore.RelocateEmbeddingFileAsync(context.CollectionName, state.ParentFileId, state.AbsolutePath, state.RelativePath, state.Fingerprint, token);
token.ThrowIfCancellationRequested();
await this.StoreDocumentStateAsync(document.Key, state, token);
await context.KeepOnlyGenerationAsync(document.Key, storedRecord.Generation, token);
this.UnchangedDocuments++;
}
/// <summary>
/// Writes the row of a document which is in the index from now on, and what the manifest knows about it.
/// </summary>
/// <param name="documentKey">The key of the document.</param>
/// <param name="state">The row.</param>
/// <param name="token">The cancellation token.</param>
private async Task StoreDocumentStateAsync(string documentKey, EmbeddingStateFile state, CancellationToken token)
{
await context.IndexStore.UpsertFileAsync(context.DataSource.Id, state, token);
context.Manifest.Files[documentKey] = new EmbeddedFileRecord(state.Fingerprint, state.FileSize, state.LastWriteUtc, state.EmbeddedAtUtc, state.ChunkCount, state.TranscriptKey, state.RelativePath, state.ParentFileId, state.ContentHash, state.Generation);
context.Manifest.UnfinishedFiles.Remove(documentKey);
await context.ForgetPermanentFailureAsync(documentKey, token);
context.Pass.Forget(documentKey);
this.readDocuments++;
}
/// <summary>
/// Records why a document could not be indexed, and removes what the attempt left behind.
/// </summary>
/// <remarks>
/// A document indexed before keeps the version it had, whatever the failure, cf.
/// IndexedRunContext.CleanupFailedDocumentAsync. So it stays in the manifest, with the
/// fingerprint of that version, and the next run tries it again once the pass allows.<br/><br/>
/// 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>
@@ -269,7 +312,6 @@ internal sealed class DocumentRunProgress(IndexedRunContext context, int totalDo
var failure = new DataSourceEmbeddingFailure(document.Key, providerFailure.UserMessage, DateTimeOffset.UtcNow, providerFailure.FailureReason, providerFailure.StatusCode, context.EmbeddingProvider.Name, DisplayName: document.DisplayName);
this.failures.Add(failure);
context.Pass.RecordFailure(document.Key, document.State.Fingerprint, failure);
context.Manifest.Files.Remove(document.Key);
await context.ForgetPermanentFailureAsync(document.Key, token);
await context.CleanupFailedDocumentAsync(document.Key, token);
@@ -297,12 +339,12 @@ internal sealed class DocumentRunProgress(IndexedRunContext context, int totalDo
// 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.
// Meanwhile, a file indexed before can still be found by its last good version.
//
this.PermanentlySkippedDocuments++;
var occurredAtUtc = DateTimeOffset.UtcNow;
var indexingMessage = this.GetFailureMessage(extractionFailure.Code, document, extractionFailure.Message);
this.failures.Add(new DataSourceEmbeddingFailure(document.Key, indexingMessage, occurredAtUtc, ExtractionCode: extractionFailure.Code, IsPermanent: true, DisplayName: document.DisplayName));
context.Manifest.Files.Remove(document.Key);
context.Pass.Forget(document.Key);
await context.CleanupFailedDocumentAsync(document.Key, token);
@@ -346,7 +388,6 @@ internal sealed class DocumentRunProgress(IndexedRunContext context, int totalDo
var failure = new DataSourceEmbeddingFailure(document.Key, failureMessage, DateTimeOffset.UtcNow, EmbeddingProviderName: context.EmbeddingProvider.Name, ExtractionCode: extractionCode, DisplayName: document.DisplayName);
this.failures.Add(failure);
context.Pass.RecordFailure(document.Key, document.State.Fingerprint, failure);
context.Manifest.Files.Remove(document.Key);
await context.ForgetPermanentFailureAsync(document.Key, token);
await context.CleanupFailedDocumentAsync(document.Key, token);
@@ -1,14 +1,27 @@
namespace AIStudio.Tools.Services.Indexing;
/// <summary>
/// One chunk as the chunking produced it, together with the page it starts on.
/// One chunk as the chunking produced it, together with where it starts and where it ends.
/// </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.
/// The locations are carried rather than read back out of the chunk text. The runtime states them,
/// and the chunking knows which source segments a chunk covers, 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);
/// <param name="Start">Where the chunk's own content starts: its page, its sheet and row, or its line.</param>
/// <param name="End">Where the chunk ends. Equal to the start when the chunk stays within one page, row or line.</param>
/// <param name="AttachmentPosition">Which attachment of a mail the chunk is from, counted from zero as the index numbers them, or null for the text of the mail itself and for anything but a mail.</param>
internal sealed record EmbeddingChunk(string Text, SourceLocation Start, SourceLocation End, int? AttachmentPosition = null)
{
/// <summary>
/// The page the chunk's own content starts on, or null when it has none.
/// </summary>
public int? PageNumber => this.Start.PageNumber;
/// <summary>
/// The page the chunk ends on, or null when it has none.
/// </summary>
public int? EndPageNumber => this.End.PageNumber;
}
@@ -29,7 +29,7 @@ internal static class ExtractedFileText
{
var normalized = TextChunker.NormalizeSegment(segment.Content);
if (!string.IsNullOrWhiteSpace(normalized))
segments.Add(new(normalized, segment.TokenCount, segment.PageNumber));
segments.Add(new(normalized, segment.TokenCount, segment.Location));
}
return new(string.Join("\n", segments.Select(segment => segment.Text)).Trim(), segments);
@@ -129,6 +129,9 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
var removedMissingFiles = await this.RemoveMissingFileEmbeddingsAsync(context, inputFiles, token);
token.ThrowIfCancellationRequested();
// Before the data source may be skipped as unchanged, since nothing else removes what is left of an unchanged file:
await context.RemoveLeftoverGenerationsAsync(key => manifest.Files.TryGetValue(key, out var record) && metadataSnapshot.FileHashes.TryGetValue(key, out var fingerprint) && string.Equals(record.Fingerprint, fingerprint, StringComparison.Ordinal), token);
logger.LogInformation(
"Compared data source hash for '{DataSourceName}' ({DataSourceId}). StoredSourceHashPrefix={StoredSourceHashPrefix}, CurrentSourceHashPrefix={CurrentSourceHashPrefix}, StoredFileRecords={StoredFileRecords}, CurrentFiles={CurrentFiles}, RemovedMissingFiles={RemovedMissingFiles}.",
dataSource.Name,
@@ -213,6 +216,29 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
try
{
//
// The fingerprint changed, but perhaps not the content: a file saved again without
// edits, synchronized or copied keeps its chunks and vectors, and only its new
// fingerprint is stored. Reading the bytes once costs far less than extracting and
// embedding them again. The transcript of an audio or video file is filed under the
// same hash, so a recording indexed before this hash was kept is recognized as well.
//
var contentHash = await FileContentHash.ComputeAsync(file.FullName, token);
document = this.CreateFileDocument(context, dataSource, file, fingerprint, null, contentHash);
if (existingRecord is { ParentFileId: not null } && string.Equals(existingRecord.ContentHash ?? existingRecord.TranscriptKey, contentHash, StringComparison.Ordinal))
{
await progress.RecordContentUnchangedAsync(document, existingRecord, token);
logger.LogInformation(
"Kept the stored chunks of file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) without reading it, because its content did not change, although {EmbeddingReason}. Progress={CompletedFiles}/{TotalFiles}.",
file.FullName,
dataSource.Name,
dataSource.Id,
GetFileEmbeddingReason(file, fingerprint, existingRecord),
progress.DoneDocuments,
totalFiles);
continue;
}
logger.LogInformation(
"Embedding file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) because {EmbeddingReason}. CurrentMetadataHashPrefix={CurrentMetadataHashPrefix}. Progress={CompletedFiles}/{TotalFiles}.",
file.FullName,
@@ -228,15 +254,15 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
// its transcript instead, and its row knows which transcript that is:
//
if (PersistentMediaTranscript.IsMediaPath(file.FullName))
document = this.CreateFileDocument(context, dataSource, file, fingerprint, await this.ReadOrCreateTranscriptAsync(dataSource, file.FullName, token));
document = this.CreateFileDocument(context, dataSource, file, fingerprint, await this.ReadOrCreateTranscriptAsync(dataSource, file.FullName, contentHash, fingerprint, token), contentHash);
var chunkCount = await context.IndexDocumentAsync(document, reportBlockProgress, token);
var indexed = 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);
await progress.RecordDocumentIndexedAsync(document, indexed, existingRecord is null, token);
// The recording changed, so its old transcript is no longer what it says:
if (existingRecord?.TranscriptKey is { } previousTranscriptKey && !string.Equals(previousTranscriptKey, document.State.TranscriptKey, StringComparison.Ordinal))
@@ -247,7 +273,7 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
file.FullName,
dataSource.Name,
dataSource.Id,
chunkCount,
indexed.ChunkCount,
(DateTimeOffset.UtcNow - startedAtUtc).TotalMilliseconds);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
@@ -552,14 +578,16 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
/// <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.
/// send the whole folder through the slow path on every run. A file which was indexed before
/// and then failed for good keeps its stored version, so it is in both lists: it counts once,
/// as settled by its failure.
/// </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)
if (manifest.Files.Keys.Union(manifest.PermanentFailures.Keys, DataSourceEmbeddingManifest.PATH_COMPARER).Count() != indexedFiles.Count)
return false;
foreach (var file in indexedFiles)
@@ -567,18 +595,9 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
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))
var isIndexed = manifest.Files.TryGetValue(file.FullName, out var existingRecord) && string.Equals(existingRecord.Fingerprint, currentHash, StringComparison.Ordinal);
var isUnreadable = manifest.PermanentFailures.TryGetValue(file.FullName, out var permanentFailure) && string.Equals(permanentFailure.Fingerprint, currentHash, StringComparison.Ordinal);
if (!isIndexed && !isUnreadable)
return false;
}
@@ -636,12 +655,17 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
/// </remarks>
/// <param name="dataSource">The data source.</param>
/// <param name="mediaPath">The audio or video file.</param>
/// <param name="contentHash">The hash of its content, cf. FileContentHash.</param>
/// <param name="fingerprint">Its fingerprint, taken before its content was hashed.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>Where the transcript is.</returns>
private async Task<MediaTranscriptSlot> ReadOrCreateTranscriptAsync(IDataSource dataSource, string mediaPath, CancellationToken token)
private async Task<MediaTranscriptSlot> ReadOrCreateTranscriptAsync(IDataSource dataSource, string mediaPath, string contentHash, string fingerprint, CancellationToken token)
{
var provider = this.GetMediaTranscriptionAvailability(dataSource).Provider;
var transcript = await PersistentMediaTranscript.ReadOrCreateAsync(DataSourceTranscriptStore.ForDataSource(dataSource.Id), mediaPath, provider, mediaTranscriptionService.TranscribeAsync, logger, token);
// The content was hashed under the fingerprint taken before, so the store need not read the whole recording again:
var location = DataSourceTranscriptStore.ForDataSource(dataSource.Id).WithKnownContentHash(contentHash, path => string.Equals(BuildFileMetadataHash(new FileInfo(path)), fingerprint, StringComparison.Ordinal));
var transcript = await PersistentMediaTranscript.ReadOrCreateAsync(location, mediaPath, provider, mediaTranscriptionService.TranscribeAsync, logger, token);
return GetTranscriptOrThrow(transcript, token);
}
@@ -1060,10 +1084,11 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
/// <param name="file">The file.</param>
/// <param name="fingerprint">The fingerprint of the file as it is about to be read.</param>
/// <param name="transcript">The transcript an audio or video file is read through, or null for every other file.</param>
/// <param name="contentHash">The hash of the content of the file, cf. FileContentHash, or null while it is not computed yet.</param>
/// <returns>The document.</returns>
private EmbeddingDocument CreateFileDocument(IndexedRunContext context, IDataSource dataSource, FileInfo file, string fingerprint, MediaTranscriptSlot? transcript) => new(
private EmbeddingDocument CreateFileDocument(IndexedRunContext context, IDataSource dataSource, FileInfo file, string fingerprint, MediaTranscriptSlot? transcript, string? contentHash = null) => new(
file.FullName,
this.CreateEmbeddingStateFile(dataSource, file, fingerprint, 0, DateTimeOffset.UtcNow) with { ParentFileId = context.GetFileParentId(file.FullName), TranscriptKey = transcript?.Key },
this.CreateEmbeddingStateFile(dataSource, file, fingerprint, 0, DateTimeOffset.UtcNow) with { ParentFileId = context.GetFileParentId(file.FullName), TranscriptKey = transcript?.Key, ContentHash = contentHash },
file.Name,
token => this.StreamEmbeddingChunksAsync(file.FullName, transcript, dataSource, context.EmbeddingProvider, context.PromptInjectionScope, token));
@@ -0,0 +1,14 @@
namespace AIStudio.Tools.Services.Indexing;
/// <summary>
/// What indexing a document stored: a new generation of its chunks, next to whatever the document had before.
/// </summary>
/// <remarks>
/// The new generation replaces the stored one only once the kind of data source records the
/// document as indexed, cf. DocumentRunProgress.RecordDocumentIndexedAsync. Until then, and for
/// good when it fails, the document keeps the version it had.
/// </remarks>
/// <param name="ChunkCount">How many chunks were stored.</param>
/// <param name="Generation">The generation of the chunks, cf. IndexedDocumentIds.CreateChunkIdSource.</param>
/// <param name="ChunkIds">The ids of the chunks, which are also the ids of their points.</param>
internal sealed record IndexedDocument(int ChunkCount, int Generation, IReadOnlyList<string> ChunkIds);
@@ -26,25 +26,35 @@ internal static class IndexedDocumentIds
CreateStableGuid($"{dataSourceId}:chunk:{source}:{chunkIndex}");
/// <summary>
/// What the chunk ids of a document derive from: nearly always its fingerprint.
/// What the chunk ids of a document derive from: nearly always its fingerprint, and the generation of the chunks.
/// </summary>
/// <remarks>
/// A moved file keeps its parent id and its chunk ids, but its fingerprint now names the new path.
/// Another file may later appear at the old path with exactly the fingerprint those chunk ids came
/// from, and the fingerprint alone would hand it the same ids. The two documents would then take
/// over each other's chunks. Hence the parent id joins whenever a document is not stored under the
/// id its key gives, which only happens after a move. Every other document keeps the ids it has.
/// id its key gives, which only happens after a move. Every other document keeps the ids it has.<br/><br/>
/// Every time a document is indexed again, its chunks get the next generation, and with it new
/// ids, even where the fingerprint stays the same, as for a mail. So the new chunks never
/// overwrite the old ones, which keep their vectors until the new chunks are stored, cf.
/// ReusableVectors. The first indexing of a document is generation 0, which gives the ids
/// documents got before generations existed.
/// </remarks>
/// <param name="dataSourceId">The data source the document belongs to.</param>
/// <param name="documentKey">The key of the document.</param>
/// <param name="parentId">The id the document is stored under.</param>
/// <param name="fingerprint">The fingerprint of the document as it is about to be read.</param>
/// <param name="generation">The generation of the chunks about to be written.</param>
/// <returns>What to pass to CreateChunkId.</returns>
public static string CreateChunkIdSource(string dataSourceId, string documentKey, string parentId, string fingerprint) =>
string.Equals(parentId, CreateParentId(dataSourceId, documentKey), StringComparison.Ordinal)
public static string CreateChunkIdSource(string dataSourceId, string documentKey, string parentId, string fingerprint, int generation = 0)
{
var source = string.Equals(parentId, CreateParentId(dataSourceId, documentKey), StringComparison.Ordinal)
? fingerprint
: $"{parentId}:{fingerprint}";
return generation is 0 ? source : $"{source}:generation:{generation}";
}
/// <summary>
/// The id of a document's row in the index store, which its chunks point at.
/// </summary>
@@ -1,3 +1,5 @@
using System.Runtime.ExceptionServices;
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Tools.Databases.IndexStore;
@@ -44,12 +46,45 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
private Dictionary<string, string>? fileParentIds;
private HashSet<string>? reservedFileParentIds;
private readonly Dictionary<string, IndexingAttempt> attempts = new(DataSourceEmbeddingManifest.PATH_COMPARER);
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.
/// One chunk on its way to the stores, with the id it is stored under and the hash of its text.
/// </summary>
private sealed record EmbeddingChunkDraft(string ChunkId, string Text, int ChunkIndex, int? PageNumber);
private sealed record EmbeddingChunkDraft(string ChunkId, EmbeddingChunk Chunk, string TextHash, int ChunkIndex)
{
public string Text => this.Chunk.Text;
}
/// <summary>
/// One attempt to index a document, from reading what it had stored until it is completed or cleaned up.
/// </summary>
/// <param name="generation">The generation the attempt writes.</param>
/// <param name="earlierGenerations">The generations the document had stored when the attempt began.</param>
private sealed class IndexingAttempt(int generation, IReadOnlyCollection<int> earlierGenerations)
{
public int Generation => generation;
public IReadOnlyCollection<int> EarlierGenerations => earlierGenerations;
/// <summary>
/// The first failure of the embedding provider, after which the attempt asks it nothing more.
/// </summary>
public Exception? ProviderFailure { get; set; }
/// <summary>
/// Whether the attempt holds every stored vector the document still needs, so it stays when it fails, cf. CleanupFailedDocumentAsync.
/// </summary>
/// <remarks>
/// True once the attempt got through the whole document, even when the embedding provider
/// failed on the way: each chunk whose text the document had stored in any generation took
/// over that vector. True as well as soon as the provider fails for a document which had
/// nothing stored: no chunk after that could take over a vector, so the attempt stops there.
/// </remarks>
public bool HoldsEveryStoredVector { get; set; }
}
public IIndexedDataSource DataSource => dataSource;
@@ -134,45 +169,81 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
public void PublishStatus(DataSourceEmbeddingStatus status) => publishStatus(status);
/// <summary>
/// Embeds one document and stores it, in place of whatever was stored for it before.
/// Embeds one document and stores its chunks as a new generation, next to whatever it had stored 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.
/// Which chunks the document has stored is read first. Each chunk whose text the document had
/// before takes over the stored vector instead of going to the embedding provider again. The
/// new chunks belong to the next generation and have ids of their own, so none of them
/// overwrites an old one.<br/><br/>
/// A document indexed before keeps its row, its chunks and its vectors meanwhile: the new
/// generation replaces them only once whoever indexes records the document as indexed, cf.
/// CompleteDocumentAsync, since only the kind of data source knows whether the document changed
/// in the meantime. Until then, no search finds the new generation, cf.
/// IndexStoreClient.GetHiddenChunksAsync. A document indexed for the first time gets a row with
/// a chunk count of zero, so its chunks have something to point at while they arrive.<br/><br/>
/// Once the embedding provider fails, it is asked nothing more for this document, and the rest
/// of the document is read only for the chunks which can take over a stored vector. So the
/// attempt ends up with every stored vector the document still needs, and stays next to the
/// stored version when it is cleaned up, cf. CleanupFailedDocumentAsync. The failure of the
/// provider is thrown at the end, as it was. A document which had nothing stored stops right
/// at the failure instead, since none of its chunks could take over a vector. Every other
/// failure is thrown right away.
/// </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)
/// <returns>What was stored for the document.</returns>
public async Task<IndexedDocument> IndexDocumentAsync(EmbeddingDocument document, Action<int, int?> reportBlockProgress, CancellationToken token)
{
var reusableVectors = await ReusableVectors.ReadAsync(indexStore, dataSource.Id, document.Key, token);
var attempt = new IndexingAttempt(reusableVectors.NextGeneration, reusableVectors.StoredGenerations);
this.attempts[document.Key] = attempt;
logger.LogDebug(
"Resetting stored embeddings for file '{FilePath}' in collection '{CollectionName}' before re-indexing.",
"Indexing file '{FilePath}' in collection '{CollectionName}' as generation {Generation}. What it had stored stays until the new chunks are complete.",
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);
this.CollectionName,
attempt.Generation);
//
// A document which was never indexed completely gets its row now, whose chunk count of zero
// hides it from every search. An attempt it kept before stays under the same row, so this one
// reuses its vectors, and a failure leaves them in place, cf. CleanupFailedDocumentAsync:
//
if (!manifest.Files.ContainsKey(document.Key))
await indexStore.UpsertFileAsync(dataSource.Id, document.State with { Generation = attempt.Generation }, token);
var embeddingBatchSize = Math.Max(1, embeddingProvider.EffectiveEmbeddingBatchSize);
var batch = new List<EmbeddingChunkDraft>(embeddingBatchSize);
var totalChunkCount = 0;
var chunkIdSource = IndexedDocumentIds.CreateChunkIdSource(dataSource.Id, document.Key, document.State.ParentFileId, document.State.Fingerprint);
var storedChunkCount = 0;
var reusedChunkCount = 0;
var storedChunkIds = new List<string>();
var chunkIdSource = IndexedDocumentIds.CreateChunkIdSource(dataSource.Id, document.Key, document.State.ParentFileId, document.State.Fingerprint, attempt.Generation);
await foreach (var chunk in document.StreamChunks(token))
{
batch.Add(new(IndexedDocumentIds.CreateChunkId(dataSource.Id, chunkIdSource, totalChunkCount), chunk.Text, totalChunkCount, chunk.PageNumber));
var chunkId = IndexedDocumentIds.CreateChunkId(dataSource.Id, chunkIdSource, totalChunkCount);
batch.Add(new(chunkId, chunk, ChunkTextHash.Of(chunk.Text), totalChunkCount));
storedChunkIds.Add(chunkId);
totalChunkCount++;
reportBlockProgress(totalChunkCount, chunk.PageNumber);
if (batch.Count >= embeddingBatchSize)
await this.FlushBatchAsync(document, batch, token);
{
var storedBatch = await this.FlushBatchAsync(document, batch, reusableVectors, attempt, token);
storedChunkCount += storedBatch.StoredCount;
reusedChunkCount += storedBatch.ReusedCount;
}
}
if (batch.Count > 0)
await this.FlushBatchAsync(document, batch, token);
{
var storedBatch = await this.FlushBatchAsync(document, batch, reusableVectors, attempt, token);
storedChunkCount += storedBatch.StoredCount;
reusedChunkCount += storedBatch.ReusedCount;
}
//
// The extraction itself did not report a failure, but nothing usable came out of it. For
@@ -182,6 +253,20 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
if (totalChunkCount == 0)
throw new FileExtractionException(FileExtractionErrorCode.NO_CONTENT, string.Format(TB("No text could be read from '{0}'."), document.Key));
attempt.HoldsEveryStoredVector = true;
if (attempt.ProviderFailure is not null)
{
logger.LogInformation(
"Read file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}) to its end without asking the failed embedding provider again. {StoredChunkCount} of {ChunkCount} chunks were stored with a vector, and stay for the next attempt.",
document.Key,
dataSource.Name,
dataSource.Id,
storedChunkCount,
totalChunkCount);
ExceptionDispatchInfo.Throw(attempt.ProviderFailure);
}
logger.LogDebug(
"Generated {ChunkCount} chunks for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
totalChunkCount,
@@ -189,7 +274,128 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
dataSource.Name,
dataSource.Id);
return totalChunkCount;
if (reusedChunkCount > 0)
logger.LogInformation(
"Reused the stored vectors of {ReusedChunkCount} of {ChunkCount} chunks for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}). Only the other chunks went to the embedding provider.",
reusedChunkCount,
totalChunkCount,
document.Key,
dataSource.Name,
dataSource.Id);
return new(totalChunkCount, attempt.Generation, storedChunkIds);
}
/// <summary>
/// Completes a document indexed again: its new generation stays, and every other one goes, in both stores.
/// </summary>
/// <remarks>
/// Called by whoever records the document as indexed, cf. DocumentRunProgress.RecordDocumentIndexedAsync,
/// right after its row names the new generation. From then on, no search finds the other
/// generations, cf. IndexStoreClient.GetHiddenChunksAsync, so a store which refuses to remove
/// them costs space, not the document. The next indexing of the document removes them then.
/// Whatever an interrupted run left behind for the document goes as well, since it belongs to
/// no generation the document keeps.
/// </remarks>
/// <param name="documentKey">The key of the document.</param>
/// <param name="indexed">What indexing the document stored.</param>
/// <param name="token">The cancellation token.</param>
public async Task CompleteDocumentAsync(string documentKey, IndexedDocument indexed, CancellationToken token)
{
this.attempts.Remove(documentKey);
try
{
await this.KeepGenerationsAsync(documentKey, [indexed.Generation], indexed.ChunkIds, token);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
throw;
}
catch (Exception exception)
{
logger.LogWarning(
exception,
"Could not remove the earlier generations of file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}) after it was indexed. No search finds them, and the next indexing of the file removes them.",
documentKey,
dataSource.Name,
dataSource.Id);
}
}
/// <summary>
/// Removes what failed attempts and interrupted runs left behind of documents which are complete and did not change since.
/// </summary>
/// <remarks>
/// The next indexing of a document removes such generations, cf. CompleteDocumentAsync. A
/// document which did not change since is not indexed again, though, e.g., a file whose edit was
/// undone, or one whose run stopped while removing them, so nothing else would ever remove
/// them. No search finds them, so they cost space only, and a failure here costs the run nothing.
/// </remarks>
/// <param name="isUnchanged">Whether the document with the given key is the same as what its row describes.</param>
/// <param name="token">The cancellation token.</param>
public async Task RemoveLeftoverGenerationsAsync(Func<string, bool> isUnchanged, CancellationToken token)
{
try
{
// The documents with hidden chunks, which are few. A document whose first indexing is not complete has none to search by, and stays as it is:
var documents = await indexStore.GetHiddenChunksAsync(dataSource.Id, token);
if (documents.Count is 0)
return;
var keysByParentId = new Dictionary<string, string>(StringComparer.Ordinal);
foreach (var (key, record) in manifest.Files)
if (record.ParentFileId is { } parentFileId)
keysByParentId.TryAdd(parentFileId, key);
foreach (var document in documents)
{
token.ThrowIfCancellationRequested();
if (document.SearchableGeneration is { } generation && keysByParentId.TryGetValue(document.ParentFileId, out var key) && manifest.Files[key].Generation == generation && isUnchanged(key))
await this.KeepOnlyGenerationAsync(key, generation, token);
}
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
throw;
}
catch (Exception exception)
{
logger.LogWarning(exception, "Could not remove what earlier runs left behind of the documents in data source '{DataSourceName}' ({DataSourceId}). No search finds it, and the next run tries again.", dataSource.Name, dataSource.Id);
}
}
/// <summary>
/// Removes every generation of a document but the given one, when it has any other.
/// </summary>
/// <param name="documentKey">The key of the document.</param>
/// <param name="generation">The generation which stays, the one its row names.</param>
/// <param name="token">The cancellation token.</param>
public async Task KeepOnlyGenerationAsync(string documentKey, int generation, CancellationToken token)
{
var storedChunks = await indexStore.GetStoredChunksAsync(dataSource.Id, documentKey, token);
if (storedChunks.All(chunk => chunk.Generation == generation))
return;
await this.KeepGenerationsAsync(documentKey, [generation], storedChunks.Where(chunk => chunk.Generation == generation).Select(chunk => chunk.ChunkId).ToList(), token);
}
/// <summary>
/// Removes every generation of a document but the given ones, from the vector store first.
/// </summary>
/// <remarks>
/// The index store forgets a chunk only once its point is gone. As long as it knows a chunk, no
/// search finds the chunk unless it belongs to the generation the document is searched by, cf.
/// IndexStoreClient.GetHiddenChunksAsync. A point the index store no longer knew of could be
/// found by the vector search.
/// </remarks>
/// <param name="documentKey">The key of the document.</param>
/// <param name="generations">The generations which stay.</param>
/// <param name="chunkIds">The ids of their chunks, which are the points that stay.</param>
/// <param name="token">The cancellation token.</param>
private async Task KeepGenerationsAsync(string documentKey, IReadOnlyCollection<int> generations, IReadOnlyCollection<string> chunkIds, CancellationToken token)
{
await this.DeleteDocumentPointsAsync(documentKey, chunkIds, token);
await indexStore.DeleteOtherGenerationsAsync(dataSource.Id, documentKey, generations, token);
}
/// <summary>
@@ -197,19 +403,40 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
/// </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)
public Task DeleteDocumentPointsAsync(string documentKey, CancellationToken token) => this.DeleteDocumentPointsAsync(documentKey, [], token);
/// <summary>
/// Removes the vectors of one document from the collection, apart from the given ones.
/// </summary>
/// <param name="documentKey">The key of the document.</param>
/// <param name="keepPointIds">The points of the document which stay.</param>
/// <param name="token">The cancellation token.</param>
private async Task DeleteDocumentPointsAsync(string documentKey, IReadOnlyCollection<string> keepPointIds, CancellationToken token)
{
// A relocated document keeps its original parent id. Delete by that identity as well as
// by path, including when a previous relocation stopped between its two database writes.
var parentId = manifest.Files.GetValueOrDefault(documentKey)?.ParentFileId ?? string.Empty;
await vectorStore.DeleteEmbeddingByFile(this.CollectionName, documentKey, parentId, token);
await vectorStore.DeleteEmbeddingByFile(this.CollectionName, documentKey, parentId, keepPointIds, token);
this.hasPendingChanges = true;
}
/// <summary>
/// Removes whatever a failed attempt left behind of one document.
/// Sorts out what a failed attempt left behind of one document.
/// </summary>
/// <remarks>
/// A document indexed before keeps the version it had. That holds for every failure, one which
/// lies in the document included: a file broken by an edit may be fixed again, and meanwhile
/// its last good version can still be found.<br/><br/>
/// An attempt which holds every stored vector the document still needs stays as well, though
/// no search finds it, cf. IndexStoreClient.GetHiddenChunksAsync: one which got through the
/// whole document, or one of a document which had nothing stored when the embedding provider
/// failed. The vectors it got from the provider before the failure are among them, so the next
/// attempt asks the provider only for the rest. Every other generation goes, so a document
/// keeps two at most. Any other attempt goes instead, and what the document had stored before
/// it stays as it was.<br/><br/>
/// A document indexed for the first time has no version to keep. It keeps an attempt which
/// holds its vectors, under a row whose chunk count stays zero. Otherwise, it keeps what an
/// earlier attempt kept, and loses everything when there is none.<br/><br/>
/// Never throws for a store which refuses: the document already failed, and that failure is the
/// one the user has to hear about.
/// </remarks>
@@ -217,27 +444,34 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
/// <param name="token">The cancellation token.</param>
public async Task CleanupFailedDocumentAsync(string documentKey, CancellationToken token)
{
try
var storedVersion = manifest.Files.GetValueOrDefault(documentKey);
this.attempts.Remove(documentKey, out var attempt);
// The attempt failed before it wrote anything, so what the document had stays as it was:
if (attempt is null)
return;
IReadOnlyCollection<int> keptGenerations = (storedVersion, attempt) switch
{
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);
}
(null, { HoldsEveryStoredVector: true } kept) => [kept.Generation],
({ } stored, { HoldsEveryStoredVector: true } kept) => [stored.Generation, kept.Generation],
_ => attempt.EarlierGenerations,
};
try
{
await indexStore.DeleteFileAsync(dataSource.Id, documentKey, token);
IReadOnlyList<string> keptChunkIds = keptGenerations.Count is 0 ? [] : await indexStore.GetChunkIdsAsync(dataSource.Id, documentKey, keptGenerations, token);
if (storedVersion is null && keptChunkIds.Count is 0)
{
await this.DeleteDocumentPointsAsync(documentKey, token);
await indexStore.DeleteFileAsync(dataSource.Id, documentKey, token);
manifest.UnfinishedFiles.Remove(documentKey);
return;
}
await this.KeepGenerationsAsync(documentKey, keptGenerations, keptChunkIds, token);
if (storedVersion is null)
manifest.UnfinishedFiles.Add(documentKey);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
@@ -247,7 +481,7 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
{
logger.LogWarning(
exception,
"Could not remove embedding state while cleaning up failed embedding for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
"Could not remove what a failed attempt left behind of file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}). No search finds it, and the next attempt removes it.",
documentKey,
dataSource.Name,
dataSource.Id);
@@ -318,16 +552,168 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
this.hasPendingChanges = false;
}
private async Task FlushBatchAsync(EmbeddingDocument document, List<EmbeddingChunkDraft> batch, CancellationToken token)
/// <summary>
/// Stores one batch of chunks, each with its vector, and leaves out the chunks without one.
/// </summary>
/// <remarks>
/// A chunk goes without a vector only once the embedding provider failed for the document, cf.
/// IndexDocumentAsync.
/// </remarks>
/// <param name="document">The document the chunks belong to.</param>
/// <param name="batch">The chunks, emptied once they are stored.</param>
/// <param name="reusableVectors">The chunks the document had stored before.</param>
/// <param name="attempt">The attempt the chunks belong to, which remembers whether the embedding provider failed.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>How many of the chunks were stored, and how many of those kept a stored vector instead of being embedded.</returns>
private async Task<StoredBatch> FlushBatchAsync(EmbeddingDocument document, List<EmbeddingChunkDraft> batch, ReusableVectors reusableVectors, IndexingAttempt attempt, CancellationToken token)
{
//
// A chunk whose text the document had stored before takes that vector over. Only the other
// chunks go to the embedding provider:
//
var storedVectors = await reusableVectors.TakeAsync(vectorStore, this.CollectionName, manifest.VectorSize, batch.Select(chunk => chunk.TextHash).ToList(), token);
var vectors = new IReadOnlyList<float>?[batch.Count];
var chunksToEmbed = new List<int>(batch.Count);
for (var index = 0; index < batch.Count; index++)
{
if (storedVectors[index] is { } storedVector)
vectors[index] = storedVector;
else
chunksToEmbed.Add(index);
}
var reusedChunkCount = batch.Count - chunksToEmbed.Count;
if (chunksToEmbed.Count > 0 && attempt.ProviderFailure is null)
{
IReadOnlyList<IReadOnlyList<float>>? embeddedVectors = null;
try
{
embeddedVectors = await this.RequestVectorsAsync(document, chunksToEmbed.Select(index => batch[index].Text).ToList(), token);
}
catch (Exception exception) when (exception is not OperationCanceledException || !token.IsCancellationRequested)
{
//
// A document which had nothing stored has no chunk which could take over a vector,
// so reading on would only cost time. What the attempt embedded so far stays:
//
if (attempt.EarlierGenerations.Count is 0)
{
attempt.HoldsEveryStoredVector = true;
throw;
}
//
// The provider is asked nothing more for this document. It may be down, and once
// the cause is fixed, e.g., a chunk too long for it, the chunks after that one may
// be cut differently, so whatever it embedded meanwhile would be paid for nothing.
// The rest of the document is still read, for the chunks with a stored vector:
//
attempt.ProviderFailure = exception;
logger.LogDebug(
"The embedding provider failed for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}). The rest of the file is read for the chunks with a stored vector only.",
document.Key,
dataSource.Name,
dataSource.Id);
}
if (embeddedVectors is not null)
{
await this.EnsureCollectionAsync(embeddedVectors[0].Count, token);
for (var position = 0; position < chunksToEmbed.Count; position++)
vectors[chunksToEmbed[position]] = embeddedVectors[position];
}
}
var chunksToStore = Enumerable.Range(0, batch.Count).Where(index => vectors[index] is not null).ToList();
if (chunksToStore.Count > 0)
{
token.ThrowIfCancellationRequested();
var embeddedAtUtc = DateTimeOffset.UtcNow;
var state = document.State;
var points = chunksToStore.Select(index => new VectorStoragePoint(
batch[index].ChunkId,
vectors[index]!,
dataSource.Id,
dataSource.Type.ToString(),
batch[index].ChunkId,
state.ParentFileId,
document.Key,
state.AbsolutePath,
state.FileName,
state.RelativePath,
state.FileType,
batch[index].Chunk.PageNumber,
batch[index].Chunk.EndPageNumber,
batch[index].ChunkIndex,
batch[index].Text,
state.Fingerprint,
state.CreationUtc,
state.LastWriteUtc,
embeddedAtUtc,
attempt.Generation)).ToList();
//
// The vector store keeps the pages only, cf. VectorStoragePoint. Everything else about
// where a chunk is goes into the index store alone:
//
var chunks = chunksToStore
.Select(index => batch[index])
.Select(draft => new EmbeddingStateChunk(
draft.ChunkId,
state.ParentFileId,
draft.Chunk.PageNumber,
draft.Chunk.EndPageNumber,
draft.ChunkIndex,
draft.Text,
embeddedAtUtc,
draft.TextHash,
attempt.Generation,
draft.Chunk.Start.SheetName,
draft.Chunk.Start.RowNumber,
draft.Chunk.End.RowNumber,
draft.Chunk.Start.LineNumber,
draft.Chunk.End.LineNumber,
draft.Chunk.AttachmentPosition))
.ToList();
//
// The index store first: as long as it knows the chunks of the new generation, no search
// finds them, cf. IndexStoreClient.GetHiddenChunksAsync. A point it did not know yet, e.g.,
// after the app stopped right here, would be found next to the stored version:
//
await indexStore.UpsertChunksAsync(dataSource.Id, chunks, token);
token.ThrowIfCancellationRequested();
await vectorStore.InsertEmbedding(this.CollectionName, points, token);
await this.RecordStoredChunksAsync(chunksToStore.Count, token);
logger.LogDebug(
"Stored {ChunkCount} embedded chunks for file '{FilePath}' in collection '{CollectionName}'.",
chunksToStore.Count,
document.Key,
this.CollectionName);
}
batch.Clear();
return new(chunksToStore.Count, reusedChunkCount);
}
/// <summary>
/// Asks the embedding provider for the vectors of some texts, and checks them before anything is stored.
/// </summary>
/// <param name="document">The document the texts belong to.</param>
/// <param name="texts">The texts.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>One vector per text, in the same order.</returns>
private async Task<IReadOnlyList<IReadOnlyList<float>>> RequestVectorsAsync(EmbeddingDocument document, List<string> texts, CancellationToken token)
{
logger.LogDebug(
"Requesting embeddings for batch of {ChunkCount} chunks from file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).",
batch.Count,
texts.Count,
document.Key,
dataSource.Name,
dataSource.Id);
var texts = batch.Select(item => item.Text).ToList();
IReadOnlyList<IReadOnlyList<float>> vectors;
try
{
@@ -353,11 +739,11 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
// 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);
throw new InvalidOperationException(string.Format(TB("The embedding provider was not able to embed {0} part(s) of '{1}'. The provider reported: {2}"), texts.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));
if (vectors.Count != texts.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, texts.Count, document.Key));
var vectorSize = vectors.FirstOrDefault()?.Count ?? 0;
if (vectorSize <= 0)
@@ -372,6 +758,16 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
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));
return vectors;
}
/// <summary>
/// Creates the store of the data source with the first vectors it gets, which decide its size.
/// </summary>
/// <param name="vectorSize">The size of the vectors the embedding provider answered with.</param>
/// <param name="token">The cancellation token.</param>
private async Task EnsureCollectionAsync(int vectorSize, CancellationToken token)
{
if (manifest.VectorSize == 0)
{
token.ThrowIfCancellationRequested();
@@ -398,46 +794,5 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
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();
}
}
@@ -73,29 +73,36 @@ internal sealed partial class MailboxIndexer
/// </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.
/// names it. A chunk of an attachment knows which one it is from, by the position of that
/// attachment among all attachments of the mail, the way the index numbers them in mail_parts,
/// so its pages, rows or lines say where in that attachment it is.
/// </remarks>
/// <param name="textChunker">Cuts the text of the mail and of its attachments into chunks.</param>
/// <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)
internal static async IAsyncEnumerable<EmbeddingChunk> StreamMailChunksAsync(TextChunker textChunker, 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)
//
// Counted over every attachment, those without text included, since the index numbers
// them that way, cf. CreateMailRecord:
//
for (var position = 0; position < attachments.Count; position++)
{
var attachment = attachments[position];
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 };
yield return chunk with { AttachmentPosition = position };
}
}
}
@@ -98,10 +98,14 @@ internal sealed partial class MailboxIndexer
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);
var indexed = await context.IndexDocumentAsync(document, reportBlockProgress, token);
await progress.RecordDocumentIndexedAsync(document, indexed, isNew, token);
// Only after the last chunk: indexing the document deleted whatever the index kept of the mail.
//
// Only after the last chunk: a mail indexed for the first time gets its row with its
// first chunk, and the rows of the mail point at it. One indexed before keeps its rows,
// which are written anew here.
//
await context.IndexStore.UpsertMailAsync(mailbox.Id, CreateMailRecord(mailId, summary, text, attachments, mailHash, location, foundAtUtc), token);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
@@ -196,10 +200,10 @@ internal sealed partial class MailboxIndexer
// 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 content = new SegmentedText(fullText, fullText.Length is 0 ? [] : [new TextSegment(fullText, null, default)]);
var chunkingOptions = DataSourceEmbeddingService.GetChunkingOptions(mailbox, context.EmbeddingProvider);
return new(key, state, displayName, chunkToken => this.StreamMailChunksAsync(content, attachments, chunkingOptions, context.EmbeddingProvider, chunkToken));
return new(key, state, displayName, chunkToken => StreamMailChunksAsync(textChunker, content, attachments, chunkingOptions, context.EmbeddingProvider, chunkToken));
}
/// <summary>
@@ -275,6 +275,7 @@ internal sealed partial class MailboxIndexer(SettingsManager settingsManager, Ru
var deletedMails = await this.ApplyRemovalsAsync(context, mailbox, plannedFolders, removals, runStartedUtc, token);
await ForgetVanishedFailuresAsync(context, encounteredKeys, token);
await ForgetVanishedAttemptsAsync(context, encounteredKeys, token);
var completedUtc = DateTimeOffset.UtcNow;
await context.IndexStore.CompleteMailboxSyncAsync(mailbox.Id, completedUtc, token);
@@ -529,6 +530,27 @@ internal sealed partial class MailboxIndexer(SettingsManager settingsManager, Ru
await context.ForgetPermanentFailureAsync(key, token);
}
/// <summary>
/// Removes what failed attempts kept of mails which did not turn up during this run.
/// </summary>
/// <remarks>
/// An attempt which failed at the embedding provider keeps its vectors for the next one, under a
/// row without chunks, cf. IndexedRunContext.CleanupFailedDocumentAsync. Such a mail has no
/// location yet, 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 what its attempt kept would stay
/// forever, hidden from every search, but asked about before each one.
/// </remarks>
private static async Task ForgetVanishedAttemptsAsync(IndexedRunContext context, IReadOnlySet<string> encounteredKeys, CancellationToken token)
{
foreach (var key in context.Manifest.UnfinishedFiles.Where(key => !encounteredKeys.Contains(key)).ToList())
{
token.ThrowIfCancellationRequested();
await context.DeleteDocumentPointsAsync(key, token);
await context.IndexStore.DeleteFileAsync(context.DataSource.Id, key, token);
context.Manifest.UnfinishedFiles.Remove(key);
}
}
/// <summary>
/// Tells the user interface how the mailbox stands, from the index alone.
/// </summary>
@@ -0,0 +1,94 @@
using AIStudio.Tools.Databases.IndexStore;
using AIStudio.Tools.Databases.VectorStore;
namespace AIStudio.Tools.Services.Indexing;
/// <summary>
/// The chunks a document had stored before it is read again, so that a chunk whose text did not
/// change keeps its vector.
/// </summary>
/// <remarks>
/// A vector depends on nothing but its text and the embedding signature, and all vectors of one
/// store share that signature: another one drops the store. So a chunk whose text is byte for byte
/// the text of a chunk the document had stored may take over that vector, instead of going to the
/// embedding provider again. That spares a file which was only saved again, or changed near its
/// end, almost every request.<br/><br/>
/// Only the id and the hash of each stored chunk are held in memory, about a hundred bytes per
/// chunk. The vectors stay in the vector store until the run is done with the document, and each
/// batch reads the few it can use, cf. TakeAsync. The new chunks belong to the next generation and
/// so have ids of their own, cf. NextGeneration: writing them never overwrites a vector a later
/// chunk may still need.
/// </remarks>
internal sealed class ReusableVectors
{
private readonly Dictionary<string, string> chunkIdsByTextHash = new(StringComparer.Ordinal);
private ReusableVectors(IReadOnlyList<StoredChunk> storedChunks)
{
foreach (var storedChunk in storedChunks)
this.chunkIdsByTextHash.TryAdd(storedChunk.TextHash, storedChunk.ChunkId);
//
// A document stored for the first time, or again after a failure removed all of it,
// starts at generation 0, which gives the ids documents got before generations existed:
//
this.NextGeneration = storedChunks.Count is 0 ? 0 : storedChunks.Max(storedChunk => storedChunk.Generation) + 1;
this.StoredGenerations = storedChunks.Select(storedChunk => storedChunk.Generation).ToHashSet();
}
/// <summary>
/// The generation of the chunks about to be written, cf. IndexedDocumentIds.CreateChunkIdSource.
/// </summary>
public int NextGeneration { get; }
/// <summary>
/// The generations the document has stored chunks of.
/// </summary>
public IReadOnlyCollection<int> StoredGenerations { get; }
/// <summary>
/// Reads which chunks a document has stored now, before anything of it is deleted.
/// </summary>
/// <param name="indexStore">The index store, which knows the chunks of the document and their texts.</param>
/// <param name="dataSourceId">The data source the document belongs to.</param>
/// <param name="documentKey">The key of the document.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The chunks of the document, which may be none.</returns>
public static async Task<ReusableVectors> ReadAsync(IndexStoreClient indexStore, string dataSourceId, string documentKey, CancellationToken token) =>
new(await indexStore.GetStoredChunksAsync(dataSourceId, documentKey, token));
/// <summary>
/// Finds the stored vector of every chunk of a batch whose text the document had before.
/// </summary>
/// <param name="vectorStore">The vector store, which still holds the vectors of the stored chunks.</param>
/// <param name="collectionName">The vector store of the data source.</param>
/// <param name="vectorSize">The size of the vectors in the store.</param>
/// <param name="textHashes">The text hash of each chunk of the batch.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>For each chunk, in the same order, its stored vector, or null when it has to be embedded.</returns>
public async Task<IReadOnlyList<float>?[]> TakeAsync(VectorStoreClient vectorStore, string collectionName, int vectorSize, IReadOnlyList<string> textHashes, CancellationToken token)
{
var vectors = new IReadOnlyList<float>?[textHashes.Count];
if (this.chunkIdsByTextHash.Count == 0 || vectorSize <= 0)
return vectors;
var sourcePointIds = textHashes.Select(textHash => this.chunkIdsByTextHash.GetValueOrDefault(textHash)).ToArray();
var pointIdsToRead = sourcePointIds.OfType<string>().Distinct(StringComparer.OrdinalIgnoreCase).ToList();
if (pointIdsToRead.Count == 0)
return vectors;
//
// A vector of another size cannot come from this store as it is now, and the store would
// refuse it next to the others:
//
var storedVectors = (await vectorStore.RetrieveVectorsAsync(collectionName, pointIdsToRead, token))
.Where(storedVector => storedVector.Vector.Count == vectorSize)
.ToDictionary(storedVector => storedVector.PointId, storedVector => storedVector.Vector, StringComparer.OrdinalIgnoreCase);
for (var index = 0; index < textHashes.Count; index++)
if (sourcePointIds[index] is { } pointId && storedVectors.TryGetValue(pointId, out var storedVector))
vectors[index] = storedVector;
return vectors;
}
}
@@ -0,0 +1,8 @@
namespace AIStudio.Tools.Services.Indexing;
/// <summary>
/// How many chunks of a batch were stored, cf. IndexedRunContext.FlushBatchAsync.
/// </summary>
/// <param name="StoredCount">How many of the chunks were stored.</param>
/// <param name="ReusedCount">How many of those kept a stored vector instead of being embedded.</param>
internal readonly record struct StoredBatch(int StoredCount, int ReusedCount);
@@ -114,15 +114,68 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
/// <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);
var requiredPrefix = heading;
foreach (var part in SplitBySheets(content))
{
var estimatedTokenCount = SumTokenCounts(part.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;
// The part starts where its first segment starts, so that is where it is until the
// splitting reaches a segment boundary:
var firstLocation = part.SourceSegments.Count > 0 ? part.SourceSegments[0].Location : default;
await foreach (var chunk in this.SplitTextByRulesAsync(part.Text, part.SourceSegments, strategy, 0, options, embeddingProvider, firstLocation, token, requiredPrefix, estimatedTokenCount))
yield return chunk;
requiredPrefix = string.Empty;
}
}
/// <summary>
/// The parts of a text which no chunk may run across: each sheet of a spreadsheet on its own,
/// and any other text as a whole.
/// </summary>
/// <remarks>
/// Every sheet is a table of its own, with columns of its own, so a chunk which ran from one
/// sheet into the next would mix what two tables mean. Neither does a chunk start with the
/// overlap of another sheet. A sheet without a single row holds nothing but its heading and its
/// fences, and is left out.
/// </remarks>
/// <param name="content">The text, in the pieces its source delivered it in.</param>
/// <returns>The parts, in the order of the text.</returns>
private static IEnumerable<SegmentedText> SplitBySheets(SegmentedText content)
{
if (content.SourceSegments.All(segment => segment.Location.SheetName is null))
{
yield return content;
yield break;
}
var sheet = new List<TextSegment>();
string? sheetName = null;
foreach (var segment in content.SourceSegments)
{
if (segment.Location.SheetName is { } segmentSheetName && sheetName is not null && !string.Equals(segmentSheetName, sheetName, StringComparison.Ordinal))
{
if (TryCreateSheet(sheet) is { } completeSheet)
yield return completeSheet;
sheet = [];
}
sheetName = segment.Location.SheetName ?? sheetName;
sheet.Add(segment);
}
if (TryCreateSheet(sheet) is { } lastSheet)
yield return lastSheet;
}
/// <summary>
/// The text of one sheet, joined the way the whole text was, or null when the sheet has no row.
/// </summary>
private static SegmentedText? TryCreateSheet(List<TextSegment> segments) => segments.Any(segment => segment.Location.RowNumber is not null)
? new SegmentedText(string.Join("\n", segments.Select(segment => segment.Text)).Trim(), segments)
: null;
private async IAsyncEnumerable<EmbeddingChunk> SplitTextByRulesAsync(
string text,
IReadOnlyList<TextSegment> sourceSegments,
@@ -130,7 +183,7 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
int ruleIndex,
ChunkingOptions options,
EmbeddingProvider embeddingProvider,
int? currentPageNumber,
SourceLocation currentLocation,
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
string requiredOverlapPrefix = "",
int? estimatedTokenCount = null)
@@ -147,14 +200,16 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
tokenCount = await this.GetEmbeddingTokenCountAsync(embeddingProvider, textWithOverlap, token);
if (tokenCount <= options.MaxChunkTokenLength)
{
yield return new(textWithOverlap, currentPageNumber);
// The whole text is one chunk, so it spans all of its segments:
var span = Locate(sourceSegments.Select(segment => segment.Location).ToList(), currentLocation);
yield return new(textWithOverlap, span.Start, span.End);
yield break;
}
}
if (ruleIndex >= strategy.Rules.Count)
{
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, currentPageNumber, token, requiredOverlapPrefix, estimatedTokenCount))
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, currentLocation, token, requiredOverlapPrefix, estimatedTokenCount))
yield return hardChunk;
yield break;
@@ -163,7 +218,7 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
var rule = strategy.Rules[ruleIndex];
if (rule.Split is null)
{
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, currentPageNumber, token, requiredOverlapPrefix, estimatedTokenCount))
await foreach (var hardChunk in this.SplitTextByHardCutAsync(text, options, embeddingProvider, currentLocation, token, requiredOverlapPrefix, estimatedTokenCount))
yield return hardChunk;
yield break;
@@ -172,7 +227,7 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
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))
await foreach (var chunk in this.SplitTextByRulesAsync(text, sourceSegments, strategy, ruleIndex + 1, options, embeddingProvider, currentLocation, token, requiredOverlapPrefix, estimatedTokenCount))
yield return chunk;
yield break;
@@ -193,12 +248,15 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
//
// 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.
// a unit is a segment and carries that segment's location. Every later rule cuts inside a
// single segment, where all units share the location they were handed. This is what ties a
// chunk to its pages, rows or lines without anybody reading the text. A segment without a
// page is on the page of the text it is part of:
//
var unitsAreSourceSegments = rule.UsesSourceSegmentCounts && sourceSegments.Count == units.Count;
int? PageOfUnit(int unitIndex) => unitsAreSourceSegments ? sourceSegments[unitIndex].PageNumber ?? currentPageNumber : currentPageNumber;
SourceLocation LocationOfUnit(int unitIndex) => unitsAreSourceSegments
? sourceSegments[unitIndex].Location with { PageNumber = sourceSegments[unitIndex].Location.PageNumber ?? currentLocation.PageNumber }
: currentLocation;
while (index < units.Count)
{
@@ -212,12 +270,15 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
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.
// From the first unit this chunk covers, not from the overlap prefix in front of
// it: the prefix repeats what the chunk before already said, while the location has
// to name where this chunk's own content begins. To its last unit.
//
if (!string.IsNullOrWhiteSpace(chunk))
yield return new(chunk, PageOfUnit(index));
{
var span = Locate(Enumerable.Range(index, unitCount).Select(LocationOfUnit).ToList(), currentLocation);
yield return new(chunk, span.Start, span.End);
}
var nextIndex = index + unitCount;
if (nextIndex >= units.Count)
@@ -249,8 +310,8 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
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))
var unitLocation = LocationOfUnit(index);
await foreach (var splitUnit in this.SplitTextByRulesAsync(units[index], [new(units[index], unitTokenCount, unitLocation)], strategy, ruleIndex + 1, options, embeddingProvider, unitLocation, token, overlapPrefix, unitTokenCount))
{
lastSplitUnit = splitUnit.Text;
yield return splitUnit;
@@ -444,15 +505,48 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
return bestStartIndex <= chunkStartIndex ? chunkEndIndex : bestStartIndex;
}
/// <summary>
/// Where a chunk starts and ends, from the locations of the units it covers, in their order.
/// </summary>
/// <remarks>
/// The chunk starts on the page of its first unit and ends on the last page among its units.
/// Its rows and lines run from the first unit which has one to the last which has one: the
/// heading of a sheet and the fences around a table are no row or line, so a chunk starting
/// with them starts at the row after them. A chunk lies within one sheet, cf. SplitBySheets.
/// </remarks>
/// <param name="locations">The locations of the units the chunk covers.</param>
/// <param name="current">The location of the text being cut, for what none of the units states.</param>
/// <returns>Where the chunk starts and where it ends.</returns>
private static ChunkSpan Locate(IReadOnlyList<SourceLocation> locations, SourceLocation current)
{
if (locations.Count is 0)
return new(current, current);
var sheetName = locations.FirstOrDefault(location => location.SheetName is not null).SheetName ?? current.SheetName;
var start = new SourceLocation(
locations[0].PageNumber ?? current.PageNumber,
sheetName,
locations.FirstOrDefault(location => location.RowNumber is not null).RowNumber ?? current.RowNumber,
locations.FirstOrDefault(location => location.LineNumber is not null).LineNumber ?? current.LineNumber);
var end = new SourceLocation(
locations.LastOrDefault(location => location.PageNumber is not null).PageNumber ?? current.PageNumber,
sheetName,
locations.LastOrDefault(location => location.RowNumber is not null).RowNumber ?? current.RowNumber,
locations.LastOrDefault(location => location.LineNumber is not null).LineNumber ?? current.LineNumber);
return new(start, end);
}
/// <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.
/// any further, so every chunk it produces sits where that piece was handed.
/// </remarks>
private async IAsyncEnumerable<EmbeddingChunk> SplitTextByHardCutAsync(
string text,
ChunkingOptions options,
EmbeddingProvider embeddingProvider,
int? currentPageNumber,
SourceLocation currentLocation,
[System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token,
string requiredOverlapPrefix = "",
int? estimatedTokenCount = null)
@@ -530,9 +624,10 @@ internal sealed partial class TextChunker(RustService rustService, ILogger logge
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));
}
// A hard cut works inside a single piece of text, so the chunk stays where it is:
var chunk = AddOverlapPrefix(text[startIndex..bestEndIndex].Trim(), overlapPrefix);
if (!string.IsNullOrWhiteSpace(chunk))
yield return new(chunk, currentPageNumber);
yield return new(chunk, currentLocation, currentLocation);
if (bestEndIndex >= text.Length)
yield break;
@@ -5,5 +5,5 @@ namespace AIStudio.Tools.Services.Indexing;
/// </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);
/// <param name="Location">Where it is: its page, its sheet and row, or its line; nothing when the source does not say.</param>
internal sealed record TextSegment(string Text, int? TokenCount, SourceLocation Location);
@@ -103,11 +103,13 @@ public sealed class LocalIndexSearchService(SettingsManager settingsManager, Rus
}
var collectionName = DataSourceEmbeddingNames.GetCollectionName(dataSource.Id);
var indexStore = await databaseClientProvider.GetIndexStoreAsync(token);
var hiddenChunks = await this.GetHiddenChunksAsync(dataSource, indexStore, token);
var results = filter is null
? await vectorStore.SearchEmbeddingAsync(collectionName, vector, maxMatches, token)
: await vectorStore.SearchEmbeddingAsync(collectionName, vector, maxMatches, filter, token);
? await vectorStore.SearchEmbeddingAsync(collectionName, vector, maxMatches, hiddenChunks, token)
: await vectorStore.SearchEmbeddingAsync(collectionName, vector, maxMatches, filter, hiddenChunks, token);
return this.LimitSearchResults(dataSource, "vector", results, maxMatches);
return await this.CompleteLocationsAsync(dataSource, indexStore, this.LimitSearchResults(dataSource, "vector", results, maxMatches), token);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
@@ -281,6 +283,67 @@ public sealed class LocalIndexSearchService(SettingsManager settingsManager, Rus
return true;
}
/// <summary>
/// The chunks of a data source which its vector search must not find, cf. IndexStoreClient.GetHiddenChunksAsync.
/// </summary>
/// <remarks>
/// Without the index store, nothing is hidden. Its keyword search fails then as well, which the
/// retrieval reports, and a hidden chunk is still text of its document, only perhaps not the
/// version the document is searched by.
/// </remarks>
/// <param name="dataSource">The data source to search.</param>
/// <param name="indexStore">The index store.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The hidden chunks, document by document.</returns>
private async Task<IReadOnlyList<HiddenChunks>> GetHiddenChunksAsync(IIndexedDataSource dataSource, IndexStoreClient indexStore, CancellationToken token)
{
if (indexStore.IsAvailable)
return await indexStore.GetHiddenChunksAsync(dataSource.Id, token);
logger.LogWarning(
"Searching the vectors of data source '{DataSourceName}' ({DataSourceId}) without hiding incomplete generations, because local RAG index '{DatabaseName}' is unavailable.",
dataSource.Name,
dataSource.Id,
indexStore.Name);
return [];
}
/// <summary>
/// Completes the hits of the vector search with where their chunks are, as the index store keeps it.
/// </summary>
/// <remarks>
/// The vector store keeps only the pages of a chunk; its sheet and rows, its lines and the
/// attachment of a mail it is from are kept in the index store alone. One query for all hits.
/// Without the index store, the hits keep their pages, which is what they said before: a hit
/// which cannot say where it is still belongs into the answer.
/// </remarks>
/// <param name="dataSource">The data source searched.</param>
/// <param name="indexStore">The index store.</param>
/// <param name="results">The hits.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The hits, in their order.</returns>
private async Task<IReadOnlyList<VectorSearchResult>> CompleteLocationsAsync(IIndexedDataSource dataSource, IndexStoreClient indexStore, IReadOnlyList<VectorSearchResult> results, CancellationToken token)
{
if (results.Count is 0 || !indexStore.IsAvailable)
return results;
try
{
var locations = await indexStore.GetChunkLocationsAsync(dataSource.Id, results.Select(result => result.ChunkId).ToList(), token);
return results.Select(result => locations.TryGetValue(result.ChunkId, out var location) ? result with { StoredLocation = location } : result).ToList();
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
throw;
}
catch (Exception exception)
{
logger.LogWarning(exception, "Could not read where the chunks found by the vector search of data source '{DataSourceName}' ({DataSourceId}) are. They name their pages only.", dataSource.Name, dataSource.Id);
return results;
}
}
private IReadOnlyList<T> LimitSearchResults<T>(IIndexedDataSource dataSource, string searchName, IReadOnlyList<T> results, int maxMatches)
{
if (results.Count <= maxMatches)
@@ -1,3 +1,5 @@
using AIStudio.Tools.Databases;
namespace AIStudio.Tools.Services;
/// <summary>
@@ -5,4 +7,5 @@ namespace AIStudio.Tools.Services;
/// </summary>
/// <param name="MailId">The id of the mail.</param>
/// <param name="Text">The text of the chunk.</param>
internal sealed record MailPassage(string MailId, string Text);
/// <param name="Location">Where the chunk is in its mail, e.g., on which page of which attachment.</param>
internal sealed record MailPassage(string MailId, string Text, ChunkLocation Location);
@@ -1,3 +1,4 @@
using AIStudio.Tools.Databases;
using AIStudio.Tools.Databases.IndexStore;
namespace AIStudio.Tools.Services;
@@ -7,4 +8,5 @@ namespace AIStudio.Tools.Services;
/// </summary>
/// <param name="Summary">What a list of mails shows about the mail.</param>
/// <param name="Passage">The passage of the mail which matched the query best, or null when the mails were listed without a query.</param>
public sealed record MailSearchHit(MailSummary Summary, string? Passage);
/// <param name="PassageLocation">Where that passage is in the mail, e.g., on which page of which attachment, or null without a passage.</param>
public sealed record MailSearchHit(MailSummary Summary, string? Passage, ChunkLocation? PassageLocation = null);
@@ -248,20 +248,20 @@ public sealed class MailboxRetrievalService(SettingsManager settingsManager, Dat
var vectorChunks = vectorTask.Result;
var keywordChunks = keywordTask.Result;
var (passages, hasMore) = RetrievalPaging.Merge(
BestPassagePerMail(vectorChunks.Select(chunk => new MailPassage(chunk.ParentFileId, chunk.Text))),
BestPassagePerMail(keywordChunks.Select(chunk => new MailPassage(chunk.ParentFileId, chunk.ChunkText))),
BestPassagePerMail(vectorChunks.Select(chunk => new MailPassage(chunk.ParentFileId, chunk.Text, chunk.Location))),
BestPassagePerMail(keywordChunks.Select(chunk => new MailPassage(chunk.ParentFileId, chunk.ChunkText, chunk.Location))),
passage => passage.MailId,
page,
pageSize);
hasMore = hasMore || MayHoldMoreMails(page, pageSize, chunkWindow, vectorChunks.Count, keywordChunks.Count);
var passageTexts = passages
var bestPassages = passages
.DistinctBy(passage => passage.MailId, StringComparer.Ordinal)
.ToDictionary(passage => passage.MailId, passage => passage.Text, StringComparer.Ordinal);
.ToDictionary(passage => passage.MailId, StringComparer.Ordinal);
var summaries = await indexStore.GetMailSummariesAsync(mailbox.Id, passages.Select(passage => passage.MailId).ToList(), token);
return new(summaries.Select(summary => new MailSearchHit(summary, passageTexts[summary.MailId])).ToList(), hasMore, []);
return new(summaries.Select(summary => new MailSearchHit(summary, bestPassages[summary.MailId].Text, bestPassages[summary.MailId].Location)).ToList(), hasMore, []);
}
/// <summary>
@@ -0,0 +1,14 @@
namespace AIStudio.Tools.Services;
/// <summary>
/// The rows or lines a passage spans, as far as it names them, cf. DataSourceLocalRetrievalService.LocatePassage.
/// </summary>
/// <param name="Start">The first row or line, or null when the passage names none.</param>
/// <param name="End">The last row or line, or null when the passage does not go on beyond its first one.</param>
internal readonly record struct PassageRange(int? Start, int? End)
{
/// <summary>
/// The passage names no rows or lines.
/// </summary>
public static readonly PassageRange NONE = new(null, null);
}
@@ -0,0 +1,87 @@
using AIStudio.Chat;
using AIStudio.Settings;
using AIStudio.Tools.DataBackup;
using AIStudio.Tools.Databases;
using AIStudio.Tools.Databases.ProvenanceStore;
namespace AIStudio.Tools.Services;
/// <summary>
/// One-shot startup service that removes from the provenance store what no stored chat refers to, see ProvenanceCleanup.
/// </summary>
public sealed class ProvenanceCleanupService(DatabaseClientProvider databaseClientProvider, DataRestoreService dataRestoreService, ILogger<ProvenanceCleanupService> logger) : BackgroundService
{
/// <summary>
/// When this run of the app started. Entries stored since belong to it and stay.
/// </summary>
private readonly DateTimeOffset startedAt = DateTimeOffset.UtcNow;
/// <summary>Waits for the data directory, performs one cleanup pass, and then exits.</summary>
/// <param name="stoppingToken">Host shutdown token.</param>
protected override async Task ExecuteAsync(CancellationToken stoppingToken)
{
while (string.IsNullOrWhiteSpace(SettingsManager.DataDirectory) && !stoppingToken.IsCancellationRequested)
await Task.Delay(TimeSpan.FromSeconds(1), stoppingToken);
if (stoppingToken.IsCancellationRequested)
return;
//
// A restore which the runtime could not finish leaves chats and store of two different states
// in the data folder, while the chats it replaced wait elsewhere to be moved back. Nothing
// tells then which chat refers to what:
//
if (SettingsManager.DataDirectory is { } dataDirectory && DataRestoreFollowUp.IsRestoreUnfinished(dataDirectory))
{
logger.LogWarning("A restore is not finished yet, so the provenance startup cleanup waits for a later start.");
return;
}
try
{
var store = await databaseClientProvider.GetProvenanceStoreAsync(stoppingToken);
if (store is NoProvenanceStoreClient)
{
logger.LogInformation("The provenance store is not available, so its startup cleanup has nothing to do.");
return;
}
//
// A restore at this start may have brought back the chats of one moment and kept the store
// of another, e.g., from a backup without the store. Then the time of the last cleanup
// tells nothing about the chats, and every chat is read:
//
await dataRestoreService.StartupStateLoaded.WaitAsync(stoppingToken);
var lastCleanup = dataRestoreService.Report is null ? await store.GetLastCleanupAsync(stoppingToken) : null;
//
// SQLite works synchronously behind its async methods, and every changed chat with entries
// is read once, so the whole pass runs on a thread of its own:
//
var result = await Task.Run(() => ProvenanceCleanup.CleanUpAsync(
store,
WorkspaceBehaviour.ReadAllChatLocations(),
(workspaceId, chatId) => WorkspaceBehaviour.LoadChatAsync(new(workspaceId, chatId)),
WorkspaceBehaviour.GetChatWrittenAt,
this.startedAt,
lastCleanup,
stoppingToken), stoppingToken);
logger.LogInformation("The provenance startup cleanup removed {OrphanedCount} entries of chats which no longer exist and {UnreferencedCount} entries their chats no longer refer to. It skipped {SkippedCount} unchanged chats.", result.Orphaned, result.Unreferenced, result.Skipped);
// A chat which could not be read is read again at the next start, as long as the time of this cleanup is not noted:
if (result.Unreadable is 0)
await store.SetLastCleanupAsync(this.startedAt, stoppingToken);
else
logger.LogWarning("The provenance startup cleanup could not read {UnreadableCount} chats, which keep all their entries. The next start reads them again.", result.Unreadable);
}
catch (OperationCanceledException)
{
// The app is shutting down; the next start cleans up.
}
catch (Exception exception)
{
logger.LogWarning(exception, "The provenance startup cleanup failed. The next start tries again.");
}
}
}
@@ -293,7 +293,7 @@ public sealed partial class RustService
{
if (segment.TokenCount is { } tokenCount)
{
yield return new(segment.Content, tokenCount, segment.PageNumber);
yield return new(segment.Content, tokenCount, segment.Location);
continue;
}
@@ -306,7 +306,7 @@ public sealed partial class RustService
var countedSegment = await this.GetTokenCount(embeddingProvider, segment.Content, token);
if (countedSegment is { Success: true } counted)
{
yield return new(segment.Content, counted.TokenCount, segment.PageNumber);
yield return new(segment.Content, counted.TokenCount, segment.Location);
continue;
}
@@ -318,7 +318,7 @@ public sealed partial class RustService
}
}
private async IAsyncEnumerable<(string Content, int? TokenCount, int? PageNumber)> StreamArbitraryFileDataCore(
private async IAsyncEnumerable<(string Content, int? TokenCount, SourceLocation Location)> StreamArbitraryFileDataCore(
string path,
string tokenizerPath,
PromptInjectionSource? reportAs,
@@ -427,13 +427,13 @@ public sealed partial class RustService
}
//
// The count and the page come from the processed event, not from the event which
// was just read: a reader may hold content back across several events, and the
// count and page of the content it releases describe that content, not the event
// that released it.
// The count and the location come from the processed event, not from the event
// which was just read: a reader may hold content back across several events, and
// the count and location of the content it releases describe that content, not the
// event that released it.
//
if (!string.IsNullOrWhiteSpace(processedEvent.Content))
yield return (processedEvent.Content, processedEvent.TokenCount, processedEvent.PageNumber);
yield return (processedEvent.Content, processedEvent.TokenCount, processedEvent.Location);
}
}
finally
@@ -442,7 +442,7 @@ public sealed partial class RustService
}
if (finalContentChunk is { } pendingContent && !string.IsNullOrWhiteSpace(pendingContent.Content))
yield return (pendingContent.Content, pendingContent.TokenCount, pendingContent.PageNumber);
yield return (pendingContent.Content, pendingContent.TokenCount, SourceLocation.OnPage(pendingContent.PageNumber));
if (promptInjectionRedactedCount is 0)
yield break;
@@ -23,17 +23,17 @@ public sealed class TemporaryChatService(ILogger<TemporaryChatService> logger, S
return;
}
this.StartMaintenance();
await this.StartMaintenanceAsync();
while (!stoppingToken.IsCancellationRequested)
{
await Task.Delay(CHECK_INTERVAL, stoppingToken);
this.StartMaintenance();
await this.StartMaintenanceAsync();
}
}
#endregion
private void StartMaintenance()
internal async Task StartMaintenanceAsync()
{
logger.LogInformation("Starting maintenance of temporary chat storage.");
var temporaryDirectories = Path.Join(SettingsManager.DataDirectory, "tempChats");
@@ -65,13 +65,42 @@ public sealed class TemporaryChatService(ILogger<TemporaryChatService> logger, S
if(deleteChat)
{
logger.LogInformation($"Deleting temporary chat storage directory '{tempChatDirPath}' due to maintenance policy.");
Directory.Delete(tempChatDirPath, true);
WorkspaceBehaviour.InvalidateWorkspaceTreeCache();
await this.DeleteTemporaryChatAsync(tempChatDirPath);
}
}
logger.LogInformation("Finished maintenance of temporary chat storage.");
}
/// <summary>
/// Deletes one temporary chat the way the user deletes a chat, so that what it received goes along.
/// </summary>
/// <remarks>
/// A chat which cannot go now, e.g., since another operation holds it or a file of it is open,
/// stays until the next maintenance. Its failure must not end this service, nor the app with it.
/// </remarks>
/// <param name="tempChatDirPath">The directory of the chat.</param>
private async Task DeleteTemporaryChatAsync(string tempChatDirPath)
{
try
{
if (Guid.TryParse(Path.GetFileName(tempChatDirPath), out var chatId))
{
if (!await WorkspaceBehaviour.DeleteChatWithoutAskingAsync(Guid.Empty, chatId))
logger.LogWarning("The temporary chat '{ChatId}' is busy, so the next maintenance deletes it.", chatId);
return;
}
// Not a chat AI Studio made, so nothing else knows about it:
Directory.Delete(tempChatDirPath, true);
WorkspaceBehaviour.InvalidateWorkspaceTreeCache();
}
catch (Exception exception)
{
logger.LogWarning(exception, "Could not delete the temporary chat storage directory '{Directory}'. The next maintenance tries again.", tempChatDirPath);
}
}
public static void Initialize()
{