AI-Studio/app/Tests/Tools/EmbeddingSignatureTests.cs

using AIStudio.Provider;
using AIStudio.Provider.HuggingFace;
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.");
    }

    [Test]
    public void ChangingTheTokenizerContentDropsTheStoredEmbeddings()
    {
        var dataSource = DataSource(ConfidenceLevel.LOW);
        var oneTokenizer = TokenizerAt("/data/tokenizers/embeddings/tokenizer.json", "AAAA");
        var anotherTokenizer = oneTokenizer with { TokenizerFingerprint = "BBBB" };

        Assert.That(
            Signature(dataSource, anotherTokenizer),
            Is.Not.EqualTo(Signature(dataSource, oneTokenizer)),
            "Another tokenizer cuts the text at other places. A tokenizer is stored under the name it came with, almost always tokenizer.json, so the path alone would not notice the swap.");
    }

    [Test]
    public void MovingTheTokenizerFileKeepsTheStoredEmbeddings()
    {
        var dataSource = DataSource(ConfidenceLevel.LOW);
        var here = TokenizerAt("/data/tokenizers/embeddings/tokenizer.json", "AAAA");
        var there = here with { TokenizerPath = "/somewhere/else/tokenizers/embeddings/tokenizer.json" };

        Assert.That(
            Signature(dataSource, there),
            Is.EqualTo(Signature(dataSource, here)),
            "It is the same tokenizer and only the data directory moved, so embedding everything again would buy nothing.");
    }

    [Test]
    public void ChangingTheHuggingFaceInferenceProviderDropsTheStoredEmbeddings()
    {
        var dataSource = DataSource(ConfidenceLevel.LOW);
        var oneBackend = EmbeddingProviderFor("text-embedding-3-small") with { HFInferenceProvider = HFInferenceProvider.GROQ };
        var anotherBackend = oneBackend with { HFInferenceProvider = HFInferenceProvider.CEREBRAS };

        Assert.That(
            Signature(dataSource, anotherBackend),
            Is.Not.EqualTo(Signature(dataSource, oneBackend)),
            "The same model name served by another backend is another vector source.");
    }

    [Test]
    public void AnIndexedDataSourceOutsideDataSourcesIsCutByItsOwnSettings()
    {
        var directory = DataSource(ConfidenceLevel.LOW);
        var elsewhere = new IndexedElsewhere
        {
            EmbeddingId = directory.EmbeddingId,
            MaxChunkTokenLength = directory.MaxChunkTokenLength,
            ChunkOverlapTokenLength = directory.ChunkOverlapTokenLength,
        };

        Assert.That(
            DataSourceEmbeddingService.BuildEmbeddingSignature(elsewhere, EmbeddingProviderFor("text-embedding-3-small")),
            Is.EqualTo(Signature(directory)),
            "The chunking follows the indexed settings alone. Asking for IInternalDataSource instead would cut such a data source at the token limit of the provider, without anybody noticing.");
    }

    [Test]
    public void TheSignatureOfAKnownConfigurationIsPinned()
    {
        Assert.That(
            Signature(DataSource(ConfidenceLevel.LOW)),
            Is.EqualTo("2|b0a4c4d2-1f3e-4f0a-8c9d-5a6b7c8d9e01|OPEN_AI|text-embedding-3-small|NONE|http://localhost:1234|NONE||8192|512|100"),
            "Reordering or extending the signature throws away every index anybody has. This test makes that a decision somebody takes rather than something which happens on the way past.");
    }

    [Test]
    public void TheSignatureOfAMailboxIsPinned()
    {
        var mailbox = new DataSourceMailbox
        {
            Id = "0c3f9b52-7d4e-4a1b-9e6f-2b8c5d7a1e40",
            EmbeddingId = "b0a4c4d2-1f3e-4f0a-8c9d-5a6b7c8d9e01",
            MaxChunkTokenLength = 512,
            ChunkOverlapTokenLength = 100,
        };

        Assert.That(
            DataSourceEmbeddingService.BuildEmbeddingSignature(mailbox, EmbeddingProviderFor("text-embedding-3-small")),
            Is.EqualTo("2|b0a4c4d2-1f3e-4f0a-8c9d-5a6b7c8d9e01|OPEN_AI|text-embedding-3-small|NONE|http://localhost:1234|NONE||8192|512|100|mail:2|attachments:10"),
            "The text of a mail has a version of its own, which rebuilds the mailboxes and nothing else. Every other data source keeps the signature pinned above.");
    }

    [Test]
    public void ReadingOtherAttachmentsRebuildsAMailbox()
    {
        var embeddingProvider = EmbeddingProviderFor("text-embedding-3-small");
        var mailbox = new DataSourceMailbox { EmbeddingId = embeddingProvider.Id, MaxAttachmentSizeMegabytes = 10 };
        var signature = DataSourceEmbeddingService.BuildEmbeddingSignature(mailbox, embeddingProvider);

        Assert.Multiple(() =>
        {
            Assert.That(DataSourceEmbeddingService.BuildEmbeddingSignature(mailbox with { MaxAttachmentSizeMegabytes = 20 }, embeddingProvider), Is.Not.EqualTo(signature), "The mails indexed before would never get their larger attachments.");
            Assert.That(DataSourceEmbeddingService.BuildEmbeddingSignature(mailbox with { IndexAttachments = false }, embeddingProvider), Does.EndWith("|attachments:none"));
            Assert.That(
                DataSourceEmbeddingService.BuildEmbeddingSignature(mailbox with { IndexAttachments = false, MaxAttachmentSizeMegabytes = 20 }, embeddingProvider),
                Is.EqualTo(DataSourceEmbeddingService.BuildEmbeddingSignature(mailbox with { IndexAttachments = false }, embeddingProvider)),
                "A limit which reads nothing threw the index away for nothing.");
        });
    }

    /// <summary>
    /// Builds the signature the way an indexing run does, working the chunking out along the way.
    /// </summary>
    /// <remarks>
    /// Handing in fixed chunking options instead would hide exactly what these tests are here for:
    /// the signature would then no longer notice a data source being cut differently.
    /// </remarks>
    /// <param name="dataSource">The data source to build the signature for.</param>
    /// <param name="embeddingProvider">The embedding provider, or the test default.</param>
    /// <returns>The signature of that pairing.</returns>
    private static string Signature(DataSourceLocalDirectory dataSource, EmbeddingProvider? embeddingProvider = null) =>
        DataSourceEmbeddingService.BuildEmbeddingSignature(
            dataSource,
            embeddingProvider ?? EmbeddingProviderFor("text-embedding-3-small"));

    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 TokenizerAt(string tokenizerPath, string tokenizerFingerprint) =>
        EmbeddingProviderFor("text-embedding-3-small") with { TokenizerPath = tokenizerPath, TokenizerFingerprint = tokenizerFingerprint };

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

    /// <summary>
    /// A data source which is embedded but not kept in DataSources, the way mailboxes are.
    /// </summary>
    private readonly record struct IndexedElsewhere() : IIndexedDataSource
    {
        public string Id => "0c3f9b52-7d4e-4a1b-9e6f-2b8c5d7a1e40";

        public uint Num => 2;

        public string Name => "Indexed elsewhere";

        public bool IsEnterpriseConfiguration => false;

        public Guid EnterpriseConfigurationPluginId => Guid.Empty;

        public DataSourceType Type { get; init; } = DataSourceType.NONE;

        public ConfidenceLevel ConfidenceLevel { get; init; } = ConfidenceLevel.LOW;

        public string EmbeddingId { get; init; } = Guid.Empty.ToString();

        public int MaxChunkTokenLength { get; init; }

        public int ChunkOverlapTokenLength { get; init; }
    }
}