Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,4 @@ public sealed record EmbeddingStateFile(
DateTimeOffset CreationUtc,
DateTimeOffset LastWriteUtc,
DateTimeOffset EmbeddedAtUtc,
int ChunkCount,
string ConfidenceLevel,
int ConfidenceLevelRank);
int ChunkCount);
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,6 @@ internal sealed class EmbeddingStateFileEntity

public int ChunkCount { get; set; }

public string ConfidenceLevel { get; set; } = string.Empty;

public int ConfidenceLevelRank { get; set; }

public EmbeddingStateDataSourceEntity? DataSource { get; set; }

public List<EmbeddingStateChunkEntity> Chunks { get; set; } = [];
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -67,13 +67,10 @@ protected override void OnModelCreating(ModelBuilder modelBuilder)
entity.Property(file => file.LastWriteUtc).HasColumnName("last_write_utc").HasConversion(utcDateTimeOffsetConverter).IsRequired();
entity.Property(file => file.EmbeddedAtUtc).HasColumnName("embedded_at_utc").HasConversion(utcDateTimeOffsetConverter).IsRequired();
entity.Property(file => file.ChunkCount).HasColumnName("chunk_count");
entity.Property(file => file.ConfidenceLevel).HasColumnName("confidence_level").IsRequired();
entity.Property(file => file.ConfidenceLevelRank).HasColumnName("confidence_level_rank");

entity.HasIndex(file => file.DataSourceId).HasDatabaseName("idx_embedded_files_data_source");
entity.HasIndex(file => file.AbsolutePath).HasDatabaseName("idx_embedded_files_absolute_path");
entity.HasIndex(file => file.FileType).HasDatabaseName("idx_embedded_files_file_type");
entity.HasIndex(file => file.ConfidenceLevelRank).HasDatabaseName("idx_embedded_files_confidence");
entity.HasIndex(file => new { file.DataSourceId, file.AbsolutePath }).HasDatabaseName("idx_embedded_files_data_source_absolute_path").IsUnique();

entity
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ internal static class IndexStoreSchemaMigrator
{
[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(Migrations.InitialRagIndex))]
[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(Migrations.PermanentIndexingFailures))]
[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(Migrations.DropFileConfidenceLevel))]
public static async Task MigrateAsync(IndexStoreDbContext context, CancellationToken token)
{
await context.Database.MigrateAsync(token);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,4 @@ public sealed record IndexStoreSearchResult(
DateTimeOffset CreationUtc,
DateTimeOffset LastWriteUtc,
DateTimeOffset EmbeddedAtUtc,
int ChunkCount,
string ConfidenceLevel,
int ConfidenceLevelRank);
int ChunkCount);
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,4 @@ internal sealed class IndexStoreSearchResultEntity
public DateTimeOffset EmbeddedAtUtc { get; set; }

public int ChunkCount { get; set; }

public string ConfidenceLevel { get; set; } = string.Empty;

public int ConfidenceLevelRank { get; set; }
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
#nullable disable

using Microsoft.EntityFrameworkCore.Infrastructure;
using Microsoft.EntityFrameworkCore.Migrations;

namespace AIStudio.Tools.Databases.IndexStore.Migrations;

