Responsible AI and LLM Security for .NET Applications

using System.Net.Http.Headers;
using System.Net.Http.Json;
using Azure.Core;

public sealed record ShieldVerdict(bool UserPromptAttack, bool DocumentAttack);

public sealed class PromptShieldsClient(HttpClient http, TokenCredential credential)
{
    private static readonly TokenRequestContext Scope =
        new(["https://cognitiveservices.azure.com/.default"]);

    public async Task<ShieldVerdict> AnalyzeAsync(
        string userPrompt, IReadOnlyList<string> documents, CancellationToken cancellationToken)
    {
        AccessToken token = await credential.GetTokenAsync(Scope, cancellationToken);
        using var request = new HttpRequestMessage(
            HttpMethod.Post, "contentsafety/text:shieldPrompt?api-version=2024-09-01")
        {
            Content = JsonContent.Create(new { userPrompt, documents }),
        };
        request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", token.Token);

        using var response = await http.SendAsync(request, cancellationToken);
        response.EnsureSuccessStatusCode();
        var body = await response.Content.ReadFromJsonAsync<ShieldResponse>(cancellationToken)
            ?? throw new InvalidOperationException("Empty Prompt Shields response.");

        return new ShieldVerdict(
            body.UserPromptAnalysis?.AttackDetected == true,
            body.DocumentsAnalysis?.Any(d => d.AttackDetected) == true);
    }

    private sealed record Analysis(bool AttackDetected);
    private sealed record ShieldResponse(
        Analysis? UserPromptAnalysis, List<Analysis>? DocumentsAnalysis);
}

using System.Runtime.CompilerServices;
using Microsoft.Extensions.AI;

public sealed class PromptAttackException(ShieldVerdict verdict)
    : Exception("Potential prompt injection detected.")
{
    public ShieldVerdict Verdict { get; } = verdict;
}

public sealed class PromptShieldChatClient(IChatClient inner, PromptShieldsClient shields)
    : DelegatingChatClient(inner)
{
    public override async Task<ChatResponse> GetResponseAsync(
        IEnumerable<ChatMessage> messages, ChatOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        List<ChatMessage> list = [.. messages];
        await EnsureSafeAsync(list, cancellationToken);
        return await base.GetResponseAsync(list, options, cancellationToken);
    }

    public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
        IEnumerable<ChatMessage> messages, ChatOptions? options = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        List<ChatMessage> list = [.. messages];
        await EnsureSafeAsync(list, cancellationToken);
        var updates = base.GetStreamingResponseAsync(list, options, cancellationToken);
        await foreach (var update in updates)
        {
            yield return update;
        }
    }

    private async Task EnsureSafeAsync(List<ChatMessage> messages, CancellationToken ct)
    {
        var prompt = messages.LastOrDefault(m => m.Role == ChatRole.User)?.Text ?? string.Empty;
        List<string> toolOutputs = [.. messages
            .SelectMany(m => m.Contents.OfType<FunctionResultContent>())
            .Select(r => r.Result?.ToString() ?? string.Empty)
            .Where(text => text.Length > 0)];

        var verdict = await shields.AnalyzeAsync(prompt, toolOutputs, ct);
        if (verdict.UserPromptAttack || verdict.DocumentAttack)
        {
            throw new PromptAttackException(verdict);
        }
    }
}