Improved the start of configuration plugins with tokenizers (#1039)

This commit is contained in:
Thorsten Sommer authored and GitHub committed 2026-10-10 12:51:20 +02:00
1 parent 5c319c45a0
commit 3d847555d7
8 files changed
+714 -28

No files matched your search

@@ -0,0 +1,191 @@
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.Services;
using Lua;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging.Abstractions;
namespace AIStudio.Tests.Tools.PluginSystem;
/// <summary>
/// Checks which tokenizer the providers of a configuration plugin point to while the plugin starts.
/// </summary>
/// <remarks>
/// A configuration plugin names its tokenizers by paths inside the plugin. The settings point to the
/// copies the runtime stored of them instead, and an embedding provider keeps the fingerprint of its
/// copy, which is part of the embedding signature. Parsing the plugin replaced both with the path
/// from the plugin and an empty fingerprint, until the tokenizers were synchronized up to a second
/// later. An indexing run in between took that for another tokenizer and reset the index of every
/// data source of the provider, without asking.<br/><br/>
/// The settings are reached through Program.SERVICE_PROVIDER, which is why this fixture does not run
/// alongside others.
/// </remarks>
[TestFixture]
[NonParallelizable]
public sealed class ManagedTokenizerStateTests
{
private static readonly Guid PLUGIN_ID = Guid.Parse("8d3e1c5a-6f2b-4a7d-9c1e-3b5f7a9d2e4c");
private const string PROVIDER_ID = "1a7c3e5f-9b2d-4f6a-8c0e-2d4f6a8c0e1b";
private const string EMBEDDING_PROVIDER_ID = "4e6a8c0e-2b4d-4f8a-9c1e-5f7a9b1d3e5f";
private const string PLUGIN_PATH = "/plugins/configuration-of-the-organization";
private const string CONFIGURED_TOKENIZER_PATH = "tokenizers/tokenizer.json";
private const string STORED_CHAT_TOKENIZER_PATH = "/data/tokenizers/chat_" + PROVIDER_ID + "/tokenizer.json";
private const string STORED_EMBEDDING_TOKENIZER_PATH = "/data/tokenizers/embedding_" + EMBEDDING_PROVIDER_ID + "/tokenizer.json";
private const string STORED_FINGERPRINT = "AAAA";
private RustService rustService = null!;
private ServiceProvider serviceProvider = null!;
private IServiceProvider previousServiceProvider = null!;
private SettingsManager settingsManager = null!;
[SetUp]
public void CreateEmptySettings()
{
// Only builds its HTTP clients. Nothing connects, because parsing a plugin does not reach the runtime:
this.rustService = new RustService("1", "unused");
this.settingsManager = new SettingsManager(NullLogger<SettingsManager>.Instance, this.rustService);
this.previousServiceProvider = Program.SERVICE_PROVIDER;
this.serviceProvider = new ServiceCollection().AddSingleton(this.settingsManager).AddSingleton(this.rustService).BuildServiceProvider();
Program.SERVICE_PROVIDER = this.serviceProvider;
}
[TearDown]
public void RestoreApplicationState()
{
Program.SERVICE_PROVIDER = this.previousServiceProvider;
this.serviceProvider.Dispose();
this.rustService.Dispose();
}
[Test]
public async Task AnEmbeddingProviderKeepsItsStoredTokenizerWhileThePluginStarts()
{
await StartPluginAsync();
this.StoreTokenizers();
var configObjects = await StartPluginAsync();
var embeddingProvider = this.settingsManager.ConfigurationData.EmbeddingProviders.Single();
Assert.Multiple(() =>
{
Assert.That(embeddingProvider.TokenizerPath, Is.EqualTo(STORED_EMBEDDING_TOKENIZER_PATH), "Until the tokenizers are synchronized, the stored copy is the tokenizer this provider uses.");
Assert.That(embeddingProvider.TokenizerFingerprint, Is.EqualTo(STORED_FINGERPRINT), "An empty fingerprint looks like another tokenizer and costs every data source of this provider its index.");
Assert.That(configObjects.Single(x => x.Type == PluginConfigurationObjectType.EMBEDDING_PROVIDER).ConfiguredTokenizerPath, Is.EqualTo(CONFIGURED_TOKENIZER_PATH), "The synchronization learns from here which tokenizer the plugin wants.");
});
}
[Test]
public async Task TheEmbeddingSignatureStaysTheSameWhileThePluginStarts()
{
await StartPluginAsync();
this.StoreTokenizers();
var before = Signature(this.settingsManager.ConfigurationData.EmbeddingProviders.Single());
await StartPluginAsync();
var after = Signature(this.settingsManager.ConfigurationData.EmbeddingProviders.Single());
Assert.That(after, Is.EqualTo(before), "An indexing run between the start of the plugin and the synchronization of its tokenizers compares exactly these two, and resets the index when they differ.");
}
[Test]
public async Task AChatProviderKeepsItsStoredTokenizerWhileThePluginStarts()
{
await StartPluginAsync();
this.StoreTokenizers();
var configObjects = await StartPluginAsync();
Assert.Multiple(() =>
{
Assert.That(this.settingsManager.ConfigurationData.Providers.Single().TokenizerPath, Is.EqualTo(STORED_CHAT_TOKENIZER_PATH), "The path inside the plugin is relative to the plugin, so no token could be counted with it.");
Assert.That(configObjects.Single(x => x.Type == PluginConfigurationObjectType.LLM_PROVIDER).ConfiguredTokenizerPath, Is.EqualTo(CONFIGURED_TOKENIZER_PATH), "The synchronization learns from here which tokenizer the plugin wants.");
});
}
[Test]
public async Task ANewProviderStartsWithoutATokenizer()
{
var configObjects = await StartPluginAsync();
var embeddingProvider = this.settingsManager.ConfigurationData.EmbeddingProviders.Single();
Assert.Multiple(() =>
{
Assert.That(this.settingsManager.ConfigurationData.Providers.Single().TokenizerPath, Is.Empty, "Nothing is stored yet, and the path inside the plugin is no tokenizer the app could load.");
Assert.That(embeddingProvider.TokenizerPath, Is.Empty);
Assert.That(embeddingProvider.TokenizerFingerprint, Is.Empty);
Assert.That(configObjects, Has.Count.EqualTo(2));
Assert.That(configObjects.Select(x => x.ConfiguredTokenizerPath), Is.All.EqualTo(CONFIGURED_TOKENIZER_PATH), "The synchronization still learns which tokenizer to store.");
});
}
/// <summary>
/// Parses the providers of the plugin into the settings, as every start of the plugin does.
/// </summary>
/// <returns>The configuration objects the plugin defined.</returns>
private static async Task<List<PluginConfigurationObject>> StartPluginAsync()
{
var state = LuaState.Create();
await state.DoStringAsync($$"""
CONFIG = {
["LLM_PROVIDERS"] = {
{
["Id"] = "{{PROVIDER_ID}}",
["InstanceName"] = "Chat of the organization",
["UsedLLMProvider"] = "SELF_HOSTED",
["Host"] = "VLLM",
["Hostname"] = "https://llm.example.org",
["AdditionalJsonApiParameters"] = "",
["TokenizerPath"] = "{{CONFIGURED_TOKENIZER_PATH}}",
["Model"] = { ["Id"] = "chat-model", ["DisplayName"] = "Chat model" },
},
},
["EMBEDDING_PROVIDERS"] = {
{
["Id"] = "{{EMBEDDING_PROVIDER_ID}}",
["Name"] = "Embeddings of the organization",
["UsedLLMProvider"] = "SELF_HOSTED",
["Host"] = "VLLM",
["Hostname"] = "https://embeddings.example.org",
["TokenizerPath"] = "{{CONFIGURED_TOKENIZER_PATH}}",
["Model"] = { ["Id"] = "embedding-model", ["DisplayName"] = "Embedding model" },
},
},
}
""");
if (!state.Environment["CONFIG"].TryRead<LuaTable>(out var mainTable))
throw new InvalidOperationException("The configuration of this test is not a Lua table.");
var configObjects = new List<PluginConfigurationObject>();
PluginConfigurationObject.TryParse(PluginConfigurationObjectType.LLM_PROVIDER, x => x.Providers, x => x.NextProviderNum, mainTable, PLUGIN_ID, ref configObjects, dryRun: false, PLUGIN_PATH);
PluginConfigurationObject.TryParse(PluginConfigurationObjectType.EMBEDDING_PROVIDER, x => x.EmbeddingProviders, x => x.NextEmbeddingNum, mainTable, PLUGIN_ID, ref configObjects, dryRun: false, PLUGIN_PATH);
return configObjects;
}
/// <summary>
/// Points the providers to stored copies of their tokenizers, as the synchronization does after
/// the plugin started.
/// </summary>
private void StoreTokenizers()
{
var data = this.settingsManager.ConfigurationData;
data.Providers[0] = data.Providers[0] with { TokenizerPath = STORED_CHAT_TOKENIZER_PATH };
data.EmbeddingProviders[0] = data.EmbeddingProviders[0] with { TokenizerPath = STORED_EMBEDDING_TOKENIZER_PATH, TokenizerFingerprint = STORED_FINGERPRINT };
}
private static string Signature(EmbeddingProvider embeddingProvider) => DataSourceEmbeddingService.BuildEmbeddingSignature(new DataSourceLocalDirectory
{
Num = 1,
Id = "7b9d1f3a-5c7e-4a9b-8d2f-4a6c8e0b2d4f",
Name = "Documents of the organization",
Type = DataSourceType.LOCAL_DIRECTORY,
EmbeddingId = embeddingProvider.Id,
MaxChunkTokenLength = 512,
Path = "/tmp/documents"
}, embeddingProvider);
}
@@ -0,0 +1,319 @@
using AIStudio.Tools;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.Rust;
using AIStudio.Tools.Services;
namespace AIStudio.Tests.Tools.PluginSystem;
/// <summary>
/// Checks what the synchronization asks the runtime to do with the tokenizer a configuration plugin
/// names for a provider.
/// </summary>
/// <remarks>
/// Whenever a configuration plugin starts, each of its providers gets the tokenizer the plugin names:
/// the runtime checks the file and stores a copy below the data directory, which the provider then
/// points to. When the plugin names none, or one which is unusable or lies outside the plugin, the
/// stored copy is deleted instead. Checking takes most of a second for a common tokenizer, so a copy
/// with the same content as the file in the plugin is kept as it is. These tests stand in for the
/// runtime, so they see every call the synchronization makes and in which order.
/// </remarks>
[TestFixture]
public sealed class TokenizerSyncTests
{
private const string CONFIGURED_TOKENIZER_PATH = "tokenizers/tokenizer.json";
private const string LOG_NAME = "embedding provider 'Embeddings of the organization'";
private static readonly string MODEL_ID = TokenizerModelId.ForEmbeddingProviderId("4e6a8c0e-2b4d-4f8a-9c1e-5f7a9b1d3e5f");
private string testDirectory = string.Empty;
private string pluginDirectory = string.Empty;
private string dataDirectory = string.Empty;
private string sourcePath = string.Empty;
[SetUp]
public void CreatePluginWithATokenizer()
{
this.testDirectory = Path.Combine(Path.GetTempPath(), $"ai-studio-tokenizers-{Guid.NewGuid():N}");
this.pluginDirectory = Path.Combine(this.testDirectory, "plugin");
this.dataDirectory = Path.Combine(this.testDirectory, "data");
this.sourcePath = Path.GetFullPath(Path.Combine(this.pluginDirectory, CONFIGURED_TOKENIZER_PATH));
Directory.CreateDirectory(Path.GetDirectoryName(this.sourcePath)!);
File.WriteAllText(this.sourcePath, """{ "version": "1.0", "model": { "type": "WordPiece" } }""");
}
[TearDown]
public void RemoveTestDirectory()
{
try
{
Directory.Delete(this.testDirectory, true);
}
catch (IOException)
{
// A temporary directory we could not remove says nothing about the code under test.
}
}
[Test]
public async Task AValidTokenizerIsCheckedAndStored()
{
var storage = new RecordingTokenizerStorage(this.dataDirectory);
var (storedPath, fingerprint) = await this.SyncAsync(storage, string.Empty);
var fingerprintOfTheCopy = await TokenizerFingerprint.ForFileAsync(storedPath);
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.EqualTo(new[] { $"validate {this.sourcePath}", $"store {MODEL_ID} {this.sourcePath}" }), "Without a stored copy, there is nothing to compare with.");
Assert.That(storedPath, Is.Not.Empty.And.Not.EqualTo(this.sourcePath), "The provider points to the copy, not to the file inside the plugin.");
Assert.That(File.ReadAllBytes(storedPath), Is.EqualTo(File.ReadAllBytes(this.sourcePath)));
Assert.That(fingerprint, Is.Not.Empty.And.EqualTo(fingerprintOfTheCopy));
});
}
[Test]
public async Task AnUnchangedTokenizerIsNeitherCheckedNorStoredAgain()
{
var firstStart = await this.SyncAsync(new RecordingTokenizerStorage(this.dataDirectory), string.Empty);
var storage = new RecordingTokenizerStorage(this.dataDirectory);
var secondStart = await this.SyncAsync(storage, firstStart.Path);
var fingerprintOfTheCopy = await TokenizerFingerprint.ForFileAsync(firstStart.Path);
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.Empty, "Checking the same file again would cost most of a second and change nothing.");
Assert.That(secondStart, Is.EqualTo(firstStart), "A fingerprint which differs from the one before would cost every data source of the provider its index.");
Assert.That(secondStart.Fingerprint, Is.EqualTo(fingerprintOfTheCopy));
});
}
[Test]
public async Task AMissingCopyIsStoredAgain()
{
var firstStart = await this.SyncAsync(new RecordingTokenizerStorage(this.dataDirectory), string.Empty);
File.Delete(firstStart.Path);
var storage = new RecordingTokenizerStorage(this.dataDirectory);
var secondStart = await this.SyncAsync(storage, firstStart.Path);
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.EqualTo(new[] { $"validate {this.sourcePath}", $"store {MODEL_ID} {this.sourcePath}" }));
Assert.That(File.Exists(secondStart.Path), Is.True);
Assert.That(secondStart.Fingerprint, Is.EqualTo(firstStart.Fingerprint), "It is the same tokenizer, so the index stays.");
});
}
[Test]
public async Task AChangedCopyIsStoredAgainEvenAtTheSameSize()
{
var firstStart = await this.SyncAsync(new RecordingTokenizerStorage(this.dataDirectory), string.Empty);
var copy = File.ReadAllBytes(firstStart.Path);
copy[^3] ^= 0x01;
File.WriteAllBytes(firstStart.Path, copy);
var storage = new RecordingTokenizerStorage(this.dataDirectory);
var secondStart = await this.SyncAsync(storage, firstStart.Path);
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.EqualTo(new[] { $"validate {this.sourcePath}", $"store {MODEL_ID} {this.sourcePath}" }), "Size and time do not tell a damaged copy apart, the content does.");
Assert.That(File.ReadAllBytes(secondStart.Path), Is.EqualTo(File.ReadAllBytes(this.sourcePath)));
Assert.That(secondStart.Fingerprint, Is.EqualTo(firstStart.Fingerprint));
});
}
[Test]
public async Task AChangedTokenizerInThePluginIsCheckedAndStored()
{
var firstStart = await this.SyncAsync(new RecordingTokenizerStorage(this.dataDirectory), string.Empty);
File.WriteAllText(this.sourcePath, """{ "version": "1.0", "model": { "type": "BPE" } }""");
var storage = new RecordingTokenizerStorage(this.dataDirectory);
var secondStart = await this.SyncAsync(storage, firstStart.Path);
var fingerprintOfThePluginFile = await TokenizerFingerprint.ForFileAsync(this.sourcePath);
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.EqualTo(new[] { $"validate {this.sourcePath}", $"store {MODEL_ID} {this.sourcePath}" }));
Assert.That(secondStart.Fingerprint, Is.Not.EqualTo(firstStart.Fingerprint), "Another tokenizer cuts the text at other places, so the index has to go.");
Assert.That(secondStart.Fingerprint, Is.EqualTo(fingerprintOfThePluginFile));
});
}
[Test]
public async Task TwoMissingFilesAreNotTheSameTokenizer()
{
var firstStart = await this.SyncAsync(new RecordingTokenizerStorage(this.dataDirectory), string.Empty);
File.Delete(firstStart.Path);
File.Delete(this.sourcePath);
var storage = new RecordingTokenizerStorage(this.dataDirectory);
var secondStart = await this.SyncAsync(storage, firstStart.Path);
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.EqualTo(new[] { $"validate {this.sourcePath}", $"delete {MODEL_ID}" }), "Neither file can be read, so nothing says they hold the same tokenizer.");
Assert.That(secondStart.Path, Is.Empty);
Assert.That(secondStart.Fingerprint, Is.Empty);
});
}
[Test]
public async Task AnUnreadableCopyTakesTheFingerprintOfThePluginFile()
{
var storage = new RecordingTokenizerStorage(this.dataDirectory) { MakesUnreadableCopies = true };
var (storedPath, fingerprint) = await this.SyncAsync(storage, string.Empty);
var fingerprintOfThePluginFile = await TokenizerFingerprint.ForFileAsync(this.sourcePath);
Assert.Multiple(() =>
{
Assert.That(storedPath, Is.Not.Empty);
Assert.That(File.Exists(storedPath), Is.False, "Otherwise this test checks nothing.");
Assert.That(fingerprint, Is.Not.Empty.And.EqualTo(fingerprintOfThePluginFile), "The runtime copies the file as it is, so the file in the plugin tells what the copy holds.");
});
}
[Test]
public async Task AnInvalidTokenizerDeletesTheStoredOne()
{
var storage = new RecordingTokenizerStorage(this.dataDirectory) { AcceptsTokenizers = false };
var (storedPath, fingerprint) = await this.SyncAsync(storage, string.Empty);
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.EqualTo(new[] { $"validate {this.sourcePath}", $"delete {MODEL_ID}" }), "A tokenizer stored earlier must not outlive the plugin naming an unusable one.");
Assert.That(storedPath, Is.Empty);
Assert.That(fingerprint, Is.Empty);
});
}
[Test]
public async Task NoTokenizerDeletesTheStoredOne()
{
var firstStart = await this.SyncAsync(new RecordingTokenizerStorage(this.dataDirectory), string.Empty);
var storage = new RecordingTokenizerStorage(this.dataDirectory);
var (storedPath, fingerprint) = await this.SyncAsync(storage, firstStart.Path, string.Empty);
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.EqualTo(new[] { $"delete {MODEL_ID}" }), "The plugin no longer names a tokenizer, so the one stored for it goes.");
Assert.That(File.Exists(firstStart.Path), Is.False);
Assert.That(storedPath, Is.Empty);
Assert.That(fingerprint, Is.Empty);
});
}
[Test]
public async Task ATokenizerOutsideThePluginIsNeitherReadNorStored()
{
var firstStart = await this.SyncAsync(new RecordingTokenizerStorage(this.dataDirectory), string.Empty);
var storage = new RecordingTokenizerStorage(this.dataDirectory);
var (storedPath, fingerprint) = await this.SyncAsync(storage, firstStart.Path, "../outside/tokenizer.json");
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.EqualTo(new[] { $"delete {MODEL_ID}" }), "A plugin must not make the app read files from anywhere else on the machine.");
Assert.That(storedPath, Is.Empty);
Assert.That(fingerprint, Is.Empty);
});
}
[Test]
public async Task ATokenizerWhichCannotBeStoredLeavesTheProviderWithoutOne()
{
var storage = new RecordingTokenizerStorage(this.dataDirectory) { StoresTokenizers = false };
var (storedPath, fingerprint) = await this.SyncAsync(storage, string.Empty);
Assert.Multiple(() =>
{
Assert.That(storage.Calls, Is.EqualTo(new[] { $"validate {this.sourcePath}", $"store {MODEL_ID} {this.sourcePath}" }));
Assert.That(storedPath, Is.Empty);
Assert.That(fingerprint, Is.Empty);
});
}
/// <summary>
/// Synchronizes the tokenizer of the provider, as one start of the plugin does.
/// </summary>
/// <param name="storage">The stand-in for the runtime.</param>
/// <param name="storedTokenizerPath">The copy the provider points to so far, or an empty string.</param>
/// <param name="configuredTokenizerPath">The tokenizer path as the plugin names it.</param>
/// <returns>The copy the provider points to afterward, and its fingerprint.</returns>
private Task<StoredTokenizer> SyncAsync(RecordingTokenizerStorage storage, string storedTokenizerPath, string configuredTokenizerPath = CONFIGURED_TOKENIZER_PATH) =>
PluginConfigurationObject.SyncTokenizerAsync(storage, configuredTokenizerPath, storedTokenizerPath, this.pluginDirectory, MODEL_ID, LOG_NAME);
/// <summary>
/// Stands in for the runtime: writes down every call, and stores a tokenizer the way the runtime
/// does, as a copy in a directory of its model below the data directory.
/// </summary>
private sealed class RecordingTokenizerStorage(string dataDirectory) : ITokenizerStorage
{
/// <summary>
/// The calls so far, each as its name and arguments.
/// </summary>
public List<string> Calls { get; } = [];
/// <summary>
/// Whether every existing file passes the check. A missing one never does.
/// </summary>
public bool AcceptsTokenizers { get; init; } = true;
/// <summary>
/// Whether storing works.
/// </summary>
public bool StoresTokenizers { get; init; } = true;
/// <summary>
/// Whether storing reports a copy which cannot be read afterward.
/// </summary>
public bool MakesUnreadableCopies { get; init; }
public Task<TokenizerResponse> ValidateTokenizer(string filePath)
{
this.Calls.Add($"validate {filePath}");
return Task.FromResult(this.AcceptsTokenizers && File.Exists(filePath)
? new TokenizerResponse(true, 0, string.Empty)
: new TokenizerResponse(false, 0, "The file is not a tokenizer."));
}
public Task<TokenizerResponse> StoreTokenizer(string modelId, string filePath)
{
this.Calls.Add($"store {modelId} {filePath}");
if (!this.StoresTokenizers)
return Task.FromResult(new TokenizerResponse(false, 0, "The disk is full."));
var modelDirectory = this.DeleteModelDirectory(modelId);
Directory.CreateDirectory(modelDirectory);
var storedPath = Path.Combine(modelDirectory, Path.GetFileName(filePath));
if (!this.MakesUnreadableCopies)
File.Copy(filePath, storedPath);
return Task.FromResult(new TokenizerResponse(true, 0, string.Empty, storedPath));
}
public Task<TokenizerResponse> DeleteTokenizer(string modelId)
{
this.Calls.Add($"delete {modelId}");
this.DeleteModelDirectory(modelId);
return Task.FromResult(new TokenizerResponse(true, 0, string.Empty));
}
private string DeleteModelDirectory(string modelId)
{
var modelDirectory = Path.Combine(dataDirectory, "tokenizers", modelId);
if (Directory.Exists(modelDirectory))
Directory.Delete(modelDirectory, true);
return modelDirectory;
}
}
}