AI-Studio/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolRegistry.cs

using System.Text.Json;

using AIStudio.Provider;
using AIStudio.Settings;

namespace AIStudio.Tools.ToolCallingSystem;


/// <summary>
/// Holds the tools AI Studio knows and decides which of them a request may use.
/// </summary>
/// <remarks>
/// Definitions arrive through tool definition sources — the app's own tools from code, later the
/// ones plugin authors write. Every definition passes the same validation regardless of where it
/// came from, which matters most for the ones AI Studio does not control.
/// </remarks>
public sealed class ToolRegistry
{
    private readonly ILogger<ToolRegistry> logger;
    private readonly SettingsManager settingsManager;
    private readonly ToolSettingsService toolSettingsService;
    private readonly Dictionary<string, ToolDefinition> definitionsById = new(StringComparer.Ordinal);
    private readonly Dictionary<string, IToolImplementation> implementationsByKey = new(StringComparer.Ordinal);

    /// <summary>
    /// What the checks of a single tool found.
    /// </summary>
    /// <param name="BlockReason">What keeps the tool from being offered, or none.</param>
    /// <param name="Implementation">The tool's implementation, once it was found.</param>
    /// <param name="MinimumConfidence">The confidence the tool requires and where that requirement came from, once it was read.</param>
    private readonly record struct ToolCheck(ToolOfferBlockReason BlockReason, IToolImplementation? Implementation, SettingsManager.ToolMinimumProviderConfidenceResolution? MinimumConfidence);

    public ToolRegistry(
        IEnumerable<IToolImplementation> implementations,
        IEnumerable<IToolDefinitionSource> definitionSources,
        SettingsManager settingsManager,
        ToolSettingsService toolSettingsService,
        ILogger<ToolRegistry> logger)
    {
        this.logger = logger;
        this.settingsManager = settingsManager;
        this.toolSettingsService = toolSettingsService;

        foreach (var implementation in implementations)
        {
            if (string.IsNullOrWhiteSpace(implementation.ImplementationKey))
            {
                this.logger.LogWarning("Skipping a tool implementation with an empty implementation key.");
                continue;
            }

            if (!this.implementationsByKey.TryAdd(implementation.ImplementationKey, implementation))
                this.logger.LogWarning("Skipping duplicate tool implementation key '{ImplementationKey}'.", implementation.ImplementationKey);
        }

        //
        // Function names are checked across all sources together: two tools offering the same
        // name would be indistinguishable to a model, no matter who defined them.
        //
        var functionNames = new HashSet<string>(StringComparer.Ordinal);
        foreach (var source in definitionSources)
        {
            foreach (var definition in source.GetDefinitions())
            {
                if (!TryValidateDefinition(definition, out var validationIssue))
                {
                    this.logger.LogWarning("Skipping tool definition '{ToolId}' from source '{SourceName}': {ValidationIssue}", definition.Id, source.SourceName, validationIssue);
                    continue;
                }

                if (!this.implementationsByKey.ContainsKey(definition.ImplementationKey))
                {
                    this.logger.LogWarning("Skipping tool definition '{ToolId}' because implementation key '{ImplementationKey}' is not registered.", definition.Id, definition.ImplementationKey);
                    continue;
                }

                if (!this.definitionsById.TryAdd(definition.Id, definition))
                {
                    this.logger.LogWarning("Skipping duplicate tool definition ID '{ToolId}' from source '{SourceName}'.", definition.Id, source.SourceName);
                    continue;
                }

                if (!functionNames.Add(definition.Function.Name))
                {
                    this.logger.LogWarning("Skipping tool definition '{ToolId}' because function name '{FunctionName}' is already registered.", definition.Id, definition.Function.Name);
                    this.definitionsById.Remove(definition.Id);
                }
            }
        }
    }

