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

using System.Diagnostics;
using System.Text.Json;
using System.Text.Json.Nodes;

using AIStudio.Chat;
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.Web;

namespace AIStudio.Tools.ToolCallingSystem;

public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogger<ToolExecutor> logger)
{
    private const string INVALID_TOOL_CALL_ERROR = "The tool call was invalid.";

    public ToolCallOutcome CreateInvalidToolCallResult(string toolCallId, int order)
    {
        logger.LogWarning(
            "Rejected invalid tool call. ToolCallId={ToolCallId}, Order={Order}, Status={Status}",
            toolCallId,
            order,
            ToolInvocationTraceStatus.ERROR);
        return new ToolCallOutcome(INVALID_TOOL_CALL_ERROR, new ToolInvocationTrace
        {
            Order = order,
            ToolName = "Invalid tool call",
            ToolCallId = toolCallId,
            Status = ToolInvocationTraceStatus.ERROR,
            StatusMessage = INVALID_TOOL_CALL_ERROR,
            Result = INVALID_TOOL_CALL_ERROR,
        });
    }

    public static bool IsValidArgumentsJson(string? argumentsJson)
    {
        if (string.IsNullOrWhiteSpace(argumentsJson))
            return false;

        try
        {
            using var document = JsonDocument.Parse(argumentsJson);
            return document.RootElement.ValueKind is JsonValueKind.Object;
        }
        catch (JsonException)
        {
            return false;
        }
    }

    public async Task<ToolCallOutcome> ExecuteAsync(
        string toolCallId,
        string toolName,
        string argumentsJson,
        IReadOnlyList<(ToolDefinition Definition, IToolImplementation Implementation)> runnableTools,
        IProvider provider,
        ChatThread chatThread,
        int order,
        CancellationToken token = default)
    {
        var runnableTool = runnableTools.FirstOrDefault(x => x.Definition.Function.Name.Equals(toolName, StringComparison.Ordinal));
        Dictionary<string, string> formattedArguments = [];
        try
        {
            using var document = JsonDocument.Parse(string.IsNullOrWhiteSpace(argumentsJson) ? "{}" : argumentsJson);
            formattedArguments = FormatArguments(document.RootElement, runnableTool.Implementation?.SensitiveTraceArgumentNames ?? EmptySensitiveTraceArgumentNames.INSTANCE);
        }
        catch (JsonException)
        {
            //
            // Only the trace loses its arguments here; the execution below parses the same JSON
            // again and reports a broken call properly. The message says which call it was, but
            // nothing about its content: arguments may carry secrets, and a parser message quotes
            // the text it stumbled over.
            //
            logger.LogWarning("Could not read the arguments of a tool call for its trace. ToolName={ToolName}, ToolCallId={ToolCallId}", toolName, toolCallId);
        }

        logger.LogInformation(
            "Starting tool execution. ToolName={ToolName}, ToolCallId={ToolCallId}",
            toolName,
            toolCallId);
        var stopwatch = Stopwatch.StartNew();
        if (runnableTool.Definition is null || runnableTool.Implementation is null)
        {
            var error = this.CreateError(toolName);
            logger.LogWarning("Completed tool execution. ToolName={ToolName}, ToolCallId={ToolCallId}, DurationMs={DurationMs}, Status={Status}", toolName, toolCallId, stopwatch.ElapsedMilliseconds, ToolInvocationTraceStatus.BLOCKED);
            return new ToolCallOutcome(error, new ToolInvocationTrace
            {
                Order = order,
                ToolId = toolName,
                ToolName = toolName,
                ToolCallId = toolCallId,
                Status = ToolInvocationTraceStatus.BLOCKED,
                StatusMessage = "Tool is not available in the current context.",
                Arguments = formattedArguments,
                Result = error,
            });
        }

        var definition = runnableTool.Definition;
        var implementation = runnableTool.Implementation;
        try
        {
            using var document = JsonDocument.Parse(string.IsNullOrWhiteSpace(argumentsJson) ? "{}" : argumentsJson);
            var settingsManager = Program.SERVICE_PROVIDER.GetRequiredService<SettingsManager>();

            //
            // Asked again here, although the request only offers what was allowed: the tools of a
            // request are chosen once, before its first round, and a tool which reads a mailbox may
            // restrict the chat in the middle of it. The model can then still call a tool it was
            // offered, with the mail content in its arguments.
            //
            var outboundDataRestriction = chatThread.RequiredOutboundDataRestriction;
            if (!ToolSelectionRules.IsOutboundDataAllowed(outboundDataRestriction.Restriction, implementation))
                throw new ToolExecutionBlockedException(outboundDataRestriction.GetToolBlockedMessage(settingsManager.ConfigurationData.Mailboxes));

            var settingsValues = await toolSettingsService.GetSettingsAsync(definition);
            var result = await implementation.ExecuteAsync(document.RootElement, new ToolExecutionContext
            {
                Definition = definition,
                ChatThread = chatThread,
                Provider = provider,
                ToolCallId = toolCallId,
                SettingsManager = settingsManager,
                SettingsValues = settingsValues,
                ProviderConfidence = provider.Provider.GetConfidence(settingsManager).Level,
            }, token);
            logger.LogInformation("Completed tool execution. ToolName={ToolName}, ToolCallId={ToolCallId}, DurationMs={DurationMs}, Status={Status}", toolName, toolCallId, stopwatch.ElapsedMilliseconds, ToolInvocationTraceStatus.SUCCESS);

            var resultModelContent = result.ToModelContent();
            var toolInvocationTrace = new ToolInvocationTrace
            {
                Order = order,
                ToolId = definition.Id,
                ToolName = implementation.GetDisplayName(),
                ToolIcon = implementation.Icon,
                ToolCallId = toolCallId,
                Status = ToolInvocationTraceStatus.SUCCESS,
                WasExecuted = true,
                Arguments = FormatArguments(document.RootElement,
                    implementation.SensitiveTraceArgumentNames),
                Result = result.TextContent ?? string.Empty,
                JsonResult = result.JsonContent,
            };

            return new ToolCallOutcome(resultModelContent, toolInvocationTrace)
            {
                RequiredProviderConfidence = result.RequiredProviderConfidence,
                RequiredDataSecurity = result.RequiredDataSecurity,
                RequiredOutboundDataRestriction = result.RequiredOutboundDataRestriction,
                Sources = result.Sources,
                ReturnedWebAddresses = FindReturnedWebAddresses(result, document.RootElement),
            };
        }
        catch (OperationCanceledException) when (token.IsCancellationRequested)
        {
            logger.LogInformation("Completed tool execution. ToolName={ToolName}, ToolCallId={ToolCallId}, DurationMs={DurationMs}, Status={Status}", toolName, toolCallId, stopwatch.ElapsedMilliseconds, "CANCELED");
            throw;
        }
        catch (ToolExecutionBlockedException exception)
        {
            logger.LogWarning("Tool execution was blocked. ToolName={ToolName}, ToolCallId={ToolCallId}, DurationMs={DurationMs}, Status={Status}, Reason={Reason}", toolName, toolCallId, stopwatch.ElapsedMilliseconds, ToolInvocationTraceStatus.BLOCKED, exception.Message);

            var toolInvocationTrace = new ToolInvocationTrace
            {
                Order = order,
                ToolId = definition.Id,
                ToolName = implementation.GetDisplayName(),
                ToolIcon = implementation.Icon,
                ToolCallId = toolCallId,
                Status = ToolInvocationTraceStatus.BLOCKED,
                StatusMessage = exception.Message,
                Arguments = formattedArguments,
                Result = exception.Message,
            };

            return new ToolCallOutcome(exception.Message, toolInvocationTrace);
        }
        catch (Exception exception)
        {
            var error = $"Tool execution failed: {exception.Message}";
            logger.LogError(exception, "Tool execution failed. ToolName={ToolName}, ToolCallId={ToolCallId}, DurationMs={DurationMs}, Status={Status}", toolName, toolCallId, stopwatch.ElapsedMilliseconds, ToolInvocationTraceStatus.ERROR);

            var toolInvocationTrace = new ToolInvocationTrace
            {
                Order = order,
                ToolId = definition.Id,
                ToolName = implementation.GetDisplayName(),
                ToolIcon = implementation.Icon,
                ToolCallId = toolCallId,
                Status = ToolInvocationTraceStatus.ERROR,
                StatusMessage = error,
                Arguments = formattedArguments,
                Result = error,
            };

            return new ToolCallOutcome(error, toolInvocationTrace);
        }
    }

    private static class EmptySensitiveTraceArgumentNames
    {
        public static readonly IReadOnlySet<string> INSTANCE = new HashSet<string>(StringComparer.Ordinal);
    }

    private string CreateError(string toolName) => $"Tool '{toolName}' is not available.";

    /// <summary>
    /// The web addresses in a result which did not come from the model.
    /// </summary>
    /// <remarks>
    /// An address which stands in one of the arguments, even as a part of one, is an echo of what
    /// the model wrote and is left out. Case does not matter for that, so a tool which writes the
    /// host in lower case does not slip one through, and neither does a tool which writes the same
    /// request in another encoding, since an address in the arguments also counts by its request
    /// key. The texts are compared as the JSON values read, so an escape in the JSON cannot hide an
    /// echo either.
    /// </remarks>
    private static HashSet<string> FindReturnedWebAddresses(ToolExecutionResult result, JsonElement arguments)
    {
        var argumentTexts = new List<string>();
        CollectStrings(arguments, argumentTexts);

        var argumentRequestKeys = new HashSet<string>(StringComparer.Ordinal);
        foreach (var address in argumentTexts.SelectMany(WebAddresses.Find))
            if (WebAddresses.TryCreateRequestKey(address, out var requestKey))
                argumentRequestKeys.Add(requestKey);

        var resultTexts = new List<string>();
        if (!string.IsNullOrWhiteSpace(result.TextContent))
            resultTexts.Add(result.TextContent);

        CollectStrings(result.JsonContent, resultTexts);

        var requestKeys = new HashSet<string>(StringComparer.Ordinal);
        foreach (var address in resultTexts.SelectMany(WebAddresses.Find))
        {
            if (argumentTexts.Any(argument => argument.Contains(address, StringComparison.OrdinalIgnoreCase)))
                continue;

            if (WebAddresses.TryCreateRequestKey(address, out var requestKey) && !argumentRequestKeys.Contains(requestKey))
                requestKeys.Add(requestKey);
        }

        return requestKeys;
    }

    private static void CollectStrings(JsonElement element, List<string> texts)
    {
        switch (element.ValueKind)
        {
            case JsonValueKind.String:
                texts.Add(element.GetString() ?? string.Empty);
                break;

            case JsonValueKind.Object:
                foreach (var property in element.EnumerateObject())
                    CollectStrings(property.Value, texts);

                break;

            case JsonValueKind.Array:
                foreach (var item in element.EnumerateArray())
                    CollectStrings(item, texts);

                break;
        }
    }

    private static void CollectStrings(JsonNode? node, List<string> texts)
    {
        switch (node)
        {
            case JsonObject jsonObject:
                foreach (var property in jsonObject)
                    CollectStrings(property.Value, texts);

                break;

            case JsonArray jsonArray:
                foreach (var item in jsonArray)
                    CollectStrings(item, texts);

                break;

            case JsonValue jsonValue when jsonValue.GetValueKind() is JsonValueKind.String:
                texts.Add(jsonValue.GetValue<string>());
                break;
        }
    }

    private static Dictionary<string, string> FormatArguments(JsonElement rootElement, IReadOnlySet<string> sensitiveNames)
    {
        if (rootElement.ValueKind is not JsonValueKind.Object)
            return [];

        var arguments = new Dictionary<string, string>(StringComparer.Ordinal);
        foreach (var property in rootElement.EnumerateObject())
        {
            arguments[property.Name] = sensitiveNames.Contains(property.Name)
                ? "*****"
                : property.Value.ValueKind switch
                {
                    JsonValueKind.String => property.Value.GetString() ?? string.Empty,
                    _ => property.Value.ToString(),
                };
        }

        return arguments;
    }
}