AI-Studio/app/MindWork AI Studio/Settings/EmbeddingProvider.cs

using System.Text.Json.Serialization;

using AIStudio.Provider;
using AIStudio.Provider.HuggingFace;
using AIStudio.Tools.PluginSystem;

using SharedTools;

using Host = AIStudio.Provider.SelfHosted.Host;
using LuaTable = Lua.LuaTable;

namespace AIStudio.Settings;

public sealed record EmbeddingProvider(
    uint Num,
    string Id,
    string Name,
    LLMProviders UsedLLMProvider,
    Model Model,
    bool IsSelfHosted = false,
    bool IsEnterpriseConfiguration = false,
    Guid EnterpriseConfigurationPluginId = default,
    string Hostname = "http://localhost:1234",
    Host Host = Host.NONE,
    string TokenizerPath = "",
    int EmbeddingBatchSize = 0,
    int TokenLimit = 0,
    bool AllowUserProvidedAPIKey = false,
    string CustomIconDataUrl = "",
    HFInferenceProvider HFInferenceProvider = HFInferenceProvider.NONE) : ConfigurationBaseObject, ISecretId, IUserProvidedAPIKey
{
    public const int DEFAULT_TOKEN_LIMIT = 8192;
    public const int DEFAULT_EMBEDDING_BATCH_SIZE = 1;

    private static readonly ILogger<EmbeddingProvider> LOGGER = Program.LOGGER_FACTORY.CreateLogger<EmbeddingProvider>();

    public static readonly EmbeddingProvider NONE = new();

    public EmbeddingProvider() : this(0, Guid.Empty.ToString(), string.Empty, LLMProviders.NONE, default, false, false, Guid.Empty)
    {
    }

    public override string ToString() => this.Name;

    #region Implementation of ISecretId

    /// <inheritdoc />
    [JsonIgnore]
    public string SecretId => this.IsEnterpriseConfiguration ? $"{ISecretId.ENTERPRISE_KEY_PREFIX}::{this.UsedLLMProvider.ToSecretId()}" : this.UsedLLMProvider.ToSecretId();

    /// <inheritdoc />
    [JsonIgnore]
    public string SecretName => this.Name;

    [JsonIgnore]
    public int EffectiveTokenLimit => this.TokenLimit > 0 ? this.TokenLimit : DEFAULT_TOKEN_LIMIT;

    [JsonIgnore]
    public int EffectiveEmbeddingBatchSize => this.EmbeddingBatchSize > 0 ? this.EmbeddingBatchSize : DEFAULT_EMBEDDING_BATCH_SIZE;

    #endregion

    public static bool TryParseEmbeddingProviderTable(int idx, LuaTable table, Guid configPluginId, string pluginPath, out ConfigurationBaseObject provider)
    {
        provider = NONE;
        if (!table.TryGetValue("Id", out var idValue) || !idValue.TryRead<string>(out var idText) || !Guid.TryParse(idText, out var id))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid ID. The ID must be a valid GUID. (Plugin ID: {configPluginId})");
            return false;
        }

        if (!table.TryGetValue("Name", out var nameValue) || !nameValue.TryRead<string>(out var name))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid name. (Plugin ID: {configPluginId})");
            return false;
        }

        if (!table.TryGetValue("UsedLLMProvider", out var usedLLMProviderValue) || !usedLLMProviderValue.TryRead<string>(out var usedLLMProviderText) || !Enum.TryParse<LLMProviders>(usedLLMProviderText, true, out var usedLLMProvider))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid LLM provider enum value. (Plugin ID: {configPluginId})");
            return false;
        }

        if (!table.TryGetValue("Host", out var hostValue) || !hostValue.TryRead<string>(out var hostText) || !Enum.TryParse<Host>(hostText, true, out var host))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid host enum value. (Plugin ID: {configPluginId})");
            return false;
        }

        if (!table.TryGetValue("Hostname", out var hostnameValue) || !hostnameValue.TryRead<string>(out var hostname))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid hostname. (Plugin ID: {configPluginId})");
            return false;
        }

        if (!table.TryGetValue("Model", out var modelValue) || !modelValue.TryRead<LuaTable>(out var modelTable))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid model table. (Plugin ID: {configPluginId})");
            return false;
        }

        if (!TryReadModelTable(idx, modelTable, configPluginId, out var model))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid model configuration. (Plugin ID: {configPluginId})");
            return false;
        }

        var tokenizerPath = string.Empty;
        if (table.TryGetValue("TokenizerPath", out var tokenizerPathValue) && !tokenizerPathValue.TryRead<string>(out tokenizerPath))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid tokenizer path. (Plugin ID: {configPluginId})");
            tokenizerPath = string.Empty;
        }

        var tokenLimit = DEFAULT_TOKEN_LIMIT;
        if (table.TryGetValue("TokenLimit", out var tokenLimitValue) && (!tokenLimitValue.TryRead<int>(out tokenLimit) || tokenLimit < 1))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid token limit. Falling back to {DEFAULT_TOKEN_LIMIT}. (Plugin ID: {configPluginId})");
            tokenLimit = DEFAULT_TOKEN_LIMIT;
        }

        var embeddingBatchSize = DEFAULT_EMBEDDING_BATCH_SIZE;
        if (table.TryGetValue("EmbeddingBatchSize", out var embeddingBatchSizeValue) && (!embeddingBatchSizeValue.TryRead<int>(out embeddingBatchSize) || embeddingBatchSize < 1))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid embedding batch size. Falling back to {DEFAULT_EMBEDDING_BATCH_SIZE}. (Plugin ID: {configPluginId})");
            embeddingBatchSize = DEFAULT_EMBEDDING_BATCH_SIZE;
        }

        var allowUserProvidedApiKey = false;
        if (table.TryGetValue("AllowUserProvidedAPIKey", out var allowUserProvidedApiKeyValue) && allowUserProvidedApiKeyValue.TryRead<bool>(out var allowUserProvidedApiKeyBool))
            allowUserProvidedApiKey = allowUserProvidedApiKeyBool;

        var hfInferenceProvider = HFInferenceProvider.NONE;
        if (table.TryGetValue("HFInferenceProvider", out var hfInferenceProviderValue) && hfInferenceProviderValue.TryRead<string>(out var hfInferenceProviderText))
        {
            if (!Enum.TryParse(hfInferenceProviderText, true, out hfInferenceProvider))
            {
                LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid Hugging Face inference provider enum value. (Plugin ID: {configPluginId})");
                hfInferenceProvider = HFInferenceProvider.NONE;
            }
        }

        var customIconDataUrl = string.Empty;
        if (table.TryGetValue("IconPath", out var iconPathValue))
        {
            if (!iconPathValue.TryRead<string>(out var iconPath))
                LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid icon path. Falling back to the built-in provider icon. (Plugin ID: {configPluginId})");
            else if (!PluginIconFile.TryLoadDataUrl(iconPath, pluginPath, out customIconDataUrl, out var iconIssue))
                LOGGER.LogWarning($"The configured embedding provider {idx} contains an invalid icon path. Falling back to the built-in provider icon. Issue: {iconIssue} (Plugin ID: {configPluginId})");
        }

        provider = new EmbeddingProvider
        {
            Num = 0, // will be set later by the PluginConfigurationObject
            Id = id.ToString(),
            Name = name,
            UsedLLMProvider = usedLLMProvider,
            Model = model,
            IsSelfHosted = usedLLMProvider is LLMProviders.SELF_HOSTED,
            IsEnterpriseConfiguration = true,
            EnterpriseConfigurationPluginId = configPluginId,
            Hostname = hostname,
            Host = host,
            TokenizerPath = tokenizerPath,
            EmbeddingBatchSize = embeddingBatchSize,
            TokenLimit = tokenLimit,
            AllowUserProvidedAPIKey = allowUserProvidedApiKey,
            CustomIconDataUrl = customIconDataUrl,
            HFInferenceProvider = hfInferenceProvider,
        };

        // Handle an encrypted API key if present. When the user manages their own key for this
        // embedding provider, we must never enqueue an embedded key: doing so would overwrite the
        // user's key in the OS keyring on every configuration reload.
        if (allowUserProvidedApiKey)
        {
            if (table.TryGetValue("APIKey", out var ignoredApiKeyValue) && ignoredApiKeyValue.TryRead<string>(out var ignoredApiKeyText) && !string.IsNullOrWhiteSpace(ignoredApiKeyText))
                LOGGER.LogWarning($"The configured embedding provider {idx} sets both AllowUserProvidedAPIKey and an embedded APIKey. Ignoring the embedded key: the user manages their own key for this provider. (Plugin ID: {configPluginId})");
        }
        else if (table.TryGetValue("APIKey", out var apiKeyValue) && apiKeyValue.TryRead<string>(out var apiKeyText) && !string.IsNullOrWhiteSpace(apiKeyText))
        {
            if (!EnterpriseEncryption.IsEncrypted(apiKeyText))
                LOGGER.LogWarning($"The configured embedding provider {idx} contains a plaintext API key. Only encrypted API keys (starting with 'ENC:v1:') are supported. (Plugin ID: {configPluginId})");
            else
            {
                var encryption = PluginFactory.EnterpriseEncryption;
                if (encryption?.IsAvailable == true)
                {
                    if (encryption.TryDecrypt(apiKeyText, out var decryptedApiKey))
                    {
                        // Queue the API key for storage in the OS keyring:
                        PendingEnterpriseApiKeys.Add(new(
                            $"{ISecretId.ENTERPRISE_KEY_PREFIX}::{usedLLMProvider.ToSecretId()}",
                            name,
                            decryptedApiKey,
                            SecretStoreType.EMBEDDING_PROVIDER));
                        LOGGER.LogDebug($"Successfully decrypted API key for embedding provider {idx}. It will be stored in the OS keyring. (Plugin ID: {configPluginId})");
                    }
                    else
                        LOGGER.LogWarning($"Failed to decrypt API key for embedding provider {idx}. The encryption secret may be incorrect. (Plugin ID: {configPluginId})");
                }
                else
                    LOGGER.LogWarning($"The configured embedding provider {idx} contains an encrypted API key, but no encryption secret is configured. (Plugin ID: {configPluginId})");
            }
        }

        return true;
    }

    private static bool TryReadModelTable(int idx, LuaTable table, Guid configPluginId, out Model model)
    {
        model = default;
        if (!table.TryGetValue("Id", out var idValue) || !idValue.TryRead<string>(out var id))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid model ID. (Plugin ID: {configPluginId})");
            return false;
        }

        if (!table.TryGetValue("DisplayName", out var displayNameValue) || !displayNameValue.TryRead<string>(out var displayName))
        {
            LOGGER.LogWarning($"The configured embedding provider {idx} does not contain a valid model display name. (Plugin ID: {configPluginId})");
            return false;
        }

        model = new(id, displayName);
        return true;
    }

    /// <summary>
    /// Exports the embedding provider configuration as a Lua configuration section.
    /// </summary>
    /// <param name="encryptedApiKey">Optional encrypted API key to include in the export.</param>
    /// <returns>A Lua configuration section string.</returns>
    public string ExportAsConfigurationSection(string? encryptedApiKey = null)
    {
        var hfInferenceProviderLine = string.Empty;
        if (this.HFInferenceProvider is not HFInferenceProvider.NONE)
        {
            hfInferenceProviderLine = $"""
                                       ["HFInferenceProvider"] = "{this.HFInferenceProvider}",
                                       """;
        }

        var apiKeyLine = string.Empty;
        if (!string.IsNullOrWhiteSpace(encryptedApiKey))
        {
            apiKeyLine = $"""
                          ["APIKey"] = "{LuaTools.EscapeLuaString(encryptedApiKey)}",
                          """;
        }

        return $$"""
                CONFIG["EMBEDDING_PROVIDERS"][#CONFIG["EMBEDDING_PROVIDERS"]+1] = {
                    ["Id"] = "{{Guid.NewGuid().ToString()}}",
                    ["Name"] = "{{LuaTools.EscapeLuaString(this.Name)}}",
                    ["UsedLLMProvider"] = "{{this.UsedLLMProvider}}",

                    ["TokenizerPath"] = "{{this.TokenizerPath}}",
                    ["TokenLimit"] = {{this.EffectiveTokenLimit}},
                    ["EmbeddingBatchSize"] = {{this.EffectiveEmbeddingBatchSize}},

                    ["Host"] = "{{this.Host}}",
                    ["Hostname"] = "{{LuaTools.EscapeLuaString(this.Hostname)}}",
                    {{hfInferenceProviderLine}}
                    {{apiKeyLine}}
                    ["Model"] = {
                        ["Id"] = "{{LuaTools.EscapeLuaString(this.Model.Id)}}",
                        ["DisplayName"] = "{{LuaTools.EscapeLuaString(this.Model.DisplayName ?? string.Empty)}}",
                    },
                }
                """;
    }
}