AI-Studio/app/MindWork AI Studio/Tools/Services/ConversationTokenCounter.cs

using System.Collections.Concurrent;
using System.Security.Cryptography;
using System.Text;

using AIStudio.Chat;
using AIStudio.Provider;
using AIStudio.Settings;

namespace AIStudio.Tools.Services;

//
// Inside the namespace on purpose. A "Provider" written here would otherwise be the namespace
// AIStudio.Provider, which every namespace below AIStudio sees before it sees a file's aliases.
//
using Provider = AIStudio.Settings.Provider;

/// <summary>
/// Counts what a conversation takes out of a model's context window.
/// </summary>
/// <remarks>
/// Asked from the chat while somebody types, so what it must not do is as important as what it
/// does. Every document is read and measured once and then remembered, because extracting a
/// thousand-page PDF on each keystroke would be unusable. The conversation so far is remembered the
/// same way, so typing measures the sentence being typed rather than the whole chat again.
///
/// The numbers are estimates and are shown as such. Unless somebody configured the model's own
/// tokenizer for their provider, the built-in one does the counting, and two tokenizers disagree by
/// a few percent on prose and by more than that on code.
/// </remarks>
public sealed class ConversationTokenCounter(RustService rustService, ILogger<ConversationTokenCounter> logger)
{
    /// <summary>
    /// How much text goes into one counting request.
    /// </summary>
    /// <remarks>
    /// The same bound the rest of the app uses when it hands text to the tokenizer. Longer
    /// conversations are counted in several pieces and added up, which costs a handful of special
    /// tokens per piece -- a rounding error against a window of hundreds of thousands.
    /// </remarks>
    private const int CHUNK_SIZE = RustService.MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH;

    private readonly ConcurrentDictionary<string, int> counted = new(StringComparer.Ordinal);

    /// <summary>
    /// Counts what the next request would carry.
    /// </summary>
    /// <param name="provider">The configured provider, which decides both the tokenizer and the window.</param>
    /// <param name="thread">The conversation so far, or null when there is none yet.</param>
    /// <param name="systemPrompt">The system prompt as it would be sent.</param>
    /// <param name="draft">What stands in the composer.</param>
    /// <param name="draftAttachments">What is attached to the composer.</param>
    /// <param name="token">Ends the counting when nobody needs the answer anymore.</param>
    /// <returns>What the conversation costs, or that nothing could be counted.</returns>
    public async Task<ConversationTokens> CountAsync(Provider provider, ChatThread? thread, string systemPrompt, string draft, IEnumerable<FileAttachment>? draftAttachments, CancellationToken token = default)
    {
        if (provider.UsedLLMProvider is LLMProviders.NONE)
            return ConversationTokens.UNAVAILABLE;

        var profile = provider.GetModelProfile();
        var parts = ConversationParts.Of(thread, systemPrompt, draft, draftAttachments, provider.SupportsImageInput());
        var tokens = 0;

        try
        {
            foreach (var text in parts.Texts)
                tokens += await this.CountTextAsync(provider, text, token);

            foreach (var document in parts.Documents)
                tokens += await this.CountDocumentAsync(provider, document, token);
        }
        catch (OperationCanceledException)
        {
            return ConversationTokens.UNAVAILABLE;
        }
        catch (Exception e)
        {
            logger.LogWarning(e, "Could not count the tokens of this conversation.");
            return ConversationTokens.UNAVAILABLE;
        }

        return new()
        {
            IsKnown = true,
            Tokens = tokens,
            IsEstimate = string.IsNullOrWhiteSpace(provider.TokenizerPath),
            Window = profile.Context,
            UncountedImages = parts.Images,
        };
    }

    /// <summary>
    /// Forgets everything counted so far.
    /// </summary>
    /// <remarks>
    /// Needed when a file changed behind our back in a way its size and time do not show, which is
    /// rare enough that nothing calls this today. It exists so that the cache has a way out other
    /// than restarting the app.
    /// </remarks>
    public void Forget() => this.counted.Clear();