    /// <summary>
    /// Whether a tool definition is complete enough to register.
    /// </summary>
    /// <remarks>
    /// What a definition cannot be is null in its parts: definitions are C# objects whose members
    /// are non-nullable and initialized, so only their content is checked here. Should definitions
    /// one day arrive from outside as data — a tool plugin, say — that assumption ends at the point
    /// where the data becomes a definition, and it is there that null has to be caught.
    /// </remarks>
    private static bool TryValidateDefinition(ToolDefinition definition, out string issue)
    {
        issue = string.Empty;
        if (definition.SchemaVersion != 1)
        {
            issue = $"unsupported schema version '{definition.SchemaVersion}'";
            return false;
        }

        if (string.IsNullOrWhiteSpace(definition.Id))
        {
            issue = "the definition ID is empty";
            return false;
        }

        if (string.IsNullOrWhiteSpace(definition.ImplementationKey))
        {
            issue = "the implementation key is empty";
            return false;
        }

        if (!IsValidFunctionName(definition.Function.Name))
        {
            issue = "the function name must contain 1-64 ASCII letters, digits, underscores, or hyphens";
            return false;
        }

        if (definition.Function.Parameters.ValueKind is not JsonValueKind.Object)
        {
            issue = "the function parameters schema must be a JSON object";
            return false;
        }

        if (definition.VisibleIn.AllowedComponents.Any(component => !Enum.IsDefined(component)) ||
            definition.VisibleIn.DeniedComponents.Any(component => !Enum.IsDefined(component)))
        {
            issue = "the visibility definition must contain valid component lists";
            return false;
        }

        if (!string.Equals(definition.SettingsSchema.Type, "object", StringComparison.OrdinalIgnoreCase))
        {
            issue = "the settings schema must have type 'object'";
            return false;
        }

        if (definition.SettingsSchema.Properties.Any(x =>
                string.IsNullOrWhiteSpace(x.Key) ||
                !string.Equals(x.Value.Type, "string", StringComparison.OrdinalIgnoreCase)))
        {
            issue = "settings properties must be named string fields";
            return false;
        }

        //
        // An empty group is how a field says it belongs to no group. Whitespace looks the
        // same in the settings file but is a different string, so it would open a second,
        // nameless group next to the ungrouped fields:
        //
        var fieldsWithBlankGroup = definition.SettingsSchema.Properties
            .Where(x => x.Value.Group.Length > 0 && string.IsNullOrWhiteSpace(x.Value.Group))
            .Select(x => x.Key)
            .ToList();
        if (fieldsWithBlankGroup.Count > 0)
        {
            issue = $"these settings declare a blank group name: {string.Join(", ", fieldsWithBlankGroup)}";
            return false;
        }

        var fieldsWithBothOptionKinds = definition.SettingsSchema.Properties
            .Where(x => !string.IsNullOrWhiteSpace(x.Value.OptionSource) && x.Value.EnumValues.Count > 0)
            .Select(x => x.Key)
            .ToList();
        if (fieldsWithBothOptionKinds.Count > 0)
        {
            issue = $"these settings declare both an option source and an enum list: {string.Join(", ", fieldsWithBothOptionKinds)}";
            return false;
        }

        var fieldsWithUnknownOptionSource = definition.SettingsSchema.Properties
            .Where(x => !string.IsNullOrWhiteSpace(x.Value.OptionSource) && !ToolSettingsOptionSources.IsKnown(x.Value.OptionSource))
            .Select(x => $"{x.Key} ('{x.Value.OptionSource}')")
            .ToList();
        if (fieldsWithUnknownOptionSource.Count > 0)
        {
            issue = $"these settings reference an unknown option source: {string.Join(", ", fieldsWithUnknownOptionSource)}";
            return false;
        }

        if (definition.SettingsSchema.Required.Any(string.IsNullOrWhiteSpace))
        {
            issue = "required setting names cannot be empty";
            return false;
        }

        var missingRequiredProperties = definition.SettingsSchema.Required
            .Where(x => !definition.SettingsSchema.Properties.ContainsKey(x))
            .ToList();
        if (missingRequiredProperties.Count > 0)
        {
            issue = $"required settings are missing definitions: {string.Join(", ", missingRequiredProperties)}";
            return false;
        }

        return true;
    }

    private static bool IsValidFunctionName(string? functionName) =>
        !string.IsNullOrWhiteSpace(functionName) &&
        functionName.Length <= 64 &&
        functionName.All(character => char.IsAsciiLetterOrDigit(character) || character is '_' or '-');

