Outbound data restrictions after searching data sources (#1106)

app/Tests/Tools/ToolCalling/SemanticSearchOutboundDataRestrictionTests.cs

@@ -0,0 +1,127 @@
+using AIStudio.Chat;
+using AIStudio.Provider;
+using AIStudio.Settings;
+using AIStudio.Settings.DataModel;
+using AIStudio.Tools.ToolCallingSystem;
+using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations;
+using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.SemanticSearch;
+using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.WebSearch;
+using AIStudio.Tools.Web;
+
+using Microsoft.Extensions.Logging.Abstractions;
+
+namespace AIStudio.Tests.Tools.ToolCalling;
+
+/// <summary>
+/// Checks where a chat may still send data once Semantic Search found something in its data sources.
+/// </summary>
+/// <remarks>
+/// A data source holds the documents of the user or of their organization. A web search or an
+/// address the model writes itself would carry them to a third party, so a search which found
+/// something restricts the chat as a mailbox or the wiki does. By default, the chat keeps the links
+/// from the chat. A data source which only self-hosted providers may use allows the configured
+/// services only: the internal pages it links to stay readable, a public page does not. The search
+/// itself needs indexed data sources and is left to a manual test; these tests start from what the
+/// passages of a search demand.
+/// </remarks>
+[TestFixture]
+[NonParallelizable]
+public sealed class SemanticSearchOutboundDataRestrictionTests : ToolRegistryTestBase
+{
+    private const string INTRANET_PAGE = "https://intranet.example.org/projects/budget";
+    private const string PUBLIC_PAGE = "https://example.org/press/2026-10";
+    private const string CHOSEN_BY_THE_MODEL = "https://attacker.example/?budget=2027";
+
+    private static readonly IDataSource PROJECT_FILES = new DataSourceLocalDirectory { Num = 1, Id = "5e2b8d4f-1c7a-4e9b-a3d6-8f4c2e7b1a59", Name = "Project files", MaxMatches = 10 };
+    private static readonly IDataSource INTRANET = new DataSourceERI_V1 { Num = 2, Id = "8a3f6c1e-4d9b-4b2a-9e7c-3d6a1f8b5c27", Name = "Intranet", MaxMatches = 10, SecurityPolicy = DataSourceSecurity.SELF_HOSTED, OutboundDataRestriction = OutboundDataRestriction.UNRESTRICTED };
+    private static readonly IDataSource PRESS_RELEASES = new DataSourceLocalFile { Num = 3, Id = "2d7c4a9e-6b1f-4c8d-b5a3-7e2f9c4d1b86", Name = "Press releases", MaxMatches = 10, OutboundDataRestriction = OutboundDataRestriction.UNRESTRICTED };
+
+    private static readonly Func<string, bool> ALLOWED_PRIVATE_HOSTS = host => string.Equals(host, "intranet.example.org", StringComparison.OrdinalIgnoreCase);
+
+    [Test]
+    public void ASearchRestrictsTheChatToLinksFromTheChatByDefault()
+    {
+        var requirements = SemanticSearchTool.GetRequirements([PROJECT_FILES, PRESS_RELEASES], [3, 0], OutboundDataRestriction.UNRESTRICTED);
+        Assert.That(requirements.OutboundData, Is.EqualTo(new OutboundDataRequirement(OutboundDataRestriction.ONLY_LINKS_FROM_CHAT, OutboundDataOrigin.DATA_SOURCE, PROJECT_FILES.Id)), "The chat names the data source, so the user learns why web search is gone.");
+    }
+
+    [Test]
+    public void ASearchInADataSourceWithoutRestrictionRestrictsNothing()
+    {
+        var requirements = SemanticSearchTool.GetRequirements([PROJECT_FILES, PRESS_RELEASES], [0, 2], OutboundDataRestriction.UNRESTRICTED);
+        Assert.That(requirements.OutboundData, Is.EqualTo(OutboundDataRequirement.NONE), "The project files found nothing, and the press releases may be read by anybody.");
+    }
+
+    [Test]
+    public void AnOrganizationMinimumTightensWhatASearchDemands()
+    {
+        var requirements = SemanticSearchTool.GetRequirements([PRESS_RELEASES], [2], OutboundDataRestriction.ONLY_CONFIGURED_SERVICES);
+        Assert.That(requirements.OutboundData.Restriction, Is.EqualTo(OutboundDataRestriction.ONLY_CONFIGURED_SERVICES));
+    }
+
+    [Test]
+    public void AfterAHitWebSearchIsNotOffered()
+    {
+        IToolImplementation webSearch = new WebSearchTool([], null!, null!, NullLogger<WebSearchTool>.Instance);
+        Assert.That(ToolSelectionRules.IsOutboundDataAllowed(ChatAfterASearch(PROJECT_FILES).RequiredOutboundDataRestriction.Restriction, webSearch), Is.False);
+    }
+
+    [Test]
+    public async Task AfterAHitAWebSearchIsTurnedDownBeforeItRuns()
+    {
+        this.SettingsManager.ConfigurationData.DataSources.Add(PROJECT_FILES);
+
+        //
+        // The request offered the web search while the chat was still unrestricted. Then Semantic
+        // Search brought in the passages, and the model calls the web search in the same request:
+        //
+        var hasRun = false;
+        var webSearch = new TestTool(Definition(), execute: _ =>
+        {
+            hasRun = true;
+            return new ToolExecutionResult { TextContent = "Results about the budget." };
+        }) { OutboundData = ToolOutboundData.THIRD_PARTY_QUERIES };
+
+        var executor = new ToolExecutor(this.CreateToolSettingsService(), NullLogger<ToolExecutor>.Instance);
+        var outcome = await executor.ExecuteAsync("call-1", TOOL_ID, "{}", [(webSearch.GetDefinition(), webSearch)], new NoProvider(), ChatAfterASearch(PROJECT_FILES), order: 1);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(hasRun, Is.False, "The content of the project files would already have left AI Studio.");
+            Assert.That(outcome.Trace.Status, Is.EqualTo(ToolInvocationTraceStatus.BLOCKED));
+            Assert.That(outcome.Content, Does.Contain("Project files"), "The user learns which data source stands in the way.");
+        });
+    }
+
+    [TestCase(INTRANET_PAGE, true)]
+    [TestCase(PUBLIC_PAGE, true)]
+    [TestCase(CHOSEN_BY_THE_MODEL, false)]
+    public void AfterAHitTheLinksOfTheHitsAreRead(string address, bool expectedAllowed)
+    {
+        var chat = ChatAfterASearch(PROJECT_FILES);
+        var allowed = ReadWebPageTool.IsAllowedByOutboundDataRestriction(new Uri(address), chat.RequiredOutboundDataRestriction.Restriction, chat, wiki: null, ALLOWED_PRIVATE_HOSTS, out _);
+        Assert.That(allowed, Is.EqualTo(expectedAllowed), "The model may follow the links of what it found, but it cannot write an address of its own.");
+    }
+
+    [TestCase(INTRANET_PAGE, true)]
+    [TestCase(PUBLIC_PAGE, false)]
+    [TestCase(CHOSEN_BY_THE_MODEL, false)]
+    public void AfterAHitInADataSourceForSelfHostedProvidersOnlyInternalPagesAreRead(string address, bool expectedAllowed)
+    {
+        var chat = ChatAfterASearch(INTRANET);
+        var allowed = ReadWebPageTool.IsAllowedByOutboundDataRestriction(new Uri(address), chat.RequiredOutboundDataRestriction.Restriction, chat, wiki: null, ALLOWED_PRIVATE_HOSTS, out _);
+        Assert.That(allowed, Is.EqualTo(expectedAllowed), "An internal page the data source links to stays readable, a public one does not, although the hit names it.");
+    }
+
+    /// <summary>
+    /// A chat in which Semantic Search found passages in the data source, linking to an internal and to a public page.
+    /// </summary>
+    private static ChatThread ChatAfterASearch(IDataSource dataSource)
+    {
+        var thread = new ChatThread();
+        thread.RuntimeWebAddressesFromTools.Add(WebAddresses.CreateRequestKey(new Uri(INTRANET_PAGE)));
+        thread.RuntimeWebAddressesFromTools.Add(WebAddresses.CreateRequestKey(new Uri(PUBLIC_PAGE)));
+        thread.RequireOutboundDataRestriction(SemanticSearchTool.GetRequirements([dataSource], [2], OutboundDataRestriction.UNRESTRICTED).OutboundData);
+        return thread;
+    }
+}

