From 4617aa5230439596099194e516ec8b3dc9c231d5 Mon Sep 17 00:00:00 2001 From: Thorsten Sommer Date: Thu, 24 Sep 2026 15:07:39 +0200 Subject: [PATCH] Let tool results restrict a chat to self-hosted providers --- app/MindWork AI Studio/Chat/ChatThread.cs | 20 +++++ .../RAG/RAGProcesses/AISrcSelWithRetCtxVal.cs | 43 ++------- .../Harness/ToolCallingLoop.cs | 4 +- .../ToolCallingSystem/ToolExecutionContext.cs | 11 +++ .../ToolCallingSystem/ToolExecutionResult.cs | 12 +++ .../Tools/ToolCallingSystem/ToolExecutor.cs | 14 +-- app/Tests/Chat/ChatThreadDataSecurityTests.cs | 52 +++++++++++ .../Tools/ToolCalling/ToolExecutorTests.cs | 90 +++++++++++++++++++ .../Tools/ToolCalling/ToolRegistryTestBase.cs | 15 ++-- 9 files changed, 209 insertions(+), 52 deletions(-) create mode 100644 app/Tests/Chat/ChatThreadDataSecurityTests.cs create mode 100644 app/Tests/Tools/ToolCalling/ToolExecutorTests.cs diff --git a/app/MindWork AI Studio/Chat/ChatThread.cs b/app/MindWork AI Studio/Chat/ChatThread.cs index 0a10b2c9..beed95ab 100644 --- a/app/MindWork AI Studio/Chat/ChatThread.cs +++ b/app/MindWork AI Studio/Chat/ChatThread.cs @@ -106,6 +106,26 @@ public sealed record ChatThread this.RequiredProviderConfidence = minimumProviderConfidence; } + /// + /// Tightens the data security of this chat to what the data brought in demands, and never + /// loosens it. + /// + /// + /// Data which may only be used with self-hosted providers keeps the chat restricted to them, + /// no matter what comes in later: the data was seen by this chat. Data which may be used with + /// any provider marks the chat as one which holds data of a data source, while a restriction set + /// earlier stays. NOT_SPECIFIED demands nothing and changes nothing.