/// <summary>
/// Drops the copy of the data source confidence level which every indexed file carried.
/// </summary>
/// <remarks>
/// The confidence level is what a data source asks of a provider. It is a property of the data
/// source, it is enforced live before anything is indexed or answered, and it changes no vector.
/// Keeping a copy per file only meant the index had to be thrown away whenever the setting changed.
/// </remarks>
[DbContext(typeof(IndexStoreDbContext))]
[Migration("20260915000000_DropFileConfidenceLevel")]
public partial class DropFileConfidenceLevel : Migration
{
/// <remarks>
/// The columns go through raw SQL instead of DropColumn on purpose. The SQLite provider answers
/// DropColumn by rebuilding the table, and a rebuild drops the table the trigger
/// embedded_files_file_name_au hangs on, which would silently stop the full-text index from
/// following a renamed file. A native ALTER TABLE ... DROP COLUMN leaves the table itself alone.
/// It does refuse a column an index names, so the index has to go first.
/// </remarks>
protected override void Up(MigrationBuilder migrationBuilder)
{
migrationBuilder.DropIndex(
name: "idx_embedded_files_confidence",
table: "embedded_files");

migrationBuilder.Sql("""
ALTER TABLE embedded_files DROP COLUMN confidence_level;
ALTER TABLE embedded_files DROP COLUMN confidence_level_rank;
""");
}

protected override void Down(MigrationBuilder migrationBuilder)
{
migrationBuilder.Sql("""
ALTER TABLE embedded_files ADD COLUMN confidence_level TEXT NOT NULL DEFAULT '';
ALTER TABLE embedded_files ADD COLUMN confidence_level_rank INTEGER NOT NULL DEFAULT 0;
""");

migrationBuilder.CreateIndex(
name: "idx_embedded_files_confidence",
table: "embedded_files",
column: "confidence_level_rank");
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -80,15 +80,6 @@ protected override void BuildModel(ModelBuilder modelBuilder)
.HasColumnType("INTEGER")
.HasColumnName("chunk_count");

entity.Property<string>("ConfidenceLevel")
.IsRequired()
.HasColumnType("TEXT")
.HasColumnName("confidence_level");

entity.Property<int>("ConfidenceLevelRank")
.HasColumnType("INTEGER")
.HasColumnName("confidence_level_rank");

entity.Property<DateTimeOffset>("CreationUtc")
.HasConversion(utcDateTimeOffsetConverter)
.HasColumnType("TEXT")
Expand Down Expand Up @@ -138,9 +129,6 @@ protected override void BuildModel(ModelBuilder modelBuilder)
entity.HasIndex("AbsolutePath")
.HasDatabaseName("idx_embedded_files_absolute_path");

entity.HasIndex("ConfidenceLevelRank")
.HasDatabaseName("idx_embedded_files_confidence");

entity.HasIndex("DataSourceId")
.HasDatabaseName("idx_embedded_files_data_source");

Expand Down Expand Up @@ -278,13 +266,6 @@ protected override void BuildModel(ModelBuilder modelBuilder)
.IsRequired()
.HasColumnType("TEXT");

entity.Property<string>("ConfidenceLevel")
.IsRequired()
.HasColumnType("TEXT");

entity.Property<int>("ConfidenceLevelRank")
.HasColumnType("INTEGER");

entity.Property<DateTimeOffset>("CreationUtc")
.HasConversion(utcDateTimeOffsetConverter)
.HasColumnType("TEXT");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -302,9 +302,7 @@ public override async Task<IReadOnlyList<IndexStoreSearchResult>> SearchChunksAs
f.creation_utc AS CreationUtc,
f.last_write_utc AS LastWriteUtc,
c.embedded_at_utc AS EmbeddedAtUtc,
f.chunk_count AS ChunkCount,
f.confidence_level AS ConfidenceLevel,
f.confidence_level_rank AS ConfidenceLevelRank
f.chunk_count AS ChunkCount
FROM embedding_chunks_fts
JOIN embedding_chunks c ON c.id = embedding_chunks_fts.rowid
JOIN embedded_files f ON f.parent_file_id = c.parent_file_id
Expand Down Expand Up @@ -394,8 +392,6 @@ private static void ApplyFile(EmbeddingStateFileEntity fileEntity, string dataSo
fileEntity.LastWriteUtc = file.LastWriteUtc;
fileEntity.EmbeddedAtUtc = file.EmbeddedAtUtc;
fileEntity.ChunkCount = file.ChunkCount;
fileEntity.ConfidenceLevel = file.ConfidenceLevel;
fileEntity.ConfidenceLevelRank = file.ConfidenceLevelRank;
}

private static void ApplyPermanentFailure(IndexingFailureEntity failureEntity, string dataSourceId, PermanentIndexingFailure failure)
Expand Down Expand Up @@ -445,9 +441,7 @@ private static void ApplyChunk(EmbeddingStateChunkEntity chunkEntity, EmbeddingS
result.CreationUtc,
result.LastWriteUtc,
result.EmbeddedAtUtc,
result.ChunkCount,
result.ConfidenceLevel,
result.ConfidenceLevelRank);
result.ChunkCount);

private static string BuildFtsQuery(string query)
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,4 @@ public sealed record VectorSearchResult(
string Fingerprint,
string CreationUtc,
string LastWriteUtc,
string EmbeddedAtUtc,
string ConfidenceLevel,
int ConfidenceLevelRank);
string EmbeddedAtUtc);
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,4 @@ public sealed record VectorStoragePoint(
string Fingerprint,
DateTimeOffset CreationUtc,
DateTimeOffset LastWriteUtc,
DateTimeOffset EmbeddedAtUtc,
string ConfidenceLevel,
int ConfidenceLevelRank);
DateTimeOffset EmbeddedAtUtc);
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ private sealed record EmbeddingChunk(string Text, int? PageNumber);

private sealed record EmbeddingChunkDraft(string ChunkId, string Text, int ChunkIndex, int? PageNumber);

private sealed record ChunkingOptions(int MaxChunkTokenLength, int OverlapTokenLength);
internal sealed record ChunkingOptions(int MaxChunkTokenLength, int OverlapTokenLength);

