mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-15 14:13:47 +00:00
Let tool results restrict a chat to self-hosted providers
This commit is contained in:
9 files changed
+209
-52
No files matched your search
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user