Fixed Semantic Search leaving out passages that did not fit into a search (#1167)

This commit is contained in:
Thorsten Sommer authored and GitHub committed 2026-11-05 16:02:59 +01:00
1 parent 0a178cb1a3
commit 7342c436dd
17 files changed
+1065 -111

No files matched your search

@@ -0,0 +1,14 @@
namespace AIStudio.Tools.RAG;
/// <summary>
/// One page of results, cut by their length, see CharacterPaging.
/// </summary>
/// <param name="Start">The index of the first result on the page.</param>
/// <param name="Count">How many results the page holds, at least one.</param>
public readonly record struct CharacterPage(int Start, int Count)
{
/// <summary>
/// The index right after the last result on the page, i.e., where the next page starts.
/// </summary>
public int End => this.Start + this.Count;
}
@@ -0,0 +1,61 @@
namespace AIStudio.Tools.RAG;
/// <summary>
/// Cuts results into pages by their length, without leaving any of them out.
/// </summary>
/// <remarks>
/// <para>
/// A page ends right before the first result which no longer fits into it, and the next page
/// starts with that result. No result makes room for a shorter one after it, so the order stays as
/// it is: every result turns up on exactly one page, and the pages before it hold every result
/// ranked above it.
/// </para>
/// <para>
/// The cut depends on the lengths alone. Whoever pages through results without keeping anything
/// between two pages, such as Semantic Search, can therefore cut them anew for every page and gets
/// the same pages each time.
/// </para>
/// </remarks>
public static class CharacterPaging
{
/// <summary>
/// Cuts results of the given lengths into pages.
/// </summary>
/// <remarks>
/// Every page holds at least one result: a result longer than the budget on its own fills a
/// page of its own, rather than holding up all results after it. Whoever must stay within the
/// budget shortens such a result before.
/// </remarks>
/// <param name="lengths">The length of each result in characters, in the order they are shown.</param>
/// <param name="budget">How many characters a page may hold.</param>
/// <returns>The pages, in order; none without results.</returns>
/// <exception cref="ArgumentOutOfRangeException">The budget is below 1, or a length is negative.</exception>
public static IReadOnlyList<CharacterPage> Cut(IReadOnlyList<int> lengths, int budget)
{
ArgumentOutOfRangeException.ThrowIfLessThan(budget, 1);
var pages = new List<CharacterPage>();
var start = 0;
var pageLength = 0L;
for (var index = 0; index < lengths.Count; index++)
{
var length = lengths[index];
ArgumentOutOfRangeException.ThrowIfNegative(length, nameof(lengths));
// The page ends before the first result which no longer fits, unless it would stay empty:
if (index > start && pageLength + length > budget)
{
pages.Add(new(start, index - start));
start = index;
pageLength = 0;
}
pageLength += length;
}
if (start < lengths.Count)
pages.Add(new(start, lengths.Count - start));
return pages;
}
}
@@ -154,7 +154,7 @@ public static class IRetrievalContextExtensions
matchedText = sanitizedTexts[nextText++];
surroundingContent = sanitizedTexts.Skip(nextText).Take(textContext.SurroundingContent.Count).ToList();
nextText += surroundingContent.Count;
AppendTextMatch(body, matchedText, surroundingContent);
AppendTextMatch(body, textContext, matchedText, surroundingContent);
break;
case RetrievalImageContext:
@@ -174,25 +174,146 @@ public static class IRetrievalContextExtensions
}
}
/// <param name="preparedContext">The prepared context.</param>
extension(PreparedRetrievalContext preparedContext)
{
/// <summary>
/// Shortens a prepared context which does not fit into a size limit even on its own.
/// </summary>
/// <remarks>
/// <para>
/// The texts the data source delivered are cut, never what the filter left of them, and the
/// shorter texts pass the filter again: the AI reads only text which the filter saw exactly
/// so. The filter may make a text longer, since its marker can be longer than what it
/// removed, so cutting repeats until the context fits or nothing is left to cut.
/// </para>
/// <para>
/// The surrounding content goes first, from its end, since the matched text is what the
/// search found. Each cut ends at a line end where possible, see CutAtLineEnd. The labels tell
/// the AI which of the texts are shorter, see AppendTextMatch, and the user is shown the
/// shorter passage, as that is what the AI read.
/// </para>
/// </remarks>
/// <param name="maxCharacters">How many characters the context may take.</param>
/// <param name="measure">How many characters the context takes as the caller writes it, e.g., with its header.</param>
/// <param name="guardService">The prompt injection filter.</param>
/// <param name="reportingScope">The action the context was retrieved for, or null when it was retrieved for none.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The context itself when it fits. Otherwise, a shorter one, which exceeds the limit only when its description alone does.</returns>
public async Task<PreparedRetrievalContext> ShortenToFitAsync(int maxCharacters, Func<PreparedRetrievalContext, int> measure, PromptInjectionGuardService guardService, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
{
var shortened = preparedContext;
var excess = measure(shortened) - maxCharacters;
while (excess > 0 && shortened.Context is RetrievalTextContext textContext && CutTexts(textContext, excess) is { } shorterContext)
{
IReadOnlyList<IRetrievalContext> shorterContexts = [shorterContext];
shortened = (await shorterContexts.PrepareForModelAsync(guardService, reportingScope, token))[0];
excess = measure(shortened) - maxCharacters;
}
return shortened;
}
}
/// <summary>
/// Cuts at least the given number of characters off the texts of a context: off its surrounding
/// content first, from the end, then off its matched text.
/// </summary>
/// <param name="textContext">The context.</param>
/// <param name="excess">How many characters to cut at least.</param>
/// <returns>The shorter context, or null when its texts are empty already.</returns>
private static RetrievalTextContext? CutTexts(RetrievalTextContext textContext, int excess)
{
var surroundingContent = textContext.SurroundingContent.ToList();
var toCut = excess;
while (toCut > 0 && surroundingContent.Count > 0)
{
var last = surroundingContent[^1];
if (last.Length <= toCut)
{
toCut -= last.Length;
surroundingContent.RemoveAt(surroundingContent.Count - 1);
continue;
}
surroundingContent[^1] = CutAtLineEnd(last, last.Length - toCut);
toCut = 0;
}
var matchedText = toCut > 0 ? CutAtLineEnd(textContext.MatchedText, textContext.MatchedText.Length - toCut) : textContext.MatchedText;
if (matchedText.Length == textContext.MatchedText.Length && surroundingContent.SequenceEqual(textContext.SurroundingContent))
return null;
return textContext.WithShorterTexts(matchedText, surroundingContent);
}
/// <summary>
/// How much of what would fit a cut gives up at most to end at a line end, as a share.
/// </summary>
/// <remarks>
/// A cut at a line end keeps a row of a table or an item of a list whole. A text whose only line
/// end lies near its beginning would lose nearly everything that way, though, so it is cut at
/// the limit instead.
/// </remarks>
private const double MAX_SHARE_GIVEN_UP_FOR_A_LINE_END = 0.5;
/// <summary>
/// Takes the beginning of a text, ending at a line end where possible.
/// </summary>
/// <remarks>
/// Never cuts between the two halves of a surrogate pair, which no JSON writer takes.
/// </remarks>
/// <param name="text">The text.</param>
/// <param name="maxLength">How many characters to keep at most.</param>
/// <returns>The text itself when it is short enough, otherwise its beginning, without the line end it was cut at.</returns>
internal static string CutAtLineEnd(string text, int maxLength)
{
if (text.Length <= maxLength)
return text;
if (maxLength <= 0)
return string.Empty;
var lineEnd = text.LastIndexOf('\n', maxLength - 1);
if (lineEnd >= maxLength * (1 - MAX_SHARE_GIVEN_UP_FOR_A_LINE_END))
return text[..lineEnd].TrimEnd('\r');
var end = char.IsHighSurrogate(text[maxLength - 1]) ? maxLength - 1 : maxLength;
return text[..end];
}
/// <summary>
/// Writes the content of a text context, after its description.
/// </summary>
/// <remarks>
/// When AI Studio shortened a text, the label says so, rather than a marker inside the text:
/// the text is what the data source wrote, and the AI must not take the label for part of it.
/// </remarks>
/// <param name="contextBuilder">The builder to write into.</param>
/// <param name="textContext">The context, which tells whether its texts were shortened.</param>
/// <param name="matchedText">The matched text, filtered.</param>
/// <param name="surroundingContent">The surrounding content, filtered.</param>
private static void AppendTextMatch(StringBuilder contextBuilder, string matchedText, IReadOnlyList<string> surroundingContent)
private static void AppendTextMatch(StringBuilder contextBuilder, RetrievalTextContext textContext, string matchedText, IReadOnlyList<string> surroundingContent)
{
contextBuilder.AppendLine();
contextBuilder.AppendLine("Matched text content:");
contextBuilder.AppendLine(textContext.IsMatchedTextShortened ? "Matched text content, shortened to fit the size limit:" : "Matched text content:");
contextBuilder.AppendLine("````");
contextBuilder.AppendLine(matchedText);
contextBuilder.AppendLine("````");
if (surroundingContent.Count is 0)
{
if (textContext.IsSurroundingContentShortened)
{
contextBuilder.AppendLine();
contextBuilder.AppendLine("Surrounding text content: left out to fit the size limit.");
}
return;
}
contextBuilder.AppendLine();
contextBuilder.AppendLine("Surrounding text content:");
contextBuilder.AppendLine(textContext.IsSurroundingContentShortened ? "Surrounding text content, shortened to fit the size limit:" : "Surrounding text content:");
foreach (var surrounding in surroundingContent)
{
contextBuilder.AppendLine();
@@ -104,4 +104,34 @@ public sealed class RetrievalTextContext : IRetrievalContext
/// such as a text file.
/// </remarks>
public int? ChunkNumber { get; init; }
/// <summary>
/// Whether AI Studio cut the matched text, since the passage did not fit into a size limit otherwise.
/// </summary>
public bool IsMatchedTextShortened { get; private set; }
/// <summary>
/// Whether AI Studio cut or left out the surrounding content, since the passage did not fit into a size limit otherwise.
/// </summary>
public bool IsSurroundingContentShortened { get; private set; }
/// <summary>
/// A copy of this context with shorter texts, made to fit a size limit.
/// </summary>
/// <remarks>
/// Everything else stays as the data source delivered it. The copy remembers which of its texts
/// are shorter than those of the data source, so that the AI is told.
/// </remarks>
/// <param name="matchedText">The matched text, possibly shortened.</param>
/// <param name="surroundingContent">The surrounding content, possibly shortened.</param>
/// <returns>The copy.</returns>
internal RetrievalTextContext WithShorterTexts(string matchedText, IReadOnlyList<string> surroundingContent)
{
var copy = (RetrievalTextContext)this.MemberwiseClone();
copy.IsMatchedTextShortened = this.IsMatchedTextShortened || matchedText.Length < this.MatchedText.Length;
copy.IsSurroundingContentShortened = this.IsSurroundingContentShortened || !surroundingContent.SequenceEqual(this.SurroundingContent);
copy.MatchedText = matchedText;
copy.SurroundingContent = surroundingContent;
return copy;
}
}
@@ -0,0 +1,14 @@
using AIStudio.Settings;
using AIStudio.Tools.RAG;
namespace AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.SemanticSearch;
/// <summary>
/// What a search tells the model about one data source, besides its passages.
/// </summary>
/// <param name="DataSource">The data source searched.</param>
/// <param name="ResultCount">How many of its passages reached the model.</param>
/// <param name="HasMore">Whether a further page of this data source holds passages.</param>
/// <param name="Gaps">What kept the search from covering the whole data source.</param>
/// <param name="LeftOutCount">How many of its passages a search of several data sources found, but had no room for.</param>
internal sealed record DataSourceSearchOutcome(IDataSource DataSource, int ResultCount, bool HasMore, IReadOnlyList<RetrievalGap> Gaps, int LeftOutCount = 0);
@@ -0,0 +1,12 @@
using AIStudio.Tools.RAG;
namespace AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.SemanticSearch;
/// <summary>
/// One page of a search, as the model and the user receive it.
/// </summary>
/// <param name="TextContent">The passages the model reads, numbered, as Markdown.</param>
/// <param name="Sources">The sources of these passages, listed below the answer.</param>
/// <param name="Passages">What the user is shown of these passages.</param>
/// <param name="Outcomes">What the model learns about each data source searched, in the order they were searched.</param>
internal sealed record SemanticSearchResult(string TextContent, IReadOnlyList<Source> Sources, IReadOnlyList<SourcePassage> Passages, IReadOnlyList<DataSourceSearchOutcome> Outcomes);
@@ -64,17 +64,18 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
private const int MAX_QUERY_CHARACTERS = 500;
/// <summary>
/// How much text one search returns at most, over all data sources searched.
/// How much text one page of a search returns at most, over all data sources searched.
/// </summary>
/// <remarks>
/// Passages are returned whole or not at all, so nothing has to be filtered again after
/// cutting it. A chunk is as long as the embedding model takes at once, by default 8,192
/// tokens, so a single passage may already fill tens of thousands of characters. The limit is
/// the one a web search has by default, and it leaves room for about three searches within the
/// budget of all tool results of an answer, see ToolSelectionRules.MAX_TOOL_RESULT_CHARACTERS:
/// A chunk is as long as the embedding model takes at once, by default 8,192 tokens, so a
/// single passage may already fill tens of thousands of characters. A passage which does not
/// fit starts the next page, see SearchAsync. Only a passage longer than a whole page is
/// shortened, and filtered again, see IRetrievalContextExtensions.ShortenToFitAsync. The limit
/// is the one a web search has by default, and it leaves room for about three searches within
/// the budget of all tool results of an answer, see ToolSelectionRules.MAX_TOOL_RESULT_CHARACTERS:
/// a question with several aspects gets one search per aspect.
/// </remarks>
private const int MAX_RESULT_CHARACTERS = 100_000;
internal const int MAX_RESULT_CHARACTERS = 100_000;
public string ImplementationKey => ToolSelectionRules.SEMANTIC_SEARCH_TOOL_ID;
@@ -162,6 +163,7 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
- When a question has several aspects, search for each of them separately.
- Leave out `data_source_ids` to search all listed data sources. Name some of them only when the question clearly concerns those.
- When the results do not fit, rephrase the query before you turn to a further page. To get a further page, name exactly one data source.
- A search of several data sources shares one page among them. When it reports passages of a data source which did not fit, search that data source alone, starting at page 1.
- A data source which reports that it could not be searched did not find nothing: its results are missing, and your answer has to say so when it matters.
- Name the documents your answer is based on.
- When your searches find nothing relevant, say so before you answer from your own knowledge.
@@ -269,8 +271,10 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
/// exactly their IDs.
/// </summary>
/// <remarks>
/// The model learns the name, the kind, and the description of each data source, and how far it
/// can page through it. Where a data source lies stays out: the model has no use for a path.
/// The model learns the name, the kind, and the description of each data source. How many
/// passages a page holds, and how many pages there are, depends on how long the passages of a
/// query are, so the result of each search tells whether a further page holds more instead.
/// Where a data source lies stays out: the model has no use for a path.
/// What the user describes and what the server of an ERI data source describes both arrive in a
/// single line, and the latter already filtered for prompt injections, see
/// DataSourceDescriptionService.
@@ -287,7 +291,7 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
description.AppendLine($"The data sources you may search, by the ID to pass in {DATA_SOURCE_IDS_ARGUMENT}:");
foreach (var (dataSource, dataSourceDescription) in dataSources)
{
description.Append($"- id={dataSource.Id}, name='{dataSource.Name}', type={GetKind(dataSource)}, results per page={dataSource.MaxMatches}, last page={dataSource.GetLastPage()}");
description.Append($"- id={dataSource.Id}, name='{dataSource.Name}', type={GetKind(dataSource)}");
if (!string.IsNullOrWhiteSpace(dataSourceDescription))
description.Append($", description='{dataSourceDescription.Trim().Shorten(MAX_DESCRIPTION_CHARACTERS)}'");
@@ -311,7 +315,7 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
private static JsonElement BuildParameters(params string[] dataSourceIds) => ToolParameterSchemaBuilder.Create()
.RequiredString(QUERY_ARGUMENT, $"What to search for: a self-contained question, statement, or a few keywords, naming the subject instead of referring to earlier messages. A single line of at most {MAX_QUERY_CHARACTERS} characters.")
.OptionalStringArray(DATA_SOURCE_IDS_ARGUMENT, "Optional IDs of the data sources to search, out of those listed in the description of this tool. Leave it out to search all of them.", dataSourceIds)
.OptionalInteger(PAGE_ARGUMENT, "Optional page of results, starting at 1. A page after the first needs exactly one data source in data_source_ids. Later pages are less relevant, so rephrase the query before you turn pages.")
.OptionalInteger(PAGE_ARGUMENT, "Optional page of results, starting at 1. A page after the first needs exactly one data source in data_source_ids, and follows its first page when searched alone, not a search of several data sources. Later pages are less relevant, so rephrase the query before you turn pages.")
.Build();
private static string GetKind(IDataSource dataSource) => dataSource switch
@@ -367,72 +371,19 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
if (thread.Blocks.Count > 0 && thread.Blocks[^1].Role is ChatRole.AI)
thread = thread with { Blocks = thread.Blocks[..^1] };
var pages = await Task.WhenAll(request.DataSources.Select(dataSource => this.SearchAsync(dataSource, request, thread, token)));
var textContent = new StringBuilder();
var sources = new List<Source>();
var passages = new List<SourcePassage>();
var resultCounts = new int[pages.Length];
var leftOutCounts = new int[pages.Length];
var passageCount = 0;
//
// Every passage goes through the same filter for prompt injections and into the same shape
// as with the classic RAG process. All of them are filtered in one request, so the user
// hears about what was filtered once for the whole search, not once per passage.
//
// The data sources take turns: first the best passage of each, then the second best of
// each, and so on. Otherwise, the data source offered first would take the budget, and the
// others would get what it left over.
//
SemanticSearchResult result;
await using (var reportingScope = guardService.BeginAction())
{
var preparedContexts = await pages.SelectMany(page => page.Contexts).ToList().PrepareForModelAsync(guardService, reportingScope, token);
// Where the prepared contexts of each data source begin, since they are prepared all together:
var firstOfPage = new int[pages.Length];
for (var index = 1; index < pages.Length; index++)
firstOfPage[index] = firstOfPage[index - 1] + pages[index - 1].Contexts.Count;
var mostPassages = pages.Select(page => page.Contexts.Count).DefaultIfEmpty(0).Max();
for (var rank = 0; rank < mostPassages; rank++)
{
for (var index = 0; index < pages.Length; index++)
{
if (rank >= pages[index].Contexts.Count)
continue;
var preparedContext = preparedContexts[firstOfPage[index] + rank];
var markdown = preparedContext.ToMarkdown(passageCount + 1);
// A passage too long for what is left makes room for shorter ones after it:
if (textContent.Length + markdown.Length > MAX_RESULT_CHARACTERS)
{
leftOutCounts[index]++;
continue;
}
passageCount++;
resultCounts[index]++;
textContent.Append(markdown);
// The user is shown what the model received, so only the passages which fit:
sources.AddRange(preparedContext.Sources);
if (preparedContext.Passage is not null)
passages.Add(preparedContext.Passage);
}
}
}
result = await SearchAsync(request, (dataSource, page) => this.RetrievePageAsync(dataSource, request.Query, page, thread, token), guardService, reportingScope, token);
var dataSourceResults = new JsonArray();
for (var index = 0; index < pages.Length; index++)
dataSourceResults.Add(DescribeResult(request.DataSources[index], pages[index], resultCounts[index], leftOutCounts[index]));
foreach (var outcome in result.Outcomes)
dataSourceResults.Add(DescribeResult(outcome));
foreach (var dataSource in unreachableDataSources)
dataSourceResults.Add(DescribeUnreachableResult(dataSource));
var leftOutCount = leftOutCounts.Sum();
logger.LogInformation("Semantic search finished. ToolCallId={ToolCallId}, DataSourceCount={DataSourceCount}, UnreachableCount={UnreachableCount}, Page={Page}, PassageCount={PassageCount}, LeftOutCount={LeftOutCount}", context.ToolCallId, request.DataSources.Count, unreachableDataSources.Count, request.Page, passageCount, leftOutCount);
var resultCounts = result.Outcomes.Select(outcome => outcome.ResultCount).ToList();
logger.LogInformation("Semantic search finished. ToolCallId={ToolCallId}, DataSourceCount={DataSourceCount}, UnreachableCount={UnreachableCount}, Page={Page}, PassageCount={PassageCount}, LeftOutCount={LeftOutCount}", context.ToolCallId, request.DataSources.Count, unreachableDataSources.Count, request.Page, resultCounts.Sum(), result.Outcomes.Sum(outcome => outcome.LeftOutCount));
// Finding nothing is no error; the model reads it from the result counts:
var requirements = GetRequirements(request.DataSources, resultCounts, settingsManager.ConfigurationData.DataSourceSecurity.MinimumOutboundDataRestriction);
@@ -443,16 +394,236 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
["query"] = request.Query,
["page"] = request.Page,
["data_sources"] = dataSourceResults,
["text_content"] = textContent.ToString(),
["text_content"] = result.TextContent,
},
Sources = sources,
Passages = passages,
Sources = result.Sources,
Passages = result.Passages,
RequiredProviderConfidence = requirements.Confidence,
RequiredDataSecurity = requirements.Security,
RequiredOutboundDataRestriction = requirements.OutboundData,
};
}
/// <summary>
/// Searches the data sources of a request, and puts together the page the model receives.
/// </summary>
/// <remarks>
/// <para>
/// Every passage goes through the same filter for prompt injections and into the same shape as
/// with the classic RAG process, before anything measures it: the filter may make a passage
/// longer, as its marker can be longer than what it removed.
/// </para>
/// <para>
/// A data source returns what it found page by page, as many matches per page as it is set to
/// return, see RetrievalPaging. A page of the search is something else, namely as many passages
/// as fit into MAX_RESULT_CHARACTERS. Each page of a data source is therefore cut into one or
/// more pages of the search, see CharacterPaging, and those are numbered on across the pages of
/// the data source. Where everything fits, both are the same.
/// </para>
/// <para>
/// A passage longer than a whole page would hold up every passage after it, so it is shortened
/// to fill a page of its own, see IRetrievalContextExtensions.ShortenToFitAsync.
/// </para>
/// </remarks>
/// <param name="request">The search to run.</param>
/// <param name="retrievePage">Retrieves a page of a data source for the query, and reports the data source as not searched when that fails.</param>
/// <param name="guardService">The prompt injection filter.</param>
/// <param name="reportingScope">The action the search belongs to.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>The page of the search.</returns>
/// <exception cref="ArgumentException">The page lies beyond the last one this query has in the data source, with a message for the model to correct it by.</exception>
internal static Task<SemanticSearchResult> SearchAsync(SemanticSearchRequest request, Func<IDataSource, int, Task<RetrievalPage>> retrievePage, PromptInjectionGuardService guardService, PromptInjectionReportingScope reportingScope, CancellationToken token) =>
request.DataSources.Count == 1
? SearchOneAsync(request.DataSources[0], request.Page, retrievePage, guardService, reportingScope, token)
: SearchSeveralAsync(request.DataSources, retrievePage, guardService, reportingScope, token);
/// <summary>
/// Searches a single data source on any page.
/// </summary>
/// <remarks>
/// Nothing is kept between two searches, since tool results do not travel into later turns. To
/// find a page, the search therefore retrieves the pages of the data source from the first one
/// on, filters and measures their passages, and cuts them again, until it reaches the page
/// asked for. A page of the search never reaches across two pages of the data source, so that
/// cutting one of them needs nothing but its own passages. Of the pages before, only their
/// number and the number of their passages are kept.
/// </remarks>
private static async Task<SemanticSearchResult> SearchOneAsync(IDataSource dataSource, int page, Func<IDataSource, int, Task<RetrievalPage>> retrievePage, PromptInjectionGuardService guardService, PromptInjectionReportingScope reportingScope, CancellationToken token)
{
var gaps = new List<RetrievalGap>();
var pagesBefore = 0;
var passagesBefore = 0;
for (var retrievalPageNumber = 1; retrievalPageNumber <= dataSource.GetLastPage(); retrievalPageNumber++)
{
var retrievalPage = await retrievePage(dataSource, retrievalPageNumber);
foreach (var gap in retrievalPage.Gaps)
{
if (!gaps.Contains(gap))
gaps.Add(gap);
}
// Without this page of the data source, nobody knows where the pages after it start:
if (retrievalPage.Gaps.Contains(RetrievalGap.NOT_SEARCHED))
return new(string.Empty, [], [], [new(dataSource, 0, false, gaps)]);
var preparedContexts = new List<PreparedRetrievalContext>();
var markdowns = new List<string>();
foreach (var preparedContext in await retrievalPage.Contexts.PrepareForModelAsync(guardService, reportingScope, token))
{
var number = passagesBefore + preparedContexts.Count + 1;
var fittingContext = await preparedContext.ShortenToFitAsync(MAX_RESULT_CHARACTERS, context => context.ToMarkdown(number).Length, guardService, reportingScope, token);
preparedContexts.Add(fittingContext);
markdowns.Add(fittingContext.ToMarkdown(number));
}
var cuts = CharacterPaging.Cut(markdowns.Select(markdown => markdown.Length).ToList(), MAX_RESULT_CHARACTERS);
if (page <= pagesBefore + cuts.Count)
{
var cut = cuts[page - pagesBefore - 1];
var hasMore = cut.End < markdowns.Count || retrievalPage.HasMore && await HoldsPassagesAsync(dataSource, retrievalPageNumber + 1, retrievePage);
var textContent = new StringBuilder();
var sources = new List<Source>();
var passages = new List<SourcePassage>();
for (var index = cut.Start; index < cut.End; index++)
Append(preparedContexts[index], markdowns[index], textContent, sources, passages);
return new(textContent.ToString(), sources, passages, [new(dataSource, cut.Count, hasMore, gaps)]);
}
pagesBefore += cuts.Count;
passagesBefore += preparedContexts.Count;
if (!retrievalPage.HasMore)
break;
}
// A search which found nothing still has its first page, an empty one:
if (page == 1)
return new(string.Empty, [], [], [new(dataSource, 0, false, gaps)]);
var lastPage = Math.Max(1, pagesBefore);
throw new ArgumentException($"Argument '{PAGE_ARGUMENT}' must be at most {lastPage} for the data source '{dataSource.Id}' and this query, but was {page}. Rephrase the query to find other passages.");
}
/// <summary>
/// Searches several data sources, on the first page only, which they share.
/// </summary>
/// <remarks>
/// <para>
/// All passages are filtered in one request, so the user hears about what was filtered once for
/// the whole search, not once per data source.
/// </para>
/// <para>
/// The data sources take turns: first the best passage of each, then the second best of each,
/// and so on. Otherwise, the data source offered first would take the budget, and the others
/// would get what it left over. A data source stops at its first passage which no longer fits,
/// while the others go on: a shorter passage after it must not take its place, or a passage
/// would go missing in the middle of what the data source found. A passage longer than a whole
/// page is shortened first, as in a search of one data source, and then has to fit like any
/// other.
/// </para>
/// <para>
/// Turning to a further page means something for a single data source only, as each of them
/// runs out at another point, see ReadRequest. So the model learns how many passages of a data
/// source did not fit, and that searching it alone from the first page on shows them. That
/// search repeats the passages read here already, which is the price of keeping nothing between
/// two searches.
/// </para>
/// </remarks>
private static async Task<SemanticSearchResult> SearchSeveralAsync(IReadOnlyList<IDataSource> dataSources, Func<IDataSource, int, Task<RetrievalPage>> retrievePage, PromptInjectionGuardService guardService, PromptInjectionReportingScope reportingScope, CancellationToken token)
{
var retrievalPages = await Task.WhenAll(dataSources.Select(dataSource => retrievePage(dataSource, 1)));
var preparedContexts = await retrievalPages.SelectMany(retrievalPage => retrievalPage.Contexts).ToList().PrepareForModelAsync(guardService, reportingScope, token);
// Where the prepared contexts of each data source begin, since they are prepared all together:
var firstOfPage = new int[retrievalPages.Length];
for (var index = 1; index < retrievalPages.Length; index++)
firstOfPage[index] = firstOfPage[index - 1] + retrievalPages[index - 1].Contexts.Count;
var textContent = new StringBuilder();
var sources = new List<Source>();
var passages = new List<SourcePassage>();
var resultCounts = new int[retrievalPages.Length];
var isFull = new bool[retrievalPages.Length];
var passageCount = 0;
var mostPassages = retrievalPages.Select(retrievalPage => retrievalPage.Contexts.Count).DefaultIfEmpty(0).Max();
for (var rank = 0; rank < mostPassages; rank++)
{
for (var index = 0; index < retrievalPages.Length; index++)
{
if (isFull[index] || rank >= retrievalPages[index].Contexts.Count)
continue;
var number = passageCount + 1;
var preparedContext = await preparedContexts[firstOfPage[index] + rank].ShortenToFitAsync(MAX_RESULT_CHARACTERS, context => context.ToMarkdown(number).Length, guardService, reportingScope, token);
var markdown = preparedContext.ToMarkdown(number);
if (textContent.Length + markdown.Length > MAX_RESULT_CHARACTERS)
{
isFull[index] = true;
continue;
}
passageCount++;
resultCounts[index]++;
Append(preparedContext, markdown, textContent, sources, passages);
}
}
var outcomes = await Task.WhenAll(retrievalPages.Select(async (retrievalPage, index) =>
{
var leftOutCount = retrievalPage.Contexts.Count - resultCounts[index];
var hasMore = leftOutCount > 0 || retrievalPage.HasMore && await HoldsPassagesAsync(dataSources[index], 2, retrievePage);
return new DataSourceSearchOutcome(dataSources[index], resultCounts[index], hasMore, retrievalPage.Gaps, leftOutCount);
}));
return new(textContent.ToString(), sources, passages, outcomes);
}
/// <summary>
/// Whether the pages of a data source from the given one on hold a passage at all.
/// </summary>
/// <remarks>
/// A data source only says that its next page may hold more: the matches there may have
/// nothing to show, e.g., a chunk without text. The model should not turn to a page which
/// turns out empty, so the search looks for itself. That a page holds a passage is enough, so
/// nothing is filtered here. When a page could not be retrieved, the data source is taken at
/// its word.
/// </remarks>
/// <param name="dataSource">The data source.</param>
/// <param name="firstRetrievalPage">The first page of the data source to look at.</param>
/// <param name="retrievePage">Retrieves a page of a data source for the query.</param>
/// <returns>True when one of these pages holds a passage, or could not be retrieved.</returns>
private static async Task<bool> HoldsPassagesAsync(IDataSource dataSource, int firstRetrievalPage, Func<IDataSource, int, Task<RetrievalPage>> retrievePage)
{
for (var retrievalPageNumber = firstRetrievalPage; retrievalPageNumber <= dataSource.GetLastPage(); retrievalPageNumber++)
{
var retrievalPage = await retrievePage(dataSource, retrievalPageNumber);
if (retrievalPage.Contexts.Count > 0 || retrievalPage.Gaps.Contains(RetrievalGap.NOT_SEARCHED))
return true;
if (!retrievalPage.HasMore)
return false;
}
return false;
}
/// <summary>
/// Adds a passage to the page of a search.
/// </summary>
/// <remarks>
/// The user is shown what the model received, so the sources and the passage go along with the
/// text, and only for the passages on the page.
/// </remarks>
private static void Append(PreparedRetrievalContext preparedContext, string markdown, StringBuilder textContent, List<Source> sources, List<SourcePassage> passages)
{
textContent.Append(markdown);
sources.AddRange(preparedContext.Sources);
if (preparedContext.Passage is not null)
passages.Add(preparedContext.Passage);
}
/// <summary>
/// Which data sources whose server could not be reached a search reports as not searched.
/// </summary>
@@ -479,7 +650,7 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
/// </remarks>
internal static JsonObject DescribeUnreachableResult(IDataSource dataSource)
{
var result = DescribeResult(dataSource, RetrievalPage.EMPTY with { Gaps = [RetrievalGap.NOT_SEARCHED] }, 0, 0);
var result = DescribeResult(new DataSourceSearchOutcome(dataSource, 0, false, [RetrievalGap.NOT_SEARCHED]));
result.Remove("id");
return result;
}
@@ -537,13 +708,25 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
if (dataSources.Count != 1)
throw new ArgumentException($"Argument '{PAGE_ARGUMENT}' may be above 1 only for exactly one data source in '{DATA_SOURCE_IDS_ARGUMENT}', but was {page} for {dataSources.Count}. Name the one data source to page through, or leave '{PAGE_ARGUMENT}' out to get the first page of each.");
var lastPage = dataSources[0].GetLastPage();
if (page > lastPage)
throw new ArgumentException($"Argument '{PAGE_ARGUMENT}' must be at most {lastPage} for the data source '{dataSources[0].Id}', but was {page}. Rephrase the query to find other passages.");
//
// How many pages a query has depends on how long its passages are, which only the search
// finds out, see SearchAsync. A page holds at least one passage, though, so no query has
// more pages than the data source returns passages for it at all:
//
var mostPages = GetMostPages(dataSources[0]);
if (page > mostPages)
throw new ArgumentException($"Argument '{PAGE_ARGUMENT}' must be at most {mostPages} for the data source '{dataSources[0].Id}', but was {page}. Rephrase the query to find other passages.");
return new(query, dataSources, page);
}
/// <summary>
/// How many pages a search in a data source can have at most, for any query.
/// </summary>
/// <param name="dataSource">The data source.</param>
/// <returns>As many pages as the data source returns passages for a query at most, since a page holds at least one of them.</returns>
internal static int GetMostPages(IDataSource dataSource) => dataSource.GetLastPage() * dataSource.MaxMatches;
/// <summary>
/// Searches one data source, and reports it as not searched when that fails.
/// </summary>
@@ -551,11 +734,11 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
/// The other data sources still answer. The failed one is reported rather than left out, so
/// that the model does not take its silence for finding nothing.
/// </remarks>
private async Task<RetrievalPage> SearchAsync(IDataSource dataSource, SemanticSearchRequest request, ChatThread thread, CancellationToken token)
private async Task<RetrievalPage> RetrievePageAsync(IDataSource dataSource, string query, int page, ChatThread thread, CancellationToken token)
{
try
{
return await dataSource.RetrieveDataAsync(request.Query, request.Page, thread, token);
return await dataSource.RetrieveDataAsync(query, page, thread, token);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
@@ -575,10 +758,11 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
/// Only AI Studio's own values: the ID and the name as configured, counts, and sentences of its
/// own. Whatever a data source returned is in the passages, which went through the filter.
/// </remarks>
internal static JsonObject DescribeResult(IDataSource dataSource, RetrievalPage page, int resultCount, int leftOutCount)
internal static JsonObject DescribeResult(DataSourceSearchOutcome outcome)
{
var (dataSource, resultCount, hasMore, gaps, leftOutCount) = outcome;
var issues = new JsonArray();
foreach (var gap in page.Gaps)
foreach (var gap in gaps)
{
issues.Add(gap switch
{
@@ -590,14 +774,14 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
}
if (leftOutCount > 0)
issues.Add($"{leftOutCount} further passages of this page were left out to keep the result within its size limit. Search this data source with a narrower query to see them.");
issues.Add($"{leftOutCount} further passages of this data source did not fit into this search of several data sources. Search this data source alone, starting at page 1, to see them.");
var result = new JsonObject
{
["id"] = dataSource.Id,
["name"] = dataSource.Name,
["result_count"] = resultCount,
["has_more"] = page.HasMore,
["has_more"] = hasMore,
};
if (issues.Count > 0)