+ /// Shared by the RAG process and by tools which search the data sources, so both tighten a chat + /// the same way. + /// + /// What the data brought in demands. + public void RequireDataSecurity(DataSourceSecurity dataSecurity) => this.DataSecurity = (this.DataSecurity, dataSecurity) switch + { + (DataSourceSecurity.SELF_HOSTED, _) or (_, DataSourceSecurity.SELF_HOSTED) => DataSourceSecurity.SELF_HOSTED, + (_, DataSourceSecurity.ALLOW_ANY) => DataSourceSecurity.ALLOW_ANY, + _ => this.DataSecurity, + }; + /// /// The name of the chat thread. Usually generated by an AI model or manually edited by the user. /// diff --git a/app/MindWork AI Studio/Tools/RAG/RAGProcesses/AISrcSelWithRetCtxVal.cs b/app/MindWork AI Studio/Tools/RAG/RAGProcesses/AISrcSelWithRetCtxVal.cs index 749a94c8..400218de 100644 --- a/app/MindWork AI Studio/Tools/RAG/RAGProcesses/AISrcSelWithRetCtxVal.cs +++ b/app/MindWork AI Studio/Tools/RAG/RAGProcesses/AISrcSelWithRetCtxVal.cs @@ -136,48 +136,15 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess // // Update the data security of the chat thread. We consider the current data security - // of the chat thread and the data security of the selected data sources: + // of the chat thread and the data security of the selected data sources: at least + // one data source with a SELF_HOSTED policy restricts the chat to self-hosted + // providers. A restriction set earlier stays either way, because the thread might + // already contain data from a data source with a SELF_HOSTED policy: // var dataSecurityRestrictedToSelfHosted = selectedDataSources .OfType() .Any(dataSource => dataSource.SecurityPolicy is DataSourceSecurity.SELF_HOSTED); - chatThread.DataSecurity = dataSecurityRestrictedToSelfHosted switch - { - // - // - // Case: the data sources which are selected have a security policy - // of SELF_HOSTED (at least one data source). - // - // When the policy was already set to ALLOW_ANY, we restrict it - // to SELF_HOSTED. - // - true => DataSourceSecurity.SELF_HOSTED, - - // - // Case: the data sources which are selected have a security policy - // of ALLOW_ANY (none of the data sources has a SELF_HOSTED policy). - // - // When the policy was already set to SELF_HOSTED, we must keep that. - // - false => chatThread.DataSecurity switch - { - // - // When the policy was not specified yet, we set it to ALLOW_ANY. - // - DataSourceSecurity.NOT_SPECIFIED => DataSourceSecurity.ALLOW_ANY, - DataSourceSecurity.ALLOW_ANY => DataSourceSecurity.ALLOW_ANY, - - // - // When the policy was already set to SELF_HOSTED, we must keep that. - // This is important since the thread might already contain data - // from a data source with a SELF_HOSTED policy. - // - DataSourceSecurity.SELF_HOSTED => DataSourceSecurity.SELF_HOSTED, - - // Default case: we use the current data security of the chat thread. - _ => chatThread.DataSecurity, - } - }; + chatThread.RequireDataSecurity(dataSecurityRestrictedToSelfHosted ? DataSourceSecurity.SELF_HOSTED : DataSourceSecurity.ALLOW_ANY); if (previousDataSecurity != chatThread.DataSecurity) LOGGER.LogInformation($"The data security of the chat thread was updated from '{previousDataSecurity}' to '{chatThread.DataSecurity}'."); diff --git a/app/MindWork AI Studio/Tools/ToolCallingSystem/Harness/ToolCallingLoop.cs b/app/MindWork AI Studio/Tools/ToolCallingSystem/Harness/ToolCallingLoop.cs index 33d33d38..433c1396 100644 --- a/app/MindWork AI Studio/Tools/ToolCallingSystem/Harness/ToolCallingLoop.cs +++ b/app/MindWork AI Studio/Tools/ToolCallingSystem/Harness/ToolCallingLoop.cs @@ -223,17 +223,19 @@ public sealed class ToolCallingLoop(ILogger logger) : IToolCall } toolCallCount++; - var (toolContent, trace, requiredProviderConfidence, sources) = await context.ToolExecutor.ExecuteAsync( + var (toolContent, trace, requiredProviderConfidence, requiredDataSecurity, sources) = await context.ToolExecutor.ExecuteAsync( call.CallId, call.ToolName, call.ArgumentsJson, context.RunnableTools, context.Provider, + context.ChatThread, toolCallCount, token); toolResultCharacterCount += toolContent.Length; context.ChatThread.RequireProviderConfidence(requiredProviderConfidence); + context.ChatThread.RequireDataSecurity(requiredDataSecurity); toolSources.MergeSources(sources); await context.AddToolInvocationAsync(trace); diff --git a/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionContext.cs b/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionContext.cs index 44c93e82..aeeffe1f 100644 --- a/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionContext.cs +++ b/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionContext.cs @@ -1,3 +1,4 @@ +using AIStudio.Chat; using AIStudio.Provider; using AIStudio.Settings; @@ -7,6 +8,16 @@ public sealed class ToolExecutionContext { public required ToolDefinition Definition { get; init; } + /// + /// The chat the call was made in. + /// + /// + /// For a tool which works with what the chat was set up with, such as Semantic Search with the + /// data sources the user picked for it. A tool reads it; what the chat has to keep because of + /// the result goes back through the ToolExecutionResult instead. + /// + public required ChatThread ChatThread { get; init; } + public string ToolCallId { get; init; } = string.Empty; public required SettingsManager SettingsManager { get; init; } diff --git a/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionResult.cs b/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionResult.cs index 6c11bf50..d3ec45ee 100644 --- a/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionResult.cs +++ b/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionResult.cs @@ -1,6 +1,7 @@ using System.Text.Json.Nodes; using AIStudio.Provider; +using AIStudio.Settings.DataModel; namespace AIStudio.Tools.ToolCallingSystem; @@ -14,6 +15,17 @@ public sealed class ToolExecutionResult public ConfidenceLevel RequiredProviderConfidence { get; init; } = ConfidenceLevel.NONE; + /// + /// The data security the chat has to keep from now on, because of what this result brings in. + /// + /// + /// The other axis next to RequiredProviderConfidence. A data source which may only be used with + /// self-hosted providers says so here, and the chat then refuses every other provider from now + /// on, see ChatThread.RequireDataSecurity. Left at NOT_SPECIFIED, the result says nothing about + /// it, and the chat stays as it was. + /// + public DataSourceSecurity RequiredDataSecurity { get; init; } = DataSourceSecurity.NOT_SPECIFIED; + public string ToModelContent() { if (this.JsonContent is not null) diff --git a/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutor.cs b/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutor.cs index a8bcf80b..5d88b09b 100644 --- a/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutor.cs +++ b/app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutor.cs @@ -1,8 +1,10 @@ using System.Diagnostics; using System.Text.Json; +using AIStudio.Chat; using AIStudio.Provider; using AIStudio.Settings; +using AIStudio.Settings.DataModel; namespace AIStudio.Tools.ToolCallingSystem; @@ -46,12 +48,13 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge } } - public async Task<(string Content, ToolInvocationTrace Trace, ConfidenceLevel RequiredProviderConfidence, IReadOnlyList Sources)> ExecuteAsync( + public async Task<(string Content, ToolInvocationTrace Trace, ConfidenceLevel RequiredProviderConfidence, DataSourceSecurity RequiredDataSecurity, IReadOnlyList Sources)> ExecuteAsync( string toolCallId, string toolName, string argumentsJson, IReadOnlyList<(ToolDefinition Definition, IToolImplementation Implementation)> runnableTools, IProvider provider, + ChatThread chatThread, int order, CancellationToken token = default) { @@ -92,7 +95,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge StatusMessage = "Tool is not available in the current context.", Arguments = formattedArguments, Result = error, - }, ConfidenceLevel.NONE, []); + }, ConfidenceLevel.NONE, DataSourceSecurity.NOT_SPECIFIED, []); } var definition = runnableTool.Definition; @@ -105,6 +108,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge var result = await implementation.ExecuteAsync(document.RootElement, new ToolExecutionContext { Definition = definition, + ChatThread = chatThread, ToolCallId = toolCallId, SettingsManager = settingsManager, SettingsValues = settingsValues, @@ -128,7 +132,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge JsonResult = result.JsonContent, }; - return (resultModelContent, toolInvocationTrace, result.RequiredProviderConfidence, result.Sources); + return (resultModelContent, toolInvocationTrace, result.RequiredProviderConfidence, result.RequiredDataSecurity, result.Sources); } catch (OperationCanceledException) when (token.IsCancellationRequested) { @@ -152,7 +156,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge Result = exception.Message, }; - return (exception.Message, toolInvocationTrace, ConfidenceLevel.NONE, []); + return (exception.Message, toolInvocationTrace, ConfidenceLevel.NONE, DataSourceSecurity.NOT_SPECIFIED, []); } catch (Exception exception) { @@ -172,7 +176,7 @@ public sealed class ToolExecutor(ToolSettingsService toolSettingsService, ILogge Result = error, }; - return (error, toolInvocationTrace, ConfidenceLevel.NONE, []); + return (error, toolInvocationTrace, ConfidenceLevel.NONE, DataSourceSecurity.NOT_SPECIFIED, []); } } diff --git a/app/Tests/Chat/ChatThreadDataSecurityTests.cs b/app/Tests/Chat/ChatThreadDataSecurityTests.cs new file mode 100644 index 00000000..8ee0b831 --- /dev/null +++ b/app/Tests/Chat/ChatThreadDataSecurityTests.cs @@ -0,0 +1,52 @@ +using AIStudio.Chat; +using AIStudio.Settings.DataModel; + +namespace AIStudio.Tests.Chat; + +/// +/// Checks how the data a chat has seen tightens the providers which may continue it. +/// +/// +/// A chat which once held data for self-hosted providers only must never be sent to any other +/// provider again, whatever it brings in afterwards. The RAG process and Semantic Search both go +/// through the same rule, so every combination of what a chat holds and what arrives is checked. +/// +[TestFixture] +public sealed class ChatThreadDataSecurityTests +{ + [TestCase(DataSourceSecurity.NOT_SPECIFIED, DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.SELF_HOSTED)] + [TestCase(DataSourceSecurity.ALLOW_ANY, DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.SELF_HOSTED)] + [TestCase(DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.SELF_HOSTED)] + public void DataForSelfHostedProvidersOnlyRestrictsTheChat(DataSourceSecurity held, DataSourceSecurity arriving, DataSourceSecurity expected) + { + Assert.That(Tightened(held, arriving), Is.EqualTo(expected)); + } + + [TestCase(DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.ALLOW_ANY)] + [TestCase(DataSourceSecurity.SELF_HOSTED, DataSourceSecurity.NOT_SPECIFIED)] + public void ARestrictionStays(DataSourceSecurity held, DataSourceSecurity arriving) + { + Assert.That(Tightened(held, arriving), Is.EqualTo(DataSourceSecurity.SELF_HOSTED), "The restricted data was seen by this chat. What arrives later cannot undo that."); + } + + [TestCase(DataSourceSecurity.NOT_SPECIFIED)] + [TestCase(DataSourceSecurity.ALLOW_ANY)] + public void DataForAnyProviderMarksTheChat(DataSourceSecurity held) + { + Assert.That(Tightened(held, DataSourceSecurity.ALLOW_ANY), Is.EqualTo(DataSourceSecurity.ALLOW_ANY)); + } + + [TestCase(DataSourceSecurity.NOT_SPECIFIED)] + [TestCase(DataSourceSecurity.ALLOW_ANY)] + public void ResultsWhichDemandNothingChangeNothing(DataSourceSecurity held) + { + Assert.That(Tightened(held, DataSourceSecurity.NOT_SPECIFIED), Is.EqualTo(held), "A web search, say, says nothing about data sources and must leave the chat as it was."); + } + + private static DataSourceSecurity Tightened(DataSourceSecurity held, DataSourceSecurity arriving) + { + var thread = new ChatThread { DataSecurity = held }; + thread.RequireDataSecurity(arriving); + return thread.DataSecurity; + } +} \ No newline at end of file diff --git a/app/Tests/Tools/ToolCalling/ToolExecutorTests.cs b/app/Tests/Tools/ToolCalling/ToolExecutorTests.cs new file mode 100644 index 00000000..6edb1333 --- /dev/null +++ b/app/Tests/Tools/ToolCalling/ToolExecutorTests.cs @@ -0,0 +1,90 @@ +using AIStudio.Chat; +using AIStudio.Provider; +using AIStudio.Settings.DataModel; +using AIStudio.Tools.ToolCallingSystem; + +using Microsoft.Extensions.Logging.Abstractions; + +namespace AIStudio.Tests.Tools.ToolCalling; + +/// +/// Checks what the tool executor hands to a tool and what it hands back to the loop. +/// +/// +/// What a result demands of the chat has to reach the loop, which tightens the chat with it: a +/// result from a data source for self-hosted providers only that got lost on the way would let the +/// next message go to a cloud provider. A call which brought nothing in must demand nothing. +/// +[TestFixture] +[NonParallelizable] +public sealed class ToolExecutorTests : ToolRegistryTestBase +{ + [Test] + public async Task WhatAResultDemandsReachesTheLoop() + { + var tool = new TestTool(Definition(), execute: _ => new ToolExecutionResult + { + TextContent = "A passage from the handbook.", + RequiredProviderConfidence = ConfidenceLevel.HIGH, + RequiredDataSecurity = DataSourceSecurity.SELF_HOSTED, + }); + + var (_, _, requiredProviderConfidence, requiredDataSecurity, _) = await this.Execute(tool, new ChatThread()); + + Assert.Multiple(() => + { + Assert.That(requiredDataSecurity, Is.EqualTo(DataSourceSecurity.SELF_HOSTED)); + Assert.That(requiredProviderConfidence, Is.EqualTo(ConfidenceLevel.HIGH)); + }); + } + + [Test] + public async Task TheToolSeesTheChatOfTheCall() + { + ChatThread? seenThread = null; + var tool = new TestTool(Definition(), execute: context => + { + seenThread = context.ChatThread; + return new ToolExecutionResult(); + }); + var thread = new ChatThread(); + + await this.Execute(tool, thread); + + Assert.That(seenThread, Is.SameAs(thread), "Semantic Search searches the data sources picked for this very chat."); + } + + [Test] + public async Task ABlockedCallDemandsNothing() + { + var tool = new TestTool(Definition(), execute: _ => throw new ToolExecutionBlockedException("The data source is not available to this provider.")); + + var (_, trace, _, requiredDataSecurity, _) = await this.Execute(tool, new ChatThread()); + + Assert.Multiple(() => + { + Assert.That(trace.Status, Is.EqualTo(ToolInvocationTraceStatus.BLOCKED)); + Assert.That(requiredDataSecurity, Is.EqualTo(DataSourceSecurity.NOT_SPECIFIED), "Nothing reached the model, so there is nothing the chat has to keep."); + }); + } + + [Test] + public async Task AFailedCallDemandsNothing() + { + var tool = new TestTool(Definition(), execute: _ => throw new InvalidOperationException("The index could not be read.")); + + var (_, trace, _, requiredDataSecurity, _) = await this.Execute(tool, new ChatThread()); + + Assert.Multiple(() => + { + Assert.That(trace.Status, Is.EqualTo(ToolInvocationTraceStatus.ERROR)); + Assert.That(requiredDataSecurity, Is.EqualTo(DataSourceSecurity.NOT_SPECIFIED), "Nothing reached the model, so there is nothing the chat has to keep."); + }); + } + + private Task<(string Content, ToolInvocationTrace Trace, ConfidenceLevel RequiredProviderConfidence, DataSourceSecurity RequiredDataSecurity, IReadOnlyList Sources)> Execute(TestTool tool, ChatThread thread) + { + var executor = new ToolExecutor(this.CreateToolSettingsService(), NullLogger.Instance); + return executor.ExecuteAsync("call-1", TOOL_ID, "{}", [(tool.GetDefinition(), tool)], new NoProvider(), thread, order: 1); + } +} \ No newline at end of file diff --git a/app/Tests/Tools/ToolCalling/ToolRegistryTestBase.cs b/app/Tests/Tools/ToolCalling/ToolRegistryTestBase.cs index c19a08ff..74ce460d 100644 --- a/app/Tests/Tools/ToolCalling/ToolRegistryTestBase.cs +++ b/app/Tests/Tools/ToolCalling/ToolRegistryTestBase.cs @@ -60,11 +60,9 @@ public abstract class ToolRegistryTestBase this.rustService.Dispose(); } - protected ToolRegistry CreateRegistry(params TestTool[] tools) - { - var toolSettingsService = new ToolSettingsService(this.SettingsManager, this.rustService, NullLogger.Instance); - return new ToolRegistry(tools, [new CodeToolDefinitionSource(tools)], this.SettingsManager, toolSettingsService, NullLogger.Instance); - } + protected ToolRegistry CreateRegistry(params TestTool[] tools) => new(tools, [new CodeToolDefinitionSource(tools)], this.SettingsManager, this.CreateToolSettingsService(), NullLogger.Instance); + + protected ToolSettingsService CreateToolSettingsService() => new(this.SettingsManager, this.rustService, NullLogger.Instance); protected ToolResolutionContext ContextFor(AIStudio.Settings.Provider provider) => new() { @@ -104,11 +102,12 @@ public abstract class ToolRegistryTestBase }; /// - /// A tool which does nothing, and offers what it is told to. + /// A tool which offers and returns what it is told to. /// /// What the tool is. /// What it offers per request; when left out, its function as defined. - protected sealed class TestTool(ToolDefinition definition, Func? resolve = null) : IToolImplementation + /// What a call returns; when left out, an empty result. + protected sealed class TestTool(ToolDefinition definition, Func? resolve = null, Func? execute = null) : IToolImplementation { public int ResolveCount { get; private set; } @@ -124,6 +123,6 @@ public abstract class ToolRegistryTestBase public IReadOnlySet SensitiveTraceArgumentNames { get; } = new HashSet(StringComparer.Ordinal); - public Task ExecuteAsync(JsonElement arguments, ToolExecutionContext context, CancellationToken token = default) => Task.FromResult(new ToolExecutionResult()); + public Task ExecuteAsync(JsonElement arguments, ToolExecutionContext context, CancellationToken token = default) => Task.FromResult(execute is null ? new ToolExecutionResult() : execute(context)); } } \ No newline at end of file