app/Tests/Tools/ToolCalling/SemanticSearchToolDescriptionTests.cs

@@ -48,7 +48,7 @@ public sealed class SemanticSearchToolDescriptionTests
     {
         var descriptions = new Dictionary<string, string> { [HANDBOOK.Id] = "Our processes.", [MINUTES.Id] = "Meetings.", [INTRANET.Id] = "Everything else." };
         ToolFunctionDefinition DescribeInOfferOrder(params IDataSource[] dataSources) =>
-            Describe(SemanticSearchTool.InOfferOrder(dataSources).Select(dataSource => (dataSource, descriptions[dataSource.Id])).ToArray());
+            Describe(dataSources.InFixedOrder().Select(dataSource => (dataSource, descriptions[dataSource.Id])).ToArray());
 
         var first = DescribeInOfferOrder(INTRANET, MINUTES, HANDBOOK);
         var second = DescribeInOfferOrder(MINUTES, HANDBOOK, INTRANET);

app/Tests/Tools/ToolCalling/SemanticSearchToolResultTests.cs

@@ -28,25 +28,29 @@ public sealed class SemanticSearchToolResultTests
     [Test]
     public void ASearchWhichFoundNothingRequiresNothing()
     {
-        var requirements = SemanticSearchTool.GetRequirements([HANDBOOK, INTRANET], [0, 0]);
+        var requirements = SemanticSearchTool.GetRequirements([HANDBOOK, INTRANET], [0, 0], OutboundDataRestriction.UNRESTRICTED);
 
-        Assert.That(requirements, Is.EqualTo((ConfidenceLevel.NONE, DataSourceSecurity.NOT_SPECIFIED)), "NOT_SPECIFIED leaves the data security of the chat as it was.");
+        Assert.That(requirements, Is.EqualTo(DataSourceRequirements.NONE), "NOT_SPECIFIED leaves the data security of the chat as it was.");
     }
 
     [Test]
     public void OnlyTheDataSourcesWithPassagesRaiseTheRequirements()
     {
-        var requirements = SemanticSearchTool.GetRequirements([HANDBOOK, MINUTES, INTRANET], [0, 2, 0]);
+        var requirements = SemanticSearchTool.GetRequirements([HANDBOOK, MINUTES, INTRANET], [0, 2, 0], OutboundDataRestriction.UNRESTRICTED);
 
-        Assert.That(requirements, Is.EqualTo((ConfidenceLevel.LOW, DataSourceSecurity.ALLOW_ANY)), "Neither the highly confidential handbook nor the self-hosted intranet brought anything into the chat.");
+        Assert.Multiple(() =>
+        {
+            Assert.That((requirements.Confidence, requirements.Security), Is.EqualTo((ConfidenceLevel.LOW, DataSourceSecurity.ALLOW_ANY)), "Neither the highly confidential handbook nor the self-hosted intranet brought anything into the chat.");
+            Assert.That(requirements.OutboundData, Is.EqualTo(new OutboundDataRequirement(OutboundDataRestriction.ONLY_LINKS_FROM_CHAT, OutboundDataOrigin.DATA_SOURCE, MINUTES.Id)), "Only the minutes restrict where the chat may still send data.");
+        });
     }
 
     [Test]
     public void ASelfHostedDataSourceWithPassagesRestrictsTheChat()
     {
-        var requirements = SemanticSearchTool.GetRequirements([MINUTES, INTRANET], [1, 1]);
+        var requirements = SemanticSearchTool.GetRequirements([MINUTES, INTRANET], [1, 1], OutboundDataRestriction.UNRESTRICTED);
 
-        Assert.That(requirements, Is.EqualTo((ConfidenceLevel.LOW, DataSourceSecurity.SELF_HOSTED)));
+        Assert.That((requirements.Confidence, requirements.Security), Is.EqualTo((ConfidenceLevel.LOW, DataSourceSecurity.SELF_HOSTED)));
     }
 
     [Test]

app/Tests/Tools/ToolCalling/ToolOrderInstructionsTests.cs

@@ -4,6 +4,7 @@ using AIStudio.Settings.DataModel;
 using AIStudio.Tools.ToolCallingSystem;
 using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations;
 using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.Mailboxes;
+using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.SemanticSearch;
 using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.WebSearch;
 
 using Microsoft.Extensions.Logging.Abstractions;
@@ -32,7 +33,7 @@ public sealed class ToolOrderInstructionsTests : ToolRegistryTestBase
 
         Assert.Multiple(() =>
         {
-            foreach (var toolId in new[] { ToolSelectionRules.SEARCH_MAILS_TOOL_ID, ToolSelectionRules.READ_MAIL_TOOL_ID, ToolSelectionRules.COUNT_MAILS_TOOL_ID, ToolSelectionRules.SEARCH_CONFLUENCE_TOOL_ID })
+            foreach (var toolId in new[] { ToolSelectionRules.SEARCH_MAILS_TOOL_ID, ToolSelectionRules.READ_MAIL_TOOL_ID, ToolSelectionRules.COUNT_MAILS_TOOL_ID, ToolSelectionRules.SEARCH_CONFLUENCE_TOOL_ID, ToolSelectionRules.SEMANTIC_SEARCH_TOOL_ID })
                 Assert.That(instructions, Does.Not.Contain($"`{toolId}`"), "The web search would name tools the request may not offer, and each of these words the rule itself.");
         });
     }
