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

@@ -305,6 +305,10 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
token.ThrowIfCancellationRequested();
// Read each attached document once, before the request sends it, see ChatThread.StoreAttachedDocumentsAsync:
await chatThread.StoreAttachedDocumentsAsync(settingsManager.ConfigurationData.Chat.ReadChangedAttachmentsAgain);
token.ThrowIfCancellationRequested();
var lastStreamingEvent = DateTimeOffset.MinValue;
if (!TrySetWaitingForRemote(state, token))
return;
@@ -371,6 +375,13 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
RemoveEmptyAIResponse(state);
//
// Before the last checkpoint, so that the stored chat refers to what the AI received instead
// of holding it. An answer which was removed for being empty keeps nothing:
//
if (request.ChatThread.Blocks.Any(block => ReferenceEquals(block.Content, aiText)))
await aiText.StoreReceivedContentAsync(request.ChatThread.ChatId);
lock (state.SyncRoot)
{
state.Snapshot = state.Snapshot with
@@ -461,12 +472,7 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
if (state.IsCompletionStarted || token.IsCancellationRequested || state.ChatGenerationRequest is null)
return false;
var aiText = state.ChatGenerationRequest.AIText;
aiText.InitialRemoteWait = false;
aiText.IsStreaming = true;
aiText.Text += contentStreamChunk;
aiText.Sources.MergeSources(contentStreamChunk.Sources);
aiText.RecordReportedUsage(contentStreamChunk.Usage, state.ChatGenerationRequest.ProviderSettings.Model.Id, blocksSent);
state.ChatGenerationRequest.AIText.ApplyStreamChunk(contentStreamChunk, state.ChatGenerationRequest.ProviderSettings.Model.Id, blocksSent);
if (state.Snapshot.Status is not AIJobStatus.RUNNING)
{
@@ -1,3 +1,4 @@
using AIStudio.Tools.Databases.ProvenanceStore;
using AIStudio.Tools.DataBackup;
namespace AIStudio.Tools.ChatArchive;
@@ -21,6 +22,11 @@ public static class ChatArchiveFormat
/// </summary>
public static readonly string TEMPORARY_CHATS_PREFIX = DataBackupFormat.DATA_ROOT + DataBackupContents.ENTRIES.First(entry => entry.Kind is DataBackupEntryKind.TEMPORARY_CHATS).RelativePath;
/// <summary>
/// Where a backup keeps the provenance store, which holds what the chats received besides their messages.
/// </summary>
public static readonly string PROVENANCE_STORE_ENTRY = DataBackupFormat.DATA_ROOT + DataBackupContents.ENTRIES.First(entry => entry.Kind is DataBackupEntryKind.CHAT_PROVENANCE).RelativePath + SqliteProvenanceStoreClientImplementation.DATABASE_FILENAME;
/// <summary>
/// The file storing the workspace or chat name.
/// </summary>
@@ -3,6 +3,7 @@ using System.Text.Json;
using AIStudio.Chat;
using AIStudio.Settings;
using AIStudio.Tools.Databases.ProvenanceStore;
using AIStudio.Tools.DataBackup;
using AIStudio.Tools.PluginSystem;
@@ -24,6 +25,11 @@ public static class ChatArchiveImporter
private const int BUFFER_SIZE = 64 * 1024;
/// <summary>
/// What SQLite appends to the name of a database for its write-ahead log.
/// </summary>
private const string WRITE_AHEAD_LOG_SUFFIX = "-wal";
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(ChatArchiveImporter).Namespace, nameof(ChatArchiveImporter));
/// <summary>
@@ -31,6 +37,13 @@ public static class ChatArchiveImporter
/// </summary>
private sealed record BackupEntry(string Name, ZipArchiveEntry Entry);
/// <summary>
/// Where the imported chats get what they received besides their messages from.
/// </summary>
/// <param name="Store">The provenance store of this computer, which receives the entries.</param>
/// <param name="BackupStorePath">The provenance store of the backup, unpacked into a file of its own.</param>
private sealed record ProvenanceImport(ProvenanceStoreClient Store, string BackupStorePath);
/// <summary>
/// What one import did so far.
/// </summary>
@@ -84,22 +97,26 @@ public static class ChatArchiveImporter
/// Adds all chats of a backup to the chat storage.
/// </summary>
/// <param name="archivePath">The backup.</param>
/// <param name="provenanceStore">The provenance store, which receives what the imported chats received besides their messages.</param>
/// <param name="collisionBehavior">What to do with chats which exist already.</param>
/// <param name="progress">Receives the progress.</param>
/// <param name="token">Cancels the import. Chats added before the cancellation are kept.</param>
/// <returns>What the import did.</returns>
public static async Task<ChatArchiveImportResult> ImportAsync(string archivePath, ChatArchiveCollisionBehavior collisionBehavior, IProgress<ChatArchiveProgress>? progress, CancellationToken token)
public static async Task<ChatArchiveImportResult> ImportAsync(string archivePath, ProvenanceStoreClient provenanceStore, ChatArchiveCollisionBehavior collisionBehavior, IProgress<ChatArchiveProgress>? progress, CancellationToken token)
{
var tally = new ImportTally();
var preview = await ReadPreviewAsync(archivePath, token);
if (!preview.Success)
return new(false, 0, 0, 0, 0, preview.Issue, false);
// Next to the copies a backup makes, which the next start clears in case AI Studio stops halfway:
var provenanceDirectory = Path.Join(DataRestoreFormat.DirectoryIn(SettingsManager.DataDirectory!, DataBackupFormat.STAGING_DIRECTORY), $"provenance-import-{Guid.NewGuid():N}");
try
{
var manifest = preview.Manifest;
using var archive = ZipFile.OpenRead(archivePath);
var entriesByChat = GroupEntriesByChat(archive);
var provenance = await ExtractProvenanceStoreAsync(archive, provenanceStore, provenanceDirectory, token);
// Which workspaces exist already is asked once per workspace of the backup, so read
// the tree once instead of building a full snapshot every time:
@@ -134,10 +151,10 @@ public static class ChatArchiveImporter
}
var workspaceEntryPrefix = $"{ChatArchiveFormat.WORKSPACES_PREFIX}{workspace.WorkspaceId}{ZipEntryNames.SEPARATOR}";
await ImportChatsAsync(tally, workspaceEntryPrefix, targetWorkspaceId, workspace.Chats, true, manifest.TotalChatCount, entriesByChat, collisionBehavior, existingChatIds, progress, token);
await ImportChatsAsync(tally, workspaceEntryPrefix, targetWorkspaceId, workspace.Chats, true, manifest.TotalChatCount, entriesByChat, collisionBehavior, existingChatIds, provenance, progress, token);
}
await ImportChatsAsync(tally, ChatArchiveFormat.TEMPORARY_CHATS_PREFIX, Guid.Empty, manifest.TemporaryChats, false, manifest.TotalChatCount, entriesByChat, collisionBehavior, existingChatIds, progress, token);
await ImportChatsAsync(tally, ChatArchiveFormat.TEMPORARY_CHATS_PREFIX, Guid.Empty, manifest.TemporaryChats, false, manifest.TotalChatCount, entriesByChat, collisionBehavior, existingChatIds, provenance, progress, token);
progress?.Report(new(tally.ProcessedChats, manifest.TotalChatCount));
LOG.LogInformation("Imported {ImportedChats} chats from the backup '{ArchivePath}'. Skipped {SkippedChats} existing chats, {FailedChats} chats failed.", tally.ImportedChats, archivePath, tally.SkippedChats, tally.FailedChats);
@@ -158,12 +175,111 @@ public static class ChatArchiveImporter
LOG.LogError(exception, "Failed to import the backup '{ArchivePath}'.", archivePath);
return new(false, tally.ImportedWorkspaces, tally.ImportedChats, tally.SkippedChats, tally.FailedChats, string.Format(TB("Unexpected error: {0}"), exception.Message), false);
}
finally
{
TryDeleteDirectory(provenanceDirectory);
}
}
/// <summary>
/// Unpacks the provenance store of a backup, so that every imported chat can take its entries from it.
/// </summary>
/// <remarks>
/// A backup of a version before the store holds none, and its chats arrive as they were: their
/// sources without the contents behind them.
/// </remarks>
/// <param name="archive">The backup.</param>
/// <param name="provenanceStore">The provenance store of this computer.</param>
/// <param name="directory">The directory to unpack the store into. The caller deletes it afterward.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>Where the entries come from, or null when the backup holds no store, or this computer cannot keep one.</returns>
private static async Task<ProvenanceImport?> ExtractProvenanceStoreAsync(ZipArchive archive, ProvenanceStoreClient provenanceStore, string directory, CancellationToken token)
{
if (provenanceStore is NoProvenanceStoreClient)
{
LOG.LogWarning("The provenance store is not available, so the chats of the backup arrive without what they received besides their messages.");
return null;
}
var entry = archive.GetEntry(ChatArchiveFormat.PROVENANCE_STORE_ENTRY);
if (entry is null)
{
LOG.LogInformation("The backup holds no provenance store, so its chats arrive without what they received besides their messages.");
return null;
}
try
{
Directory.CreateDirectory(directory);
var path = Path.Join(directory, SqliteProvenanceStoreClientImplementation.DATABASE_FILENAME);
await ExtractEntryAsync(entry, path, token);
//
// A backup which took the store as it lay in the data folder may hold its newest entries
// in the write-ahead log next to it, which SQLite reads along:
//
if (archive.GetEntry(ChatArchiveFormat.PROVENANCE_STORE_ENTRY + WRITE_AHEAD_LOG_SUFFIX) is { } logEntry)
await ExtractEntryAsync(logEntry, path + WRITE_AHEAD_LOG_SUFFIX, token);
return new(provenanceStore, path);
}
catch (OperationCanceledException)
{
throw;
}
catch (Exception exception)
{
// As in ImportProvenanceAsync: without the store, the chats lose the contents behind their sources, not themselves.
LOG.LogWarning(exception, "The provenance store of the backup could not be unpacked, so its chats arrive without what they received besides their messages.");
return null;
}
}
/// <summary>
/// Writes one entry of the backup into a file, as a stream.
/// </summary>
/// <param name="entry">The entry.</param>
/// <param name="destinationPath">The file, which is replaced when it exists.</param>
/// <param name="token">The cancellation token.</param>
private static async Task ExtractEntryAsync(ZipArchiveEntry entry, string destinationPath, CancellationToken token)
{
await using var entryStream = entry.Open();
await using var fileStream = new FileStream(destinationPath, FileMode.Create, FileAccess.Write, FileShare.None, BUFFER_SIZE, FileOptions.Asynchronous | FileOptions.SequentialScan);
await entryStream.CopyToAsync(fileStream, token);
}
/// <summary>
/// Adds what an imported chat received besides its messages, from the provenance store of the backup.
/// </summary>
/// <remarks>
/// The chat is stored already, so this ends neither by a cancellation nor by a failure: the next
/// chat checks for the cancellation. A chat without its entries still has its messages and
/// sources; it only lacks the contents behind them, and the store of the backup may be damaged
/// or of a newer form.
/// </remarks>
/// <param name="provenance">Where the entries come from, or null when they come from nowhere.</param>
/// <param name="sourceChatId">The chat in the backup.</param>
/// <param name="targetChatId">The chat as it was imported, possibly under a new id.</param>
private static async Task ImportProvenanceAsync(ProvenanceImport? provenance, Guid sourceChatId, Guid targetChatId)
{
if (provenance is null)
return;
try
{
var count = await provenance.Store.ImportChatAsync(provenance.BackupStorePath, sourceChatId, targetChatId, CancellationToken.None);
LOG.LogDebug("Imported {Count} provenance entries of the chat '{ChatId}'.", count, targetChatId);
}
catch (Exception exception)
{
LOG.LogWarning(exception, "The chat '{ChatId}' was imported without what it received besides its messages, because the provenance store of the backup could not be read.", targetChatId);
}
}
/// <summary>
/// Adds the chats of one workspace of the backup, or its temporary chats.
/// </summary>
private static async Task ImportChatsAsync(ImportTally tally, string entryPrefix, Guid targetWorkspaceId, IReadOnlyList<ChatArchiveManifestChat> chats, bool countsAsWorkspace, int totalChats, Dictionary<string, List<BackupEntry>> entriesByChat, ChatArchiveCollisionBehavior collisionBehavior, HashSet<Guid> existingChatIds, IProgress<ChatArchiveProgress>? progress, CancellationToken token)
private static async Task ImportChatsAsync(ImportTally tally, string entryPrefix, Guid targetWorkspaceId, IReadOnlyList<ChatArchiveManifestChat> chats, bool countsAsWorkspace, int totalChats, Dictionary<string, List<BackupEntry>> entriesByChat, ChatArchiveCollisionBehavior collisionBehavior, HashSet<Guid> existingChatIds, ProvenanceImport? provenance, IProgress<ChatArchiveProgress>? progress, CancellationToken token)
{
var isFirstImportOfWorkspace = true;
foreach (var chat in chats)
@@ -175,7 +291,7 @@ public static class ChatArchiveImporter
var chatEntryPrefix = $"{entryPrefix}{chat.ChatId}{ZipEntryNames.SEPARATOR}";
var chatEntries = entriesByChat.GetValueOrDefault(chatEntryPrefix, []);
switch (await ImportChatAsync(chatEntries, chatEntryPrefix, targetWorkspaceId, chat, collisionBehavior, existingChatIds, token))
switch (await ImportChatAsync(chatEntries, chatEntryPrefix, targetWorkspaceId, chat, collisionBehavior, existingChatIds, provenance, token))
{
case ChatArchiveImportOutcome.IMPORTED:
tally.ImportedChats++;
@@ -383,8 +499,9 @@ public static class ChatArchiveImporter
/// <param name="chat">The chat as the preview lists it.</param>
/// <param name="collisionBehavior">What to do when a chat with the same ID exists already.</param>
/// <param name="existingChatIds">The IDs of all chats in the storage. An imported chat is added.</param>
/// <param name="provenance">Where the chat gets what it received besides its messages from, or null when from nowhere.</param>
/// <param name="token">Cancels the import.</param>
private static async Task<ChatArchiveImportOutcome> ImportChatAsync(IReadOnlyList<BackupEntry> chatEntries, string chatEntryPrefix, Guid targetWorkspaceId, ChatArchiveManifestChat chat, ChatArchiveCollisionBehavior collisionBehavior, HashSet<Guid> existingChatIds, CancellationToken token)
private static async Task<ChatArchiveImportOutcome> ImportChatAsync(IReadOnlyList<BackupEntry> chatEntries, string chatEntryPrefix, Guid targetWorkspaceId, ChatArchiveManifestChat chat, ChatArchiveCollisionBehavior collisionBehavior, HashSet<Guid> existingChatIds, ProvenanceImport? provenance, CancellationToken token)
{
var threadEntry = chatEntries.FirstOrDefault(entry => entry.Name.Equals($"{chatEntryPrefix}{ChatArchiveFormat.THREAD_FILE_NAME}", StringComparison.OrdinalIgnoreCase));
if (threadEntry is null)
@@ -427,10 +544,7 @@ public static class ChatArchiveImporter
var destinationPath = ZipEntryNames.GetDestinationPath(chatDirectory, relativePath);
Directory.CreateDirectory(Path.GetDirectoryName(destinationPath)!);
await using var entryStream = entry.Entry.Open();
await using var fileStream = new FileStream(destinationPath, FileMode.Create, FileAccess.Write, FileShare.None, BUFFER_SIZE, FileOptions.Asynchronous | FileOptions.SequentialScan);
await entryStream.CopyToAsync(fileStream, token);
await ExtractEntryAsync(entry.Entry, destinationPath, token);
extractedPaths.Add(relativePath);
}
@@ -477,6 +591,9 @@ public static class ChatArchiveImporter
}
existingChatIds.Add(targetChatId);
// Under the id of the chat in the backup, which is what the store of the backup knows it by:
await ImportProvenanceAsync(provenance, chat.ChatId, targetChatId);
return ChatArchiveImportOutcome.IMPORTED;
}
catch (OperationCanceledException)
@@ -533,7 +650,7 @@ public static class ChatArchiveImporter
}
catch (Exception exception)
{
LOG.LogWarning(exception, "Failed to clean up the incomplete chat directory '{ChatDirectory}'.", directory);
LOG.LogWarning(exception, "Failed to clean up the directory '{Directory}'.", directory);
}
}
}
@@ -12,8 +12,8 @@ namespace AIStudio.Tools;
/// <param name="Error">The reported failure, or null when the event was processed successfully.</param>
/// <param name="PromptInjection">What the runtime filtered out of the content, or null when it filtered nothing.</param>
/// <param name="TokenCount">The number of tokens of the content, or null when it is unknown.</param>
/// <param name="PageNumber">The page the content came from, or null when it has none.</param>
public readonly record struct ContentStreamProcessedEvent(string? Content, ContentStreamErrorDetails? Error, ContentStreamPromptInjectionDetails? PromptInjection = null, int? TokenCount = null, int? PageNumber = null)
/// <param name="Location">Where the content came from: its page, its sheet and row, or its line.</param>
public readonly record struct ContentStreamProcessedEvent(string? Content, ContentStreamErrorDetails? Error, ContentStreamPromptInjectionDetails? PromptInjection = null, int? TokenCount = null, SourceLocation Location = default)
{
/// <summary>
/// An event which neither produced content nor reported a failure.
@@ -21,18 +21,18 @@ public readonly record struct ContentStreamProcessedEvent(string? Content, Conte
public static readonly ContentStreamProcessedEvent NOTHING = new(null, null);
/// <summary>
/// An event which produced content, with the token count and the page of that very content.
/// An event which produced content, with the token count and the location of that very content.
/// </summary>
/// <remarks>
/// The count travels with the content because a reader may hold content back across several
/// events: pairing it with the count of the event which released it would size it by the
/// wrong text. The page travels along for the same reason, and so that whoever indexes the
/// wrong text. The location travels along for the same reason, and so that whoever indexes the
/// content is told where it came from instead of having to read it back out of the text.
/// </remarks>
/// <param name="content">The content to append.</param>
/// <param name="tokenCount">The number of tokens of that content, or null when it is unknown.</param>
/// <param name="pageNumber">The page that content came from, or null when it has none.</param>
public static ContentStreamProcessedEvent FromContent(string? content, int? tokenCount = null, int? pageNumber = null) => new(content, null, TokenCount: tokenCount, PageNumber: pageNumber);
/// <param name="location">Where that content came from, or nothing when the runtime stated nothing.</param>
public static ContentStreamProcessedEvent FromContent(string? content, int? tokenCount = null, SourceLocation location = default) => new(content, null, TokenCount: tokenCount, Location: location);
public static ContentStreamProcessedEvent FromError(ContentStreamErrorDetails? error) => new(null, error);
@@ -16,8 +16,8 @@ public static class ContentStreamSseHandler
case { Content: not null, Metadata: not null }:
switch (sseEvent.Metadata)
{
case ContentStreamTextMetadata:
return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount);
case ContentStreamTextMetadata textMetadata:
return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount, SourceLocation.OnLine(textMetadata.Text?.LineNumber));
//
// The heading tells the AI which page it is reading. The number is handed on
@@ -31,8 +31,12 @@ public static class ContentStreamSseHandler
# Page {pageNumber}
{sseEvent.Content}
""", sseEvent.TokenCount, pageNumber > 0 ? pageNumber : null);
""", sseEvent.TokenCount, SourceLocation.OnPage(pageNumber));
//
// Row 0 is the heading of a sheet together with the start of its table, which
// is no row of the sheet, cf. SourceLocation.InSheet:
//
case ContentStreamSpreadsheetMetadata spreadsheetMetadata:
var sheetName = spreadsheetMetadata.Spreadsheet?.SheetName;
var rowNumber = spreadsheetMetadata.Spreadsheet?.RowNumber;
@@ -44,7 +48,7 @@ public static class ContentStreamSseHandler
}
spreadSheetResult.Append(sseEvent.Content);
return ContentStreamProcessedEvent.FromContent(spreadSheetResult.ToString(), sseEvent.TokenCount);
return ContentStreamProcessedEvent.FromContent(spreadSheetResult.ToString(), sseEvent.TokenCount, SourceLocation.InSheet(sheetName, rowNumber));
//
// Documents which the runtime reads page by page are buffered, so the images of
@@ -62,7 +66,7 @@ public static class ContentStreamSseHandler
var documentManager = DOCUMENT_MANAGERS.GetOrAdd(sseEvent.StreamId!, _ => new());
var documentContent = documentManager.AddPage(documentMetadata, sseEvent.Content, sseEvent.TokenCount, extractImages);
return documentContent is null ? ContentStreamProcessedEvent.NOTHING : ContentStreamProcessedEvent.FromContent(documentContent.Value.Content, documentContent.Value.TokenCount, documentContent.Value.PageNumber);
return documentContent is null ? ContentStreamProcessedEvent.NOTHING : ContentStreamProcessedEvent.FromContent(documentContent.Value.Content, documentContent.Value.TokenCount, SourceLocation.OnPage(documentContent.Value.PageNumber));
case ContentStreamImageMetadata:
return ContentStreamProcessedEvent.FromContent(sseEvent.Content, sseEvent.TokenCount);
@@ -41,6 +41,7 @@ public static class DataBackupContents
new(DataBackupLocation.CONFIG_DIRECTORY, "settings.json", DataBackupEntryKind.SETTINGS),
new(DataBackupLocation.DATA_DIRECTORY, "workspaces/", DataBackupEntryKind.WORKSPACES),
new(DataBackupLocation.DATA_DIRECTORY, "tempChats/", DataBackupEntryKind.TEMPORARY_CHATS),
new(DataBackupLocation.DATA_DIRECTORY, "databases/provenance/", DataBackupEntryKind.CHAT_PROVENANCE),
new(DataBackupLocation.DATA_DIRECTORY, PLUGINS_FOLDER, DataBackupEntryKind.PLUGINS),
new(DataBackupLocation.DATA_DIRECTORY, "tokenizers/", DataBackupEntryKind.TOKENIZERS),
new(DataBackupLocation.DATA_DIRECTORY, "visualBriefings/", DataBackupEntryKind.VISUAL_BRIEFINGS),
@@ -14,6 +14,7 @@ public sealed record DataBackupEntry(DataBackupLocation Location, string Relativ
public bool IsIncluded => this.Kind is DataBackupEntryKind.SETTINGS
or DataBackupEntryKind.WORKSPACES
or DataBackupEntryKind.TEMPORARY_CHATS
or DataBackupEntryKind.CHAT_PROVENANCE
or DataBackupEntryKind.PLUGINS
or DataBackupEntryKind.TOKENIZERS
or DataBackupEntryKind.VISUAL_BRIEFINGS
@@ -33,6 +34,16 @@ public sealed record DataBackupEntry(DataBackupLocation Location, string Relativ
or DataBackupEntryKind.VECTOR_STORES
or DataBackupEntryKind.MEDIA_TRANSCRIPTS;
/// <summary>
/// Whether a backup takes the entry from the copy it makes of the provenance store, rather than from the data folder.
/// </summary>
/// <remarks>
/// The chats write to the store at any time, so a file-by-file copy could catch it halfway
/// through a write. Unlike the indexes, it is copied for every backup, since it belongs to the
/// chats, which every backup holds. Cf. DataBackupSources.CopyProvenanceStoreAsync.
/// </remarks>
public bool ComesFromProvenanceCopy => this.Kind is DataBackupEntryKind.CHAT_PROVENANCE;
/// <summary>
/// Whether the entry is a folder rather than a single file.
/// </summary>
@@ -20,6 +20,12 @@ public enum DataBackupEntryKind
/// </summary>
TEMPORARY_CHATS,
/// <summary>
/// What the chats received besides their messages, such as the content of attached files and the
/// results of tools. It belongs to the chats, so every backup takes it along, like them.
/// </summary>
CHAT_PROVENANCE,
/// <summary>
/// The plugins of the user. A backup leaves out the plugins of the organization.
/// </summary>
@@ -8,4 +8,5 @@ namespace AIStudio.Tools.DataBackup;
/// <param name="SettingsJson">The settings as the backup takes them along, cf. DataBackupSettings. They come from memory rather than from the disk, since the backup leaves out the objects of configuration plugins.</param>
/// <param name="IndexCopiesDirectory">The copies of the indexes, in the layout of the data folder, or an empty string when the backup leaves the indexes out. The indexes in the data folder change while AI Studio runs, so the backup never reads them there.</param>
/// <param name="OrganizationPluginFolders">The plugins an organization deployed among the user's own ones, cf. DataBackupContents.FindOrganizationPluginsAmongUserPlugins.</param>
public sealed record DataBackupSources(string DataDirectory, string ConfigDirectory, string SettingsJson, string IndexCopiesDirectory, IReadOnlyList<string> OrganizationPluginFolders);
/// <param name="CopyProvenanceStoreAsync">Copies the provenance store once the chats are written, and answers the directory of the copy, in the layout of the data folder, or an empty string when there is none. Copied after the chats, the store holds everything the chats of the backup refer to, also what an answer completed meanwhile received; what no chat of the backup refers to is removed after a restore. Without a copy, the backup takes the store from the data folder, which is consistent only while nothing writes to it, i.e., when AI Studio could not open the store at all.</param>
public sealed record DataBackupSources(string DataDirectory, string ConfigDirectory, string SettingsJson, string IndexCopiesDirectory, IReadOnlyList<string> OrganizationPluginFolders, Func<CancellationToken, Task<string>>? CopyProvenanceStoreAsync = null);
@@ -43,8 +43,9 @@ public static class DataBackupWriter
/// <returns>The manifest as the backup holds it.</returns>
public static async Task<DataBackupManifest> WriteAsync(Stream target, DataBackupSources sources, DataBackupManifest manifest, IProgress<DataBackupProgress>? progress, CancellationToken token)
{
var files = CollectFiles(sources, token);
var totalFiles = files.Count + 1;
// The provenance store follows the chats, cf. DataBackupSources.CopyProvenanceStoreAsync. Until it is copied, it counts as one file:
var files = CollectFiles(sources, entry => !entry.ComesFromProvenanceCopy, string.Empty, token);
var totalFiles = files.Count + 2;
var entries = new List<DataBackupManifestEntry>(totalFiles);
var buffer = new byte[BUFFER_SIZE];
@@ -56,25 +57,12 @@ public static class DataBackupWriter
entries.Add(await WriteEntryAsync(archive, DataBackupFormat.EntryName(settingsEntry.Location, settingsEntry.RelativePath), settingsStream, DateTimeOffset.Now, buffer, token));
progress?.Report(new(DataBackupPhase.WRITING, entries.Count, totalFiles));
foreach (var file in files)
{
token.ThrowIfCancellationRequested();
await WriteFilesAsync(archive, files, entries, totalFiles, progress, buffer, token);
FileStream source;
try
{
source = new FileStream(file.SourcePath, FileMode.Open, FileAccess.Read, FileShare.ReadWrite | FileShare.Delete, BUFFER_SIZE, FileOptions.Asynchronous | FileOptions.SequentialScan);
}
catch (Exception exception) when (exception is FileNotFoundException or DirectoryNotFoundException)
{
continue;
}
await using (source)
entries.Add(await WriteEntryAsync(archive, file.EntryName, source, File.GetLastWriteTime(file.SourcePath), buffer, token));
progress?.Report(new(DataBackupPhase.WRITING, entries.Count, totalFiles));
}
var provenanceCopyDirectory = sources.CopyProvenanceStoreAsync is null ? string.Empty : await sources.CopyProvenanceStoreAsync(token);
var provenanceFiles = CollectFiles(sources, entry => entry.ComesFromProvenanceCopy, provenanceCopyDirectory, token);
totalFiles = files.Count + 1 + provenanceFiles.Count;
await WriteFilesAsync(archive, provenanceFiles, entries, totalFiles, progress, buffer, token);
manifest = manifest with { Entries = entries };
await using (var manifestStream = new MemoryStream(JsonSerializer.SerializeToUtf8Bytes(manifest, DataBackupFormat.JSON_OPTIONS), writable: false))
@@ -88,9 +76,14 @@ public static class DataBackupWriter
/// Collects the files of the data folder which belong into a backup, cf. DataBackupContents.
/// </summary>
/// <remarks>
/// The settings are no part of it: they come from memory. The indexes come from their copies.
/// The settings are no part of it: they come from memory. The indexes and the provenance store
/// come from their copies.
/// </remarks>
private static List<BackupFile> CollectFiles(DataBackupSources sources, CancellationToken token)
/// <param name="sources">Where the files come from.</param>
/// <param name="include">Which entries of the list to collect.</param>
/// <param name="provenanceCopyDirectory">The copy of the provenance store, or an empty string to take it from the data folder.</param>
/// <param name="token">Cancels collecting.</param>
private static List<BackupFile> CollectFiles(DataBackupSources sources, Func<DataBackupEntry, bool> include, string provenanceCopyDirectory, CancellationToken token)
{
var excludedPaths = DataBackupContents.ENTRIES
.Where(entry => !entry.IsIncluded && entry.Location is DataBackupLocation.DATA_DIRECTORY)
@@ -99,11 +92,14 @@ public static class DataBackupWriter
.ToList();
var files = new List<BackupFile>();
foreach (var entry in DataBackupContents.ENTRIES.Where(entry => entry.IsIncluded && entry.Location is DataBackupLocation.DATA_DIRECTORY))
foreach (var entry in DataBackupContents.ENTRIES.Where(entry => entry.IsIncluded && entry.Location is DataBackupLocation.DATA_DIRECTORY && include(entry)))
{
var root = entry.ComesFromIndexCopies
? sources.IndexCopiesDirectory
: sources.DataDirectory;
var root = entry switch
{
{ ComesFromIndexCopies: true } => sources.IndexCopiesDirectory,
{ ComesFromProvenanceCopy: true } when !string.IsNullOrWhiteSpace(provenanceCopyDirectory) => provenanceCopyDirectory,
_ => sources.DataDirectory,
};
if (string.IsNullOrWhiteSpace(root))
continue;
@@ -141,6 +137,32 @@ public static class DataBackupWriter
return files;
}
/// <summary>
/// Writes files into the backup, and notes each one in the entries of the manifest.
/// </summary>
private static async Task WriteFilesAsync(ZipArchive archive, IReadOnlyList<BackupFile> files, List<DataBackupManifestEntry> entries, int totalFiles, IProgress<DataBackupProgress>? progress, byte[] buffer, CancellationToken token)
{
foreach (var file in files)
{
token.ThrowIfCancellationRequested();
FileStream source;
try
{
source = new FileStream(file.SourcePath, FileMode.Open, FileAccess.Read, FileShare.ReadWrite | FileShare.Delete, BUFFER_SIZE, FileOptions.Asynchronous | FileOptions.SequentialScan);
}
catch (Exception exception) when (exception is FileNotFoundException or DirectoryNotFoundException)
{
continue;
}
await using (source)
entries.Add(await WriteEntryAsync(archive, file.EntryName, source, File.GetLastWriteTime(file.SourcePath), buffer, token));
progress?.Report(new(DataBackupPhase.WRITING, entries.Count, totalFiles));
}
}
/// <summary>
/// Writes one entry, and notes its size and hash while it streams through.
/// </summary>
@@ -5,7 +5,7 @@ These instructions cover the databases in `app/MindWork AI Studio/Tools/Database
## Databases
Local RAG runs on two databases, addressed through `DatabaseRole`:
AI Studio keeps three databases, addressed through `DatabaseRole`. Local RAG runs on the first two:
- **`VECTOR_STORE`** — Qdrant Edge through the `qdrant-edge` crate, running **in-process inside the
Rust runtime**. There is no sidecar process, no port 6333 and no Qdrant API key; .NET reaches it
@@ -15,6 +15,13 @@ Local RAG runs on two databases, addressed through `DatabaseRole`:
- **`INDEX_STORE`** — SQLite at `<data directory>/databases/sqlite/rag-index.sqlite3`, reached
through EF Core. It holds the data sources, the file fingerprints, the chunk texts and an FTS5
index over them, plus the files which permanently failed to index.
- **`PROVENANCE_STORE`** — SQLite at `<data directory>/databases/provenance/provenance.sqlite3`,
reached through EF Core with migrations of its own. It holds what chats received besides their
messages, exactly as the AI received it: attached files, tool results, and passages of data
sources. A chat keeps only the ids. Never put this into the index store: the index can be
rebuilt from the files and is treated that way, while what a chat received cannot. Its
connections run with `secure_delete`, see `SecureDeleteInterceptor`, so that the text of a
deleted chat does not stay readable in the file; a new way to connect to it keeps that.
When working on these, keep in mind:
@@ -32,7 +39,8 @@ When working on these, keep in mind:
- **Raw SQL against SQLite goes through `context.Database.GetDbConnection()`**, not through
`SqlQueryRaw<T>`: that one expects a column named `Value` and wraps the statement, so a `PRAGMA`
never works with it.
- **A new EF Core migration needs a `[DynamicDependency]`** in `IndexStoreSchemaMigrator`, because
`PublishTrimmed` is on and the migration type would otherwise be trimmed away. The "Schema version"
- **A new EF Core migration needs a `[DynamicDependency]`** in the schema migrator of its store,
`IndexStoreSchemaMigrator` or `ProvenanceStoreSchemaMigrator`, because `PublishTrimmed` is on and
the migration type would otherwise be trimmed away. The "Schema version"
line on the information page shows the applied and pending counts, so a forgotten entry becomes
visible there.
@@ -0,0 +1,27 @@
namespace AIStudio.Tools.Databases;
/// <summary>
/// Where a stored chunk is in its document, as the index store keeps it.
/// </summary>
/// <remarks>
/// Each value is null when the document has no such thing, and for a chunk stored before the value
/// was kept: it reads as unknown then. The vector store keeps only the pages, so a hit of the vector
/// search learns the rest from the index store, cf. LocalIndexSearchService.
/// </remarks>
/// <param name="PageNumber">The page the chunk starts on.</param>
/// <param name="EndPageNumber">The page the chunk ends on.</param>
/// <param name="SheetName">The sheet of a spreadsheet the chunk is from.</param>
/// <param name="StartRow">The row of its sheet the chunk starts at.</param>
/// <param name="EndRow">The row of its sheet the chunk ends at.</param>
/// <param name="StartLine">The line the chunk starts at, for text, CSV and source code.</param>
/// <param name="EndLine">The line the chunk ends at.</param>
/// <param name="AttachmentPosition">Which attachment of a mail the chunk is from, as mail_parts numbers them, or null for the text of the mail itself and for any file.</param>
public sealed record ChunkLocation(
int? PageNumber = null,
int? EndPageNumber = null,
string? SheetName = null,
int? StartRow = null,
int? EndRow = null,
int? StartLine = null,
int? EndLine = null,
int? AttachmentPosition = null);
@@ -1,4 +1,5 @@
using AIStudio.Tools.Databases.IndexStore;
using AIStudio.Tools.Databases.ProvenanceStore;
using AIStudio.Tools.Databases.VectorStore;
using AIStudio.Tools.Services;
@@ -69,6 +70,18 @@ public sealed class DatabaseClientProvider(RustService rustService, ILoggerFacto
client.Status);
}
public async Task<ProvenanceStoreClient> GetProvenanceStoreAsync(CancellationToken cancellationToken = default)
{
var client = await this.GetClientAsync(DatabaseRole.PROVENANCE_STORE, cancellationToken);
if (client is ProvenanceStoreClient provenanceStore)
return provenanceStore;
return new NoProvenanceStoreClient(
client.Name,
"The configured database client does not support provenance operations.",
client.Status);
}
/// <summary>
/// Builds the client which stands in for a database role that cannot serve right now.
/// </summary>
@@ -85,6 +98,7 @@ public sealed class DatabaseClientProvider(RustService rustService, ILoggerFacto
{
DatabaseRole.VECTOR_STORE => new NoVectorStoreClient(name, reason, status),
DatabaseRole.INDEX_STORE => new NoIndexStoreClient(name, reason, status),
DatabaseRole.PROVENANCE_STORE => new NoProvenanceStoreClient(name, reason, status),
_ => new NoDatabaseClient(name, reason, status)
};
@@ -125,6 +139,7 @@ public sealed class DatabaseClientProvider(RustService rustService, ILoggerFacto
{
DatabaseRole.VECTOR_STORE => await QdrantEdgeClientImplementation.CreateAsync(rustService, this.GetIndexStoreAsync, this.logger, this.databaseClientLogger, cancellationToken),
DatabaseRole.INDEX_STORE => await SqliteIndexStoreClientImplementation.CreateAsync(this.logger, this.databaseClientLogger, cancellationToken),
DatabaseRole.PROVENANCE_STORE => await SqliteProvenanceStoreClientImplementation.CreateAsync(this.logger, this.databaseClientLogger, cancellationToken),
_ => new NoDatabaseClient(databaseRole.ToString(), "The requested database role is not supported.")
};
@@ -4,4 +4,5 @@ public enum DatabaseRole
{
VECTOR_STORE,
INDEX_STORE,
PROVENANCE_STORE,
}
@@ -0,0 +1,15 @@
namespace AIStudio.Tools.Databases;
/// <summary>
/// The chunks of one document which no search may find.
/// </summary>
/// <remarks>
/// A document is searched by the generation its index row names, and only once that generation is
/// complete. Indexing it again writes the next generation next to it, and an attempt which failed
/// may stay there afterward, so the next one can reuse its vectors, cf. IndexedRunContext. A
/// document indexed for the first time has no searchable generation until its first one is
/// complete. The index store knows all of this, and the vector store leaves out what it names.
/// </remarks>
/// <param name="ParentFileId">The parent id of the document.</param>
/// <param name="SearchableGeneration">The generation which can be searched, so every other one is hidden. Null while the document has none, so all of its chunks are hidden.</param>
public sealed record HiddenChunks(string ParentFileId, int? SearchableGeneration);
@@ -0,0 +1,23 @@
using System.Security.Cryptography;
using System.Text;
namespace AIStudio.Tools.Databases.IndexStore;
/// <summary>
/// The hash of the text of a chunk, exactly as it was embedded.
/// </summary>
/// <remarks>
/// A vector depends on nothing but the text it was made from and the embedding signature, and all
/// vectors of one store share that signature. So two chunks of a store with the same hash may share
/// a vector, which is what spares a document the embedding provider for every chunk that did not
/// change. The text is hashed with its overlap and its heading line, as the provider received it.
/// </remarks>
internal static class ChunkTextHash
{
/// <summary>
/// SHA-256 over the UTF-8 bytes of the text, as lowercase hex.
/// </summary>
/// <param name="text">The text of the chunk.</param>
/// <returns>The hash.</returns>
public static string Of(string text) => Convert.ToHexStringLower(SHA256.HashData(Encoding.UTF8.GetBytes(text)));
}
@@ -1,3 +1,41 @@
namespace AIStudio.Tools.Databases.IndexStore;
public sealed record EmbeddingStateChunk(string ChunkId, string ParentFileId, int? PageNumber, int ChunkIndex, string ChunkText, DateTimeOffset EmbeddedAtUtc);
/// <summary>
/// One chunk as the index stores it.
/// </summary>
/// <remarks>
/// Where a chunk is goes into the typed columns its document has: pages, sheet and rows, or lines.
/// The vector store keeps only the pages, cf. VectorStoragePoint: a search completes its hits from
/// here by their chunk id.
/// </remarks>
/// <param name="ChunkId">The id of the chunk, which is the id of its point in the vector store.</param>
/// <param name="ParentFileId">The document the chunk belongs to.</param>
/// <param name="PageNumber">The page the chunk starts on, or null.</param>
/// <param name="EndPageNumber">The page the chunk ends on, or null.</param>
/// <param name="ChunkIndex">The place of the chunk in its document, counted from zero.</param>
/// <param name="ChunkText">The text which was embedded.</param>
/// <param name="EmbeddedAtUtc">When the chunk was stored.</param>
/// <param name="TextHash">The hash of the text, cf. ChunkTextHash.</param>
/// <param name="Generation">Which indexing of the document wrote the chunk.</param>
/// <param name="SheetName">The sheet the chunk is from, or null. A chunk never runs from one sheet into the next.</param>
/// <param name="StartRow">The row of its sheet the chunk starts at, or null.</param>
/// <param name="EndRow">The row of its sheet the chunk ends at, or null.</param>
/// <param name="StartLine">The line the chunk starts at, or null.</param>
/// <param name="EndLine">The line the chunk ends at, or null.</param>
/// <param name="AttachmentPosition">Which attachment of a mail the chunk is from, as mail_parts numbers it, or null for anything else.</param>
public sealed record EmbeddingStateChunk(
string ChunkId,
string ParentFileId,
int? PageNumber,
int? EndPageNumber,
int ChunkIndex,
string ChunkText,
DateTimeOffset EmbeddedAtUtc,
string? TextHash = null,
int Generation = 0,
string? SheetName = null,
int? StartRow = null,
int? EndRow = null,
int? StartLine = null,
int? EndLine = null,
int? AttachmentPosition = null);
@@ -10,10 +10,68 @@ internal sealed class EmbeddingStateChunkEntity
public int? PageNumber { get; set; }
public int? EndPageNumber { get; set; }
public int ChunkIndex { get; set; }
public string ChunkText { get; set; } = string.Empty;
/// <summary>
/// The hash of the chunk text, cf. ChunkTextHash, or null for a chunk stored before v26.10.1.
/// </summary>
/// <remarks>
/// Stored so that a document read again need not hash every old chunk text before it can tell
/// which vectors to keep. A chunk without it is hashed from its text instead.
/// </remarks>
public string? TextHash { get; set; }
/// <summary>
/// How many times the document was indexed before this chunk was written, cf. IndexedDocumentIds.CreateChunkIdSource.
/// </summary>
/// <remarks>
/// The chunk id includes it, so a document indexed again never overwrites its old chunks with
/// new ones, even where its fingerprint stays the same, as a mail's does. Chunks stored before
/// v26.10.1 are generation 0.
/// </remarks>
public int Generation { get; set; }
/// <summary>
/// The sheet the chunk is from, or null when its document has no sheets or the chunk was stored before v26.10.1.
/// </summary>
/// <remarks>
/// A chunk never runs from one sheet into the next, cf. TextChunker.SplitBySheets. The location
/// columns hold what the runtime states about the pieces of a document. A chunk stored before
/// they existed keeps NULL in all of them, which reads as unknown, and gains them once its
/// document is indexed again.
/// </remarks>
public string? SheetName { get; set; }
/// <summary>
/// The row of its sheet the chunk starts at, counted from one as a spreadsheet program counts it.
/// </summary>
public int? StartRow { get; set; }
/// <summary>
/// The row of its sheet the chunk ends at.
/// </summary>
public int? EndRow { get; set; }
/// <summary>
/// The line the chunk starts at, for text, CSV and source code, counted from one.
/// </summary>
public int? StartLine { get; set; }
/// <summary>
/// The line the chunk ends at.
/// </summary>
public int? EndLine { get; set; }
/// <summary>
/// Which attachment of a mail the chunk is from, as mail_parts numbers the attachments of a
/// mail, or null for the text of the mail itself and for every chunk of a file.
/// </summary>
public int? AttachmentPosition { get; set; }
public DateTimeOffset EmbeddedAtUtc { get; set; }
public EmbeddingStateFileEntity? File { get; set; }
@@ -12,4 +12,6 @@ public sealed record EmbeddingStateFile(
DateTimeOffset LastWriteUtc,
DateTimeOffset EmbeddedAtUtc,
int ChunkCount,
string? TranscriptKey = null);
string? TranscriptKey = null,
string? ContentHash = null,
int Generation = 0);
@@ -28,6 +28,28 @@ internal sealed class EmbeddingStateFileEntity
public string? TranscriptKey { get; set; }
/// <summary>
/// The SHA-256 hash of the raw bytes of the file when it was indexed, cf. FileContentHash, or
/// null for a file indexed before v26.10.1 and for every mail.
/// </summary>
/// <remarks>
/// The fingerprint says only that path, size or write time changed. When it does, the content
/// hash tells whether the content changed as well, or whether the file was merely saved again,
/// synchronized or copied, which needs neither extraction nor embedding.
/// </remarks>
public string? ContentHash { get; set; }
/// <summary>
/// The generation of the chunks the document is searched by, cf. IndexedDocumentIds.CreateChunkIdSource.
/// </summary>
/// <remarks>
/// A document read again writes its chunks under the next generation, while the chunks of this
/// one stay. Only once all new chunks are stored does this switch to the new generation, and
/// the old chunks go. An attempt which fails leaves this as it was, and only its own chunks go,
/// so the document keeps the version it had.
/// </remarks>
public int Generation { get; set; }
public EmbeddingStateDataSourceEntity? DataSource { get; set; }
public List<EmbeddingStateChunkEntity> Chunks { get; set; } = [];
@@ -42,6 +42,69 @@ public abstract partial class IndexStoreClient(string name, string path) : Datab
public abstract Task UpsertChunksAsync(string dataSourceId, IReadOnlyList<EmbeddingStateChunk> chunks, CancellationToken token);
/// <summary>
/// The chunks a document has stored now, each with the hash of its text and its generation.
/// </summary>
/// <remarks>
/// The same chunks DeleteFileAsync removes for the document. A chunk stored before its hash was
/// kept is hashed from its text here, and only such a chunk has its text read at all.
/// </remarks>
/// <param name="dataSourceId">The data source the document belongs to.</param>
/// <param name="filePath">The key of the document, as DeleteFileAsync takes it.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The chunks, in no particular order.</returns>
public abstract Task<IReadOnlyList<StoredChunk>> GetStoredChunksAsync(string dataSourceId, string filePath, CancellationToken token);
/// <summary>
/// The ids of the chunks of some generations of a document.
/// </summary>
/// <param name="dataSourceId">The data source the document belongs to.</param>
/// <param name="filePath">The key of the document, as DeleteFileAsync takes it.</param>
/// <param name="generations">The generations.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The chunk ids, in no particular order.</returns>
public abstract Task<IReadOnlyList<string>> GetChunkIdsAsync(string dataSourceId, string filePath, IReadOnlyCollection<int> generations, CancellationToken token);
/// <summary>
/// Removes the chunks of every generation of a document but the given ones.
/// </summary>
/// <remarks>
/// Keeping the new generation completes a document read again. Keeping the stored one, and an
/// attempt which failed but holds the vectors the document needs, makes room for the next attempt.
/// </remarks>
/// <param name="dataSourceId">The data source the document belongs to.</param>
/// <param name="filePath">The key of the document, as DeleteFileAsync takes it.</param>
/// <param name="generations">The generations which stay.</param>
/// <param name="token">The cancellation token.</param>
public abstract Task DeleteOtherGenerationsAsync(string dataSourceId, string filePath, IReadOnlyCollection<int> generations, CancellationToken token);
/// <summary>
/// Where some chunks of a data source are in their documents, by their chunk ids.
/// </summary>
/// <remarks>
/// The vector store keeps only the pages of a chunk, so its hits learn the rest from here.
/// </remarks>
/// <param name="dataSourceId">The data source.</param>
/// <param name="chunkIds">The ids of the chunks.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The location of every chunk the index store knows, by its id.</returns>
public abstract Task<IReadOnlyDictionary<string, ChunkLocation>> GetChunkLocationsAsync(string dataSourceId, IReadOnlyCollection<string> chunkIds, CancellationToken token);
/// <summary>
/// The chunks of a data source which no search may find, document by document.
/// </summary>
/// <remarks>
/// A document is searched by the generation its row names, once its chunk count says that
/// generation is complete. Every other chunk of it is hidden: a generation which is still being
/// written, an attempt which failed and stays for the vectors of the next one, and all chunks of
/// a document whose first indexing is not complete. The keyword search applies the same rule
/// itself, cf. SearchChunksAsync. Only documents with such chunks are listed, which are few.
/// </remarks>
/// <param name="dataSourceId">The data source.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>One entry per document with hidden chunks, in no particular order.</returns>
public abstract Task<IReadOnlyList<HiddenChunks>> GetHiddenChunksAsync(string dataSourceId, CancellationToken token);
public abstract Task<IReadOnlyList<IndexStoreSearchResult>> SearchChunksAsync(string dataSourceId, string query, int maxMatches, CancellationToken token);
public abstract Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token);
@@ -81,6 +81,8 @@ internal sealed class IndexStoreDbContext(DbContextOptions<IndexStoreDbContext>
entity.Property(file => file.EmbeddedAtUtc).HasColumnName("embedded_at_utc").HasConversion(utcDateTimeOffsetConverter).IsRequired();
entity.Property(file => file.ChunkCount).HasColumnName("chunk_count");
entity.Property(file => file.TranscriptKey).HasColumnName("transcript_key");
entity.Property(file => file.ContentHash).HasColumnName("content_hash");
entity.Property(file => file.Generation).HasColumnName("generation");
entity.HasIndex(file => file.DataSourceId).HasDatabaseName("idx_embedded_files_data_source");
entity.HasIndex(file => file.AbsolutePath).HasDatabaseName("idx_embedded_files_absolute_path");
@@ -103,14 +105,23 @@ internal sealed class IndexStoreDbContext(DbContextOptions<IndexStoreDbContext>
entity.Property(chunk => chunk.ChunkId).HasColumnName("chunk_id").IsRequired();
entity.Property(chunk => chunk.ParentFileId).HasColumnName("parent_file_id").IsRequired();
entity.Property(chunk => chunk.PageNumber).HasColumnName("page_number");
entity.Property(chunk => chunk.EndPageNumber).HasColumnName("end_page_number");
entity.Property(chunk => chunk.ChunkIndex).HasColumnName("chunk_index");
entity.Property(chunk => chunk.ChunkText).HasColumnName("chunk_text").IsRequired();
entity.Property(chunk => chunk.TextHash).HasColumnName("text_hash");
entity.Property(chunk => chunk.Generation).HasColumnName("generation");
entity.Property(chunk => chunk.SheetName).HasColumnName("sheet_name");
entity.Property(chunk => chunk.StartRow).HasColumnName("start_row");
entity.Property(chunk => chunk.EndRow).HasColumnName("end_row");
entity.Property(chunk => chunk.StartLine).HasColumnName("start_line");
entity.Property(chunk => chunk.EndLine).HasColumnName("end_line");
entity.Property(chunk => chunk.AttachmentPosition).HasColumnName("attachment_position");
entity.Property(chunk => chunk.EmbeddedAtUtc).HasColumnName("embedded_at_utc").HasConversion(utcDateTimeOffsetConverter).IsRequired();
entity.HasIndex(chunk => chunk.ChunkId).HasDatabaseName("idx_embedding_chunks_chunk_id").IsUnique();
entity.HasIndex(chunk => chunk.ParentFileId).HasDatabaseName("idx_embedding_chunks_parent_file");
entity.HasIndex(chunk => chunk.PageNumber).HasDatabaseName("idx_embedding_chunks_page");
entity.HasIndex(chunk => new { chunk.ParentFileId, chunk.ChunkIndex }).HasDatabaseName("idx_embedding_chunks_parent_file_chunk_index").IsUnique();
entity.HasIndex(chunk => new { chunk.ParentFileId, chunk.Generation, chunk.ChunkIndex }).HasDatabaseName("idx_embedding_chunks_parent_file_generation_chunk_index").IsUnique();
});
modelBuilder.Entity<IndexingFailureEntity>(entity =>
@@ -12,6 +12,7 @@ internal static class IndexStoreSchemaMigrator
[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(Migrations.DropDataSourceName))]
[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(Migrations.MailboxIndex))]
[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(Migrations.MediaTranscriptKey))]
[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(Migrations.ChunkMetadata))]
public static async Task MigrateAsync(IndexStoreDbContext context, CancellationToken token)
{
await context.Database.MigrateAsync(token);
@@ -10,6 +10,7 @@ public sealed record IndexStoreSearchResult(
string RelativePath,
string FileType,
int? PageNumber,
int? EndPageNumber,
int ChunkIndex,
string ChunkText,
double Score,
@@ -18,4 +19,16 @@ public sealed record IndexStoreSearchResult(
DateTimeOffset CreationUtc,
DateTimeOffset LastWriteUtc,
DateTimeOffset EmbeddedAtUtc,
int ChunkCount);
int ChunkCount,
string? SheetName = null,
int? StartRow = null,
int? EndRow = null,
int? StartLine = null,
int? EndLine = null,
int? AttachmentPosition = null)
{
/// <summary>
/// Where the chunk is in its document.
/// </summary>
public ChunkLocation Location => new(this.PageNumber, this.EndPageNumber, this.SheetName, this.StartRow, this.EndRow, this.StartLine, this.EndLine, this.AttachmentPosition);
}
@@ -20,6 +20,20 @@ internal sealed class IndexStoreSearchResultEntity
public int? PageNumber { get; set; }
public int? EndPageNumber { get; set; }
public string? SheetName { get; set; }
public int? StartRow { get; set; }
public int? EndRow { get; set; }
public int? StartLine { get; set; }
public int? EndLine { get; set; }
public int? AttachmentPosition { get; set; }
public int ChunkIndex { get; set; }
public string ChunkText { get; set; } = string.Empty;
@@ -0,0 +1,83 @@
#nullable disable
using Microsoft.EntityFrameworkCore.Infrastructure;
using Microsoft.EntityFrameworkCore.Migrations;
namespace AIStudio.Tools.Databases.IndexStore.Migrations;
/// <summary>
/// Adds what a chunk and its file know beyond their text: where a chunk is, the hash of its text,
/// its generation, and the hash of the content of its file.
/// </summary>
/// <remarks>
/// Each chunk gains:
/// <list type="bullet">
/// <item>the page it ends on, next to the page it starts on, and its sheet and rows, its lines, and
/// the attachment of a mail it is from, cf. SourceLocation;</item>
/// <item>the hash of its text, so a document read again keeps the vector of every chunk whose text
/// did not change, cf. ChunkTextHash;</item>
/// <item>its generation, so a document read again gets its new chunks next to the old ones and
/// keeps the old ones until the new ones are complete, cf. IndexedDocumentIds.CreateChunkIdSource.
/// Two generations may hold the same position, so the unique index over parent and position
/// gains the generation as well.</item>
/// </list>
/// Each file gains the hash of its content, so a file saved again without changes needs neither
/// extraction nor embedding, and the generation it is searched by.<br/><br/>
/// Chunks and files stored before keep NULL, which reads as unknown, and generation 0, which is
/// the generation every one of them has: an index built before has one generation per document.
/// They gain the rest once their document is indexed again, so no index has to be rebuilt, cf.
/// DataSourceEmbeddingService.CHUNK_METADATA_VERSION.
/// </remarks>
[DbContext(typeof(IndexStoreDbContext))]
[Migration("20261022000000_ChunkMetadata")]
public partial class ChunkMetadata : Migration
{
/// <remarks>
/// Both directions go through raw SQL, like MediaTranscriptKey: the SQLite provider answers
/// DropColumn by rebuilding the table, which would drop the triggers that keep the full-text
/// index of the chunks in step, and dropping the old embedded_files table would let the cascade
/// of the foreign key in embedding_chunks take every chunk with it. Adding a column and
/// replacing an index are native in SQLite and leave the rows alone.
/// </remarks>
protected override void Up(MigrationBuilder migrationBuilder)
{
migrationBuilder.Sql("ALTER TABLE embedding_chunks ADD COLUMN end_page_number INTEGER NULL;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks ADD COLUMN sheet_name TEXT NULL;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks ADD COLUMN start_row INTEGER NULL;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks ADD COLUMN end_row INTEGER NULL;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks ADD COLUMN start_line INTEGER NULL;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks ADD COLUMN end_line INTEGER NULL;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks ADD COLUMN attachment_position INTEGER NULL;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks ADD COLUMN text_hash TEXT NULL;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks ADD COLUMN generation INTEGER NOT NULL DEFAULT 0;");
migrationBuilder.Sql("ALTER TABLE embedded_files ADD COLUMN content_hash TEXT NULL;");
migrationBuilder.Sql("ALTER TABLE embedded_files ADD COLUMN generation INTEGER NOT NULL DEFAULT 0;");
migrationBuilder.Sql("DROP INDEX IF EXISTS idx_embedding_chunks_parent_file_chunk_index;");
migrationBuilder.Sql("CREATE UNIQUE INDEX idx_embedding_chunks_parent_file_generation_chunk_index ON embedding_chunks (parent_file_id, generation, chunk_index);");
}
/// <remarks>
/// Going back needs every document to have a single generation again, as it has whenever no
/// indexing run is under way and no failed attempt is kept.
/// </remarks>
protected override void Down(MigrationBuilder migrationBuilder)
{
migrationBuilder.Sql("DROP INDEX IF EXISTS idx_embedding_chunks_parent_file_generation_chunk_index;");
migrationBuilder.Sql("CREATE UNIQUE INDEX idx_embedding_chunks_parent_file_chunk_index ON embedding_chunks (parent_file_id, chunk_index);");
migrationBuilder.Sql("ALTER TABLE embedded_files DROP COLUMN generation;");
migrationBuilder.Sql("ALTER TABLE embedded_files DROP COLUMN content_hash;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks DROP COLUMN generation;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks DROP COLUMN text_hash;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks DROP COLUMN attachment_position;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks DROP COLUMN end_line;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks DROP COLUMN start_line;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks DROP COLUMN end_row;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks DROP COLUMN start_row;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks DROP COLUMN sheet_name;");
migrationBuilder.Sql("ALTER TABLE embedding_chunks DROP COLUMN end_page_number;");
}
}
@@ -75,6 +75,10 @@ partial class IndexStoreDbContextModelSnapshot : ModelSnapshot
.HasColumnType("INTEGER")
.HasColumnName("chunk_count");
entity.Property<string>("ContentHash")
.HasColumnType("TEXT")
.HasColumnName("content_hash");
entity.Property<DateTimeOffset>("CreationUtc")
.HasConversion(utcDateTimeOffsetConverter)
.HasColumnType("TEXT")
@@ -109,6 +113,10 @@ partial class IndexStoreDbContextModelSnapshot : ModelSnapshot
.HasColumnType("TEXT")
.HasColumnName("fingerprint");
entity.Property<int>("Generation")
.HasColumnType("INTEGER")
.HasColumnName("generation");
entity.Property<DateTimeOffset>("LastWriteUtc")
.HasConversion(utcDateTimeOffsetConverter)
.HasColumnType("TEXT")
@@ -149,6 +157,10 @@ partial class IndexStoreDbContextModelSnapshot : ModelSnapshot
.HasColumnName("id")
.HasAnnotation("Sqlite:Autoincrement", true);
entity.Property<int?>("AttachmentPosition")
.HasColumnType("INTEGER")
.HasColumnName("attachment_position");
entity.Property<string>("ChunkId")
.IsRequired()
.HasColumnType("TEXT")
@@ -168,6 +180,22 @@ partial class IndexStoreDbContextModelSnapshot : ModelSnapshot
.HasColumnType("TEXT")
.HasColumnName("embedded_at_utc");
entity.Property<int?>("EndLine")
.HasColumnType("INTEGER")
.HasColumnName("end_line");
entity.Property<int?>("EndPageNumber")
.HasColumnType("INTEGER")
.HasColumnName("end_page_number");
entity.Property<int?>("EndRow")
.HasColumnType("INTEGER")
.HasColumnName("end_row");
entity.Property<int>("Generation")
.HasColumnType("INTEGER")
.HasColumnName("generation");
entity.Property<int?>("PageNumber")
.HasColumnType("INTEGER")
.HasColumnName("page_number");
@@ -177,6 +205,22 @@ partial class IndexStoreDbContextModelSnapshot : ModelSnapshot
.HasColumnType("TEXT")
.HasColumnName("parent_file_id");
entity.Property<string>("SheetName")
.HasColumnType("TEXT")
.HasColumnName("sheet_name");
entity.Property<int?>("StartLine")
.HasColumnType("INTEGER")
.HasColumnName("start_line");
entity.Property<int?>("StartRow")
.HasColumnType("INTEGER")
.HasColumnName("start_row");
entity.Property<string>("TextHash")
.HasColumnType("TEXT")
.HasColumnName("text_hash");
entity.HasKey("Id");
entity.HasIndex("ChunkId")
@@ -189,9 +233,9 @@ partial class IndexStoreDbContextModelSnapshot : ModelSnapshot
entity.HasIndex("ParentFileId")
.HasDatabaseName("idx_embedding_chunks_parent_file");
entity.HasIndex("ParentFileId", "ChunkIndex")
entity.HasIndex("ParentFileId", "Generation", "ChunkIndex")
.IsUnique()
.HasDatabaseName("idx_embedding_chunks_parent_file_chunk_index");
.HasDatabaseName("idx_embedding_chunks_parent_file_generation_chunk_index");
entity.ToTable("embedding_chunks");
});
@@ -251,6 +295,9 @@ partial class IndexStoreDbContextModelSnapshot : ModelSnapshot
.IsRequired()
.HasColumnType("TEXT");
entity.Property<int?>("AttachmentPosition")
.HasColumnType("INTEGER");
entity.Property<int>("ChunkCount")
.HasColumnType("INTEGER");
@@ -281,6 +328,15 @@ partial class IndexStoreDbContextModelSnapshot : ModelSnapshot
.HasConversion(utcDateTimeOffsetConverter)
.HasColumnType("TEXT");
entity.Property<int?>("EndLine")
.HasColumnType("INTEGER");
entity.Property<int?>("EndPageNumber")
.HasColumnType("INTEGER");
entity.Property<int?>("EndRow")
.HasColumnType("INTEGER");
entity.Property<string>("FileName")
.IsRequired()
.HasColumnType("TEXT");
@@ -314,6 +370,15 @@ partial class IndexStoreDbContextModelSnapshot : ModelSnapshot
entity.Property<double>("Score")
.HasColumnType("REAL");
entity.Property<string>("SheetName")
.HasColumnType("TEXT");
entity.Property<int?>("StartLine")
.HasColumnType("INTEGER");
entity.Property<int?>("StartRow")
.HasColumnType("INTEGER");
entity.HasNoKey();
entity.ToView("embedding_chunk_search_results");
@@ -63,6 +63,16 @@ public sealed partial class NoIndexStoreClient(string name, string? unavailableR
public override Task UpsertChunksAsync(string dataSourceId, IReadOnlyList<EmbeddingStateChunk> chunks, CancellationToken token) => Task.CompletedTask;
public override Task<IReadOnlyList<StoredChunk>> GetStoredChunksAsync(string dataSourceId, string filePath, CancellationToken token) => Task.FromResult<IReadOnlyList<StoredChunk>>([]);
public override Task<IReadOnlyList<string>> GetChunkIdsAsync(string dataSourceId, string filePath, IReadOnlyCollection<int> generations, CancellationToken token) => Task.FromResult<IReadOnlyList<string>>([]);
public override Task DeleteOtherGenerationsAsync(string dataSourceId, string filePath, IReadOnlyCollection<int> generations, CancellationToken token) => Task.CompletedTask;
public override Task<IReadOnlyDictionary<string, ChunkLocation>> GetChunkLocationsAsync(string dataSourceId, IReadOnlyCollection<string> chunkIds, CancellationToken token) => Task.FromResult<IReadOnlyDictionary<string, ChunkLocation>>(new Dictionary<string, ChunkLocation>(StringComparer.Ordinal));
public override Task<IReadOnlyList<HiddenChunks>> GetHiddenChunksAsync(string dataSourceId, CancellationToken token) => Task.FromResult<IReadOnlyList<HiddenChunks>>([]);
public override Task<IReadOnlyList<IndexStoreSearchResult>> SearchChunksAsync(string dataSourceId, string query, int maxMatches, CancellationToken token) =>
Task.FromResult<IReadOnlyList<IndexStoreSearchResult>>([]);
@@ -35,8 +35,8 @@ public sealed partial class SqliteIndexStoreClientImplementation
{
await using var context = this.CreateContext();
var mailIds = FilterMails(context, dataSourceId, filter).Select(mail => mail.ParentFileId);
return await context.EmbeddingChunks
.AsNoTracking()
return await SearchableChunks(context)
.Where(chunk => mailIds.Contains(chunk.ParentFileId))
.Select(chunk => chunk.ChunkId)
.ToListAsync(token);
@@ -177,6 +177,13 @@ public sealed partial class SqliteIndexStoreClientImplementation
.FirstOrDefaultAsync(token);
}
/// <summary>
/// The chunks a search may find: those of the generation their document is searched by, cf. GetHiddenChunksAsync.
/// </summary>
private static IQueryable<EmbeddingStateChunkEntity> SearchableChunks(IndexStoreDbContext context) => context.EmbeddingChunks
.AsNoTracking()
.Where(chunk => context.EmbeddedFiles.Any(file => file.ParentFileId == chunk.ParentFileId && file.ChunkCount > 0 && file.Generation == chunk.Generation));
/// <summary>
/// The mails of a mailbox which meet the conditions, as a query to build on.
/// </summary>
@@ -7,7 +7,6 @@ using AIStudio.Tools.Mail;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.Services;
using Microsoft.Data.Sqlite;
using Microsoft.EntityFrameworkCore;
namespace AIStudio.Tools.Databases.IndexStore;
@@ -151,7 +150,9 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
file.ChunkCount,
file.TranscriptKey,
file.RelativePath,
file.ParentFileId);
file.ParentFileId,
file.ContentHash,
file.Generation);
}
var permanentFailures = await context.PermanentIndexingFailures
@@ -336,6 +337,110 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
await transaction.CommitAsync(token);
}
public override async Task<IReadOnlyList<StoredChunk>> GetStoredChunksAsync(string dataSourceId, string filePath, CancellationToken token)
{
await using var context = this.CreateContext();
var parentFileIds = context.EmbeddedFiles
.Where(file => file.DataSourceId == dataSourceId && file.AbsolutePath == filePath)
.Select(file => file.ParentFileId);
//
// Only a description of the query: the rows are read in the loop below, one at a time. A
// chunk without a stored hash brings its whole text along, and a large document has
// thousands of them, so they must never be in memory together. A chunk with a hash brings
// no text at all.
//
var chunkQuery = context.EmbeddingChunks
.AsNoTracking()
.Where(chunk => parentFileIds.Contains(chunk.ParentFileId))
.Select(chunk => new { chunk.ChunkId, chunk.TextHash, chunk.Generation, ChunkText = chunk.TextHash == null ? chunk.ChunkText : null })
.AsAsyncEnumerable();
var storedChunks = new List<StoredChunk>();
await foreach (var chunk in chunkQuery.WithCancellation(token))
storedChunks.Add(new(chunk.ChunkId, chunk.TextHash ?? ChunkTextHash.Of(chunk.ChunkText ?? string.Empty), chunk.Generation));
return storedChunks;
}
public override async Task<IReadOnlyList<string>> GetChunkIdsAsync(string dataSourceId, string filePath, IReadOnlyCollection<int> generations, CancellationToken token)
{
await using var context = this.CreateContext();
var parentFileIds = context.EmbeddedFiles
.Where(file => file.DataSourceId == dataSourceId && file.AbsolutePath == filePath)
.Select(file => file.ParentFileId);
return await context.EmbeddingChunks
.Where(chunk => parentFileIds.Contains(chunk.ParentFileId) && generations.Contains(chunk.Generation))
.Select(chunk => chunk.ChunkId)
.ToListAsync(token);
}
public override async Task DeleteOtherGenerationsAsync(string dataSourceId, string filePath, IReadOnlyCollection<int> generations, CancellationToken token)
{
await using var context = this.CreateContext();
var parentFileIds = context.EmbeddedFiles
.Where(file => file.DataSourceId == dataSourceId && file.AbsolutePath == filePath)
.Select(file => file.ParentFileId);
await context.EmbeddingChunks
.Where(chunk => parentFileIds.Contains(chunk.ParentFileId) && !generations.Contains(chunk.Generation))
.ExecuteDeleteAsync(token);
}
public override async Task<IReadOnlyDictionary<string, ChunkLocation>> GetChunkLocationsAsync(string dataSourceId, IReadOnlyCollection<string> chunkIds, CancellationToken token)
{
if (chunkIds.Count is 0)
return new Dictionary<string, ChunkLocation>(StringComparer.Ordinal);
await using var context = this.CreateContext();
var parentFileIds = context.EmbeddedFiles
.Where(file => file.DataSourceId == dataSourceId)
.Select(file => file.ParentFileId);
// The locations only, never the texts of the chunks:
return await context.EmbeddingChunks
.AsNoTracking()
.Where(chunk => chunkIds.Contains(chunk.ChunkId) && parentFileIds.Contains(chunk.ParentFileId))
.Select(chunk => new
{
chunk.ChunkId,
Location = new ChunkLocation(chunk.PageNumber, chunk.EndPageNumber, chunk.SheetName, chunk.StartRow, chunk.EndRow, chunk.StartLine, chunk.EndLine, chunk.AttachmentPosition),
})
.ToDictionaryAsync(chunk => chunk.ChunkId, chunk => chunk.Location, StringComparer.Ordinal, token);
}
/// <inheritdoc />
/// <remarks>
/// On a thread of its own, cf. the Databases AGENTS.md: it runs before every vector search,
/// which a message in the chat may start, and it looks at every document of the data source.
/// </remarks>
public override Task<IReadOnlyList<HiddenChunks>> GetHiddenChunksAsync(string dataSourceId, CancellationToken token) => Task.Run<IReadOnlyList<HiddenChunks>>(async () =>
{
await using var context = this.CreateContext();
var documents = await DocumentsWithHiddenChunks(context, dataSourceId)
.Select(file => new { file.ParentFileId, file.ChunkCount, file.Generation })
.ToListAsync(token);
return documents.Select(document => new HiddenChunks(document.ParentFileId, document.ChunkCount > 0 ? document.Generation : null)).ToList();
}, token);
/// <summary>
/// The documents of a data source with chunks a search must not find: those whose row names no
/// generation yet, and those with chunks of another generation than the one their row names.
/// </summary>
/// <remarks>
/// Asked before every vector search. For a document without hidden chunks, which is nearly
/// every document, this costs two lookups in the index over parent, generation and position:
/// one below the generation of its row and one above it.
/// </remarks>
private static IQueryable<EmbeddingStateFileEntity> DocumentsWithHiddenChunks(IndexStoreDbContext context, string dataSourceId) => context.EmbeddedFiles
.AsNoTracking()
.Where(file => file.DataSourceId == dataSourceId)
.Where(file => file.ChunkCount <= 0
|| context.EmbeddingChunks.Any(chunk => chunk.ParentFileId == file.ParentFileId && chunk.Generation < file.Generation)
|| context.EmbeddingChunks.Any(chunk => chunk.ParentFileId == file.ParentFileId && chunk.Generation > file.Generation));
public override async Task<IReadOnlyList<IndexStoreSearchResult>> SearchChunksAsync(string dataSourceId, string query, int maxMatches, CancellationToken token)
{
if (maxMatches <= 0)
@@ -358,7 +463,9 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
/// </summary>
/// <remarks>
/// Kept free of ORDER BY and LIMIT, so EF Core can wrap it and filter it further before either
/// of them applies, cf. InSearchOrder.
/// of them applies, cf. InSearchOrder.<br/><br/>
/// Only the generation the row of a document names can be found, and only once the row says it
/// is complete. These are the same chunks the vector search finds, cf. GetHiddenChunksAsync.
/// </remarks>
private static IQueryable<IndexStoreSearchResultEntity> MatchChunks(IndexStoreDbContext context, string dataSourceId, string ftsQuery) => context.SearchResults
.FromSqlInterpolated($"""
@@ -372,6 +479,13 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
f.relative_path AS RelativePath,
f.file_type AS FileType,
c.page_number AS PageNumber,
c.end_page_number AS EndPageNumber,
c.sheet_name AS SheetName,
c.start_row AS StartRow,
c.end_row AS EndRow,
c.start_line AS StartLine,
c.end_line AS EndLine,
c.attachment_position AS AttachmentPosition,
c.chunk_index AS ChunkIndex,
c.chunk_text AS ChunkText,
bm25(embedding_chunks_fts) AS Score,
@@ -386,6 +500,8 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
JOIN embedded_files f ON f.parent_file_id = c.parent_file_id
JOIN data_sources ds ON ds.data_source_id = f.data_source_id
WHERE ds.data_source_id = {dataSourceId}
AND f.chunk_count > 0
AND c.generation = f.generation
AND embedding_chunks_fts MATCH {ftsQuery}
""")
.AsNoTracking();
@@ -609,70 +725,14 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
/// <inheritdoc />
/// <remarks>
/// VACUUM INTO writes a consistent snapshot of the database, including what the write-ahead log
/// still holds, and in WAL mode, reading it holds back no writer.<br/><br/>
///
/// It runs on connections of their own, although the databases guide has raw SQL go through
/// GetDbConnection. The connections of the context share one cache, and a shared cache locks
/// whole tables: reading the whole database for the copy would hold back every writer until the
/// copy is done. Without pooling, nothing keeps the copy open afterward, so the backup can
/// delete it on every OS.<br/><br/>
///
/// SQLite runs its async methods synchronously, and copying and checking a large index takes a
/// while. The copy therefore runs on a thread of its own, so that it never holds up its caller.
/// A consistent snapshot which holds back no writer, made on a thread of its own, cf. SqliteDatabaseCopy.
/// </remarks>
public override Task<string> BackupToDirectoryAsync(string targetDirectory, CancellationToken token) => Task.Run(() => this.CopyAndCheckAsync(targetDirectory, token), token);
private async Task<string> CopyAndCheckAsync(string targetDirectory, CancellationToken token)
public override Task<string> BackupToDirectoryAsync(string targetDirectory, CancellationToken token) => Task.Run(async () =>
{
Directory.CreateDirectory(targetDirectory);
var targetPath = Path.Join(targetDirectory, DATABASE_FILENAME);
// VACUUM INTO refuses to write into a file which holds anything:
if (File.Exists(targetPath))
File.Delete(targetPath);
try
{
await using (var connection = new SqliteConnection(BuildBackupConnectionString(this.databasePath, SqliteOpenMode.ReadWrite)))
{
await connection.OpenAsync(token);
await using var command = connection.CreateCommand();
command.CommandText = "VACUUM INTO $targetPath;";
command.Parameters.AddWithValue("$targetPath", targetPath);
await command.ExecuteNonQueryAsync(token);
}
await using (var connection = new SqliteConnection(BuildBackupConnectionString(targetPath, SqliteOpenMode.ReadOnly)))
{
await connection.OpenAsync(token);
await using var command = connection.CreateCommand();
command.CommandText = "PRAGMA quick_check;";
var checkResult = (await command.ExecuteScalarAsync(token))?.ToString();
if (checkResult is not "ok")
throw new InvalidOperationException($"The copy of the index at '{targetPath}' failed its check: '{checkResult}'.");
}
return targetPath;
}
catch
{
if (File.Exists(targetPath))
File.Delete(targetPath);
throw;
}
}
private static string BuildBackupConnectionString(string path, SqliteOpenMode mode) => new SqliteConnectionStringBuilder
{
DataSource = path,
Mode = mode,
Cache = SqliteCacheMode.Private,
Pooling = false,
DefaultTimeout = 30,
}.ToString();
await SqliteDatabaseCopy.CopyAndCheckAsync(this.databasePath, targetPath, token);
return targetPath;
}, token);
private IndexStoreDbContext CreateContext() => new(this.dbContextOptions);
@@ -706,6 +766,8 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
fileEntity.EmbeddedAtUtc = file.EmbeddedAtUtc;
fileEntity.ChunkCount = file.ChunkCount;
fileEntity.TranscriptKey = file.TranscriptKey;
fileEntity.ContentHash = file.ContentHash;
fileEntity.Generation = file.Generation;
}
private static void ApplyPermanentFailure(IndexingFailureEntity failureEntity, string dataSourceId, PermanentIndexingFailure failure)
@@ -731,8 +793,17 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
chunkEntity.ChunkId = chunk.ChunkId;
chunkEntity.ParentFileId = chunk.ParentFileId;
chunkEntity.PageNumber = chunk.PageNumber;
chunkEntity.EndPageNumber = chunk.EndPageNumber;
chunkEntity.ChunkIndex = chunk.ChunkIndex;
chunkEntity.ChunkText = chunk.ChunkText;
chunkEntity.TextHash = chunk.TextHash;
chunkEntity.Generation = chunk.Generation;
chunkEntity.SheetName = chunk.SheetName;
chunkEntity.StartRow = chunk.StartRow;
chunkEntity.EndRow = chunk.EndRow;
chunkEntity.StartLine = chunk.StartLine;
chunkEntity.EndLine = chunk.EndLine;
chunkEntity.AttachmentPosition = chunk.AttachmentPosition;
chunkEntity.EmbeddedAtUtc = chunk.EmbeddedAtUtc;
}
@@ -746,6 +817,7 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
result.RelativePath,
result.FileType,
result.PageNumber,
result.EndPageNumber,
result.ChunkIndex,
result.ChunkText,
result.Score,
@@ -754,7 +826,13 @@ public sealed partial class SqliteIndexStoreClientImplementation(string name, st
result.CreationUtc,
result.LastWriteUtc,
result.EmbeddedAtUtc,
result.ChunkCount);
result.ChunkCount,
result.SheetName,
result.StartRow,
result.EndRow,
result.StartLine,
result.EndLine,
result.AttachmentPosition);
private static string BuildFtsQuery(string query)
{
@@ -0,0 +1,9 @@
namespace AIStudio.Tools.Databases.IndexStore;
/// <summary>
/// A chunk the index store holds, as far as indexing its document again needs to know it.
/// </summary>
/// <param name="ChunkId">The id of the chunk, which is also the id of its point in the vector store.</param>
/// <param name="TextHash">The hash of its text, cf. ChunkTextHash.</param>
/// <param name="Generation">How many times its document was indexed before it, cf. IndexedDocumentIds.CreateChunkIdSource.</param>
public sealed record StoredChunk(string ChunkId, string TextHash, int Generation);
@@ -0,0 +1,43 @@
#nullable disable
using Microsoft.EntityFrameworkCore.Infrastructure;
using Microsoft.EntityFrameworkCore.Migrations;
namespace AIStudio.Tools.Databases.ProvenanceStore.Migrations;
/// <summary>
/// Creates the table provenance, which holds what chats received besides their messages.
/// </summary>
[DbContext(typeof(ProvenanceStoreDbContext))]
[Migration("20261024000000_InitialProvenanceStore")]
public partial class InitialProvenanceStore : Migration
{
protected override void Up(MigrationBuilder migrationBuilder)
{
migrationBuilder.CreateTable(
name: "provenance",
columns: table => new
{
chat_id = table.Column<string>(type: "TEXT", nullable: false),
id = table.Column<string>(type: "TEXT", nullable: false),
kind = table.Column<string>(type: "TEXT", nullable: false),
content = table.Column<string>(type: "TEXT", nullable: false),
surrounding_content = table.Column<string>(type: "TEXT", nullable: true),
created_at_utc = table.Column<string>(type: "TEXT", nullable: false),
},
constraints: table =>
{
table.PrimaryKey("PK_provenance", entry => new { entry.chat_id, entry.id });
});
migrationBuilder.CreateIndex(
name: "idx_provenance_id",
table: "provenance",
column: "id");
}
protected override void Down(MigrationBuilder migrationBuilder)
{
migrationBuilder.DropTable(name: "provenance");
}
}
@@ -0,0 +1,60 @@
#nullable disable
using AIStudio.Tools.Databases.IndexStore;
using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Infrastructure;
namespace AIStudio.Tools.Databases.ProvenanceStore.Migrations;
// EF Core writes this file and declares the type as partial. Dropping the keyword would only
// last until the next migration is added, so the inspection is silenced instead.
// ReSharper disable once PartialTypeWithSinglePart
[DbContext(typeof(ProvenanceStoreDbContext))]
partial class ProvenanceStoreDbContextModelSnapshot : ModelSnapshot
{
protected override void BuildModel(ModelBuilder modelBuilder)
{
#pragma warning disable 612, 618
modelBuilder.HasAnnotation("ProductVersion", "9.0.18");
var utcDateTimeOffsetConverter = new IndexStoreDateTimeOffsetConverter();
modelBuilder.Entity("AIStudio.Tools.Databases.ProvenanceStore.ProvenanceEntity", entity =>
{
entity.Property<string>("ChatId")
.HasColumnType("TEXT")
.HasColumnName("chat_id");
entity.Property<string>("Id")
.HasColumnType("TEXT")
.HasColumnName("id");
entity.Property<string>("Content")
.IsRequired()
.HasColumnType("TEXT")
.HasColumnName("content");
entity.Property<DateTimeOffset>("CreatedAtUtc")
.HasConversion(utcDateTimeOffsetConverter)
.HasColumnType("TEXT")
.HasColumnName("created_at_utc");
entity.Property<string>("Kind")
.IsRequired()
.HasColumnType("TEXT")
.HasColumnName("kind");
entity.Property<string>("SurroundingContent")
.HasColumnType("TEXT")
.HasColumnName("surrounding_content");
entity.HasKey("ChatId", "Id");
entity.HasIndex("Id")
.HasDatabaseName("idx_provenance_id");
entity.ToTable("provenance");
});
#pragma warning restore 612, 618
}
}
@@ -0,0 +1,69 @@
using AIStudio.Tools.Databases.IndexStore;
using AIStudio.Tools.PluginSystem;
namespace AIStudio.Tools.Databases.ProvenanceStore;
/// <summary>
/// Stands in for the provenance store while it cannot serve: it keeps nothing and finds nothing.
/// </summary>
/// <remarks>
/// Storing answers with no id, which tells the caller to keep the content in the chat itself.
/// </remarks>
public sealed class NoProvenanceStoreClient(string name, string? unavailableReason, DatabaseClientStatus status = DatabaseClientStatus.UNAVAILABLE) : ProvenanceStoreClient(name, string.Empty)
{
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoProvenanceStoreClient).Namespace, nameof(NoProvenanceStoreClient));
public override DatabaseClientStatus Status => status;
public override async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo()
{
yield return (TB("Status"), status switch
{
DatabaseClientStatus.STARTING => TB("Starting"),
_ => TB("Unavailable")
});
if (!string.IsNullOrWhiteSpace(unavailableReason))
yield return (TB("Reason"), unavailableReason);
// As in NoIndexStoreClient: which library this process bound to matters most when SQLite cannot be loaded at all.
yield return (TB("Native library"), OrUnknown(SqliteRuntimeInfo.GetNativeLibraryName()));
yield return (TB("Process architecture"), OrUnknown(SqliteRuntimeInfo.GetProcessArchitecture()));
await Task.CompletedTask;
}
private static string OrUnknown(string value) => string.IsNullOrWhiteSpace(value) ? TB("unknown") : value;
public override Task<Guid?> StoreAsync(Guid chatId, ProvenanceKind kind, string content, IReadOnlyList<string>? surroundingContent, CancellationToken token) => Task.FromResult<Guid?>(null);
public override Task<IReadOnlyDictionary<Guid, ProvenanceEntry>> GetEntriesAsync(IReadOnlyCollection<Guid> ids, CancellationToken token) => Task.FromResult<IReadOnlyDictionary<Guid, ProvenanceEntry>>(new Dictionary<Guid, ProvenanceEntry>());
public override Task CopyChatAsync(Guid sourceChatId, Guid targetChatId, CancellationToken token) => Task.CompletedTask;
public override Task<int> DeleteChatsAsync(IReadOnlyCollection<Guid> chatIds, CancellationToken token, DateTimeOffset? createdBefore = null) => Task.FromResult(0);
public override Task<int> DeleteChatEntriesExceptAsync(Guid chatId, IReadOnlyCollection<Guid> keptIds, CancellationToken token, DateTimeOffset? createdBefore = null) => Task.FromResult(0);
public override Task<int> DeleteEntriesAsync(Guid chatId, IReadOnlyCollection<Guid> ids, CancellationToken token) => Task.FromResult(0);
public override Task<IReadOnlyList<Guid>> GetChatIdsAsync(CancellationToken token, DateTimeOffset? createdSince = null) => Task.FromResult<IReadOnlyList<Guid>>([]);
public override Task<DateTimeOffset?> GetLastCleanupAsync(CancellationToken token) => Task.FromResult<DateTimeOffset?>(null);
public override Task SetLastCleanupAsync(DateTimeOffset startedAt, CancellationToken token) => Task.CompletedTask;
/// <inheritdoc />
/// <remarks>
/// Unlike the other methods, this one does not pretend success: a backup would take an empty
/// copy for the real one. Whoever makes the backup decides what to do instead.
/// </remarks>
public override Task<string> BackupToDirectoryAsync(string targetDirectory, CancellationToken token) =>
Task.FromException<string>(new InvalidOperationException(unavailableReason ?? "The provenance store is not available."));
public override Task<int> ImportChatAsync(string sourceDatabasePath, Guid sourceChatId, Guid targetChatId, CancellationToken token) => Task.FromResult(0);
public override void Dispose()
{
}
}
@@ -0,0 +1,30 @@
namespace AIStudio.Tools.Databases.ProvenanceStore;
/// <summary>
/// One row of the table provenance.
/// </summary>
/// <remarks>
/// The key is the chat together with the id. A chat which is copied, e.g., an assistant's
/// conversation which the user carries over into a chat, takes its rows along under the same ids,
/// so its references stay valid, and deleting one of the two chats leaves the other one whole.
/// </remarks>
internal sealed class ProvenanceEntity
{
public string ChatId { get; set; } = string.Empty;
public string Id { get; set; } = string.Empty;
/// <summary>
/// The ProvenanceKind, stored by name.
/// </summary>
public string Kind { get; set; } = string.Empty;
public string Content { get; set; } = string.Empty;
/// <summary>
/// The content around an ERI passage as a JSON array of strings, or null when there is none.
/// </summary>
public string? SurroundingContent { get; set; }
public DateTimeOffset CreatedAtUtc { get; set; }
}
@@ -0,0 +1,16 @@
namespace AIStudio.Tools.Databases.ProvenanceStore;
/// <summary>
/// One piece of content a chat received besides its messages, exactly as the AI received it.
/// </summary>
/// <remarks>
/// Where the content comes from, such as the file, the page, or the tool, is kept by the chat
/// itself, next to the id of this entry. The chat needs it to list its sources even without this
/// store, and it would only be stored twice here.
/// </remarks>
/// <param name="Id">The id the chat refers to the entry by.</param>
/// <param name="Kind">What sort of content this is.</param>
/// <param name="Content">The content, after the prompt injection filter.</param>
/// <param name="SurroundingContent">The content an ERI server sent along around its passage, empty for everything else.</param>
/// <param name="CreatedAtUtc">When the entry was stored.</param>
public sealed record ProvenanceEntry(Guid Id, ProvenanceKind Kind, string Content, IReadOnlyList<string> SurroundingContent, DateTimeOffset CreatedAtUtc);
@@ -0,0 +1,32 @@
namespace AIStudio.Tools.Databases.ProvenanceStore;
/// <summary>
/// What sort of content a chat received, apart from its messages.
/// </summary>
/// <remarks>
/// Stored by name, like the kinds of mail parts in the index: a new member may then go anywhere
/// in this list, and the database reads on its own.
/// </remarks>
public enum ProvenanceKind
{
/// <summary>
/// The content of a file the user attached to a message.
/// </summary>
ATTACHMENT,
/// <summary>
/// A passage a data source handed to the AI: a local one or an ERI server, through the classic
/// process or the semantic search.
/// </summary>
DATA_SOURCE_PASSAGE,
/// <summary>
/// A mail, or a part of one, which the AI received from a mailbox.
/// </summary>
MAIL,
/// <summary>
/// A web page the AI received, e.g., which a tool read for it.
/// </summary>
WEB_PAGE,
}
@@ -0,0 +1,129 @@
namespace AIStudio.Tools.Databases.ProvenanceStore;
/// <summary>
/// Keeps what chats received besides their messages: the content of attached files, the results
/// of tools, and the passages of data sources, each exactly as the AI received it.
/// </summary>
/// <remarks>
/// A chat keeps only the ids of its entries, next to where each one came from. The content lives
/// here, so that a chat file stays small and an attached file is read only once.<br/><br/>
///
/// A database of its own rather than a table of the index: the index of the data sources can be
/// rebuilt from the files and is treated that way, e.g., a backup may leave it out. What a chat
/// received cannot be rebuilt, since the files, pages and mails may have changed since.<br/><br/>
///
/// Every entry belongs to one chat. The same id may stand for an entry of several chats, after a
/// chat was copied, and then it holds the same content in each.
/// </remarks>
public abstract class ProvenanceStoreClient(string name, string path) : DatabaseClient(name, path)
{
/// <summary>
/// Stores one piece of content a chat received.
/// </summary>
/// <param name="chatId">The chat which received the content.</param>
/// <param name="kind">What sort of content it is.</param>
/// <param name="content">The content, exactly as the AI received it.</param>
/// <param name="surroundingContent">The content an ERI server sent along around its passage, or nothing.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The id of the new entry, or null when this store cannot keep anything right now.</returns>
public abstract Task<Guid?> StoreAsync(Guid chatId, ProvenanceKind kind, string content, IReadOnlyList<string>? surroundingContent, CancellationToken token);
/// <summary>
/// Reads entries by their ids.
/// </summary>
/// <remarks>
/// By the ids alone, without the chat: the place which shows an entry seldom knows its chat,
/// and every copy of an entry holds the same content.
/// </remarks>
/// <param name="ids">The ids to read.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The entries found, by their ids. An id without an entry is missing from it.</returns>
public abstract Task<IReadOnlyDictionary<Guid, ProvenanceEntry>> GetEntriesAsync(IReadOnlyCollection<Guid> ids, CancellationToken token);
/// <summary>
/// Copies every entry of a chat to another chat, keeping their ids.
/// </summary>
/// <remarks>
/// The copy is what a chat made from another one refers to. Entries the target holds already
/// stay as they are.
/// </remarks>
/// <param name="sourceChatId">The chat to copy from.</param>
/// <param name="targetChatId">The chat to copy to.</param>
/// <param name="token">The cancellation token.</param>
public abstract Task CopyChatAsync(Guid sourceChatId, Guid targetChatId, CancellationToken token);
/// <summary>
/// Deletes every entry of some chats.
/// </summary>
/// <remarks>
/// All of them at once: deleting a workspace or cleaning up at the start may concern many chats.
/// </remarks>
/// <param name="chatIds">The chats.</param>
/// <param name="token">The cancellation token.</param>
/// <param name="createdBefore">Deletes only the entries stored before this moment, or all of them when null.</param>
/// <returns>How many entries were deleted.</returns>
public abstract Task<int> DeleteChatsAsync(IReadOnlyCollection<Guid> chatIds, CancellationToken token, DateTimeOffset? createdBefore = null);
/// <summary>
/// Deletes the entries of a chat which it no longer refers to.
/// </summary>
/// <param name="chatId">The chat.</param>
/// <param name="keptIds">The ids the chat still refers to.</param>
/// <param name="token">The cancellation token.</param>
/// <param name="createdBefore">Deletes only the entries stored before this moment, or all of them when null.</param>
/// <returns>How many entries were deleted.</returns>
public abstract Task<int> DeleteChatEntriesExceptAsync(Guid chatId, IReadOnlyCollection<Guid> keptIds, CancellationToken token, DateTimeOffset? createdBefore = null);
/// <summary>
/// Deletes some entries of a chat, e.g., texts which a preview kept for a message that went without them.
/// </summary>
/// <param name="chatId">The chat.</param>
/// <param name="ids">The ids of the entries.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>How many entries were deleted.</returns>
public abstract Task<int> DeleteEntriesAsync(Guid chatId, IReadOnlyCollection<Guid> ids, CancellationToken token);
/// <summary>
/// Lists the chats this store holds entries of.
/// </summary>
/// <param name="token">The cancellation token.</param>
/// <param name="createdSince">Lists only the chats with an entry stored at or after this moment, or all of them when null.</param>
/// <returns>The ids of the chats, Guid.Empty included when entries were stored for a chat without an id.</returns>
public abstract Task<IReadOnlyList<Guid>> GetChatIdsAsync(CancellationToken token, DateTimeOffset? createdSince = null);
/// <summary>
/// Reads when the last complete cleanup of this store began, see ProvenanceCleanup.
/// </summary>
/// <param name="token">The cancellation token.</param>
/// <returns>The moment, or null when no cleanup was noted, e.g., for a store a restore brought back.</returns>
public abstract Task<DateTimeOffset?> GetLastCleanupAsync(CancellationToken token);
/// <summary>
/// Notes when a complete cleanup of this store began, see ProvenanceCleanup.
/// </summary>
/// <param name="startedAt">The moment the cleanup began.</param>
/// <param name="token">The cancellation token.</param>
public abstract Task SetLastCleanupAsync(DateTimeOffset startedAt, CancellationToken token);
/// <summary>
/// Writes a copy of this store into a directory, for a backup.
/// </summary>
/// <param name="targetDirectory">The directory of the copy. A copy written there before is replaced.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The path of the copy.</returns>
public abstract Task<string> BackupToDirectoryAsync(string targetDirectory, CancellationToken token);
/// <summary>
/// Adds the entries of one chat from a copy of this store, such as the one a backup holds.
/// </summary>
/// <remarks>
/// The entries keep their ids, so the chat's references stay valid, also when the chat is added
/// under another id. Entries the target chat holds already stay as they are.
/// </remarks>
/// <param name="sourceDatabasePath">The copy to read from. It is only read.</param>
/// <param name="sourceChatId">The chat in the copy.</param>
/// <param name="targetChatId">The chat in this store.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>How many entries were added.</returns>
public abstract Task<int> ImportChatAsync(string sourceDatabasePath, Guid sourceChatId, Guid targetChatId, CancellationToken token);
}
@@ -0,0 +1,45 @@
using AIStudio.Tools.Databases.IndexStore;
using Microsoft.Data.Sqlite;
using Microsoft.EntityFrameworkCore;
namespace AIStudio.Tools.Databases.ProvenanceStore;
internal sealed class ProvenanceStoreDbContext(DbContextOptions<ProvenanceStoreDbContext> options) : DbContext(options)
{
private static readonly SecureDeleteInterceptor SECURE_DELETE = new();
public static DbContextOptions<ProvenanceStoreDbContext> CreateOptions(string databasePath) => new DbContextOptionsBuilder<ProvenanceStoreDbContext>()
.UseSqlite(BuildConnectionString(databasePath))
.AddInterceptors(SECURE_DELETE)
.Options;
public DbSet<ProvenanceEntity> Entries => this.Set<ProvenanceEntity>();
protected override void OnModelCreating(ModelBuilder modelBuilder)
{
modelBuilder.Entity<ProvenanceEntity>(entity =>
{
entity.ToTable("provenance");
entity.HasKey(entry => new { entry.ChatId, entry.Id });
entity.Property(entry => entry.ChatId).HasColumnName("chat_id");
entity.Property(entry => entry.Id).HasColumnName("id");
entity.Property(entry => entry.Kind).HasColumnName("kind").IsRequired();
entity.Property(entry => entry.Content).HasColumnName("content").IsRequired();
entity.Property(entry => entry.SurroundingContent).HasColumnName("surrounding_content");
entity.Property(entry => entry.CreatedAtUtc).HasColumnName("created_at_utc").HasConversion(new IndexStoreDateTimeOffsetConverter()).IsRequired();
// A chat reads its entries by their ids alone, see ProvenanceStoreClient.GetEntriesAsync:
entity.HasIndex(entry => entry.Id).HasDatabaseName("idx_provenance_id");
});
}
private static string BuildConnectionString(string databasePath) => new SqliteConnectionStringBuilder
{
DataSource = databasePath,
Mode = SqliteOpenMode.ReadWriteCreate,
Cache = SqliteCacheMode.Shared,
DefaultTimeout = 30,
}.ToString();
}
@@ -0,0 +1,15 @@
using Microsoft.EntityFrameworkCore.Design;
namespace AIStudio.Tools.Databases.ProvenanceStore;
internal sealed class ProvenanceStoreDesignTimeDbContextFactory : IDesignTimeDbContextFactory<ProvenanceStoreDbContext>
{
public ProvenanceStoreDbContext CreateDbContext(string[] args)
{
var databasePath = args.FirstOrDefault(argument => argument.EndsWith(".sqlite3", StringComparison.OrdinalIgnoreCase));
if (string.IsNullOrWhiteSpace(databasePath))
databasePath = Path.Combine(Path.GetTempPath(), "mindwork-ai-studio-provenance-design.sqlite3");
return new ProvenanceStoreDbContext(ProvenanceStoreDbContext.CreateOptions(databasePath));
}
}
@@ -0,0 +1,15 @@
using System.Diagnostics.CodeAnalysis;
using Microsoft.EntityFrameworkCore;
namespace AIStudio.Tools.Databases.ProvenanceStore;
internal static class ProvenanceStoreSchemaMigrator
{
[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(Migrations.InitialProvenanceStore))]
public static async Task MigrateAsync(ProvenanceStoreDbContext context, CancellationToken token)
{
await context.Database.MigrateAsync(token);
await context.Database.ExecuteSqlRawAsync("PRAGMA journal_mode=WAL;", token);
}
}
@@ -0,0 +1,36 @@
using System.Data.Common;
using Microsoft.EntityFrameworkCore.Diagnostics;
namespace AIStudio.Tools.Databases.ProvenanceStore;
/// <summary>
/// Turns on secure_delete for every connection to the provenance store.
/// </summary>
/// <remarks>
/// SQLite only marks the space of a deleted row as free, and its bytes stay in the file until
/// another row takes their place. The provenance store holds whole documents, mails, and web pages,
/// so the text of a deleted chat would stay readable in the file, perhaps for months. With
/// secure_delete, SQLite overwrites it with zeros while deleting. The write-ahead log keeps the old
/// pages until its next checkpoint, which SQLite makes by itself.<br/><br/>
///
/// The setting lasts as long as its connection, hence it is made for every connection anew.
/// </remarks>
internal sealed class SecureDeleteInterceptor : DbConnectionInterceptor
{
private const string PRAGMA = "PRAGMA secure_delete = ON;";
public override void ConnectionOpened(DbConnection connection, ConnectionEndEventData eventData)
{
using var command = connection.CreateCommand();
command.CommandText = PRAGMA;
command.ExecuteNonQuery();
}
public override async Task ConnectionOpenedAsync(DbConnection connection, ConnectionEndEventData eventData, CancellationToken cancellationToken = default)
{
await using var command = connection.CreateCommand();
command.CommandText = PRAGMA;
await command.ExecuteNonQueryAsync(cancellationToken);
}
}
@@ -0,0 +1,459 @@
using System.Data;
using System.Globalization;
using System.Text.Encodings.Web;
using System.Text.Json;
using AIStudio.Settings;
using AIStudio.Tools.Databases.IndexStore;
using AIStudio.Tools.PluginSystem;
using Microsoft.Data.Sqlite;
using Microsoft.EntityFrameworkCore;
namespace AIStudio.Tools.Databases.ProvenanceStore;
/// <summary>
/// The provenance store in SQLite, see ProvenanceStoreClient.
/// </summary>
/// <remarks>
/// Every reading and writing of entries runs on a thread of its own. SQLite works synchronously
/// behind its async methods, cf. the Databases AGENTS.md, and nearly everything this store holds is
/// a whole document, a web page, or a mail. Many callers sit in the user interface, e.g., a preview
/// which keeps its text, or a click which deletes a chat, and would freeze the window meanwhile.
/// </remarks>
public sealed class SqliteProvenanceStoreClientImplementation(string name, string databasePath, string basePath, string version) : ProvenanceStoreClient(name, basePath)
{
private const string DATABASE_NAME = "SQLite";
/// <summary>
/// The file of the store, in its folder in the data folder and in a backup alike.
/// </summary>
public const string DATABASE_FILENAME = "provenance.sqlite3";
/// <summary>
/// The file which notes when the last complete cleanup began, see GetLastCleanupAsync.
/// </summary>
private const string LAST_CLEANUP_FILENAME = "last-cleanup.txt";
/// <summary>
/// Writes the content around a passage the way it reads: the default encoder would turn every
/// character outside ASCII into an escape, cf. ToolExecutionResult.MODEL_CONTENT_OPTIONS.
/// </summary>
private static readonly JsonSerializerOptions JSON_OPTIONS = new() { Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping };
private readonly string databasePath = databasePath;
private readonly string lastCleanupPath = Path.Join(basePath, LAST_CLEANUP_FILENAME);
private readonly DbContextOptions<ProvenanceStoreDbContext> dbContextOptions = ProvenanceStoreDbContext.CreateOptions(databasePath);
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(SqliteProvenanceStoreClientImplementation).Namespace, nameof(SqliteProvenanceStoreClientImplementation));
public override string CacheKey => $"{this.Name}:{this.databasePath}:{version}";
public override string Version => version;
public static async Task<DatabaseClient> CreateAsync(
ILogger logger,
ILogger<DatabaseClient> databaseClientLogger,
CancellationToken cancellationToken)
{
if (string.IsNullOrWhiteSpace(SettingsManager.DataDirectory))
return CreateNoProvenanceStoreClient(DATABASE_NAME, "The application data directory is not available yet.", DatabaseClientStatus.STARTING, databaseClientLogger);
try
{
SQLitePCL.Batteries_V2.Init();
var basePath = Path.Combine(SettingsManager.DataDirectory, "databases", "provenance");
Directory.CreateDirectory(basePath);
var databasePath = Path.Combine(basePath, DATABASE_FILENAME);
var client = new SqliteProvenanceStoreClientImplementation(DATABASE_NAME, databasePath, basePath, string.Empty);
await client.InitializeAsync(cancellationToken);
var version = await client.GetSqliteVersionAsync(cancellationToken);
client = new SqliteProvenanceStoreClientImplementation(DATABASE_NAME, databasePath, basePath, version);
client.SetLogger(databaseClientLogger);
return client;
}
catch (Exception exception)
{
logger.LogWarning(exception, "The provenance store ({DatabaseName}) is not available. Chats keep what they received in their own files.", DATABASE_NAME);
return CreateNoProvenanceStoreClient(DATABASE_NAME, exception.Message, DatabaseClientStatus.UNAVAILABLE, databaseClientLogger);
}
}
public override async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo()
{
//
// Everything is read before the first line, since a yield return cannot stand in a try
// block. Each probe catches its own failure, cf. the Databases AGENTS.md:
//
var schemaVersion = string.Empty;
var journalMode = string.Empty;
IReadOnlyDictionary<string, int>? countsByKind = null;
int? chatCount = null;
try
{
await using var context = this.CreateContext();
schemaVersion = await GetSchemaVersionAsync(context);
journalMode = (await QueryScalarTextAsync(context, "PRAGMA journal_mode;")).ToUpperInvariant();
countsByKind = await TryCountByKindAsync(context);
chatCount = await TryCountAsync(context.Entries.Select(entry => entry.ChatId).Distinct());
}
catch (Exception exception)
{
this.Logger?.LogWarning(exception, "Failed to read the display details of the provenance store.");
}
yield return (TB("Reported version"), version);
yield return (TB("Journal mode"), OrUnknown(journalMode));
yield return (TB("Schema version"), OrUnknown(schemaVersion));
yield return (TB("Storage size"), this.GetStorageSize());
// The total counts an entry of a kind this version does not know as well, e.g., one written by a newer version:
yield return (TB("Stored items across all chats"), countsByKind?.Values.Sum().CompactCount() ?? TB("unknown"));
yield return (TB("Chats with stored items"), chatCount?.CompactCount() ?? TB("unknown"));
// One line per kind which any chat received, in the order of the kinds. A kind no chat received would only say zero:
if (countsByKind is null)
yield break;
foreach (var kind in Enum.GetValues<ProvenanceKind>())
if (countsByKind.TryGetValue(kind.ToString(), out var count) && count > 0)
yield return (GetKindLabel(kind), count.CompactCount());
}
private static string GetKindLabel(ProvenanceKind kind) => kind switch
{
ProvenanceKind.ATTACHMENT => TB("Attached files"),
ProvenanceKind.DATA_SOURCE_PASSAGE => TB("Passages of data sources"),
ProvenanceKind.MAIL => TB("Mails"),
ProvenanceKind.WEB_PAGE => TB("Web pages"),
_ => kind.ToString(),
};
public override Task<Guid?> StoreAsync(Guid chatId, ProvenanceKind kind, string content, IReadOnlyList<string>? surroundingContent, CancellationToken token) => Task.Run<Guid?>(async () =>
{
var id = Guid.NewGuid();
await using var context = this.CreateContext();
context.Entries.Add(new ProvenanceEntity
{
ChatId = ToKey(chatId),
Id = ToKey(id),
Kind = kind.ToString(),
Content = content,
SurroundingContent = surroundingContent is { Count: > 0 } ? JsonSerializer.Serialize(surroundingContent, JSON_OPTIONS) : null,
CreatedAtUtc = DateTimeOffset.UtcNow,
});
await context.SaveChangesAsync(token);
return id;
}, token);
public override Task<IReadOnlyDictionary<Guid, ProvenanceEntry>> GetEntriesAsync(IReadOnlyCollection<Guid> ids, CancellationToken token) => Task.Run<IReadOnlyDictionary<Guid, ProvenanceEntry>>(async () =>
{
var entries = new Dictionary<Guid, ProvenanceEntry>();
if (ids.Count is 0)
return entries;
var keys = ids.Select(ToKey).Distinct().ToList();
await using var context = this.CreateContext();
var rows = await context.Entries
.AsNoTracking()
.Where(entry => keys.Contains(entry.Id))
.ToListAsync(token);
// Copies of an entry hold the same content, so the first one found stands for all of them:
foreach (var row in rows)
{
if (!Guid.TryParse(row.Id, out var id) || entries.ContainsKey(id))
continue;
if (!Enum.TryParse<ProvenanceKind>(row.Kind, out var kind))
{
this.Logger?.LogWarning("The provenance entry '{EntryId}' has the unknown kind '{Kind}', e.g., from a newer version of AI Studio, and is left out.", row.Id, row.Kind);
continue;
}
entries[id] = new ProvenanceEntry(id, kind, row.Content, this.ReadSurroundingContent(row), row.CreatedAtUtc);
}
return entries;
}, token);
public override Task CopyChatAsync(Guid sourceChatId, Guid targetChatId, CancellationToken token) => Task.Run(async () =>
{
var sourceKey = ToKey(sourceChatId);
var targetKey = ToKey(targetChatId);
await using var context = this.CreateContext();
//
// In one statement, so no row passes through this process. The time of each row stays,
// since it says when the AI received the content, not when the chat was copied:
//
await context.Database.ExecuteSqlAsync(
$"""
INSERT OR IGNORE INTO provenance (chat_id, id, kind, content, surrounding_content, created_at_utc)
SELECT {targetKey}, id, kind, content, surrounding_content, created_at_utc
FROM provenance
WHERE chat_id = {sourceKey}
""",
token);
}, token);
public override Task<int> DeleteChatsAsync(IReadOnlyCollection<Guid> chatIds, CancellationToken token, DateTimeOffset? createdBefore = null) => Task.Run(async () =>
{
if (chatIds.Count is 0)
return 0;
var chatKeys = chatIds.Select(ToKey).Distinct().ToList();
await using var context = this.CreateContext();
return await CreatedBefore(context.Entries.Where(entry => chatKeys.Contains(entry.ChatId)), createdBefore).ExecuteDeleteAsync(token);
}, token);
public override Task<int> DeleteChatEntriesExceptAsync(Guid chatId, IReadOnlyCollection<Guid> keptIds, CancellationToken token, DateTimeOffset? createdBefore = null) => Task.Run(async () =>
{
var chatKey = ToKey(chatId);
var keptKeys = keptIds.Select(ToKey).Distinct().ToList();
await using var context = this.CreateContext();
return await CreatedBefore(context.Entries.Where(entry => entry.ChatId == chatKey && !keptKeys.Contains(entry.Id)), createdBefore).ExecuteDeleteAsync(token);
}, token);
public override Task<int> DeleteEntriesAsync(Guid chatId, IReadOnlyCollection<Guid> ids, CancellationToken token) => Task.Run(async () =>
{
if (ids.Count is 0)
return 0;
var chatKey = ToKey(chatId);
var keys = ids.Select(ToKey).Distinct().ToList();
await using var context = this.CreateContext();
return await context.Entries.Where(entry => entry.ChatId == chatKey && keys.Contains(entry.Id)).ExecuteDeleteAsync(token);
}, token);
/// <summary>
/// Narrows entries down to those stored before a moment.
/// </summary>
/// <remarks>
/// The time is stored as text in the round-trip format of UTC, whose order is the order of time,
/// so SQLite compares it as text, after the parameter went through the same conversion.
/// </remarks>
private static IQueryable<ProvenanceEntity> CreatedBefore(IQueryable<ProvenanceEntity> entries, DateTimeOffset? createdBefore) => createdBefore is { } moment
? entries.Where(entry => entry.CreatedAtUtc < moment)
: entries;
public override Task<IReadOnlyList<Guid>> GetChatIdsAsync(CancellationToken token, DateTimeOffset? createdSince = null) => Task.Run<IReadOnlyList<Guid>>(async () =>
{
await using var context = this.CreateContext();
var entries = createdSince is { } moment ? context.Entries.Where(entry => entry.CreatedAtUtc >= moment) : context.Entries;
var chatKeys = await entries.Select(entry => entry.ChatId).Distinct().ToListAsync(token);
// Guid.Empty included: entries stored for a chat without an id belong to no chat on the disk, so the cleanup takes them:
return chatKeys
.Select(chatKey => Guid.TryParse(chatKey, out var chatId) ? chatId : (Guid?)null)
.OfType<Guid>()
.ToList();
}, token);
/// <inheritdoc />
/// <remarks>
/// Kept in a file next to the database rather than in it: the copy of the store which a backup
/// holds comes without it, so the first cleanup of a store brought back from a backup reads every chat.
/// </remarks>
public override async Task<DateTimeOffset?> GetLastCleanupAsync(CancellationToken token)
{
if (!File.Exists(this.lastCleanupPath))
return null;
// A file which cannot be read costs one cleanup which reads every chat, and then it is written anew:
try
{
var text = await File.ReadAllTextAsync(this.lastCleanupPath, token);
return DateTimeOffset.TryParse(text.Trim(), CultureInfo.InvariantCulture, DateTimeStyles.RoundtripKind, out var moment) ? moment : null;
}
catch (Exception exception) when (exception is IOException or UnauthorizedAccessException)
{
this.Logger?.LogWarning(exception, "Failed to read when the provenance store was cleaned up last.");
return null;
}
}
public override Task SetLastCleanupAsync(DateTimeOffset startedAt, CancellationToken token) => File.WriteAllTextAsync(this.lastCleanupPath, IndexStoreDateTimeOffset.ToUtcText(startedAt), token);
/// <inheritdoc />
/// <remarks>
/// A consistent snapshot which holds back no writer, made on a thread of its own, cf. SqliteDatabaseCopy.
/// </remarks>
public override Task<string> BackupToDirectoryAsync(string targetDirectory, CancellationToken token) => Task.Run(async () =>
{
var targetPath = Path.Join(targetDirectory, DATABASE_FILENAME);
await SqliteDatabaseCopy.CopyAndCheckAsync(this.databasePath, targetPath, token);
return targetPath;
}, token);
/// <inheritdoc />
/// <remarks>
/// The copy is read on a connection of its own, read-only and without a pool, so nothing keeps it
/// open afterward and the caller can delete it. Only the columns this version knows are read: a
/// copy made by a newer version may hold more.<br/><br/>
///
/// The entries pass one at a time, since each one may be a whole document and a chat may hold
/// many of them. All of them go in one transaction, so a chat arrives with all of its entries
/// or with none. The time of each entry stays, as in CopyChatAsync.
/// </remarks>
public override Task<int> ImportChatAsync(string sourceDatabasePath, Guid sourceChatId, Guid targetChatId, CancellationToken token) => Task.Run(async () =>
{
var targetKey = ToKey(targetChatId);
await using var source = new SqliteConnection(SqliteDatabaseCopy.BuildConnectionString(sourceDatabasePath, SqliteOpenMode.ReadOnly));
await source.OpenAsync(token);
await using var command = source.CreateCommand();
command.CommandText = "SELECT id, kind, content, surrounding_content, created_at_utc FROM provenance WHERE chat_id = $chatId;";
command.Parameters.AddWithValue("$chatId", ToKey(sourceChatId));
await using var context = this.CreateContext();
await using var transaction = await context.Database.BeginTransactionAsync(token);
await using var reader = await command.ExecuteReaderAsync(token);
var added = 0;
while (await reader.ReadAsync(token))
{
var id = reader.GetString(0);
var kind = reader.GetString(1);
var content = reader.GetString(2);
var surroundingContent = reader.IsDBNull(3) ? null : reader.GetString(3);
var createdAtUtc = reader.GetString(4);
// Entries the chat holds already stay as they are:
added += await context.Database.ExecuteSqlAsync(
$"""
INSERT OR IGNORE INTO provenance (chat_id, id, kind, content, surrounding_content, created_at_utc)
VALUES ({targetKey}, {id}, {kind}, {content}, {surroundingContent}, {createdAtUtc})
""",
token);
}
await transaction.CommitAsync(token);
return added;
}, token);
public override void Dispose()
{
}
/// <summary>
/// Writes an id the way this store keeps it: as text, lowercase and with hyphens.
/// </summary>
private static string ToKey(Guid id) => id.ToString("D");
private IReadOnlyList<string> ReadSurroundingContent(ProvenanceEntity row)
{
if (string.IsNullOrWhiteSpace(row.SurroundingContent))
return [];
try
{
return JsonSerializer.Deserialize<List<string>>(row.SurroundingContent) ?? [];
}
catch (JsonException exception)
{
this.Logger?.LogWarning(exception, "The content around the provenance entry '{EntryId}' cannot be read and is left out.", row.Id);
return [];
}
}
private static string OrUnknown(string value) => string.IsNullOrWhiteSpace(value) ? TB("unknown") : value;
private static async Task<string> QueryScalarTextAsync(ProvenanceStoreDbContext context, string sql)
{
try
{
// The raw connection rather than SqlQueryRaw, which cannot run a PRAGMA, cf. the Databases AGENTS.md:
var connection = context.Database.GetDbConnection();
if (connection.State is not ConnectionState.Open)
await connection.OpenAsync();
await using var command = connection.CreateCommand();
command.CommandText = sql;
var result = await command.ExecuteScalarAsync();
return result?.ToString() ?? string.Empty;
}
catch
{
return string.Empty;
}
}
private static async Task<string> GetSchemaVersionAsync(ProvenanceStoreDbContext context)
{
try
{
var appliedMigrations = (await context.Database.GetAppliedMigrationsAsync()).ToList();
if (appliedMigrations.Count == 0)
return TB("no migration applied");
var pendingMigrations = (await context.Database.GetPendingMigrationsAsync()).ToList();
return pendingMigrations.Count == 0
? string.Format(I18N.I.Culture, TB("{0} ({1} applied)"), appliedMigrations[^1], appliedMigrations.Count.CompactCount())
: string.Format(I18N.I.Culture, TB("{0} ({1} applied, {2} pending)"), appliedMigrations[^1], appliedMigrations.Count.CompactCount(), pendingMigrations.Count.CompactCount());
}
catch
{
return string.Empty;
}
}
/// <summary>
/// Counts the entries of every kind, by the name each kind is stored with.
/// </summary>
/// <returns>The counts by kind, or null when the count failed.</returns>
private static async Task<IReadOnlyDictionary<string, int>?> TryCountByKindAsync(ProvenanceStoreDbContext context)
{
try
{
return await context.Entries
.GroupBy(entry => entry.Kind)
.Select(group => new { Kind = group.Key, Count = group.Count() })
.ToDictionaryAsync(group => group.Kind, group => group.Count);
}
catch
{
return null;
}
}
private static async Task<int?> TryCountAsync<T>(IQueryable<T> query)
{
try
{
return await query.CountAsync();
}
catch
{
return null;
}
}
private async Task InitializeAsync(CancellationToken token)
{
await using var context = this.CreateContext();
await ProvenanceStoreSchemaMigrator.MigrateAsync(context, token);
}
private async Task<string> GetSqliteVersionAsync(CancellationToken token)
{
await using var context = this.CreateContext();
var versions = await context.Database
.SqlQueryRaw<string>("SELECT sqlite_version() AS Value")
.ToListAsync(token);
return versions.FirstOrDefault() ?? string.Empty;
}
private ProvenanceStoreDbContext CreateContext() => new(this.dbContextOptions);
private static NoProvenanceStoreClient CreateNoProvenanceStoreClient(string name, string? unavailableReason, DatabaseClientStatus status, ILogger<DatabaseClient> databaseClientLogger)
{
var client = new NoProvenanceStoreClient(name, unavailableReason, status);
client.SetLogger(databaseClientLogger);
return client;
}
}
@@ -0,0 +1,85 @@
using Microsoft.Data.Sqlite;
namespace AIStudio.Tools.Databases;
/// <summary>
/// Copies a SQLite database while AI Studio keeps using it, e.g., for a backup.
/// </summary>
/// <remarks>
/// VACUUM INTO writes a consistent snapshot of the database, including what the write-ahead log
/// still holds, and in WAL mode, reading it holds back no writer.<br/><br/>
///
/// It runs on connections of their own, although the databases guide has raw SQL go through
/// GetDbConnection. The connections of a context share one cache, and a shared cache locks whole
/// tables: reading the whole database for the copy would hold back every writer until the copy is
/// done. Without pooling, nothing keeps the copy open afterward, so the caller can delete it on
/// every OS.<br/><br/>
///
/// SQLite runs its async methods synchronously, and copying and checking a large database takes a
/// while. The callers therefore run this on a thread of their own, so that it never holds up theirs.
/// </remarks>
internal static class SqliteDatabaseCopy
{
/// <summary>
/// Copies a database into a file and checks the copy.
/// </summary>
/// <remarks>
/// A copy which fails its check is removed, and an InvalidOperationException says so.
/// </remarks>
/// <param name="databasePath">The database to copy.</param>
/// <param name="targetPath">The file of the copy. A file there is replaced.</param>
/// <param name="token">The cancellation token.</param>
public static async Task CopyAndCheckAsync(string databasePath, string targetPath, CancellationToken token)
{
Directory.CreateDirectory(Path.GetDirectoryName(targetPath)!);
// VACUUM INTO refuses to write into a file which holds anything:
if (File.Exists(targetPath))
File.Delete(targetPath);
try
{
await using (var connection = new SqliteConnection(BuildConnectionString(databasePath, SqliteOpenMode.ReadWrite)))
{
await connection.OpenAsync(token);
await using var command = connection.CreateCommand();
command.CommandText = "VACUUM INTO $targetPath;";
command.Parameters.AddWithValue("$targetPath", targetPath);
await command.ExecuteNonQueryAsync(token);
}
await using (var connection = new SqliteConnection(BuildConnectionString(targetPath, SqliteOpenMode.ReadOnly)))
{
await connection.OpenAsync(token);
await using var command = connection.CreateCommand();
command.CommandText = "PRAGMA quick_check;";
var checkResult = (await command.ExecuteScalarAsync(token))?.ToString();
if (checkResult is not "ok")
throw new InvalidOperationException($"The copy of the database at '{targetPath}' failed its check: '{checkResult}'.");
}
}
catch
{
if (File.Exists(targetPath))
File.Delete(targetPath);
throw;
}
}
/// <summary>
/// Builds the connection string for a connection of its own, which shares no cache and stays in no pool.
/// </summary>
/// <param name="path">The database.</param>
/// <param name="mode">How to open it.</param>
/// <returns>The connection string.</returns>
public static string BuildConnectionString(string path, SqliteOpenMode mode) => new SqliteConnectionStringBuilder
{
DataSource = path,
Mode = mode,
Cache = SqliteCacheMode.Private,
Pooling = false,
DefaultTimeout = 30,
}.ToString();
}
@@ -31,13 +31,16 @@ public sealed class NoVectorStoreClient(string name, string? unavailableReason,
public override Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>
Task.FromException(this.CreateUnavailableException());
public override Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, CancellationToken token) =>
public override Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, IReadOnlyList<HiddenChunks> hiddenChunks, CancellationToken token) =>
Task.FromException<IReadOnlyList<VectorSearchResult>>(this.CreateUnavailableException());
public override Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, VectorSearchFilter filter, CancellationToken token) =>
public override Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, VectorSearchFilter filter, IReadOnlyList<HiddenChunks> hiddenChunks, CancellationToken token) =>
Task.FromException<IReadOnlyList<VectorSearchResult>>(this.CreateUnavailableException());
public override Task DeleteEmbeddingByFile(string storeName, string filePath, string parentFileId, CancellationToken token) =>
public override Task<IReadOnlyList<StoredVector>> RetrieveVectorsAsync(string storeName, IReadOnlyList<string> pointIds, CancellationToken token) =>
Task.FromException<IReadOnlyList<StoredVector>>(this.CreateUnavailableException());
public override Task DeleteEmbeddingByFile(string storeName, string filePath, string parentFileId, IReadOnlyCollection<string> keepPointIds, CancellationToken token) =>
Task.FromException(this.CreateUnavailableException());
public override Task RelocateEmbeddingFileAsync(string storeName, string parentFileId, string absolutePath, string relativePath, string fingerprint, CancellationToken token) =>
@@ -26,6 +26,17 @@ public sealed class QdrantEdgeClientImplementation(
private const string EXISTS_PATH = "/system/qdrant-edge/exists";
private const string INSERT_PATH = "/system/qdrant-edge/insert";
private const string SEARCH_PATH = "/system/qdrant-edge/search";
private const string RETRIEVE_PATH = "/system/qdrant-edge/retrieve";
/// <summary>
/// How many vectors one request to the runtime reads at most.
/// </summary>
/// <remarks>
/// A vector travels as JSON text, some 15 KB at 1,536 dimensions. A document of thousands of
/// chunks in one answer would be hundreds of megabytes on the connection all other calls of the
/// app share.
/// </remarks>
private const int RETRIEVE_BATCH_SIZE = 256;
private const string DELETE_FILE_PATH = "/system/qdrant-edge/delete-file";
private const string RELOCATE_FILE_PATH = "/system/qdrant-edge/relocate-file";
private const string OPTIMIZE_PATH = "/system/qdrant-edge/optimize";
@@ -133,19 +144,19 @@ public sealed class QdrantEdgeClientImplementation(
public override Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, INSERT_PATH, new InsertEmbeddingRequest(storeName, points), token);
public override Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, CancellationToken token) =>
this.QueryEmbeddingAsync(storeName, vector, maxMatches, null, token);
public override Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, IReadOnlyList<HiddenChunks> hiddenChunks, CancellationToken token) =>
this.QueryEmbeddingAsync(storeName, vector, maxMatches, null, hiddenChunks, token);
public override async Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, VectorSearchFilter filter, CancellationToken token)
public override async Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, VectorSearchFilter filter, IReadOnlyList<HiddenChunks> hiddenChunks, CancellationToken token)
{
if (maxMatches <= 0 || filter.MatchesNothing)
return [];
var candidates = await this.QueryEmbeddingAsync(storeName, vector, filter.GetCandidateCount(maxMatches), filter.GetRequestPointIds(), token);
var candidates = await this.QueryEmbeddingAsync(storeName, vector, filter.GetCandidateCount(maxMatches), filter.GetRequestPointIds(), hiddenChunks, token);
return filter.Apply(candidates, maxMatches);
}
private async Task<IReadOnlyList<VectorSearchResult>> QueryEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, IReadOnlyList<string>? pointIds, CancellationToken token)
private async Task<IReadOnlyList<VectorSearchResult>> QueryEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, IReadOnlyList<string>? pointIds, IReadOnlyList<HiddenChunks> hiddenChunks, CancellationToken token)
{
if (maxMatches <= 0)
return [];
@@ -153,12 +164,28 @@ public sealed class QdrantEdgeClientImplementation(
return await rustService.ExecuteDatabaseQuery<SearchEmbeddingRequest, List<VectorSearchResult>>(
DATABASE_NAME,
SEARCH_PATH,
new SearchEmbeddingRequest(storeName, vector, maxMatches, pointIds),
new SearchEmbeddingRequest(storeName, vector, maxMatches, pointIds, hiddenChunks),
token) ?? [];
}
public override Task DeleteEmbeddingByFile(string storeName, string filePath, string parentFileId, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, DELETE_FILE_PATH, new DeleteEmbeddingByFileRequest(storeName, filePath, parentFileId), token);
public override async Task<IReadOnlyList<StoredVector>> RetrieveVectorsAsync(string storeName, IReadOnlyList<string> pointIds, CancellationToken token)
{
var vectors = new List<StoredVector>(pointIds.Count);
foreach (var batch in pointIds.Chunk(RETRIEVE_BATCH_SIZE))
{
token.ThrowIfCancellationRequested();
var batchVectors = await rustService.ExecuteDatabaseQuery<RetrieveVectorsRequest, List<StoredVector>>(DATABASE_NAME, RETRIEVE_PATH,
new RetrieveVectorsRequest(storeName, batch), token);
if (batchVectors is not null)
vectors.AddRange(batchVectors);
}
return vectors;
}
public override Task DeleteEmbeddingByFile(string storeName, string filePath, string parentFileId, IReadOnlyCollection<string> keepPointIds, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, DELETE_FILE_PATH, new DeleteEmbeddingByFileRequest(storeName, filePath, parentFileId, keepPointIds), token);
public override Task RelocateEmbeddingFileAsync(string storeName, string parentFileId, string absolutePath, string relativePath, string fingerprint, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, RELOCATE_FILE_PATH, new RelocateEmbeddingFileRequest(storeName, parentFileId, absolutePath, Path.GetFileName(absolutePath), relativePath, fingerprint), token);
@@ -196,9 +223,11 @@ public sealed class QdrantEdgeClientImplementation(
private sealed record InsertEmbeddingRequest(string StoreName, IReadOnlyList<VectorStoragePoint> Points);
/// <param name="PointIds">The only points the search may return, or null for the whole store.</param>
private sealed record SearchEmbeddingRequest(string StoreName, IReadOnlyList<float> Vector, int MaxMatches, IReadOnlyList<string>? PointIds);
private sealed record SearchEmbeddingRequest(string StoreName, IReadOnlyList<float> Vector, int MaxMatches, IReadOnlyList<string>? PointIds, IReadOnlyList<HiddenChunks> HiddenChunks);
private sealed record DeleteEmbeddingByFileRequest(string StoreName, string FilePath, string ParentFileId);
private sealed record RetrieveVectorsRequest(string StoreName, IReadOnlyList<string> PointIds);
private sealed record DeleteEmbeddingByFileRequest(string StoreName, string FilePath, string ParentFileId, IReadOnlyCollection<string> KeepPointIds);
private sealed record RelocateEmbeddingFileRequest(string StoreName, string ParentFileId, string AbsolutePath, string FileName, string RelativePath, string Fingerprint);
@@ -0,0 +1,8 @@
namespace AIStudio.Tools.Databases.VectorStore;
/// <summary>
/// The vector of one point, as the vector store holds it.
/// </summary>
/// <param name="PointId">The id of the point, which is also the id of its chunk.</param>
/// <param name="Vector">The vector.</param>
public sealed record StoredVector(string PointId, IReadOnlyList<float> Vector);
@@ -13,9 +13,24 @@ public sealed record VectorSearchResult(
string RelativePath,
string FileType,
int? PageNumber,
int? EndPageNumber,
int ChunkIndex,
string Text,
string Fingerprint,
string CreationUtc,
string LastWriteUtc,
string EmbeddedAtUtc);
string EmbeddedAtUtc)
{
/// <summary>
/// Where the chunk is, as the index store keeps it, or null until the search completed the hit from there.
/// </summary>
/// <remarks>
/// The vector store keeps only the pages of a chunk, cf. LocalIndexSearchService.SearchVectorsAsync.
/// </remarks>
public ChunkLocation? StoredLocation { get; init; }
/// <summary>
/// Where the chunk is: what the index store keeps, or the pages alone when it could not say.
/// </summary>
public ChunkLocation Location => this.StoredLocation ?? new ChunkLocation(this.PageNumber, this.EndPageNumber);
}
@@ -13,9 +13,11 @@ public sealed record VectorStoragePoint(
string RelativePath,
string FileType,
int? PageNumber,
int? EndPageNumber,
int ChunkIndex,
string Text,
string Fingerprint,
DateTimeOffset CreationUtc,
DateTimeOffset LastWriteUtc,
DateTimeOffset EmbeddedAtUtc);
DateTimeOffset EmbeddedAtUtc,
int Generation);
@@ -14,7 +14,16 @@ public abstract class VectorStoreClient(string name, string path): DatabaseClien
public abstract Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token);
public abstract Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, CancellationToken token);
/// <summary>
/// Searches the points of a store.
/// </summary>
/// <param name="storeName">The name of the store.</param>
/// <param name="vector">The vector to search for.</param>
/// <param name="maxMatches">How many matches to return at most.</param>
/// <param name="hiddenChunks">The chunks of documents the search must not find, as the index store names them.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The matches, best first.</returns>
public abstract Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, IReadOnlyList<HiddenChunks> hiddenChunks, CancellationToken token);
/// <summary>
/// Searches only those points of a store which pass the filter.
@@ -23,11 +32,29 @@ public abstract class VectorStoreClient(string name, string path): DatabaseClien
/// <param name="vector">The vector to search for.</param>
/// <param name="maxMatches">How many matches to return at most.</param>
/// <param name="filter">The points the search may return. Without any point, it finds nothing.</param>
/// <param name="hiddenChunks">The chunks of documents the search must not find, as the index store names them.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The matches, best first, all of them inside the filter.</returns>
public abstract Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, VectorSearchFilter filter, CancellationToken token);
public abstract Task<IReadOnlyList<VectorSearchResult>> SearchEmbeddingAsync(string storeName, IReadOnlyList<float> vector, int maxMatches, VectorSearchFilter filter, IReadOnlyList<HiddenChunks> hiddenChunks, CancellationToken token);
public abstract Task DeleteEmbeddingByFile(string storeName, string filePath, string parentFileId, CancellationToken token);
/// <summary>
/// Reads the vectors of the given points.
/// </summary>
/// <param name="storeName">The name of the store.</param>
/// <param name="pointIds">The points.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The vectors of the points which exist, in no particular order. None when the store does not exist.</returns>
public abstract Task<IReadOnlyList<StoredVector>> RetrieveVectorsAsync(string storeName, IReadOnlyList<string> pointIds, CancellationToken token);
/// <summary>
/// Removes the points of one document, found by its path and by its parent id.
/// </summary>
/// <param name="storeName">The name of the store.</param>
/// <param name="filePath">The key of the document.</param>
/// <param name="parentFileId">The parent id of the document, or empty when it has none yet.</param>
/// <param name="keepPointIds">The points of the document which stay, e.g., the ones an indexing run just wrote.</param>
/// <param name="token">The cancellation token.</param>
public abstract Task DeleteEmbeddingByFile(string storeName, string filePath, string parentFileId, IReadOnlyCollection<string> keepPointIds, CancellationToken token);
/// <summary>
/// Updates the location of a document's points without changing their ids or vectors.
@@ -10,6 +10,11 @@ internal static class FileContentHash
/// <summary>
/// Computes a lowercase SHA-256 hash for a file without loading it fully into memory.
/// </summary>
/// <remarks>
/// The file is read while other programs may keep it open, even for writing, as an office
/// program does with a document being edited. The runtime reads files for extraction the same
/// way, so a file which can be indexed can always be hashed.
/// </remarks>
/// <param name="path">The file path.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The lowercase hexadecimal hash.</returns>
@@ -19,7 +24,7 @@ internal static class FileContentHash
path,
FileMode.Open,
FileAccess.Read,
FileShare.Read,
FileShare.ReadWrite | FileShare.Delete,
65_536,
FileOptions.Asynchronous | FileOptions.SequentialScan);
+19
View File
@@ -1,4 +1,5 @@
using Markdig;
using Markdig.Extensions.Mathematics;
using Markdig.Syntax;
using System.Text;
@@ -17,6 +18,24 @@ public static class Markdown
.DisableHtml()
.Build();
/// <summary>
/// The pipeline for text AI Studio read from somewhere else, such as a document or a passage of
/// a data source: the safe pipeline without math.
/// </summary>
/// <remarks>
/// Such text uses the dollar sign for prices, and a PDF puts it among the characters it could
/// not read in a figure. Taken for math, it makes MudMarkdown typeset the whole page with
/// MathJax, which replaces elements Blazor renders, so that the next render of them fails.
/// </remarks>
public static readonly MarkdownPipeline CONTENT_MARKDOWN_PIPELINE = CreateContentPipeline();
private static MarkdownPipeline CreateContentPipeline()
{
var builder = new MarkdownPipelineBuilder().UseAdvancedExtensions().DisableHtml();
builder.Extensions.TryRemove<MathExtension>();
return builder.Build();
}
public static MudMarkdownProps DefaultConfig => new()
{
Heading =
@@ -55,6 +55,27 @@ internal sealed partial class DataSourceTranscriptStore(string dataDirectory, st
return new(this.GetPath(key), key);
}
/// <summary>
/// This store, for a media file whose content the caller hashed already, so the file is not read twice.
/// </summary>
/// <remarks>
/// Hashing a recording reads all of it, which takes long for a video. The hash counts only while
/// the file is as it was when it was hashed: the location is asked after its caller took its own
/// stamp of the file, cf. PersistentMediaTranscript, so a change in between keeps a transcript
/// out of the place of the content before. Then the file is hashed anew.
/// </remarks>
/// <param name="contentHash">The hash of the content, cf. FileContentHash.</param>
/// <param name="isUnchangedSinceHashed">Whether the file at the given path is still as it was when it was hashed.</param>
/// <returns>The location.</returns>
public IMediaTranscriptLocation WithKnownContentHash(string contentHash, Func<string, bool> isUnchangedSinceHashed) => new KnownContentHashLocation(this, contentHash, isUnchangedSinceHashed);
private sealed class KnownContentHashLocation(DataSourceTranscriptStore store, string contentHash, Func<string, bool> isUnchangedSinceHashed) : IMediaTranscriptLocation
{
public Task<MediaTranscriptSlot> ResolveAsync(string mediaPath, CancellationToken token) => isUnchangedSinceHashed(mediaPath)
? Task.FromResult(new MediaTranscriptSlot(store.GetPath(contentHash), contentHash))
: store.ResolveAsync(mediaPath, token);
}
/// <summary>
/// Whether a text is a key of this store, i.e. a lowercase SHA-256 hash.
/// </summary>
@@ -436,6 +436,9 @@ public sealed class PluginConfiguration(bool isInternal, LuaState state, PluginT
ManagedConfiguration.TryProcessConfiguration(x => x.Chat, x => x.PreselectedDataSourcesRetrievalMode, this.Id, settingsTable, dryRun);
ManagedConfiguration.TryProcessConfiguration(x => x.Chat, x => x.SendToChatDataSourceBehavior, this.Id, settingsTable, dryRun);
// Config: read changed attachments of a chat again?
ManagedConfiguration.TryProcessConfiguration(x => x.Chat, x => x.ReadChangedAttachmentsAgain, this.Id, settingsTable, dryRun);
// Config: Slide Assistant template and defaults?
ManagedConfiguration.TryProcessConfiguration(x => x.SlideBuilder, x => x.PreselectOptions, this.Id, settingsTable, dryRun);
ManagedConfiguration.TryProcessConfiguration(x => x.SlideBuilder, x => x.TemplatePath, this.Id, settingsTable, dryRun, convert: this.ResolvePluginFile);
@@ -27,24 +27,25 @@ public sealed class AugmentationOne : IAugmentationProcess
public string Description => TB("This is the standard augmentation process, which uses all retrieval contexts to augment the chat thread.");
/// <inheritdoc />
public async Task<ChatThread> ProcessAsync(IProvider provider, IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IRetrievalContext> retrievalContexts, CancellationToken token = default)
public async Task<AugmentationResult> ProcessAsync(IProvider provider, IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IRetrievalContext> retrievalContexts, CancellationToken token = default)
{
var settings = Program.SERVICE_PROVIDER.GetService<SettingsManager>()!;
if(retrievalContexts.Count == 0)
{
LOGGER.LogWarning("No retrieval contexts were issued. Skipping the augmentation process.");
return chatThread;
return new(chatThread, []);
}
var numTotalRetrievalContexts = retrievalContexts.Count;
//
// One report for the whole answer. With the check of which passages fit, every passage is
// filtered twice: once for that check, and once more for the model.
// One report for the whole answer. Every passage is filtered once, before the check of
// which passages fit: the check then judges exactly what the model would receive, and the
// model receives exactly what the check judged.
//
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
await using var reportingScope = guardService.BeginAction();
var preparedContexts = await retrievalContexts.PrepareForModelAsync(guardService, reportingScope, token);
var usedContexts = preparedContexts;
// Want the user to validate all retrieval contexts?
if (settings.ConfigurationData.AgentRetrievalContextValidation.EnableRetrievalContextValidation && chatThread.DataSourceOptions.AutomaticValidation)
@@ -56,7 +57,7 @@ public sealed class AugmentationOne : IAugmentationProcess
try
{
// Let's validate all retrieval contexts:
var validationResults = await validationAgent.ValidateRetrievalContextsAsync(lastUserPrompt, chatThread, retrievalContexts, reportingScope, token);
var validationResults = await validationAgent.ValidateRetrievalContextsAsync(lastUserPrompt, chatThread, preparedContexts, token);
if (validationResults.Count == 0)
LOGGER.LogWarning("Retrieval context validation returned no results. Continuing augmentation with all retrieved contexts.");
else
@@ -68,7 +69,7 @@ public sealed class AugmentationOne : IAugmentationProcess
var threshold = validationResults.GetConfidenceThreshold(targetWindow);
// Filter the retrieval contexts:
retrievalContexts = validationResults.Where(x => x.RetrievalContext is not null && x.Confidence >= threshold).Select(x => x.RetrievalContext!).ToList();
usedContexts = validationResults.Where(x => x.RetrievalContext is not null && x.Confidence >= threshold).Select(x => x.RetrievalContext!).ToList();
}
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
@@ -97,7 +98,7 @@ public sealed class AugmentationOne : IAugmentationProcess
}
}
LOGGER.LogInformation($"Starting the augmentation process over {retrievalContexts.Count:###,###,###,###} of {numTotalRetrievalContexts:###,###,###,###} retrieved contexts.");
LOGGER.LogInformation($"Starting the augmentation process over {usedContexts.Count:###,###,###,###} of {retrievalContexts.Count:###,###,###,###} retrieved contexts.");
//
// We build a huge prompt from all retrieval contexts:
@@ -107,11 +108,12 @@ public sealed class AugmentationOne : IAugmentationProcess
sb.AppendLine();
// Let's convert all retrieval contexts to Markdown:
await retrievalContexts.AsMarkdown(sb, reportingScope, token);
for (var index = 0; index < usedContexts.Count; index++)
sb.Append(usedContexts[index].ToMarkdown(index + 1, usedContexts.Count));
// Add the augmented data to the chat thread:
chatThread.AugmentedData = sb.ToString();
return chatThread;
return new(chatThread, usedContexts);
}
#endregion
@@ -0,0 +1,15 @@
using AIStudio.Chat;
namespace AIStudio.Tools.RAG;
/// <summary>
/// What an augmentation process did: the chat thread it augmented, and what it handed to the AI.
/// </summary>
/// <remarks>
/// The contexts the AI received can be fewer than the ones retrieved, e.g., when a check of which
/// passages fit the question left some out. Whoever lists the sources and passages of the answer
/// takes them from here, so the user sees what the AI read and nothing else.
/// </remarks>
/// <param name="ChatThread">The augmented chat thread.</param>
/// <param name="UsedContexts">The retrieval contexts handed to the AI, in the order it received them.</param>
public sealed record AugmentationResult(ChatThread ChatThread, IReadOnlyList<PreparedRetrievalContext> UsedContexts);
@@ -28,6 +28,6 @@ public interface IAugmentationProcess
/// <param name="chatThread">The chat thread.</param>
/// <param name="retrievalContexts">The retrieval contexts that were issued by the retrieval process.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The altered chat thread.</returns>
public Task<ChatThread> ProcessAsync(IProvider provider, IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IRetrievalContext> retrievalContexts, CancellationToken token = default);
/// <returns>The altered chat thread, together with the retrieval contexts handed to the AI.</returns>
public Task<AugmentationResult> ProcessAsync(IProvider provider, IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IRetrievalContext> retrievalContexts, CancellationToken token = default);
}
@@ -14,8 +14,10 @@ public static class IRetrievalContextExtensions
/// <remarks>
/// The location is what lets the AI say where an answer comes from. Naming only the file is
/// not enough in a document of two hundred pages, and we know the page: it travels from the
/// runtime through the index into the context. A slide or a sheet has no page, and then
/// nothing is claimed rather than something made up.
/// runtime through the index into the context, together with the page a passage ends on when it
/// goes on beyond its first one. A document without pages names its sheet and rows, or its
/// lines, the same way. A slide has none of them, and then nothing is claimed rather than
/// something made up.
/// </remarks>
/// <param name="contextBuilder">The builder to write into.</param>
/// <param name="retrievalContext">The context to describe.</param>
@@ -29,8 +31,8 @@ public static class IRetrievalContextExtensions
contextBuilder.AppendLine($"Content type: {retrievalContext.Type}");
contextBuilder.AppendLine($"Content path: {retrievalContext.Path}");
if(retrievalContext is RetrievalTextContext { PageNumber: > 0 } locatedContext)
contextBuilder.AppendLine($"Content location: page {locatedContext.PageNumber}");
if(retrievalContext is RetrievalTextContext textContext && DescribeLocation(textContext.PageNumber, textContext.EndPageNumber, textContext.SheetName, textContext.StartRow, textContext.EndRow, textContext.StartLine, textContext.EndLine) is { } location)
contextBuilder.AppendLine($"Content location: {location}");
if(retrievalContext.Links.Count is 0)
return;
@@ -40,6 +42,38 @@ public static class IRetrievalContextExtensions
contextBuilder.AppendLine($"- {link}");
}
/// <summary>
/// Says where a passage is, the way the AI is told, e.g., "pages 12 to 14", "sheet "Budget", rows 3 to 9" or "lines 40 to 85".
/// </summary>
/// <remarks>
/// In English, like everything else the AI is told about a passage. Pages come first, then
/// rows, then lines: a passage names one kind only. An end only names a range when it lies
/// beyond the start.
/// </remarks>
/// <param name="pageNumber">The page the passage starts on.</param>
/// <param name="endPageNumber">The page it ends on.</param>
/// <param name="sheetName">The sheet its rows are of, or null for a table without sheets.</param>
/// <param name="startRow">The row it starts at.</param>
/// <param name="endRow">The row it ends at.</param>
/// <param name="startLine">The line it starts at.</param>
/// <param name="endLine">The line it ends at.</param>
/// <returns>Where the passage is, or null when it names nothing.</returns>
internal static string? DescribeLocation(int? pageNumber, int? endPageNumber, string? sheetName, int? startRow, int? endRow, int? startLine, int? endLine)
{
if (pageNumber is > 0)
return DescribeRange("page", "pages", pageNumber.Value, endPageNumber);
if (startRow is > 0)
{
var rows = DescribeRange("row", "rows", startRow.Value, endRow);
return string.IsNullOrEmpty(sheetName) ? rows : $"sheet \"{sheetName}\", {rows}";
}
return startLine is > 0 ? DescribeRange("line", "lines", startLine.Value, endLine) : null;
}
private static string DescribeRange(string one, string several, int start, int? end) => end > start ? $"{several} {start} to {end}" : $"{one} {start}";
/// <summary>
/// Writes what the AI is told about an image which matched, in place of the image itself.
/// </summary>
@@ -57,140 +91,184 @@ public static class IRetrievalContextExtensions
}
/// <summary>
/// Writes what the AI is told about several retrieval contexts, filtered for prompt injections.
/// Prepares retrieval contexts for the AI: filters them for prompt injections, writes out what
/// the AI is told about each, and keeps what the user is shown of it.
/// </summary>
/// <param name="retrievalContexts">The contexts to describe.</param>
/// <param name="sb">The builder to write into, or null for a new one.</param>
/// <remarks>
/// A data source is as untrusted as any other external content: it may serve text written to
/// steer the model rather than to answer the query. So everything a data source wrote passes
/// the filter: the description of each context, its matched text, and each part of its
/// surrounding content. What the app writes around them, such as the labels and the fences,
/// does not need to.<br/><br/>
/// All contexts go to the runtime in one request, which also gives the user one report for all
/// of them. The passages are made from the filtered texts, so the user reads what the AI read.
/// </remarks>
/// <param name="retrievalContexts">The contexts to prepare.</param>
/// <param name="guardService">The prompt injection filter.</param>
/// <param name="reportingScope">The action the contexts are retrieved for, or null when they are retrieved for none.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>Everything written into the builder so far.</returns>
public static async Task<string> AsMarkdown(this IReadOnlyList<IRetrievalContext> retrievalContexts, StringBuilder? sb = null, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
/// <returns>One prepared context per context, in the order they came in.</returns>
public static async Task<IReadOnlyList<PreparedRetrievalContext>> PrepareForModelAsync(this IReadOnlyList<IRetrievalContext> retrievalContexts, PromptInjectionGuardService guardService, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
{
sb ??= new StringBuilder();
var index = 0;
if (retrievalContexts.Count is 0)
return [];
//
// One report for the whole retrieval run: a query may pull in dozens of contexts, and
// the user wants to know that something was filtered, not to acknowledge it per context.
// The texts of all contexts in one list, each context with its description first, then its
// matched text, then its surrounding content. The answer comes back in the same order, and
// is read back below in the same way:
//
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
await using var retrievalScope = guardService.BeginAction(within: reportingScope);
foreach(var retrievalContext in retrievalContexts)
var texts = new List<PromptInjectionText>();
foreach (var retrievalContext in retrievalContexts)
{
index++;
await retrievalContext.AsMarkdown(sb, index, retrievalContexts.Count, retrievalScope, token);
var source = PromptInjectionSource.RetrievalContext(retrievalContext.DataSourceName, retrievalContext.Path);
var description = new StringBuilder();
AppendContextDescription(description, retrievalContext);
texts.Add(new(description.ToString(), source));
if (retrievalContext is not RetrievalTextContext textContext)
continue;
texts.Add(new(textContext.MatchedText, source));
foreach (var surrounding in textContext.SurroundingContent)
texts.Add(new(surrounding, source));
}
return sb.ToString();
token.ThrowIfCancellationRequested();
var sanitizedTexts = await guardService.SanitizeAsync(texts, reportingScope);
var preparedContexts = new List<PreparedRetrievalContext>(retrievalContexts.Count);
var nextText = 0;
foreach (var retrievalContext in retrievalContexts)
{
var body = new StringBuilder();
body.Append(sanitizedTexts[nextText++]);
var matchedText = string.Empty;
IReadOnlyList<string> surroundingContent = [];
switch (retrievalContext)
{
case RetrievalTextContext textContext:
matchedText = sanitizedTexts[nextText++];
surroundingContent = sanitizedTexts.Skip(nextText).Take(textContext.SurroundingContent.Count).ToList();
nextText += surroundingContent.Count;
AppendTextMatch(body, matchedText, surroundingContent);
break;
case RetrievalImageContext:
AppendImageMatch(body);
break;
default:
LOGGER.LogWarning($"The retrieval content type '{retrievalContext.Type}' of data source '{retrievalContext.DataSourceName}' at location '{retrievalContext.Path}' is not supported yet.");
break;
}
var sources = retrievalContext.ToSources();
preparedContexts.Add(new(retrievalContext, body.ToString(), sources, CreatePassage(retrievalContext, sources, matchedText, surroundingContent)));
}
return preparedContexts;
}
/// <summary>
/// Writes what the AI is told about one retrieval context, filtered for prompt injections.
/// Writes the content of a text context, after its description.
/// </summary>
/// <param name="retrievalContext">The context to describe.</param>
/// <param name="sb">The builder to write into, or null for a new one.</param>
/// <param name="index">The number of the context among all of them, or -1 when it stands alone.</param>
/// <param name="numTotalRetrievalContexts">How many contexts there are, or -1 when that is unknown.</param>
/// <param name="reportingScope">The action the context is retrieved for, or null when it is retrieved for none.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>Everything written into the builder so far.</returns>
public static async Task<string> AsMarkdown(this IRetrievalContext retrievalContext, StringBuilder? sb = null, int index = -1, int numTotalRetrievalContexts = -1, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
/// <param name="contextBuilder">The builder to write into.</param>
/// <param name="matchedText">The matched text, filtered.</param>
/// <param name="surroundingContent">The surrounding content, filtered.</param>
private static void AppendTextMatch(StringBuilder contextBuilder, string matchedText, IReadOnlyList<string> surroundingContent)
{
sb ??= new StringBuilder();
var contextBuilder = new StringBuilder();
switch (index)
{
case > 0 when numTotalRetrievalContexts is -1:
contextBuilder.AppendLine($"# Retrieval context {index}");
break;
case > 0 when numTotalRetrievalContexts > 0:
contextBuilder.AppendLine($"# Retrieval context {index} of {numTotalRetrievalContexts}");
break;
default:
contextBuilder.AppendLine("# Retrieval context");
break;
}
AppendContextDescription(contextBuilder, retrievalContext);
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
var source = PromptInjectionSource.RetrievalContext(retrievalContext.DataSourceName, retrievalContext.Path);
switch(retrievalContext)
{
case RetrievalTextContext textContext:
contextBuilder.AppendLine();
contextBuilder.AppendLine("Matched text content:");
contextBuilder.AppendLine("````");
contextBuilder.AppendLine(textContext.MatchedText);
contextBuilder.AppendLine("````");
if(textContext.SurroundingContent.Count > 0)
{
contextBuilder.AppendLine();
contextBuilder.AppendLine("Surrounding text content:");
foreach(var surrounding in textContext.SurroundingContent)
{
contextBuilder.AppendLine();
contextBuilder.AppendLine("````");
contextBuilder.AppendLine(surrounding);
contextBuilder.AppendLine("````");
}
}
await FilterWhatWeHaveSoFar();
break;
case RetrievalImageContext:
AppendImageMatch(contextBuilder);
await FilterWhatWeHaveSoFar();
break;
default:
await FilterWhatWeHaveSoFar();
LOGGER.LogWarning($"The retrieval content type '{retrievalContext.Type}' of data source '{retrievalContext.DataSourceName}' at location '{retrievalContext.Path}' is not supported yet.");
break;
}
contextBuilder.AppendLine();
sb.Append(contextBuilder);
return sb.ToString();
contextBuilder.AppendLine("Matched text content:");
contextBuilder.AppendLine("````");
contextBuilder.AppendLine(matchedText);
contextBuilder.AppendLine("````");
//
// Replaces what has been built so far with its filtered version. A data source is as
// untrusted as any other external content: it may serve text written to steer the model
// rather than to answer the query.
//
async Task FilterWhatWeHaveSoFar()
if (surroundingContent.Count is 0)
return;
contextBuilder.AppendLine();
contextBuilder.AppendLine("Surrounding text content:");
foreach (var surrounding in surroundingContent)
{
var sanitized = await guardService.SanitizeAsync(contextBuilder.ToString(), source, reportingScope);
contextBuilder.Clear();
contextBuilder.Append(sanitized);
contextBuilder.AppendLine();
contextBuilder.AppendLine("````");
contextBuilder.AppendLine(surrounding);
contextBuilder.AppendLine("````");
}
}
/// <summary>
/// Keeps what the user is shown of a context the AI received.
/// </summary>
/// <remarks>
/// The name and the path stay as the data source delivered them: they tell the user where the
/// passage is, and the AI read them in the filtered description anyway.
/// </remarks>
/// <param name="retrievalContext">The context.</param>
/// <param name="sources">The sources of the context.</param>
/// <param name="matchedText">The matched text, filtered; empty for anything but text.</param>
/// <param name="surroundingContent">The surrounding content, filtered.</param>
/// <returns>The passage, or null when the context has no source.</returns>
private static SourcePassage? CreatePassage(IRetrievalContext retrievalContext, IReadOnlyList<Source> sources, string matchedText, IReadOnlyList<string> surroundingContent)
{
//
// The user finds a passage under the source it belongs to. A context without a source, i.e.,
// one which names neither a document nor its data source, is listed nowhere, so its passage
// could never be shown and would only take up room:
//
if (sources.Count is 0)
return null;
var textContext = retrievalContext as RetrievalTextContext;
return new SourcePassage
{
SourceUrl = sources[0].URL,
DataSourceName = retrievalContext.DataSourceName,
Path = retrievalContext.Path,
Category = retrievalContext.Category,
Type = retrievalContext.Type,
Text = matchedText,
SurroundingContent = surroundingContent,
PageNumber = textContext?.PageNumber is > 0 ? textContext.PageNumber : null,
EndPageNumber = textContext?.EndPageNumber is > 0 ? textContext.EndPageNumber : null,
SheetName = textContext?.SheetName is { Length: > 0 } sheetName ? sheetName : null,
StartRow = textContext?.StartRow is > 0 ? textContext.StartRow : null,
EndRow = textContext?.EndRow is > 0 ? textContext.EndRow : null,
StartLine = textContext?.StartLine is > 0 ? textContext.StartLine : null,
EndLine = textContext?.EndLine is > 0 ? textContext.EndLine : null,
ChunkNumber = textContext?.ChunkNumber is > 0 ? textContext.ChunkNumber : null,
};
}
/// <summary>
/// The sources a retrieval context lends to an answer, as they are listed below it.
/// </summary>
/// <remarks>
/// The reference comes first: the title and link of the passage itself where the data source
/// names them, e.g., a local file with its page, and otherwise the data source and the path.
/// The further links of the context follow. Only what can be opened becomes a source, i.e., a
/// The further links of the context follow. Only what can be opened becomes a link, i.e., a
/// web address or a file with an absolute path. A relative path would point elsewhere depending
/// on where it is opened from.
/// on where it is opened from.<br/><br/>
///
/// A context which names nothing to open at all, e.g., by a path which only its ERI server knows,
/// still lends one source: the AI received it, and its passage needs a place to be read at. That
/// source is named as text, see SourceExtensions.CreateRetrievalSourceUrl.
/// </remarks>
/// <param name="retrievalContext">The retrieval context.</param>
/// <returns>The sources, which may be none.</returns>
/// <returns>The sources; none only for a context which names neither a document nor its data source.</returns>
public static IReadOnlyList<Source> ToSources(this IRetrievalContext retrievalContext)
{
var sources = new List<Source>();
AddSource(sources, GetReferenceTitle(retrievalContext), GetReferenceLink(retrievalContext));
var referenceTitle = GetReferenceTitle(retrievalContext);
AddSource(sources, referenceTitle, GetReferenceLink(retrievalContext));
foreach (var link in retrievalContext.Links)
AddSource(sources, retrievalContext.DataSourceName, link);
if (sources.Count is 0 && !string.IsNullOrWhiteSpace(referenceTitle))
sources.Add(new Source(referenceTitle, SourceExtensions.CreateRetrievalSourceUrl(referenceTitle, retrievalContext.Path), SourceOrigin.RAG));
return sources;
}
@@ -0,0 +1,15 @@
namespace AIStudio.Tools.RAG;
/// <summary>
/// Where a passage is in its document, counted from one: the pages it spans, or the part it is of a
/// document without pages.
/// </summary>
/// <param name="First">The page the passage starts on, or its part.</param>
/// <param name="Last">The page the passage ends on; the first one again when it stays on one, and always for a part.</param>
public readonly record struct LocationRange(int First, int Last)
{
/// <summary>
/// The range as it is written, e.g., "4" or "4–5".
/// </summary>
public override string ToString() => this.First == this.Last ? $"{this.First}" : $"{this.First}–{this.Last}";
}
@@ -0,0 +1,26 @@
namespace AIStudio.Tools.RAG;
/// <summary>
/// Where passages are in their document: their pages or parts, and the words around them.
/// </summary>
/// <remarks>
/// Kept in parts rather than as one text, so that each page can be shown as a link of its own,
/// while the words around the numbers still come from one translation, e.g., "Pages " and nothing
/// for "Pages 1, 11". Written out as a whole, see ToString, it reads the same as the parts.
/// </remarks>
/// <param name="Prefix">What stands before the numbers, e.g., "Pages ".</param>
/// <param name="Ranges">The pages or parts, in the order of the document, each once.</param>
/// <param name="Suffix">What stands after the numbers, often nothing.</param>
/// <param name="ArePages">Whether the numbers are pages, at which a document can be opened, or the parts of a document without pages.</param>
public sealed record PassageLocation(string Prefix, IReadOnlyList<LocationRange> Ranges, string Suffix, bool ArePages)
{
/// <summary>
/// What stands between two ranges.
/// </summary>
public const string RANGE_SEPARATOR = ", ";
/// <summary>
/// The location as one text, e.g., "Pages 1, 11" or "Part 3".
/// </summary>
public override string ToString() => $"{this.Prefix}{string.Join(RANGE_SEPARATOR, this.Ranges)}{this.Suffix}";
}
@@ -0,0 +1,48 @@
using System.Text;
namespace AIStudio.Tools.RAG;
/// <summary>
/// A retrieval context as the AI receives it: filtered for prompt injections, written out, and
/// ready to be numbered.
/// </summary>
/// <remarks>
/// Made by IRetrievalContextExtensions.PrepareForModelAsync. Which number a context gets depends on
/// which contexts the caller hands on in the end, so the header is written only then, see
/// ToMarkdown. The app writes it itself, which is why it needs no filter.
/// </remarks>
/// <param name="Context">The context as the data source delivered it.</param>
/// <param name="Body">What the AI is told about the context below its header: the filtered description, then the filtered content.</param>
/// <param name="Sources">The sources the context lends to the answer, which may be none.</param>
/// <param name="Passage">What the user is shown of the context, or null when the context has no source to list it under.</param>
public sealed record PreparedRetrievalContext(IRetrievalContext Context, string Body, IReadOnlyList<Source> Sources, SourcePassage? Passage)
{
/// <summary>
/// Writes what the AI is told about the context, headed by its number.
/// </summary>
/// <param name="index">The number of the context among all of them, or null when it stands alone.</param>
/// <param name="numTotalRetrievalContexts">How many contexts there are, or null when that is unknown.</param>
/// <returns>The context as Markdown, ending with an empty line to set it apart from the next one.</returns>
public string ToMarkdown(int? index = null, int? numTotalRetrievalContexts = null)
{
var sb = new StringBuilder();
switch (index)
{
case > 0 when numTotalRetrievalContexts is null:
sb.AppendLine($"# Retrieval context {index}");
break;
case > 0 when numTotalRetrievalContexts > 0:
sb.AppendLine($"# Retrieval context {index} of {numTotalRetrievalContexts}");
break;
default:
sb.AppendLine("# Retrieval context");
break;
}
sb.Append(this.Body);
sb.AppendLine();
return sb.ToString();
}
}
@@ -225,10 +225,13 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
//
// Perform the augmentation of the R(A)G process:
//
IReadOnlyList<PreparedRetrievalContext> usedContexts = [];
if (proceedWithRAG)
{
var augmentationProcess = new AugmentationOne();
chatThread = await augmentationProcess.ProcessAsync(provider, lastUserPrompt, chatThread, dataContexts, token);
var augmentation = await augmentationProcess.ProcessAsync(provider, lastUserPrompt, chatThread, dataContexts, token);
chatThread = augmentation.ChatThread;
usedContexts = augmentation.UsedContexts;
}
//
@@ -240,24 +243,19 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
chatThread.AugmentedData = WithUnreachableDataSources(chatThread.AugmentedData, dataSources.GetUnreachableOfTheChat(chatThread.DataSourceOptions));
//
// Add sources from the selected data
// Add the sources and passages of what the AI received:
//
// We know that the last block is the AI answer block (cf. check above):
var aiAnswerBlock = chatThread.Blocks.Last();
var aiAnswerSources = aiAnswerBlock.Content?.Sources;
var aiAnswerContent = aiAnswerBlock.Content;
// It should never happen that the AI answer block does not contain a content part.
// Just in case, we check this:
if(aiAnswerSources is null)
if(aiAnswerContent is null)
return chatThread;
var ragSources = new List<ISource>();
foreach (var retrievalContext in dataContexts)
ragSources.AddRange(retrievalContext.ToSources());
// Merge the sources, avoiding duplicates:
aiAnswerSources.MergeSources(ragSources);
AddWhatTheAIReceived(aiAnswerContent, usedContexts);
}
return chatThread;
@@ -265,6 +263,26 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
#endregion
/// <summary>
/// Lists below the answer what the AI received: the sources and the passages of the retrieval
/// contexts handed to it.
/// </summary>
/// <remarks>
/// Only those: a passage which the check of which passages fit left out never reached the AI,
/// and the answer does not rest on it. Only a text answer keeps passages; every answer keeps
/// its sources.
/// </remarks>
/// <param name="aiAnswerContent">The content of the answer.</param>
/// <param name="usedContexts">The retrieval contexts handed to the AI.</param>
internal static void AddWhatTheAIReceived(IContent aiAnswerContent, IReadOnlyList<PreparedRetrievalContext> usedContexts)
{
var sources = usedContexts.SelectMany(context => context.Sources);
if (aiAnswerContent is ContentText aiAnswerText)
aiAnswerText.AddSources(sources, usedContexts.Select(context => context.Passage).OfType<SourcePassage>());
else
aiAnswerContent.Sources.MergeSources(sources);
}
/// <summary>
/// Restricts where the chat may still send data, because of the passages the retrieval brought in.
/// </summary>
@@ -63,4 +63,45 @@ public sealed class RetrievalTextContext : IRetrievalContext
/// to. A slide or a sheet has no page and leaves this empty.
/// </remarks>
public int? PageNumber { get; init; }
/// <summary>
/// The page this passage ends on when it goes on beyond the page it starts on, or null when it
/// stays on that page or the end is not known.
/// </summary>
public int? EndPageNumber { get; init; }
/// <summary>
/// The sheet of a spreadsheet this passage is from, or null when it is from none.
/// </summary>
public string? SheetName { get; init; }
/// <summary>
/// The row this passage starts at, counted from one: a row of its sheet, or of a table such as a CSV file.
/// </summary>
public int? StartRow { get; init; }
/// <summary>
/// The row this passage ends at when it goes on beyond the row it starts at, or null.
/// </summary>
public int? EndRow { get; init; }
/// <summary>
/// The line this passage starts at, counted from one, for text and source code, or null.
/// </summary>
public int? StartLine { get; init; }
/// <summary>
/// The line this passage ends at when it goes on beyond the line it starts at, or null.
/// </summary>
public int? EndLine { get; init; }
/// <summary>
/// Which part of its document this passage is, counted from one, or null when the data source does not say.
/// </summary>
/// <remarks>
/// A document is split into parts to be searched, and each part is one passage. The page says
/// better where a passage is, so the part only names the place in a document without pages,
/// such as a text file.
/// </remarks>
public int? ChunkNumber { get; init; }
}
@@ -0,0 +1,130 @@
using System.Text.Json.Serialization;
namespace AIStudio.Tools.RAG;
/// <summary>
/// One passage a data source handed to the AI for an answer, as the AI received it.
/// </summary>
/// <remarks>
/// Stored with the answer, so the user can still read what the answer rests on after the chat was
/// reloaded: where it is with the answer, its texts in the provenance store, see ProvenanceId. The
/// texts are the filtered ones: whatever the prompt injection filter removed did not reach the AI,
/// so it is not shown as if it had.<br/><br/>
/// A passage belongs to the source it is listed under, which is why it is only made for a context
/// with a source, see SourceUrl. A context which names nothing to open, e.g., by a path which only
/// its ERI server knows, gets a source named as text, see IRetrievalContextExtensions.ToSources.<br/><br/>
/// Plain stored values only: a computed property would be written into every chat file as well.
/// </remarks>
public sealed record SourcePassage
{
/// <summary>
/// The address of the source this passage is listed under.
/// </summary>
/// <remarks>
/// The first source of the context, i.e., its own reference, such as the file with the page of
/// the passage. The list of sources tells sources apart by their address without the fragment,
/// and so does the search for the passages of a source, see GetPassagesOf.
/// </remarks>
public string SourceUrl { get; init; } = string.Empty;
/// <summary>
/// The name of where the passage was found, e.g., a document or a database.
/// </summary>
public string DataSourceName { get; init; } = string.Empty;
/// <summary>
/// The path to the content, e.g., a file path or a URL.
/// </summary>
public string Path { get; init; } = string.Empty;
/// <summary>
/// The category of the content, like text or image.
/// </summary>
public RetrievalContentCategory Category { get; init; }
/// <summary>
/// What type of content the passage is, like a document or a transcript.
/// </summary>
public RetrievalContentType Type { get; init; }
/// <summary>
/// The entry of the provenance store which holds the matched text and the content around it,
/// or null when the passage holds them itself.
/// </summary>
/// <remarks>
/// Set once the answer is complete, see ContentText.StoreReceivedContentAsync. Until then, and
/// whenever the store could not take them, the passage holds both texts itself; once they are
/// stored, Text and SurroundingContent are empty. Left out of the chat file while null.
/// </remarks>
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
public Guid? ProvenanceId { get; init; }
/// <summary>
/// The matched text, filtered for prompt injections, as long as it is not stored, see ProvenanceId.
/// </summary>
/// <remarks>
/// Empty for an image which matched: the AI is only told that it did, never shown the image.
/// </remarks>
public string Text { get; init; } = string.Empty;
/// <summary>
/// The content around the matched text, filtered for prompt injections, as long as it is not stored, see ProvenanceId.
/// </summary>
public IReadOnlyList<string> SurroundingContent { get; init; } = [];
/// <summary>
/// The page the passage starts on, or null when it has none.
/// </summary>
public int? PageNumber { get; init; }
/// <summary>
/// The page the passage ends on, or null when it is not known.
/// </summary>
public int? EndPageNumber { get; init; }
//
// Where a passage without pages is. Left out of the chat file while unknown, which they are for
// every passage of a document with pages, and for every passage stored before they existed:
//
/// <summary>
/// The sheet of a spreadsheet the passage is from, or null.
/// </summary>
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
public string? SheetName { get; init; }
/// <summary>
/// The row the passage starts at, counted from one: a row of its sheet, or of a table such as a CSV file.
/// </summary>
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
public int? StartRow { get; init; }
/// <summary>
/// The row the passage ends at, or null when it is not known.
/// </summary>
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
public int? EndRow { get; init; }
/// <summary>
/// The line the passage starts at, counted from one, for text and source code, or null.
/// </summary>
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
public int? StartLine { get; init; }
/// <summary>
/// The line the passage ends at, or null when it is not known.
/// </summary>
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
public int? EndLine { get; init; }
/// <summary>
/// Which part of its document the passage is, counted from one, or null when the data source
/// does not say.
/// </summary>
/// <remarks>
/// Names the place of a passage which says neither pages, nor rows, nor lines, e.g., one
/// indexed before these were kept. Wherever they are known, they say it better, see
/// SourcePassageExtensions.GetLocation.
/// </remarks>
public int? ChunkNumber { get; init; }
}
@@ -0,0 +1,318 @@
using AIStudio.Tools.Databases.ProvenanceStore;
using AIStudio.Tools.PluginSystem;
namespace AIStudio.Tools.RAG;
public static class SourcePassageExtensions
{
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(SourcePassageExtensions).Namespace, nameof(SourcePassageExtensions));
/// <summary>
/// Whether the passage comes from a transcript, e.g., of an audio or a video file.
/// </summary>
/// <param name="passage">The passage.</param>
/// <returns>True for a transcript.</returns>
public static bool IsTranscript(this SourcePassage passage) => passage.Type is RetrievalContentType.TEXT_TRANSCRIPT;
/// <summary>
/// What sort of content the passage is, as the provenance store files it.
/// </summary>
/// <remarks>
/// Told by what the content is rather than by where it came from: a mail is a mail, whether a
/// mailbox data source or a tool handed it over.
/// </remarks>
/// <param name="passage">The passage.</param>
/// <returns>The kind of its entry.</returns>
public static ProvenanceKind GetProvenanceKind(this SourcePassage passage) => passage.Type switch
{
RetrievalContentType.TEXT_EMAIL => ProvenanceKind.MAIL,
RetrievalContentType.TEXT_WEBSITE => ProvenanceKind.WEB_PAGE,
_ => ProvenanceKind.DATA_SOURCE_PASSAGE,
};
/// <summary>
/// Makes the passage of a web page which a tool read for the AI, so that the user can read under
/// the source of the page what the AI received from it.
/// </summary>
/// <remarks>
/// Listed under the source by its address, see GetPassagesOf, and named by its title in the
/// dialog of its passages. A page says nothing about pages or lines, so its source keeps the
/// title the tool gave it, see NamePassagesInTitles.
/// </remarks>
/// <param name="source">The source of the page, as the tool lists it.</param>
/// <param name="content">The content of the page as the model received it: filtered for prompt injections, and cut where the tool cut it.</param>
/// <returns>The passage of the page.</returns>
public static SourcePassage ToWebPagePassage(this Source source, string content) => new()
{
SourceUrl = source.URL,
DataSourceName = source.Title,
Path = source.URL,
Category = RetrievalContentCategory.TEXT,
Type = RetrievalContentType.TEXT_WEBSITE,
Text = content,
};
/// <summary>
/// Makes the passage of a mail which a tool handed to the AI, so that the user can read under
/// the source of the mail what the AI received of it.
/// </summary>
/// <remarks>
/// Listed and named like the passage of a web page, see ToWebPagePassage. A mail has no path:
/// its address names it for AI Studio only, so there is nothing to copy or to open.
/// </remarks>
/// <param name="source">The source of the mail, as the tool lists it.</param>
/// <param name="content">What the model received of the mail, filtered for prompt injections.</param>
/// <returns>The passage of the mail.</returns>
public static SourcePassage ToMailPassage(this Source source, string content) => new()
{
SourceUrl = source.URL,
DataSourceName = source.Title,
Category = RetrievalContentCategory.TEXT,
Type = RetrievalContentType.TEXT_EMAIL,
Text = content,
};
/// <summary>
/// Whether the passage holds texts which belong into the provenance store, see SourcePassage.ProvenanceId.
/// </summary>
/// <param name="passage">The passage.</param>
/// <returns>False once the texts are stored, and for a passage without any, such as an image.</returns>
public static bool HoldsTextsToStore(this SourcePassage passage) => passage.ProvenanceId is null && (passage.Text.Length > 0 || passage.SurroundingContent.Count > 0);
/// <summary>
/// What stands between the sheets of a spreadsheet in a location, e.g., "Budget: rows 3–9 · Summary: row 2".
/// </summary>
private const string SHEET_SEPARATOR = " · ";
/// <summary>
/// Names where the passage is in its document, e.g., "Page 4", "Pages 4–5", "Budget: rows 3–9", "Lines 40–85", or "Part 3".
/// </summary>
/// <param name="passage">The passage.</param>
/// <returns>The location, or an empty string when the data source does not say.</returns>
public static string GetLocationLabel(this SourcePassage passage) => GetLocationLabel([passage]);
/// <summary>
/// Where the passage is in its document, as numbers only, e.g., its pages 4–5, its lines 40–85, or its part 3.
/// </summary>
/// <remarks>
/// Orders the passages of one document, in the order of the document. The passages of the
/// sheets of a spreadsheet give their parts, since rows start anew in every sheet.
/// </remarks>
/// <param name="passage">The passage.</param>
/// <returns>The numbers, or null when the data source does not say.</returns>
public static LocationRange? GetLocationRange(this SourcePassage passage) => GetLocation([passage])?.Ranges[0];
/// <summary>
/// Names where several passages are in their document, e.g., "Pages 1, 11", "Budget: rows 3–9 · Summary: row 2", "Lines 40–85, 120" or "Parts 2, 5".
/// </summary>
/// <param name="passages">The passages.</param>
/// <returns>The location, or an empty string when the data source does not say.</returns>
public static string GetLocationLabel(this IEnumerable<SourcePassage> passages)
{
var passagesOfDocument = passages.ToList();
if (passagesOfDocument.All(passage => passage.PageNumber is not > 0) && GetSheetsLabel(passagesOfDocument) is { } sheets)
return sheets;
return passagesOfDocument.GetLocation()?.ToString() ?? string.Empty;
}
/// <summary>
/// Finds where several passages are in their document, in parts, so that each page can be shown
/// as a link.
/// </summary>
/// <remarks>
/// A document with pages is named by its pages: they are what the user finds again in the
/// document, and what a program can open it at. A table without sheets, such as a CSV file, is
/// named by its rows, and a text by its lines. Rows and lines are nothing a program can open a
/// document at, so they become no links. A document which says none of them is named by the
/// parts it was split into for the search. Either way, the numbers follow the order of the
/// document, each once: two passages on the same page name that page once. An end which is not
/// after the start names no range.<br/><br/>
/// The rows of a spreadsheet start anew in every sheet, so they need the name of their sheet
/// and are named by GetLocationLabel, sheet by sheet. Here, its passages give their parts.
/// </remarks>
/// <param name="passages">The passages.</param>
/// <returns>The location, or null when the data source says nothing about it.</returns>
public static PassageLocation? GetLocation(this IEnumerable<SourcePassage> passages)
{
var passagesOfDocument = passages.ToList();
var pages = passagesOfDocument
.Where(passage => passage.PageNumber is > 0)
.Select(passage => RangeOf(passage.PageNumber!.Value, passage.EndPageNumber))
.ToList();
if (pages.Count > 0)
return CreateLocation(pages, TB("Page {0}"), TB("Pages {0}"), arePages: true);
var rows = passagesOfDocument
.Where(passage => passage is { StartRow: > 0, SheetName: null or "" })
.Select(passage => RangeOf(passage.StartRow!.Value, passage.EndRow))
.ToList();
if (rows.Count > 0)
return CreateLocation(rows, TB("Row {0}"), TB("Rows {0}"), arePages: false);
var lines = passagesOfDocument
.Where(passage => passage.StartLine is > 0)
.Select(passage => RangeOf(passage.StartLine!.Value, passage.EndLine))
.ToList();
if (lines.Count > 0)
return CreateLocation(lines, TB("Line {0}"), TB("Lines {0}"), arePages: false);
var parts = passagesOfDocument
.Where(passage => passage.ChunkNumber is > 0)
.Select(passage => new LocationRange(passage.ChunkNumber!.Value, passage.ChunkNumber!.Value))
.ToList();
return parts.Count > 0 ? CreateLocation(parts, TB("Part {0}"), TB("Parts {0}"), arePages: false) : null;
}
/// <summary>
/// Names the rows of passages from the sheets of a spreadsheet, sheet by sheet, e.g., "Budget: rows 3–9, 20 · Summary: row 2".
/// </summary>
/// <remarks>
/// The sheets follow the order of the document, which is the order of the parts the passages
/// are. A passage never runs from one sheet into the next, cf. TextChunker.SplitBySheets.
/// </remarks>
/// <param name="passages">The passages of one document.</param>
/// <returns>The location, or null when no passage names a sheet.</returns>
private static string? GetSheetsLabel(IReadOnlyList<SourcePassage> passages)
{
var sheets = passages
.Where(passage => passage is { StartRow: > 0, SheetName.Length: > 0 })
.OrderBy(passage => passage.ChunkNumber ?? int.MaxValue)
.GroupBy(passage => passage.SheetName!, StringComparer.Ordinal)
.Select(sheet => GetSheetLabel(sheet.Key, sheet.Select(passage => RangeOf(passage.StartRow!.Value, passage.EndRow))))
.ToList();
return sheets.Count > 0 ? string.Join(SHEET_SEPARATOR, sheets) : null;
}
private static string GetSheetLabel(string sheetName, IEnumerable<LocationRange> rows)
{
var orderedRows = rows
.Distinct()
.OrderBy(range => range.First)
.ThenBy(range => range.Last)
.ToList();
var format = orderedRows is [var single] && single.First == single.Last ? TB("{0}: row {1}") : TB("{0}: rows {1}");
return string.Format(format, sheetName, string.Join(PassageLocation.RANGE_SEPARATOR, orderedRows));
}
/// <summary>
/// The numbers from a start to an end; an end which is not after the start names no range.
/// </summary>
private static LocationRange RangeOf(int first, int? last) => new(first, Math.Max(first, last ?? 0));
private static PassageLocation CreateLocation(IEnumerable<LocationRange> ranges, string oneFormat, string severalFormat, bool arePages)
{
var orderedRanges = ranges
.Distinct()
.OrderBy(range => range.First)
.ThenBy(range => range.Last)
.ToList();
//
// The words around the numbers come from the translation, which may put them before or after
// the numbers. A translation which lost its placeholder still gets the numbers, after its words:
//
const string PLACEHOLDER = "{0}";
var format = orderedRanges is [var single] && single.First == single.Last ? oneFormat : severalFormat;
var placeholderAt = format.IndexOf(PLACEHOLDER, StringComparison.Ordinal);
return placeholderAt < 0
? new PassageLocation($"{format} ", orderedRanges, string.Empty, arePages)
: new PassageLocation(format[..placeholderAt], orderedRanges, format[(placeholderAt + PLACEHOLDER.Length)..], arePages);
}
/// <summary>
/// Names in the title of each source where its passages are, e.g., "handbook.pdf (Pages 1, 11)".
/// </summary>
/// <remarks>
/// The list of sources keeps one entry per document, see SourceExtensions.MergeSources, and its
/// title came from the passage found first, e.g., "handbook.pdf (Page 1)". A second passage on
/// page 11 would go unmentioned, so the title is made anew from all passages of the source. A
/// spreadsheet is named by its sheets and rows the same way, a text file by its lines, see
/// GetLocationLabel. Passages whose data source
/// says neither, e.g., an external one, are counted instead as soon as there are several of
/// them. Otherwise, a source keeps the title its data source gave it, and so does every source
/// without passages, such as a web page.
/// </remarks>
/// <param name="sources">The sources of an answer.</param>
/// <param name="passages">The passages of the same answer.</param>
public static void NamePassagesInTitles(this IList<Source> sources, IReadOnlyCollection<SourcePassage> passages)
{
if (passages.Count is 0)
return;
var passagesBySource = passages.SortBySource();
for (var index = 0; index < sources.Count; index++)
{
var passagesOfSource = passagesBySource.GetPassagesOf(sources[index]);
if (passagesOfSource.Count is 0 || string.IsNullOrWhiteSpace(passagesOfSource[0].DataSourceName))
continue;
var location = passagesOfSource.GetLocationLabel();
if (location.Length is 0 && passagesOfSource.Count > 1)
location = string.Format(TB("{0} passages"), passagesOfSource.Count.CompactCount());
if (location.Length is 0)
continue;
sources[index] = sources[index] with { Title = $"{passagesOfSource[0].DataSourceName} ({location})" };
}
}
/// <summary>
/// Merges added passages into an existing list of passages, leaving out those it holds already.
/// </summary>
/// <remarks>
/// A passage is known when one with the same address and the same text is, e.g., the same chunk
/// found by two searches for one answer. Unlike the sources, the address counts with its page:
/// the same sentence on two pages of a document is two passages. A stored passage is told
/// apart by its entry, since its text lies in the provenance store.
/// </remarks>
/// <param name="passages">The existing list of passages to merge into.</param>
/// <param name="addedPassages">The passages to add.</param>
public static void MergePassages(this IList<SourcePassage> passages, IEnumerable<SourcePassage> addedPassages)
{
var knownPassages = passages
.Select(passage => (passage.SourceUrl, passage.Text, passage.ProvenanceId))
.ToHashSet();
foreach (var addedPassage in addedPassages)
{
if (knownPassages.Add((addedPassage.SourceUrl, addedPassage.Text, addedPassage.ProvenanceId)))
passages.Add(addedPassage);
}
}
/// <summary>
/// Finds the passages listed under a source.
/// </summary>
/// <remarks>
/// The list of sources keeps one entry per document, whichever page its link names, see
/// SourceExtensions.MergeSources. The passages are matched the same way, so the entry of a
/// document shows the passages of all its pages, and every passage shows up under its source.
/// </remarks>
/// <param name="passages">The passages of an answer.</param>
/// <param name="source">The source.</param>
/// <returns>The passages of the source, in their order; empty when it has none.</returns>
public static IReadOnlyList<SourcePassage> GetPassagesOf(this IEnumerable<SourcePassage> passages, ISource source) => passages.SortBySource().GetPassagesOf(source);
/// <summary>
/// Sorts passages by the source they are listed under, to find the passages of many sources.
/// </summary>
/// <param name="passages">The passages of an answer.</param>
/// <returns>The passages by the identity of their source, see SourceExtensions.GetSourceIdentity.</returns>
public static ILookup<string, SourcePassage> SortBySource(this IEnumerable<SourcePassage> passages) => passages.ToLookup(passage => SourceExtensions.GetSourceIdentity(passage.SourceUrl), StringComparer.Ordinal);
/// <summary>
/// Finds the passages listed under a source, among passages sorted by their sources.
/// </summary>
/// <param name="passagesBySource">The passages of an answer, see SortBySource.</param>
/// <param name="source">The source.</param>
/// <returns>The passages of the source, in their order; empty when it has none.</returns>
public static IReadOnlyList<SourcePassage> GetPassagesOf(this ILookup<string, SourcePassage> passagesBySource, ISource source) => [..passagesBySource[SourceExtensions.GetSourceIdentity(source.URL)]];
}
@@ -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()
{
@@ -0,0 +1,81 @@
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.Rust;
using AIStudio.Tools.Services;
namespace AIStudio.Tools;
/// <summary>
/// What the user can do with the document a source names: open it, or show it in the file manager.
/// </summary>
/// <remarks>
/// The list of sources and the passages of a source both offer it, and a failure reaches the user
/// the same way from either of them.
/// </remarks>
public static class SourceDocumentLocationExtensions
{
private static readonly ILogger LOGGER = Program.LOGGER_FACTORY.CreateLogger(nameof(SourceDocumentLocationExtensions));
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(SourceDocumentLocationExtensions).Namespace, nameof(SourceDocumentLocationExtensions));
/// <summary>
/// Opens the document in the program the system uses for it, at its page.
/// </summary>
/// <remarks>
/// Whether the program can be sent to a page is the runtime's business, and it says afterwards
/// whether it managed to. Nothing is shown about that: the document is open, and the page is
/// named where the user asked to open it.
/// </remarks>
/// <param name="document">The document to open, and the page to show.</param>
/// <param name="rustService">The runtime, which opens the document.</param>
public static async Task OpenAsync(this SourceDocumentLocation document, RustService rustService)
{
OpenDocumentResponse response;
try
{
response = await rustService.TryOpenDocumentInSystemViewer(document.Path, document.PageNumber);
}
catch (Exception e)
{
LOGGER.LogWarning(e, "Could not open a source document.");
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.Description, TB("Could not open the document.")));
return;
}
if (response.Success)
return;
var issue = string.IsNullOrWhiteSpace(response.Issue) ? TB("Unknown error") : response.Issue;
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.Description, string.Format(TB("Could not open the document: {0}"), issue)));
}
/// <summary>
/// Opens the file browser of the system and selects the document in it.
/// </summary>
/// <remarks>
/// The second way out: a document which the system opens in the wrong program, or which the user
/// wants to move or send on instead of read, is reached from here without being opened. This is
/// the same way out the embeddings page offers for a file it could not read.
/// </remarks>
/// <param name="document">The document to show.</param>
/// <param name="rustService">The runtime, which opens the file browser.</param>
public static async Task ShowInFileManagerAsync(this SourceDocumentLocation document, RustService rustService)
{
OpenPathResponse response;
try
{
response = await rustService.TryOpenPathInRuntimeFileManager(document.Path);
}
catch (Exception e)
{
LOGGER.LogWarning(e, "Could not show a source document in the file manager.");
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.FolderOpen, TB("Could not open the file location.")));
return;
}
if (response.Success)
return;
var issue = string.IsNullOrWhiteSpace(response.Issue) ? TB("Unknown error") : response.Issue;
await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.FolderOpen, string.Format(TB("Could not open the file location: {0}"), issue)));
}
}
@@ -36,6 +36,36 @@ public static partial class SourceExtensions
public static bool IsMailSource(this ISource source) =>
Uri.TryCreate(source.URL.Trim(), UriKind.Absolute, out var address) && address.Scheme.Equals(MAIL_SOURCE_SCHEME, StringComparison.OrdinalIgnoreCase);
/// <summary>
/// The scheme of the address a passage of a data source leaves as a source, when its path is nothing a program can open.
/// </summary>
private const string RETRIEVAL_SOURCE_SCHEME = "retrieval";
/// <summary>
/// The address of the source a passage of a data source leaves in a chat, when its path is
/// nothing a program can open, e.g., a path which only the ERI server knows.
/// </summary>
/// <remarks>
/// Made from the name and the path the data source gave, so that the passages of one document
/// gather under one source, see MergeSources. The host is the same for all of them, since the
/// identity of an address writes its host in lower case. Like the address of a mail, nothing can
/// open it, see HasNoLink.
/// </remarks>
/// <param name="name">The name of the content, e.g., the document.</param>
/// <param name="path">The path the data source gave, which may be empty.</param>
/// <returns>The address.</returns>
public static string CreateRetrievalSourceUrl(string name, string path) => $"{RETRIEVAL_SOURCE_SCHEME}://context/{Uri.EscapeDataString(name)}/{Uri.EscapeDataString(path)}";
/// <summary>
/// Whether nothing can open a source by its address: a mail, or the passage of a data source
/// whose path names nothing to open, see CreateRetrievalSourceUrl.
/// </summary>
/// <param name="source">The source.</param>
/// <returns>True for a source which is named as text, without a link.</returns>
public static bool HasNoLink(this ISource source) =>
Uri.TryCreate(source.URL.Trim(), UriKind.Absolute, out var address)
&& (address.Scheme.Equals(MAIL_SOURCE_SCHEME, StringComparison.OrdinalIgnoreCase) || address.Scheme.Equals(RETRIEVAL_SOURCE_SCHEME, StringComparison.OrdinalIgnoreCase));
private static void AppendMarkdownLink(StringBuilder sb, string title, string url)
{
sb.Append('[');
@@ -181,8 +211,8 @@ public static partial class SourceExtensions
{
sb.Append($"- [{numberedSource.Number}] ");
// A link to a mail would lead nowhere in a document or in the clipboard, so a mail is named only:
if (numberedSource.Source.IsMailSource())
// A link to a mail, or to what only a data source knows, would lead nowhere in a document or in the clipboard, so it is named only:
if (numberedSource.Source.HasNoLink())
sb.Append(EscapeMarkdownLinkText(numberedSource.Source.Title));
else
AppendMarkdownLink(sb, numberedSource.Source.Title, keepPageAnchors ? numberedSource.Source.URL : WithoutPageAnchor(numberedSource.Source.URL));
@@ -305,7 +335,16 @@ public static partial class SourceExtensions
}
}
private static string GetSourceIdentity(string url)
/// <summary>
/// What tells two sources apart: the address without its fragment.
/// </summary>
/// <remarks>
/// The passages of an answer are matched to their sources by the same identity, see
/// SourcePassageExtensions.GetPassagesOf, so a source and its passages always agree.
/// </remarks>
/// <param name="url">The address of a source.</param>
/// <returns>The identity of the source.</returns>
internal static string GetSourceIdentity(string url)
{
var cleanedUrl = url.Trim().Replace("\r", string.Empty).Replace("\n", string.Empty);
if (!Uri.TryCreate(cleanedUrl, UriKind.Absolute, out var absoluteUri))
@@ -0,0 +1,38 @@
namespace AIStudio.Tools;
/// <summary>
/// Where a piece of an extracted document is, as the runtime stated it.
/// </summary>
/// <remarks>
/// Which values a piece has depends on its document: a page for PDF and Word files, a sheet and a
/// row for spreadsheets, a line for text, CSV and source code. Slides are left out while the
/// runtime may deliver the slides of a deck out of order: a slide number would point at the wrong
/// slide.<br/><br/>
/// What the app writes around the content, such as the heading of a sheet or the fence around a
/// table, has no row or line of its own: the runtime numbers it 0, and it is stated as none here.
/// </remarks>
/// <param name="PageNumber">The page, counted from one, or null.</param>
/// <param name="SheetName">The sheet, or null.</param>
/// <param name="RowNumber">The row of the sheet, counted from one as a spreadsheet program counts it, or null.</param>
/// <param name="LineNumber">The line, counted from one, or null.</param>
public readonly record struct SourceLocation(int? PageNumber = null, string? SheetName = null, int? RowNumber = null, int? LineNumber = null)
{
/// <summary>
/// A piece on a page.
/// </summary>
/// <param name="pageNumber">The page, or null when the runtime stated none.</param>
public static SourceLocation OnPage(int? pageNumber) => new(PageNumber: pageNumber is > 0 ? pageNumber : null);
/// <summary>
/// A row of a sheet.
/// </summary>
/// <param name="sheetName">The sheet.</param>
/// <param name="rowNumber">The row, or 0 for what stands before the first row.</param>
public static SourceLocation InSheet(string? sheetName, int? rowNumber) => new(SheetName: string.IsNullOrEmpty(sheetName) ? null : sheetName, RowNumber: rowNumber is > 0 ? rowNumber : null);
/// <summary>
/// A line of a text.
/// </summary>
/// <param name="lineNumber">The line, or 0 for what stands before the first line.</param>
public static SourceLocation OnLine(int? lineNumber) => new(LineNumber: lineNumber is > 0 ? lineNumber : null);
}
@@ -1,6 +1,7 @@
using System.Runtime.CompilerServices;
using AIStudio.Provider;
using AIStudio.Tools.RAG;
namespace AIStudio.Tools.ToolCallingSystem.Harness;
@@ -38,6 +39,7 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
var toolCallCount = 0;
var toolResultCharacterCount = 0L;
var toolSources = new List<Source>();
var toolPassages = new List<SourcePassage>();
var hasStreamedTextBefore = false;
var isFirstRound = true;
@@ -111,11 +113,15 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
//
// No outcome means the round failed: the request errored out, or the stream ended
// mid-sentence. Either way the adapter has already reported it.
// mid-sentence. Either way the adapter has already reported it. The answer keeps what
// was written so far, and with it what the tools of the earlier rounds contributed:
//
if (round is null)
{
await context.ResetToolRuntimeStatusAsync();
if (toolSources.Count > 0 || toolPassages.Count > 0)
yield return WithWhatTheToolsContributed(string.Empty);
yield break;
}
@@ -133,7 +139,7 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
var unanswerable = context.ToolExecutor.CreateInvalidToolCallResult(string.Empty, toolCallCount);
await context.AddToolInvocationAsync(unanswerable.Trace);
await context.ResetToolRuntimeStatusAsync();
yield return new ContentStreamChunk(unanswerable.Content, [..toolSources]);
yield return WithWhatTheToolsContributed(unanswerable.Content);
yield break;
}
@@ -146,10 +152,7 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
// the tools contributed. An empty chunk is how sources travel on their own; the
// streaming paths of the providers attach their annotations the same way.
//
yield return new ContentStreamChunk(
roundAnswered ? string.Empty : NO_ANSWER_AFTER_LIMIT,
[..toolSources]);
yield return WithWhatTheToolsContributed(roundAnswered ? string.Empty : NO_ANSWER_AFTER_LIMIT);
yield break;
}
@@ -158,13 +161,13 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
await context.ResetToolRuntimeStatusAsync();
if (roundAnswered)
{
yield return new ContentStreamChunk(string.Empty, [..toolSources]);
yield return WithWhatTheToolsContributed(string.Empty);
yield break;
}
if (toolCallCount > 0)
{
yield return new ContentStreamChunk(NO_ANSWER_AFTER_TOOL_CALL, [..toolSources]);
yield return WithWhatTheToolsContributed(NO_ANSWER_AFTER_TOOL_CALL);
yield break;
}
@@ -239,6 +242,7 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
context.ChatThread.RequireOutboundDataRestriction(outcome.RequiredOutboundDataRestriction);
context.ChatThread.RuntimeWebAddressesFromTools.UnionWith(outcome.ReturnedWebAddresses);
toolSources.MergeSources(outcome.Sources);
toolPassages.MergePassages(outcome.Passages);
await context.AddToolInvocationAsync(outcome.Trace);
// A blocked call counts as a failure towards the model as much as an errored
@@ -252,6 +256,13 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> logger) : IToolCall
await context.ResetToolRuntimeStatusAsync();
}
}
//
// What the tools contributed goes out with the last chunk of the conversation: the sources,
// and the passages the model received from them. One way out for every way the conversation
// ends, so none of them hands over the sources and forgets the passages.
//
ContentStreamChunk WithWhatTheToolsContributed(string content) => new(content, [..toolSources]) { Passages = [..toolPassages] };
}
private static string GetDisplayName(ToolCallingLoopContext context, string toolName) => context.RunnableTools
@@ -1,5 +1,6 @@
using AIStudio.Provider;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.RAG;
namespace AIStudio.Tools.ToolCallingSystem;
@@ -35,6 +36,11 @@ public sealed record ToolCallOutcome(string Content, ToolInvocationTrace Trace)
/// </summary>
public IReadOnlyList<Source> Sources { get; init; } = [];
/// <summary>
/// The passages of the sources, as the model received them in the result.
/// </summary>
public IReadOnlyList<SourcePassage> Passages { get; init; } = [];
/// <summary>
/// The web addresses in the result, as request keys, except those the model wrote into the call.
/// </summary>
Loaded 100 of 229 files, more files were not shown because too many files have changed in this diff. Show more