private sealed record ChunkingStrategy(string Name, IReadOnlyList<ChunkingRule> Rules);

Expand Down Expand Up @@ -986,7 +986,23 @@ private bool IsSkippedRagDirectory(string path)
}
}

private string BuildEmbeddingSignature(IDataSource dataSource, EmbeddingProvider embeddingProvider, ChunkingOptions chunkingOptions)
/// <summary>
/// Describes how the vectors of a data source were made.
/// </summary>
/// <remarks>
/// What appears here decides when stored embeddings are thrown away: a signature differing from
/// the persisted one drops the whole index and builds it again. So it names the embedding model,
/// where it runs, how the text was cut for it, and the chunk metadata version — the things a
/// vector actually depends on.
///
/// The confidence level a data source asks of a provider is deliberately not among them. It
/// changes no vector, and it is enforced live on every request anyway: DataSourceService checks
/// it against the participating chat providers and against the embedding provider, and this
/// service checks it again before each indexing run. It was part of this signature once, which
/// re-embedded every file of a data source whenever somebody raised or lowered it — real money
/// at a cloud embedding provider, for nothing.
/// </remarks>
internal static string BuildEmbeddingSignature(IDataSource dataSource, EmbeddingProvider embeddingProvider, ChunkingOptions chunkingOptions)
{
return string.Join('|',
CHUNK_METADATA_VERSION,
Expand All @@ -997,7 +1013,6 @@ private string BuildEmbeddingSignature(IDataSource dataSource, EmbeddingProvider
embeddingProvider.Hostname,
embeddingProvider.TokenizerPath,
embeddingProvider.EffectiveTokenLimit,
GetDataSourceConfidenceLevel(dataSource).ToString(),
dataSource is IInternalDataSource internalDataSource ? internalDataSource.MaxChunkTokenLength : 0,
dataSource is IInternalDataSource overlapDataSource ? overlapDataSource.ChunkOverlapTokenLength : DEFAULT_CHUNK_OVERLAP_TOKEN_LENGTH,
chunkingOptions.MaxChunkTokenLength,
Expand Down Expand Up @@ -1101,7 +1116,6 @@ private EmbeddingStateFile CreateEmbeddingStateFile(IDataSource dataSource, File
{
file.Refresh();
var absolutePath = Path.GetFullPath(file.FullName);
var confidenceLevel = GetDataSourceConfidenceLevel(dataSource);
return new(
this.CreateParentFileId(dataSource.Id, absolutePath),
absolutePath,
Expand All @@ -1113,9 +1127,7 @@ private EmbeddingStateFile CreateEmbeddingStateFile(IDataSource dataSource, File
file.Exists ? new DateTimeOffset(file.CreationTimeUtc) : DateTimeOffset.UnixEpoch,
file.Exists ? new DateTimeOffset(file.LastWriteTimeUtc) : DateTimeOffset.UnixEpoch,
embeddedAtUtc,
chunkCount,
confidenceLevel.ToString(),
(int)confidenceLevel);
chunkCount);
}

private IReadOnlyList<EmbeddingStateChunk> CreateEmbeddingStateChunks(EmbeddingStateFile parentFile, IReadOnlyList<EmbeddingChunkDraft> batch, DateTimeOffset embeddedAtUtc)
Expand All @@ -1131,11 +1143,6 @@ private IReadOnlyList<EmbeddingStateChunk> CreateEmbeddingStateChunks(EmbeddingS
.ToList();
}

private static ConfidenceLevel GetDataSourceConfidenceLevel(IDataSource dataSource) =>
dataSource is not IInternalDataSource internalDataSource || internalDataSource.ConfidenceLevel is ConfidenceLevel.NONE
? ConfidenceLevel.UNKNOWN
: internalDataSource.ConfidenceLevel;

private static string GetFileType(FileInfo file)
{
var extension = file.Extension.TrimStart('.').ToLowerInvariant();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1030,9 +1030,7 @@ private async Task UpsertPointsAsync(
fingerprint,
parentFile.CreationUtc,
parentFile.LastWriteUtc,
embeddedAtUtc,
parentFile.ConfidenceLevel,
parentFile.ConfidenceLevelRank)).ToList();
embeddedAtUtc)).ToList();

await vectorStore.InsertEmbedding(collectionName, points, token);
}
Expand Down Expand Up @@ -1237,7 +1235,7 @@ private async Task<DataSourceEmbeddingManifest> EnsureCompatibleManifestAsync(
CancellationToken token)
{
var chunkingOptions = this.GetChunkingOptions(dataSource, embeddingProvider);
var embeddingSignature = this.BuildEmbeddingSignature(dataSource, embeddingProvider, chunkingOptions);
var embeddingSignature = BuildEmbeddingSignature(dataSource, embeddingProvider, chunkingOptions);
var manifest = await indexStore.GetManifestAsync(dataSource.Id, token);

logger.LogInformation(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,9 +52,7 @@ private sealed record LocalRetrievalHit(
int ChunkIndex,
string Text,
double Score,
int Rank,
string ConfidenceLevel,
int ConfidenceLevelRank);
int Rank);
// ReSharper restore NotAccessedPositionalProperty.Local

public Task<IReadOnlyList<IRetrievalContext>> RetrieveDataAsync(DataSourceLocalFile dataSource, IContent lastUserPrompt, ChatThread thread, CancellationToken token = default) =>
Expand Down Expand Up @@ -354,9 +352,7 @@ private static LocalRetrievalHit FromVectorResult(VectorSearchResult result, int
result.ChunkIndex,
result.Text,
result.Score,
rank,
result.ConfidenceLevel,
result.ConfidenceLevelRank);
rank);

private static LocalRetrievalHit FromBm25Result(IndexStoreSearchResult result, int rank) =>
new(
Expand All @@ -374,9 +370,7 @@ private static LocalRetrievalHit FromBm25Result(IndexStoreSearchResult result, i
result.ChunkIndex,
result.ChunkText,
result.Score,
rank,
result.ConfidenceLevel,
result.ConfidenceLevelRank);
rank);

private static RetrievalTextContext ToRetrievalContext(LocalRetrievalHit hit)
{
Expand Down
72 changes: 72 additions & 0 deletions app/Tests/Tools/EmbeddingSignatureTests.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.Services;

namespace AIStudio.Tests.Tools;

/// <summary>
/// Checks what makes the stored embeddings of a data source invalid.
/// </summary>
/// <remarks>
/// The embedding signature decides whether an index survives: when it differs from the one persisted
/// for a data source, everything stored is thrown away and embedded again. That is the right answer
/// for anything a vector depends on, and an expensive mistake for everything else. The confidence
/// level a data source asks of a provider used to be part of it, so changing that one setting
/// re-embedded every file of the source — at a cloud embedding provider, for real money and no gain.
/// </remarks>
[TestFixture]
public sealed class EmbeddingSignatureTests
{
[Test]
public void ChangingTheConfidenceLevelKeepsTheStoredEmbeddings()
{
var low = DataSource(ConfidenceLevel.LOW);
var high = DataSource(ConfidenceLevel.HIGH);

Assert.That(Signature(high), Is.EqualTo(Signature(low)), "The confidence level changes no vector, so the stored index stays valid and nothing is embedded again.");
}

[Test]
public void ChangingTheChunkSizeDropsTheStoredEmbeddings()
{
var small = DataSource(ConfidenceLevel.LOW) with { MaxChunkTokenLength = 512 };
var large = DataSource(ConfidenceLevel.LOW) with { MaxChunkTokenLength = 1024 };

Assert.That(Signature(large), Is.Not.EqualTo(Signature(small)), "Other chunk boundaries mean other vectors, so the index has to be built again.");
}

[Test]
public void ChangingTheEmbeddingModelDropsTheStoredEmbeddings()
{
var dataSource = DataSource(ConfidenceLevel.LOW);

Assert.That(
Signature(dataSource, EmbeddingProviderFor("text-embedding-3-large")),
Is.Not.EqualTo(Signature(dataSource, EmbeddingProviderFor("text-embedding-3-small"))),
"Another model means another vector space, so nothing stored may be kept.");
}

private static string Signature(DataSourceLocalDirectory dataSource, EmbeddingProvider? embeddingProvider = null) =>
DataSourceEmbeddingService.BuildEmbeddingSignature(
dataSource,
embeddingProvider ?? EmbeddingProviderFor("text-embedding-3-small"),
new(512, 100));

private static DataSourceLocalDirectory DataSource(ConfidenceLevel confidenceLevel) => new()
{
Num = 1,
Id = "6f1d6a4e-6a5e-4c62-9a4f-0f2d2c8b7a11",
Name = "Test data",
Description = "Documents used by the tests.",
Type = DataSourceType.LOCAL_DIRECTORY,
EmbeddingId = "b0a4c4d2-1f3e-4f0a-8c9d-5a6b7c8d9e01",
MaxChunkTokenLength = 512,
ChunkOverlapTokenLength = 100,
ConfidenceLevel = confidenceLevel,
Path = "/tmp/test-data",
};

private static EmbeddingProvider EmbeddingProviderFor(string modelId) =>
new(1, "b0a4c4d2-1f3e-4f0a-8c9d-5a6b7c8d9e01", "Test embeddings", LLMProviders.OPEN_AI, new(modelId, modelId));
}
Loading
Loading