@@ -82,6 +83,46 @@ public sealed class ToolOrderInstructionsTests : ToolRegistryTestBase
         Assert.That(instructions, Does.Not.Contain("web_search"), "The request does not offer the web search, so a rule about it would only confuse the model.");
     }
 
+    [TestCase(OutboundDataRestriction.ONLY_LINKS_FROM_CHAT, true, true)]
+    [TestCase(OutboundDataRestriction.ONLY_CONFIGURED_SERVICES, true, true)]
+    [TestCase(OutboundDataRestriction.UNRESTRICTED, true, false)]
+    [TestCase(OutboundDataRestriction.ONLY_LINKS_FROM_CHAT, false, false)]
+    public async Task TheSemanticSearchSaysSoWhenItEndsTheWebSearch(OutboundDataRestriction restriction, bool offersWebSearch, bool expected)
+    {
+        var dataSource = new DataSourceLocalDirectory { Id = "6f1a8c3e-9d2b-4e7a-b4c5-1e8d3f6a9c72", Name = "Project files", OutboundDataRestriction = restriction };
+        this.SettingsManager.ConfigurationData.DataSources.Add(dataSource);
+        var chat = new ChatThread { DataSourceOptions = new() { PreselectedDataSourceIds = [dataSource.Id] } };
+        var tool = new SemanticSearchTool(this.SettingsManager, null!, null!, null!, NullLogger<SemanticSearchTool>.Instance);
+        string[] candidates = offersWebSearch ? [ToolSelectionRules.SEMANTIC_SEARCH_TOOL_ID, ToolSelectionRules.WEB_SEARCH_TOOL_ID] : [ToolSelectionRules.SEMANTIC_SEARCH_TOOL_ID];
+
+        var instructions = await tool.ResolveSystemPromptInstructionsAsync(tool.GetDefinition(), this.ResolutionContext(chat, candidates));
+
+        Assert.That(instructions.Contains("search the web before you search the data sources", StringComparison.Ordinal), Is.EqualTo(expected), "Without a restriction, a search leaves the web search in place, and the rule would only cost the model a detour.");
+    }
+
+    [Test]
+    public void OnlyTheDataSourcesTheModelMaySearchCount()
+    {
+        var open = new DataSourceLocalFile { Id = "3c9e1a7f-2b5d-4f8c-a6e1-9d4b7c2a5f13", Name = "Press releases", OutboundDataRestriction = OutboundDataRestriction.UNRESTRICTED };
+        var closed = new DataSourceLocalDirectory { Id = "7b4d2f9a-5e1c-4a6b-8d3f-2c9a6e1b4d58", Name = "Project files" };
+        var settings = new Data { DataSources = [open, closed] };
+        var selected = new ChatThread { DataSourceOptions = new() { PreselectedDataSourceIds = [open.Id] } };
+        var automatic = new ChatThread { DataSourceOptions = new() { AutomaticDataSourceSelection = true } };
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(SemanticSearchTool.EndsWebSearch(this.ResolutionContext(selected), settings), Is.False, "The project files are not selected, so the model cannot find anything in them.");
+            Assert.That(SemanticSearchTool.EndsWebSearch(this.ResolutionContext(automatic), settings), Is.True, "With the automatic selection, the model may search the project files as well.");
+        });
+    }
+
+    [Test]
+    public void TheSemanticSearchCountsWithTheRule()
+    {
+        var instructions = new SemanticSearchTool(null!, null!, null!, null!, null!).GetDefinition().SystemPromptInstructions;
+        Assert.That(instructions, Does.Contain("search the web before you search the data sources"), "The token count assumes a request which offers the web search and a data source which ends it.");
+    }
+
     [TestCase(true)]
     [TestCase(false)]
     public void ReadWebPageSaysSoWhenAnInternalPageEndsTheWebSearch(bool internalPageEndsWebSearch)

