mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-17 14:58:00 +00:00
Make the embedding signature comparable from outside the service
This commit is contained in:
4 files changed
+260
-4
No files matched your search
@@ -67,7 +67,7 @@ public sealed partial class DataSourceEmbeddingService
|
|||||||
|
|
||||||
private async IAsyncEnumerable<EmbeddingChunk> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
private async IAsyncEnumerable<EmbeddingChunk> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
||||||
{
|
{
|
||||||
var options = this.GetChunkingOptions(dataSource, embeddingProvider);
|
var options = GetChunkingOptions(dataSource, embeddingProvider);
|
||||||
var strategy = this.GetChunkingStrategy(filePath);
|
var strategy = this.GetChunkingStrategy(filePath);
|
||||||
var content = await this.ReadExtractedFileContentAsync(filePath, embeddingProvider, token);
|
var content = await this.ReadExtractedFileContentAsync(filePath, embeddingProvider, token);
|
||||||
|
|
||||||
@@ -582,7 +582,18 @@ public sealed partial class DataSourceEmbeddingService
|
|||||||
throw new InvalidOperationException(string.Format(TB("The tokens of the text could not be counted for the embedding provider '{0}'. {1}"), embeddingProvider.Name, message));
|
throw new InvalidOperationException(string.Format(TB("The tokens of the text could not be counted for the embedding provider '{0}'. {1}"), embeddingProvider.Name, message));
|
||||||
}
|
}
|
||||||
|
|
||||||
private ChunkingOptions GetChunkingOptions(IDataSource dataSource, EmbeddingProvider embeddingProvider)
|
/// <summary>
|
||||||
|
/// Works out how the text of a data source is cut for a given embedding provider.
|
||||||
|
/// </summary>
|
||||||
|
/// <remarks>
|
||||||
|
/// Static, because the answer follows from its two arguments alone. That lets the embedding
|
||||||
|
/// signature be built for a configuration which is not stored yet, which is what the dialogs ask
|
||||||
|
/// before they save a change.
|
||||||
|
/// </remarks>
|
||||||
|
/// <param name="dataSource">The data source whose own chunk settings apply.</param>
|
||||||
|
/// <param name="embeddingProvider">The embedding provider whose token limit caps them.</param>
|
||||||
|
/// <returns>The chunk size and overlap which are actually used.</returns>
|
||||||
|
internal static ChunkingOptions GetChunkingOptions(IDataSource dataSource, EmbeddingProvider embeddingProvider)
|
||||||
{
|
{
|
||||||
var providerMaxChunkTokenLength = Math.Max(1, embeddingProvider.EffectiveTokenLimit);
|
var providerMaxChunkTokenLength = Math.Max(1, embeddingProvider.EffectiveTokenLimit);
|
||||||
var dataSourceMaxChunkTokenLength = dataSource is IInternalDataSource { MaxChunkTokenLength: > 0 } internalDataSource
|
var dataSourceMaxChunkTokenLength = dataSource is IInternalDataSource { MaxChunkTokenLength: > 0 } internalDataSource
|
||||||
@@ -1026,6 +1037,15 @@ public sealed partial class DataSourceEmbeddingService
|
|||||||
chunkingOptions.OverlapTokenLength);
|
chunkingOptions.OverlapTokenLength);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// <summary>
|
||||||
|
/// Describes how the vectors of a data source were made, working the chunking out along the way.
|
||||||
|
/// </summary>
|
||||||
|
/// <param name="dataSource">The data source the vectors belong to.</param>
|
||||||
|
/// <param name="embeddingProvider">The embedding provider which makes them.</param>
|
||||||
|
/// <returns>The signature of this pairing.</returns>
|
||||||
|
internal static string BuildEmbeddingSignature(IDataSource dataSource, EmbeddingProvider embeddingProvider) =>
|
||||||
|
BuildEmbeddingSignature(dataSource, embeddingProvider, GetChunkingOptions(dataSource, embeddingProvider));
|
||||||
|
|
||||||
private DataSourceMetadataSnapshot BuildDataSourceMetadataSnapshot(IDataSource dataSource, IReadOnlyList<FileInfo> indexedFiles)
|
private DataSourceMetadataSnapshot BuildDataSourceMetadataSnapshot(IDataSource dataSource, IReadOnlyList<FileInfo> indexedFiles)
|
||||||
{
|
{
|
||||||
var fileHashes = indexedFiles
|
var fileHashes = indexedFiles
|
||||||
|
|||||||
@@ -209,6 +209,16 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
|||||||
}
|
}
|
||||||
|
|
||||||
var manifest = await indexStore.GetManifestAsync(dataSourceId, token);
|
var manifest = await indexStore.GetManifestAsync(dataSourceId, token);
|
||||||
|
return HasStoredIndexState(manifest);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// <summary>
|
||||||
|
/// Whether the index holds anything at all about a data source.
|
||||||
|
/// </summary>
|
||||||
|
/// <param name="manifest">What the index store returned for it.</param>
|
||||||
|
/// <returns>True when there is stored index state.</returns>
|
||||||
|
private static bool HasStoredIndexState(DataSourceEmbeddingManifest manifest)
|
||||||
|
{
|
||||||
return !string.IsNullOrWhiteSpace(manifest.EmbeddingProviderId)
|
return !string.IsNullOrWhiteSpace(manifest.EmbeddingProviderId)
|
||||||
|| !string.IsNullOrWhiteSpace(manifest.EmbeddingSignature)
|
|| !string.IsNullOrWhiteSpace(manifest.EmbeddingSignature)
|
||||||
|| !string.IsNullOrWhiteSpace(manifest.SourceHash)
|
|| !string.IsNullOrWhiteSpace(manifest.SourceHash)
|
||||||
@@ -220,6 +230,59 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
|||||||
|| manifest.PermanentFailures.Count > 0;
|
|| manifest.PermanentFailures.Count > 0;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// <summary>
|
||||||
|
/// Picks the data sources which already hold something in the index.
|
||||||
|
/// </summary>
|
||||||
|
/// <remarks>
|
||||||
|
/// Asked before a setting is saved which would throw those indexes away, so the question can be
|
||||||
|
/// put to the user with the names in it. Anything unclear counts as holding something — the
|
||||||
|
/// opposite of IsAwaitingReindexAsync, and for the opposite reason: there, a wrongly greyed-out
|
||||||
|
/// row would stay wrong for good, while a question asked once too often costs a click, and one
|
||||||
|
/// skipped costs whatever a cloud provider charges for embedding everything again.
|
||||||
|
/// </remarks>
|
||||||
|
/// <param name="dataSources">The data sources to ask about.</param>
|
||||||
|
/// <param name="token">The cancellation token.</param>
|
||||||
|
/// <returns>Those of them which have stored index state.</returns>
|
||||||
|
public async Task<IReadOnlyList<IDataSource>> GetDataSourcesWithStoredIndexAsync(IReadOnlyCollection<IDataSource> dataSources, CancellationToken token = default)
|
||||||
|
{
|
||||||
|
//
|
||||||
|
// Filtering first also keeps the index database from being created while local RAG is off:
|
||||||
|
// asking for the store runs its migrations on the first call, which must not happen because
|
||||||
|
// somebody opened a dialog.
|
||||||
|
//
|
||||||
|
var candidates = dataSources.Where(this.IsSupportedInternalDataSource).ToList();
|
||||||
|
if (candidates.Count == 0)
|
||||||
|
return [];
|
||||||
|
|
||||||
|
try
|
||||||
|
{
|
||||||
|
using var timeout = CancellationTokenSource.CreateLinkedTokenSource(token);
|
||||||
|
timeout.CancelAfter(REINDEX_CHECK_TIMEOUT);
|
||||||
|
|
||||||
|
var indexStore = await databaseClientProvider.GetIndexStoreAsync(timeout.Token);
|
||||||
|
if (!indexStore.IsAvailable)
|
||||||
|
{
|
||||||
|
logger.LogWarning("Could not tell which data sources hold a stored index because the local RAG index database '{DatabaseName}' is unavailable. Treating all {DataSourceCount} of them as affected.", indexStore.Name, candidates.Count);
|
||||||
|
return candidates;
|
||||||
|
}
|
||||||
|
|
||||||
|
var affected = new List<IDataSource>(candidates.Count);
|
||||||
|
foreach (var dataSource in candidates)
|
||||||
|
{
|
||||||
|
var manifest = await indexStore.GetManifestAsync(dataSource.Id, timeout.Token);
|
||||||
|
if (HasStoredIndexState(manifest))
|
||||||
|
affected.Add(dataSource);
|
||||||
|
}
|
||||||
|
|
||||||
|
return affected;
|
||||||
|
}
|
||||||
|
catch (Exception exception)
|
||||||
|
{
|
||||||
|
logger.LogWarning(exception, "Could not tell which of {DataSourceCount} data source(s) hold a stored index. Treating all of them as affected.", candidates.Count);
|
||||||
|
return candidates;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// <summary>
|
/// <summary>
|
||||||
/// Whether a data source cannot answer a search right now because its index has to be built anew.
|
/// Whether a data source cannot answer a search right now because its index has to be built anew.
|
||||||
/// </summary>
|
/// </summary>
|
||||||
@@ -269,7 +332,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
|||||||
return false;
|
return false;
|
||||||
|
|
||||||
var indexState = await indexStore.GetDataSourceStateAsync(dataSource.Id, timeout.Token);
|
var indexState = await indexStore.GetDataSourceStateAsync(dataSource.Id, timeout.Token);
|
||||||
var chunkingOptions = this.GetChunkingOptions(dataSource, embeddingProvider);
|
var chunkingOptions = GetChunkingOptions(dataSource, embeddingProvider);
|
||||||
var embeddingSignature = BuildEmbeddingSignature(dataSource, embeddingProvider, chunkingOptions);
|
var embeddingSignature = BuildEmbeddingSignature(dataSource, embeddingProvider, chunkingOptions);
|
||||||
var runState = this.statuses.TryGetValue(dataSource.Id, out var status) ? status.State : (DataSourceEmbeddingState?)null;
|
var runState = this.statuses.TryGetValue(dataSource.Id, out var status) ? status.State : (DataSourceEmbeddingState?)null;
|
||||||
|
|
||||||
@@ -1444,7 +1507,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
|||||||
IndexStoreClient indexStore,
|
IndexStoreClient indexStore,
|
||||||
CancellationToken token)
|
CancellationToken token)
|
||||||
{
|
{
|
||||||
var chunkingOptions = this.GetChunkingOptions(dataSource, embeddingProvider);
|
var chunkingOptions = GetChunkingOptions(dataSource, embeddingProvider);
|
||||||
var embeddingSignature = BuildEmbeddingSignature(dataSource, embeddingProvider, chunkingOptions);
|
var embeddingSignature = BuildEmbeddingSignature(dataSource, embeddingProvider, chunkingOptions);
|
||||||
var manifest = await indexStore.GetManifestAsync(dataSource.Id, token);
|
var manifest = await indexStore.GetManifestAsync(dataSource.Id, token);
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,40 @@
|
|||||||
|
using AIStudio.Settings;
|
||||||
|
|
||||||
|
namespace AIStudio.Tools.Services;
|
||||||
|
|
||||||
|
/// <summary>
|
||||||
|
/// Answers whether an edit throws the stored index of a data source away.
|
||||||
|
/// </summary>
|
||||||
|
/// <remarks>
|
||||||
|
/// Nothing here knows which settings matter. Both questions are answered by building the embedding
|
||||||
|
/// signature twice and comparing the two, so the single place which decides stays
|
||||||
|
/// BuildEmbeddingSignature and this cannot drift away from what an indexing run then does.
|
||||||
|
/// </remarks>
|
||||||
|
internal static class EmbeddingChangeImpact
|
||||||
|
{
|
||||||
|
/// <summary>
|
||||||
|
/// Whether an edited embedding provider invalidates what is stored for one of its data sources.
|
||||||
|
/// </summary>
|
||||||
|
/// <param name="dataSource">The data source, which the edit leaves alone.</param>
|
||||||
|
/// <param name="before">The embedding provider as it is stored.</param>
|
||||||
|
/// <param name="after">The embedding provider as it would be stored.</param>
|
||||||
|
/// <returns>True when the stored index would be discarded.</returns>
|
||||||
|
public static bool AffectsStoredIndex(IDataSource dataSource, EmbeddingProvider before, EmbeddingProvider after) =>
|
||||||
|
!string.Equals(
|
||||||
|
DataSourceEmbeddingService.BuildEmbeddingSignature(dataSource, before),
|
||||||
|
DataSourceEmbeddingService.BuildEmbeddingSignature(dataSource, after),
|
||||||
|
StringComparison.Ordinal);
|
||||||
|
|
||||||
|
/// <summary>
|
||||||
|
/// Whether an edited data source invalidates what is stored for it.
|
||||||
|
/// </summary>
|
||||||
|
/// <param name="embeddingProvider">The embedding provider, which the edit leaves alone.</param>
|
||||||
|
/// <param name="before">The data source as it is stored.</param>
|
||||||
|
/// <param name="after">The data source as it would be stored.</param>
|
||||||
|
/// <returns>True when the stored index would be discarded.</returns>
|
||||||
|
public static bool AffectsStoredIndex(EmbeddingProvider embeddingProvider, IDataSource before, IDataSource after) =>
|
||||||
|
!string.Equals(
|
||||||
|
DataSourceEmbeddingService.BuildEmbeddingSignature(before, embeddingProvider),
|
||||||
|
DataSourceEmbeddingService.BuildEmbeddingSignature(after, embeddingProvider),
|
||||||
|
StringComparison.Ordinal);
|
||||||
|
}
|
||||||
@@ -0,0 +1,133 @@
|
|||||||
|
using AIStudio.Provider;
|
||||||
|
using AIStudio.Provider.HuggingFace;
|
||||||
|
using AIStudio.Settings;
|
||||||
|
using AIStudio.Settings.DataModel;
|
||||||
|
using AIStudio.Tools.Services;
|
||||||
|
|
||||||
|
using Host = AIStudio.Provider.SelfHosted.Host;
|
||||||
|
|
||||||
|
namespace AIStudio.Tests.Tools;
|
||||||
|
|
||||||
|
/// <summary>
|
||||||
|
/// Checks which edits have to be asked about before they are saved.
|
||||||
|
/// </summary>
|
||||||
|
/// <remarks>
|
||||||
|
/// An edit which changes the embedding signature throws away everything indexed for the data sources
|
||||||
|
/// behind it, and sends every one of their documents to the embedding provider again. Asking about an
|
||||||
|
/// edit which costs nothing trains people to click the question away; not asking about one which does
|
||||||
|
/// costs them money at a cloud provider. So both directions are pinned down here.
|
||||||
|
/// </remarks>
|
||||||
|
[TestFixture]
|
||||||
|
public sealed class EmbeddingChangeImpactTests
|
||||||
|
{
|
||||||
|
[Test]
|
||||||
|
public void HarmlessEmbeddingProviderEditsKeepTheStoredIndex()
|
||||||
|
{
|
||||||
|
var dataSource = StoredDataSource();
|
||||||
|
var stored = StoredEmbeddingProvider();
|
||||||
|
|
||||||
|
Assert.Multiple(() =>
|
||||||
|
{
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { Name = "Another name" }), Is.False, "The name of an embedding provider reaches no vector.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { Num = 42 }), Is.False, "The number is there to sort the list with.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { EmbeddingBatchSize = 16 }), Is.False, "How many chunks travel in one request says nothing about the vectors which come back.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { CustomIconDataUrl = "data:image/png;base64,AAAA" }), Is.False, "An icon is an icon.");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
[Test]
|
||||||
|
public void ChangingTheModelDropsTheStoredIndex()
|
||||||
|
{
|
||||||
|
var dataSource = StoredDataSource();
|
||||||
|
var stored = StoredEmbeddingProvider();
|
||||||
|
|
||||||
|
Assert.That(
|
||||||
|
EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { Model = new("text-embedding-3-large", "text-embedding-3-large") }),
|
||||||
|
Is.True,
|
||||||
|
"Another model means another vector space.");
|
||||||
|
}
|
||||||
|
|
||||||
|
[Test]
|
||||||
|
public void ChangingTheTokenLimitDropsTheStoredIndex()
|
||||||
|
{
|
||||||
|
var dataSource = StoredDataSource();
|
||||||
|
var stored = StoredEmbeddingProvider();
|
||||||
|
|
||||||
|
Assert.That(
|
||||||
|
EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { TokenLimit = 4096 }),
|
||||||
|
Is.True,
|
||||||
|
"The token limit decides where the text is cut, and other chunks are other vectors.");
|
||||||
|
}
|
||||||
|
|
||||||
|
[Test]
|
||||||
|
public void ChangingWhereTheProviderRunsDropsTheStoredIndex()
|
||||||
|
{
|
||||||
|
var dataSource = StoredDataSource();
|
||||||
|
var stored = StoredEmbeddingProvider();
|
||||||
|
|
||||||
|
Assert.Multiple(() =>
|
||||||
|
{
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { Hostname = "http://localhost:9999" }), Is.True, "Another server can serve another model under the same name.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { Host = Host.LM_STUDIO }), Is.True, "Another kind of host speaks another API.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { HFInferenceProvider = HFInferenceProvider.GROQ }), Is.True, "The same model name served by another backend is another vector source.");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
[Test]
|
||||||
|
public void TheTokenizerIsComparedByItsContentNotItsPath()
|
||||||
|
{
|
||||||
|
var dataSource = StoredDataSource();
|
||||||
|
var stored = StoredEmbeddingProvider() with { TokenizerPath = "/data/tokenizers/embeddings/tokenizer.json", TokenizerFingerprint = "AAAA" };
|
||||||
|
|
||||||
|
Assert.Multiple(() =>
|
||||||
|
{
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { TokenizerFingerprint = "BBBB" }), Is.True, "Another tokenizer counts tokens differently, so the text is cut elsewhere.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(dataSource, stored, stored with { TokenizerPath = "/somewhere/else/tokenizer.json" }), Is.False, "It is the same tokenizer under another path.");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
[Test]
|
||||||
|
public void ChangingTheChunkSettingsOfADataSourceDropsItsStoredIndex()
|
||||||
|
{
|
||||||
|
var embeddingProvider = StoredEmbeddingProvider();
|
||||||
|
var stored = StoredDataSource();
|
||||||
|
|
||||||
|
Assert.Multiple(() =>
|
||||||
|
{
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(embeddingProvider, stored, stored with { MaxChunkTokenLength = 256 }), Is.True, "Other chunk boundaries mean other vectors.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(embeddingProvider, stored, stored with { ChunkOverlapTokenLength = 50 }), Is.True, "Another overlap changes what every chunk starts with.");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
[Test]
|
||||||
|
public void HarmlessDataSourceEditsKeepTheStoredIndex()
|
||||||
|
{
|
||||||
|
var embeddingProvider = StoredEmbeddingProvider();
|
||||||
|
var stored = StoredDataSource();
|
||||||
|
|
||||||
|
Assert.Multiple(() =>
|
||||||
|
{
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(embeddingProvider, stored, stored with { Name = "Another name" }), Is.False, "The name is how the data source is offered, not how it was read.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(embeddingProvider, stored, stored with { Description = "Another description" }), Is.False, "The description is there for the agent which picks data sources.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(embeddingProvider, stored, stored with { MaxMatches = 42 }), Is.False, "How many matches an answer may use is decided per query.");
|
||||||
|
Assert.That(EmbeddingChangeImpact.AffectsStoredIndex(embeddingProvider, stored, stored with { ConfidenceLevel = ConfidenceLevel.HIGH }), Is.False, "The confidence level is enforced live on every request and changes no vector.");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
private static DataSourceLocalDirectory StoredDataSource() => 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.LOW,
|
||||||
|
Path = "/tmp/test-data",
|
||||||
|
};
|
||||||
|
|
||||||
|
private static EmbeddingProvider StoredEmbeddingProvider() =>
|
||||||
|
new(1, "b0a4c4d2-1f3e-4f0a-8c9d-5a6b7c8d9e01", "Test embeddings", LLMProviders.OPEN_AI, new("text-embedding-3-small", "text-embedding-3-small"));
|
||||||
|
}
|
||||||
Reference in new issue
Block a user