MindWorkAI/AI-Studio
1
0 Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2026-10-16 06:13:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

made tokenizer concurrent and threadsafe

Browse Source
This commit is contained in:
Paul Koudelka committed 2026-08-10 21:36:01 +02:00
1 parent 6985678d2f
commit c8c1e38475
11 files changed
+110 -306
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs
+9 -106
View File
@@ -1,4 +1,4 @@
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Tools.Rust;
namespace AIStudio.Tools.Services;
@@ -7,35 +7,12 @@ public sealed partial class RustService
{
internal const int MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH = 200_000;
private readonly SemaphoreSlim tokenizerLock = new(1, 1);
private string currentTokenizerPath = string.Empty;
private bool hasInitializedTokenizer;
private static TokenizerResponse CreateUnavailableTokenizerResponse(string message) => new(
false,
0,
message,
TokenizerStatus.UNAVAILABLE,
string.Empty);
public async Task<TokenizerResponse> GetTokenizerInfo(CancellationToken cancellationToken = default)
{
try
{
return await this.http.GetFromJsonAsync<TokenizerResponse>("/system/tokenizer/info", this.jsonRustSerializerOptions, cancellationToken);
}
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
{
this.logger?.LogWarning("Fetching tokenizer info from Rust service was cancelled by caller.");
return CreateUnavailableTokenizerResponse("Operation cancelled by caller.");
}
catch (Exception e)
{
this.logger?.LogError(e, "Error while fetching tokenizer info from Rust service.");
return CreateUnavailableTokenizerResponse(e.Message);
}
}
public async Task<TokenizerResponse> ValidateTokenizer(string filePath)
{
var result = await this.http.PostAsJsonAsync("/tokenizer/validate", new {
@@ -86,99 +63,25 @@ public sealed partial class RustService
return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
}
public Task<TokenizerResponse?> GetTokenCount(string text)
{
return this.GetTokenCountCoreAsync(text);
}
public Task<TokenizerResponse?> GetTokenCount(AIStudio.Settings.Provider provider, string text, CancellationToken cancellationToken = default) =>
this.GetTokenCount(provider.InstanceName, provider.TokenizerPath, text, cancellationToken);
public async Task<TokenizerResponse?> GetTokenCount(string providerName, string path, string text, CancellationToken cancellationToken = default)
{
await this.tokenizerLock.WaitAsync(cancellationToken);
try
{
var tokenizerResponse = await this.EnsureTokenizerCoreAsync(providerName, path);
if (tokenizerResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE })
return tokenizerResponse;
public Task<TokenizerResponse?> GetTokenCount(EmbeddingProvider provider, string text, CancellationToken cancellationToken = default) =>
this.GetTokenCount(provider.Name, provider.TokenizerPath, text, cancellationToken);
return await this.GetTokenCountCoreAsync(text, cancellationToken);
}
finally
{
this.tokenizerLock.Release();
}
}
private async Task<TokenizerResponse?> GetTokenCountCoreAsync(string text, CancellationToken cancellationToken = default)
private async Task<TokenizerResponse?> GetTokenCount(string providerName, string tokenizerPath, string text, CancellationToken cancellationToken)
{
var result = await this.http.PostAsJsonAsync("/tokenizer/count", new {
text = text,
tokenizer_path = tokenizerPath,
}, this.jsonRustSerializerOptions, cancellationToken);
if (!result.IsSuccessStatusCode)
{
this.logger!.LogError($"Failed to get the token count '{result.StatusCode}'");
this.hasInitializedTokenizer = false;
this.logger!.LogError("Failed to get the token count for provider '{ProviderName}': {StatusCode}", providerName, result.StatusCode);
return CreateUnavailableTokenizerResponse("Error while getting token count from Rust service: "+result.StatusCode);
}
var response = await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
if (response is not { Status: TokenizerStatus.AVAILABLE })
this.hasInitializedTokenizer = false;
return response;
}
public async Task<TokenizerResponse?> SetTokenizer(string providerName, string path)
{
this.logger!.LogInformation($"Setting a new tokenizer for '{providerName}'");
var result = await this.http.PostAsJsonAsync("/tokenizer/set", new {
file_path = path,
}, this.jsonRustSerializerOptions);
if (!result.IsSuccessStatusCode)
{
this.logger!.LogError($"Failed to set the tokenizer '{result.StatusCode}'");
this.hasInitializedTokenizer = false;
return CreateUnavailableTokenizerResponse("An error occured while sending the path to the Rust framework for setting a tokenizer: "+result.StatusCode);
}
var response = await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
if (response is not { Success: true, Status: TokenizerStatus.AVAILABLE })
this.hasInitializedTokenizer = false;
return response;
}
public async Task<TokenizerResponse?> EnsureTokenizer(string providerName, string path)
{
await this.tokenizerLock.WaitAsync();
try
{
return await this.EnsureTokenizerCoreAsync(providerName, path);
}
finally
{
this.tokenizerLock.Release();
}
}
private async Task<TokenizerResponse?> EnsureTokenizerCoreAsync(string providerName, string path)
{
if (this.hasInitializedTokenizer && this.currentTokenizerPath == path)
return new TokenizerResponse(true, 0, string.Empty, TokenizerStatus.AVAILABLE);
var response = await this.SetTokenizer(providerName, path);
if (response is { Success: true, Status: TokenizerStatus.AVAILABLE })
{
this.currentTokenizerPath = path;
this.hasInitializedTokenizer = true;
}
else
{
this.currentTokenizerPath = string.Empty;
this.hasInitializedTokenizer = false;
}
return response;
return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions, cancellationToken);
}
}