app/Tests/Tools/ToolCalling/ToolOutboundDataTests.cs

@@ -23,6 +23,7 @@ namespace AIStudio.Tests.Tools.ToolCalling;
 public sealed class ToolOutboundDataTests : ToolRegistryTestBase
 {
     private const string MAILBOX_ID = "6e3a9d2f-1b7c-4f8e-a5d4-9c2b7e1f3a6d";
+    private const string DATA_SOURCE_ID = "1c7e4a9f-3b6d-4e2a-8f5c-9d2b6e1a4c38";
 
     [TestCase(OutboundDataRestriction.UNRESTRICTED, ToolOutboundData.NONE, true)]
     [TestCase(OutboundDataRestriction.UNRESTRICTED, ToolOutboundData.CONFIGURED_SERVICE, true)]
@@ -185,6 +186,34 @@ public sealed class ToolOutboundDataTests : ToolRegistryTestBase
         });
     }
 
+    [Test]
+    public void ADataSourceIsNamedByItsName()
+    {
+        var requirement = new OutboundDataRequirement(OutboundDataRestriction.ONLY_LINKS_FROM_CHAT, OutboundDataOrigin.DATA_SOURCE, DATA_SOURCE_ID);
+        var message = requirement.GetToolBlockedMessage(new Data { DataSources = [new DataSourceLocalDirectory { Id = DATA_SOURCE_ID, Name = "Project files" }] });
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(message, Does.Contain("Project files"), "The user learns which data source stands in the way.");
+            Assert.That(message, Does.Not.Contain("mailbox"), "No mailbox restricted the chat.");
+            Assert.That(message, Does.Contain(OutboundDataRestriction.ONLY_LINKS_FROM_CHAT.GetName()));
+        });
+    }
+
+    [Test]
+    public void ADataSourceRemovedSinceIsNotNamed()
+    {
+        var requirement = new OutboundDataRequirement(OutboundDataRestriction.ONLY_CONFIGURED_SERVICES, OutboundDataOrigin.DATA_SOURCE, DATA_SOURCE_ID);
+        var message = requirement.GetToolBlockedMessage(new Data { DataSources = [new DataSourceLocalFile { Id = "another-data-source", Name = "Minutes" }] });
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(message, Does.Contain("removed since"));
+            Assert.That(message, Does.Not.Contain("Minutes"), "Another data source did not restrict the chat.");
+            Assert.That(message, Does.Contain(OutboundDataRestriction.ONLY_CONFIGURED_SERVICES.GetName()), "The restriction stays, so the user still learns why.");
+        });
+    }
+
     [Test]
     public async Task ACallTurnedDownByTheWikiNamesTheWiki()
     {