    public IReadOnlyList<ToolDefinition> GetDefinitionsForComponent(Components component)
    {
        return this.definitionsById.Values
            .Where(x => x.VisibleIn.IsVisibleIn(component))
            .OrderBy(x => this.implementationsByKey.GetValueOrDefault(x.ImplementationKey)?.GetDisplayName(), StringComparer.OrdinalIgnoreCase)
            .ToList();
    }

    public IReadOnlyList<ToolDefinition> GetAllDefinitions() => this.definitionsById.Values
        .OrderBy(x => this.implementationsByKey.GetValueOrDefault(x.ImplementationKey)?.GetDisplayName(), StringComparer.OrdinalIgnoreCase)
        .ToList();

    public ToolDefinition? GetDefinition(string toolId) => this.definitionsById.GetValueOrDefault(toolId);

    public IToolImplementation? GetImplementation(string implementationKey) => this.implementationsByKey.GetValueOrDefault(implementationKey);

    /// <summary>
    /// The provider confidence a tool needs: its own minimum, unless the user or an administrator
    /// raised or lowered it.
    /// </summary>
    /// <remarks>
    /// This is the place that knows both halves — the definition's own minimum and the stored
    /// overrides — so callers holding only a tool ID come here instead of to the settings.
    /// </remarks>
    public ConfidenceLevel GetMinimumProviderConfidence(string toolId) => this.GetDefinition(toolId) is { } definition
        ? this.GetMinimumProviderConfidence(definition)
        : ConfidenceLevel.NONE;

    public ConfidenceLevel GetMinimumProviderConfidence(ToolDefinition definition) =>
        this.settingsManager.GetMinimumProviderConfidenceForTool(definition.Id, definition.MinimumProviderConfidence);

    /// <summary>
    /// Narrows a selection of tool IDs to those the given provider may actually use.
    /// </summary>
    /// <remarks>
    /// Used before a request is sent, so the chat records what will really be available rather
    /// than what the user once ticked. Lives here because judging a tool needs its definition:
    /// the settings know the overrides, the definition knows the tool's own minimum.
    /// </remarks>
    /// <param name="provider">The provider the request goes to.</param>
    /// <param name="selectedToolIds">The tools the user selected.</param>
    /// <returns>The subset that is enabled, active, and allowed by the provider's confidence.</returns>
    public HashSet<string> FilterToolIdsForProvider(AIStudio.Settings.Provider provider, IEnumerable<string> selectedToolIds)
    {
        if (!this.settingsManager.AreToolsEnabled())
            return [];

        if (!provider.GetToolCallingAvailability().IsAvailable)
            return [];

        var providerConfidence = provider.UsedLLMProvider.GetConfidence(this.settingsManager).Level;
        var filtered = ToolSelectionRules.NormalizeSelection(selectedToolIds);
        foreach (var toolId in filtered.ToList())
        {
            if (!this.settingsManager.IsToolActive(toolId))
            {
                filtered.Remove(toolId);
                continue;
            }

            if (!ToolSelectionRules.IsProviderConfidenceAllowed(providerConfidence, this.GetMinimumProviderConfidence(toolId)))
                filtered.Remove(toolId);
        }

        return filtered;
    }

    /// <summary>
    /// The tools somebody can select in this component.
    /// </summary>
    /// <remarks>
    /// Every selection in the app is built from this list: the one below the message field, the
    /// defaults, the templates, and the tools the AI picks for a new assistant. A tool which offers
    /// itself from the context of a chat is left out, because selecting it would change nothing.
    /// The tool list of the app settings asks for all definitions instead, so an organization can
    /// still switch such a tool off or set the trust it requires.
    /// </remarks>
    public async Task<IReadOnlyList<ToolCatalogItem>> GetCatalogAsync(Components component)
    {
        var definitions = this.GetDefinitionsForComponent(component).Where(x => x.Activation is ToolActivation.SELECTION);
        return await this.GetCatalogAsync(definitions);
    }