    /// <summary>
    /// Counts one text, remembering the answer under a fingerprint of it.
    /// </summary>
    /// <remarks>
    /// This is what makes typing affordable. A message which was sent an hour ago says exactly what
    /// it said then, and its tokens are the same number every time -- so the whole conversation is
    /// measured once and every keystroke afterwards measures the sentence being written.
    ///
    /// Keyed by a hash rather than by the text, because the key of the cache would otherwise be a
    /// second copy of the whole conversation in memory. Hashing is not free, but it is two orders of
    /// magnitude cheaper than tokenizing the same bytes, so the trade pays for itself on the first
    /// repeat.
    /// </remarks>
    private async Task<int> CountTextAsync(Provider provider, string text, CancellationToken token)
    {
        if (string.IsNullOrWhiteSpace(text))
            return 0;

        var key = $"{provider.TokenizerPath} {Fingerprint(text)}";
        if (this.counted.TryGetValue(key, out var known))
            return known;

        var tokens = 0;

        //
        // A text longer than one request is split. Cutting between characters rather than between
        // words costs a token or two where the cut falls, which is the cheapest way to stay inside
        // the bound without pretending to know the language.
        //
        for (var start = 0; start < text.Length; start += CHUNK_SIZE)
            tokens += await this.AskTokenizerAsync(provider, text.Substring(start, Math.Min(CHUNK_SIZE, text.Length - start)), token);

        this.counted[key] = tokens;
        return tokens;
    }

    /// <summary>
    /// A short, stable name for a piece of text.
    /// </summary>
    /// <remarks>
    /// The length travels along with the hash. Two texts colliding on the hash and agreeing on
    /// their length as well is not something which happens by accident, and nothing here is a
    /// security decision: the worst a collision could do is show a number which is a few tokens off.
    /// </remarks>
    private static string Fingerprint(string text) => $"{text.Length}:{Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(text)))}";

    /// <summary>
    /// Counts one document, reading it the first time and remembering it afterwards.
    /// </summary>
    /// <remarks>
    /// The key carries the tokenizer as well as the file: the same document counted for two
    /// providers is two different numbers, and handing one of them to the other would be wrong in
    /// exactly the case somebody switches providers to see whether their chat fits.
    /// </remarks>
    private async Task<int> CountDocumentAsync(Provider provider, FileAttachment document, CancellationToken token)
    {
        var file = new FileInfo(document.FilePath);
        if (!file.Exists)
            return 0;

        var key = $"{provider.TokenizerPath}{file.FullName}{file.Length}{file.LastWriteTimeUtc.Ticks}";
        if (this.counted.TryGetValue(key, out var known))
            return known;

        //
        // Read without telling the user about filtered passages. Nothing here is sent anywhere: the
        // text is measured and dropped, and the warning belongs to the moment the document actually
        // travels -- where it is still given.
        //
        var extraction = await rustService.ReadArbitraryFileData(document.FilePath, int.MaxValue, reportPromptInjections: false, token: token);
        if (!extraction.HasUsableContent)
        {
            //
            // A document which cannot be read is not sent either, so it costs nothing. Remembered
            // as zero so that a broken file is not read again on every keystroke.
            //
            logger.LogInformation("The attachment '{FilePath}' could not be read and is therefore counted as nothing.", document.FilePath);
            this.counted[key] = 0;
            return 0;
        }

        var tokens = await this.CountTextAsync(provider, extraction.Content, token);
        this.counted[key] = tokens;
        return tokens;
    }

    private async Task<int> AskTokenizerAsync(Provider provider, string text, CancellationToken token)
    {
        if (string.IsNullOrWhiteSpace(text))
            return 0;

        var response = await rustService.GetTokenCount(provider, text, token);
        if (response is null || !response.Value.Success)
            throw new InvalidOperationException($"The tokenizer did not answer: {response?.Message}");

        return response.Value.TokenCount;
    }
}