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

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

No files matched your search

@@ -5,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.