    /// <summary>
    /// Reduces a set of tool IDs to the tools a user could switch on themselves in this component.
    /// </summary>
    /// <remarks>
    /// For preselecting tools on someone's behalf, such as when a launcher opens a chat. A tool
    /// this installation does not know, one an organization switched off, or one whose settings are
    /// incomplete cannot be enabled by hand either, so handing it over as enabled would show the
    /// user a state they could not have produced and could not fix from where they are. The
    /// provider confidence stays out of this: it belongs to the moment a message is sent, not to
    /// the selection, and it may well be a different provider by then.
    /// </remarks>
    public async Task<HashSet<string>> FilterSelectableToolIdsAsync(Components component, IEnumerable<string> toolIds)
    {
        var wantedToolIds = ToolSelectionRules.NormalizeSelection(toolIds);
        if (wantedToolIds.Count is 0 || !this.settingsManager.AreToolsEnabled())
            return [];

        var catalog = await this.GetCatalogAsync(component);
        return catalog
            .Where(x => wantedToolIds.Contains(x.Definition.Id) && x is { IsActive: true, ConfigurationState.IsConfigured: true })
            .Select(x => x.Definition.Id)
            .ToHashSet(StringComparer.Ordinal);
    }

    public async Task<IReadOnlyList<ToolCatalogItem>> GetCatalogAsync(IEnumerable<ToolDefinition> definitions)
    {
        var definitionList = definitions.ToList();
        var items = new List<ToolCatalogItem>(definitionList.Count);
        foreach (var definition in definitionList)
        {
            if (!this.implementationsByKey.TryGetValue(definition.ImplementationKey, out var implementation))
                continue;

            items.Add(new ToolCatalogItem
            {
                Definition = definition,
                Implementation = implementation,
                ConfigurationState = await this.toolSettingsService.GetConfigurationStateAsync(definition, implementation),
                IsActive = this.settingsManager.IsToolActive(definition.Id),
                MinimumProviderConfidence = this.GetMinimumProviderConfidence(definition),
            });
        }

        return items;
    }

