mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-09-13 22:53:37 +00:00
80 lines
3.2 KiB
C#
80 lines
3.2 KiB
C#
|
using AIStudio.Tools.Security;
|
||
|
|
|
||
|
|
namespace AIStudio.Tools.Web;
|
||
|
|
|
||
|
|
/// <summary>
|
||
|
|
/// Filters prompt injections out of the web page content a tool returns to a model.
|
||
|
|
/// </summary>
|
||
|
|
/// <remarks>
|
||
|
|
/// Web search and reading a single page share this: both hand the model text they fetched from
|
||
|
|
/// the public web, and neither may pass it on unchecked. All pages of one tool call are filtered
|
||
|
|
/// in a single runtime request, so a search across five pages costs one round trip and produces
|
||
|
|
/// one report for the user.
|
||
|
|
/// </remarks>
|
||
|
|
public static class WebPageContentSanitizer
|
||
|
|
{
|
||
|
|
/// <summary>
|
||
|
|
/// How many single-value fields each page contributes, in the order they are collected. The
|
||
|
|
/// author list follows them and varies in length, so rebuilding depends on this being right.
|
||
|
|
/// </summary>
|
||
|
|
private const int SINGLE_VALUE_FIELD_COUNT = 6;
|
||
|
|
|
||
|
|
/// <summary>
|
||
|
|
/// Filters every model-facing text of the given pages.
|
||
|
|
/// </summary>
|
||
|
|
/// <param name="guardService">The guard service performing the filtering.</param>
|
||
|
|
/// <param name="pages">The page contents to filter, each with the source it came from.</param>
|
||
|
|
/// <returns>The filtered contents, in the order they came in.</returns>
|
||
|
|
public static async Task<IReadOnlyList<WebPageModelContent>> SanitizeAsync(PromptInjectionGuardService guardService,
|
||
|
|
IReadOnlyList<(WebPageModelContent Content, PromptInjectionSource Source)> pages)
|
||
|
|
{
|
||
|
|
if (pages.Count is 0)
|
||
|
|
return [];
|
||
|
|
|
||
|
|
List<PromptInjectionText> texts = [];
|
||
|
|
foreach (var (content, source) in pages)
|
||
|
|
{
|
||
|
|
texts.Add(new(content.Markdown, source));
|
||
|
|
texts.Add(new(content.Title, source));
|
||
|
|
texts.Add(new(content.Description, source));
|
||
|
|
texts.Add(new(content.Language, source));
|
||
|
|
texts.Add(new(content.PublishedTime, source));
|
||
|
|
texts.Add(new(content.ModifiedTime, source));
|
||
|
|
|
||
|
|
foreach (var author in content.Authors)
|
||
|
|
texts.Add(new(author, source));
|
||
|
|
}
|
||
|
|
|
||
|
|
var sanitizedTexts = await guardService.SanitizeAsync(texts);
|
||
|
|
var sanitizedPages = new List<WebPageModelContent>(pages.Count);
|
||
|
|
var offset = 0;
|
||
|
|
|
||
|
|
foreach (var (content, _) in pages)
|
||
|
|
{
|
||
|
|
var authors = new List<string>(content.Authors.Count);
|
||
|
|
for (var authorIndex = 0; authorIndex < content.Authors.Count; authorIndex++)
|
||
|
|
authors.Add(sanitizedTexts[offset + SINGLE_VALUE_FIELD_COUNT + authorIndex]);
|
||
|
|
|
||
|
|
sanitizedPages.Add(new(
|
||
|
|
sanitizedTexts[offset],
|
||
|
|
sanitizedTexts[offset + 1],
|
||
|
|
sanitizedTexts[offset + 2],
|
||
|
|
authors,
|
||
|
|
sanitizedTexts[offset + 3],
|
||
|
|
sanitizedTexts[offset + 4],
|
||
|
|
sanitizedTexts[offset + 5]));
|
||
|
|
|
||
|
|
offset += SINGLE_VALUE_FIELD_COUNT + content.Authors.Count;
|
||
|
|
}
|
||
|
|
|
||
|
|
return sanitizedPages;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// <summary>
|
||
|
|
/// Filters every model-facing text of a single page.
|
||
|
|
/// </summary>
|
||
|
|
public static async Task<WebPageModelContent> SanitizeAsync(
|
||
|
|
PromptInjectionGuardService guardService,
|
||
|
|
WebPageModelContent content,
|
||
|
|
PromptInjectionSource source) => (await SanitizeAsync(guardService, [(content, source)]))[0];
|
||
|
|
}
|