MindWorkAI/AI-Studio
MindWorkAI/AI-Studio
1
0
Fork 0
mirror of https://github.com/MindWorkAI/AI-Studio.git synced 2026-09-27 01:33:37 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Let tool results restrict a chat to self-hosted providers

Browse Source
This commit is contained in:
Thorsten Sommer 2026-09-24 15:07:39 +02:00
parent e9460a9cde
commit 4617aa5230
Signed by untrusted user who does not match committer: tsommer
GPG Key ID: 371BBA77A02C0108
9 changed files with 209 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
app/MindWork AI Studio/Chat/ChatThread.cs
View File

@ -106,6 +106,26 @@ public sealed record ChatThread
this.RequiredProviderConfidence = minimumProviderConfidence;
}
/// <summary>
/// Tightens the data security of this chat to what the data brought in demands, and never
/// loosens it.
/// </summary>
/// <remarks>
/// 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.<br/><br/>
/// Shared by the RAG process and by tools which search the data sources, so both tighten a chat
/// the same way.
/// </remarks>
/// <param name="dataSecurity">What the data brought in demands.</param>
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,
};
/// <summary>
/// The name of the chat thread. Usually generated by an AI model or manually edited by the user.
/// </summary>

43
app/MindWork AI Studio/Tools/RAG/RAGProcesses/AISrcSelWithRetCtxVal.cs
View File

@ -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<IExternalDataSource>()
.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}'.");

4
app/MindWork AI Studio/Tools/ToolCallingSystem/Harness/ToolCallingLoop.cs
View File

@ -223,17 +223,19 @@ public sealed class ToolCallingLoop(ILogger<ToolCallingLoop> 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);

11
app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionContext.cs
View File

@ -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; }
/// <summary>
/// The chat the call was made in.
/// </summary>
/// <remarks>
/// 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.
/// </remarks>
public required ChatThread ChatThread { get; init; }
public string ToolCallId { get; init; } = string.Empty;
public required SettingsManager SettingsManager { get; init; }

12
app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutionResult.cs
View File

@ -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;
/// <summary>
/// The data security the chat has to keep from now on, because of what this result brings in.
/// </summary>
/// <remarks>
/// 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.
/// </remarks>
public DataSourceSecurity RequiredDataSecurity { get; init; } = DataSourceSecurity.NOT_SPECIFIED;
public string ToModelContent()
{
if (this.JsonContent is not null)

14
app/MindWork AI Studio/Tools/ToolCallingSystem/ToolExecutor.cs
View File

@ -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<Source> Sources)> ExecuteAsync(
public async Task<(string Content, ToolInvocationTrace Trace, ConfidenceLevel RequiredProviderConfidence, DataSourceSecurity RequiredDataSecurity, IReadOnlyList<Source> 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, []);
}
}

52
app/Tests/Chat/ChatThreadDataSecurityTests.cs Normal file
View File

@ -0,0 +1,52 @@
using AIStudio.Chat;
using AIStudio.Settings.DataModel;
namespace AIStudio.Tests.Chat;
/// <summary>
/// Checks how the data a chat has seen tightens the providers which may continue it.
/// </summary>
/// <remarks>
/// 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.
/// </remarks>
[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;
}
}

90
app/Tests/Tools/ToolCalling/ToolExecutorTests.cs Normal file
View File

@ -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;
/// <summary>
/// Checks what the tool executor hands to a tool and what it hands back to the loop.
/// </summary>
/// <remarks>
/// 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.
/// </remarks>
[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<AIStudio.Tools.Source> Sources)> Execute(TestTool tool, ChatThread thread)
{
var executor = new ToolExecutor(this.CreateToolSettingsService(), NullLogger<ToolExecutor>.Instance);
return executor.ExecuteAsync("call-1", TOOL_ID, "{}", [(tool.GetDefinition(), tool)], new NoProvider(), thread, order: 1);
}
}

15
app/Tests/Tools/ToolCalling/ToolRegistryTestBase.cs
View File

@ -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<ToolSettingsService>.Instance);
return new ToolRegistry(tools, [new CodeToolDefinitionSource(tools)], this.SettingsManager, toolSettingsService, NullLogger<ToolRegistry>.Instance);
}
protected ToolRegistry CreateRegistry(params TestTool[] tools) => new(tools, [new CodeToolDefinitionSource(tools)], this.SettingsManager, this.CreateToolSettingsService(), NullLogger<ToolRegistry>.Instance);
protected ToolSettingsService CreateToolSettingsService() => new(this.SettingsManager, this.rustService, NullLogger<ToolSettingsService>.Instance);
protected ToolResolutionContext ContextFor(AIStudio.Settings.Provider provider) => new()
{
@ -104,11 +102,12 @@ public abstract class ToolRegistryTestBase
};
/// <summary>
/// A tool which does nothing, and offers what it is told to.
/// A tool which offers and returns what it is told to.
/// </summary>
/// <param name="definition">What the tool is.</param>
/// <param name="resolve">What it offers per request; when left out, its function as defined.</param>
protected sealed class TestTool(ToolDefinition definition, Func<ToolDefinition, ToolFunctionDefinition?>? resolve = null) : IToolImplementation
/// <param name="execute">What a call returns; when left out, an empty result.</param>
protected sealed class TestTool(ToolDefinition definition, Func<ToolDefinition, ToolFunctionDefinition?>? resolve = null, Func<ToolExecutionContext, ToolExecutionResult>? execute = null) : IToolImplementation
{
public int ResolveCount { get; private set; }
@ -124,6 +123,6 @@ public abstract class ToolRegistryTestBase
public IReadOnlySet<string> SensitiveTraceArgumentNames { get; } = new HashSet<string>(StringComparer.Ordinal);
public Task<ToolExecutionResult> ExecuteAsync(JsonElement arguments, ToolExecutionContext context, CancellationToken token = default) => Task.FromResult(new ToolExecutionResult());
public Task<ToolExecutionResult> ExecuteAsync(JsonElement arguments, ToolExecutionContext context, CancellationToken token = default) => Task.FromResult(execute is null ? new ToolExecutionResult() : execute(context));
}
}