    /// <summary>
    /// The tools a request offers the model, each with the function it offers in this request.
    /// </summary>
    /// <remarks>
    /// Model capabilities are not a parameter on purpose: they are read from the given provider,
    /// which carries the user's expert capability overrides. Passing them in separately allowed a
    /// caller to gate tools on capabilities that differed from the ones the availability check saw.<br/><br/>
    /// The candidates are the selected tools and every tool which offers itself from the context of
    /// the chat, see ToolActivation. Each one passes the same checks, and only then is it asked what
    /// it offers in this request, see IToolImplementation.ResolveFunctionAsync.
    /// </remarks>
    /// <param name="context">The request being prepared.</param>
    /// <param name="selectedToolIds">The tools selected for the request.</param>
    /// <param name="mayRunTools">Whether the request may run tools at all, as its caller decides.</param>
    /// <param name="token">The cancellation token of the request.</param>
    /// <returns>The runnable tools, with their definitions as offered in this request.</returns>
    public async Task<IReadOnlyList<(ToolDefinition Definition, IToolImplementation Implementation)>> GetRunnableToolsAsync(ToolResolutionContext context, IEnumerable<string> selectedToolIds, bool mayRunTools, CancellationToken token = default)
    {
        var provider = context.Provider;
        var component = context.Component;
        var providerConfidence = context.ProviderConfidence;
        if (!this.settingsManager.AreToolsEnabled())
        {
            this.logger.LogDebug("Tool calling is skipped because tools are disabled by managed configuration.");
            return [];
        }

        //
        // Where the user selects the tools, they must be able to see that selection; where the
        // assistant's own rules name them, there is nothing to see. Which of the two applies is
        // decided by the caller, because only it knows where its tools came from:
        //
        if (!mayRunTools)
        {
            this.logger.LogDebug("Tool calling is skipped for component '{Component}' because its tool selection is hidden and no assistant rule names the tools.", component);
            return [];
        }

        var toolCallingAvailability = provider.GetToolCallingAvailability();
        if (!toolCallingAvailability.IsAvailable)
        {
            this.logger.LogDebug("Tool calling is unavailable for provider '{Provider}' with model '{ModelId}': {Reason}", provider.InstanceName, provider.Model.Id, toolCallingAvailability.Message);
            return [];
        }

        var selectedToolIdSet = ToolSelectionRules.NormalizeSelection(selectedToolIds);
        this.logger.LogDebug("Resolving runnable tools for provider '{Provider}' with model '{ModelId}'. Selected tool IDs: [{ToolIds}].", provider.InstanceName, provider.Model.Id, string.Join(", ", selectedToolIdSet.OrderBy(x => x, StringComparer.Ordinal)));

        var definitions = this.GetDefinitionsForComponent(component)
            .Where(x => x.Activation is ToolActivation.CONTEXT || selectedToolIdSet.Contains(x.Id))
            .ToList();

        var result = new List<(ToolDefinition, IToolImplementation)>(definitions.Count);
        foreach (var definition in definitions)
        {
            var check = await this.CheckToolAsync(definition, providerConfidence);
            if (check.MinimumConfidence is { } minimumConfidence)
                this.logger.LogDebug("Tool '{ToolId}' uses minimum provider confidence '{ConfidenceLevel}' from {Source}.", definition.Id, minimumConfidence.ConfidenceLevel, minimumConfidence.Source);

            switch (check)
            {
                case { BlockReason: ToolOfferBlockReason.NONE, Implementation: { } implementation }:
                    if (await this.ResolveAsync(definition, implementation, context, token) is { } offeredDefinition)
                        result.Add((offeredDefinition, implementation));

                    break;

                case { BlockReason: ToolOfferBlockReason.TOOL_SWITCHED_OFF }:
                    this.logger.LogDebug("Skipping tool '{ToolId}' because it is disabled by managed configuration.", definition.Id);
                    break;

                case { BlockReason: ToolOfferBlockReason.NOT_CONFIGURED }:
                    this.logger.LogDebug("Skipping tool '{ToolId}' because it is not configured.", definition.Id);
                    break;

                case { BlockReason: ToolOfferBlockReason.PROVIDER_CONFIDENCE_TOO_LOW }:
                    this.logger.LogInformation("Skipping tool '{ToolId}' because provider confidence '{ProviderConfidence}' is below the required minimum '{MinimumConfidence}'.", definition.Id, providerConfidence, check.MinimumConfidence?.ConfidenceLevel);
                    break;

                case { BlockReason: ToolOfferBlockReason.NOT_AVAILABLE_HERE }:
                    this.logger.LogWarning("Skipping tool '{ToolId}' because no implementation is registered.", definition.Id);
                    break;
            }
        }

        foreach (var selectedToolId in selectedToolIdSet.Where(selectedToolId => definitions.All(definition => !definition.Id.Equals(selectedToolId, StringComparison.Ordinal))))
            this.logger.LogDebug("Skipping tool '{ToolId}' because it is not selected in this component or not available in this context.", selectedToolId);

        return result;
    }

    /// <summary>
    /// Whether a tool can be offered to a provider in this component, and if not, what is in the way.
    /// </summary>
    /// <remarks>
    /// Asks the same questions, in the same order, as the preparation of a request does, because
    /// whoever decides something on the tool's behalf must not come to another answer than the
    /// request will. The RAG process, for instance, leaves the searching of the data sources to
    /// Semantic Search only when this says it can be offered; checks of its own which forgot one
    /// of these would leave a chat without its data sources.<br/><br/>
    /// Two questions stay out. Whether the tool is selected is the caller's business, and whether
    /// the tool has anything to offer right now depends on the chat, so only the preparation of a
    /// request can answer it.
    /// </remarks>
    /// <param name="toolId">The tool to check.</param>
    /// <param name="provider">The provider the request would go to.</param>
    /// <param name="component">Where the request would come from.</param>
    /// <returns>ToolOfferBlockReason.NONE when nothing is in the way, otherwise the first obstacle found.</returns>
    public async Task<ToolOfferBlockReason> GetOfferBlockReasonAsync(string toolId, AIStudio.Settings.Provider provider, Components component)
    {
        if (!this.settingsManager.AreToolsEnabled())
            return ToolOfferBlockReason.TOOLS_SWITCHED_OFF;

        if (!provider.GetToolCallingAvailability().IsAvailable)
            return ToolOfferBlockReason.MODEL_CANNOT_USE_TOOLS;

        if (this.GetDefinition(toolId) is not { } definition || !definition.VisibleIn.IsVisibleIn(component))
            return ToolOfferBlockReason.NOT_AVAILABLE_HERE;

        var providerConfidence = provider.UsedLLMProvider.GetConfidence(this.settingsManager).Level;
        return (await this.CheckToolAsync(definition, providerConfidence)).BlockReason;
    }

    /// <summary>
    /// Checks one tool on its own, apart from what applies to all tools of a request.
    /// </summary>
    /// <remarks>
    /// Shared by the preparation of a request and by GetOfferBlockReasonAsync, so the two cannot
    /// drift apart. It reports rather than logs: the preparation of a request writes down why a
    /// tool was left out, while a question asked by the user interface on every render must not.
    /// </remarks>
    private async Task<ToolCheck> CheckToolAsync(ToolDefinition definition, ConfidenceLevel providerConfidence)
    {
        if (!this.settingsManager.IsToolActive(definition.Id))
            return new(ToolOfferBlockReason.TOOL_SWITCHED_OFF, null, null);

        if (!this.implementationsByKey.TryGetValue(definition.ImplementationKey, out var implementation))
            return new(ToolOfferBlockReason.NOT_AVAILABLE_HERE, null, null);

        var configurationState = await this.toolSettingsService.GetConfigurationStateAsync(definition, implementation);
        if (!configurationState.IsConfigured)
            return new(ToolOfferBlockReason.NOT_CONFIGURED, implementation, null);

        var minimumConfidence = this.settingsManager.GetMinimumProviderConfidenceResolutionForTool(definition.Id, definition.MinimumProviderConfidence);
        if (!ToolSelectionRules.IsProviderConfidenceAllowed(providerConfidence, minimumConfidence.ConfidenceLevel))
            return new(ToolOfferBlockReason.PROVIDER_CONFIDENCE_TOO_LOW, implementation, minimumConfidence);

        return new(ToolOfferBlockReason.NONE, implementation, minimumConfidence);
    }

    /// <summary>
    /// Asks a tool which passed every check what it offers in this request.
    /// </summary>
    /// <remarks>
    /// Only the description and the parameters of the answer are taken. The name and the strict
    /// mode stay as registered, because the model's calls find their tool by that name, and the
    /// rest of the definition was checked a moment ago and must not change after that.
    /// </remarks>
    /// <returns>The definition as offered in this request, or null when the tool has nothing to offer or could not say what.</returns>
    private async Task<ToolDefinition?> ResolveAsync(ToolDefinition definition, IToolImplementation implementation, ToolResolutionContext context, CancellationToken token)
    {
        ToolFunctionDefinition? function;
        try
        {
            function = await implementation.ResolveFunctionAsync(definition, context, token);
        }
        catch (OperationCanceledException) when (token.IsCancellationRequested)
        {
            throw;
        }
        catch (Exception exception)
        {
            this.logger.LogError(exception, "Skipping tool '{ToolId}' because it could not say what it offers in this request.", definition.Id);
            return null;
        }

        if (function is null)
        {
            this.logger.LogDebug("Skipping tool '{ToolId}' because it has nothing to offer in this request.", definition.Id);
            return null;
        }

        if (ReferenceEquals(function, definition.Function))
            return definition;

        if (function.Parameters.ValueKind is not JsonValueKind.Object)
        {
            this.logger.LogWarning("Tool '{ToolId}' offered parameters which are not a JSON object schema. It is offered as registered instead.", definition.Id);
            return definition;
        }

        if (!string.Equals(function.Name, definition.Function.Name, StringComparison.Ordinal) || function.Strict != definition.Function.Strict)
            this.logger.LogWarning("Tool '{ToolId}' changed the name or the strict mode of its function for a request. Both stay as registered.", definition.Id);

        return definition with
        {
            Function = function with
            {
                Name = definition.Function.Name,
                Strict = definition.Function.Strict,
            },
        };
    }
}