Chat.razor

@page "/"
@using System.ComponentModel
@using System.Runtime.InteropServices
@using System.Text.RegularExpressions
@using System.Drawing
@using Microsoft.Extensions.AI
@using imageGeneratorSample.Web.Services.Images
@inject IChatClient ChatClient
@inject NavigationManager Nav
@inject SemanticSearch Search
@inject IImageCacheService ImageCache
@implements IDisposable

<PageTitle>Chat</PageTitle>

<ChatHeader OnNewChat="@ResetConversationAsync" />

<ChatMessageList Messages="@messages" InProgressMessage="@currentResponseMessage" >
    <NoMessagesContent>
        <div>To get started, try asking about these example documents. You can replace these with your own data and replace this message.</div>
        <ChatCitation File="Example_Emergency_Survival_Kit.pdf"/>
        <ChatCitation File="Example_GPS_Watch.pdf"/>
    </NoMessagesContent>
</ChatMessageList>

<div class="chat-container">
    <ChatSuggestions OnSelected="@AddUserMessageAsync" @ref="@chatSuggestions" />
    <ChatInput OnSend="@AddUserMessageAsync" @ref="@chatInput" />
    <SurveyPrompt /> @* Remove this line to eliminate the template survey message *@
</div>

@code {
    private const string SystemPrompt = @"
        You are an assistant who answers questions about information you retrieve.
        Do not answer questions about anything else.
        Use only simple markdown to format your responses.

        Use the search tool to find relevant information. When you do this, end your
        reply with citations in the special XML format:

        <citation filename='string' page_number='number'>exact quote here</citation>

        Always include the citation in your response if there are results.

        The quote must be max 5 words, taken word-for-word from the search result, and is the basis for why the citation is relevant.
        Don't refer to the presence of citations; just emit these tags right at the end, with no surrounding text.
        ";

    private readonly ChatOptions chatOptions = new();
    private readonly List<ChatMessage> messages = new();
    private CancellationTokenSource? currentResponseCancellation;
    private ChatMessage? currentResponseMessage;
    private ChatInput? chatInput;
    private ChatSuggestions? chatSuggestions;

    protected override void OnInitialized()
    {
        messages.Add(new(ChatRole.System, SystemPrompt));
        chatOptions.Tools = 
        [
            AIFunctionFactory.Create(SearchAsync),
            #pragma warning disable MEAI001 // Type is for evaluation purposes only
            new HostedImageGenerationTool()
            {
                Options = new()
                {
                    ModelId = "gpt-image-1",
                    StreamingCount = 3
                }
            }
        ];
    }

    private async Task AddUserMessageAsync(ChatMessage userMessage)
    {
        CancelAnyCurrentResponse();

        // Add the user message to the conversation
        messages.Add(userMessage);
        chatSuggestions?.Clear();
        await chatInput!.FocusAsync();


        var chatMessages = chatOptions.ConversationId is null
            ? messages
            : [userMessage];

        // Stream and display a new response from the IChatClient
        var responseText = new TextContent("");
        currentResponseMessage = new ChatMessage(ChatRole.Assistant, [responseText]);

        // initialize the new currentResponseMessage
        await InvokeAsync(StateHasChanged);

        currentResponseCancellation = new();
        await foreach (var update in ChatClient.GetStreamingResponseAsync(chatMessages, chatOptions, currentResponseCancellation.Token))
        {
            chatOptions.ConversationId = update.ConversationId;
            messages.AddMessages(update, filter: c => c is not TextContent);
            await CacheImages(update.Contents, responseText);
            responseText.Text += update.Text;
            ChatMessageItem.NotifyChanged(currentResponseMessage);
        }


        // Store the final response in the conversation, and begin getting suggestions
        messages.Add(currentResponseMessage!);
        currentResponseMessage = null;
        chatSuggestions?.Update(messages);
    }

    private async Task CacheImages(IList<AIContent> contents, TextContent textContent)
    {
        Dictionary<string, string> cachedUris = new();
        for (int i = 0; i < contents.Count; i++)
        {
            var content = contents[i];
            if (content is not ImageGenerationToolResultContent imageResult)
            {
                continue;
            }

            contents.RemoveAt(i--);

            foreach (var dataContent in imageResult.Outputs?. OfType<DataContent>() ?? Array.Empty<DataContent>())
            {
                string? fileName = dataContent.Name;
                fileName ??= imageResult.ImageId;
                if (String.IsNullOrEmpty(fileName))
                {
                    _ = dataContent.AdditionalProperties?.TryGetValue("ItemId", out fileName);
                }
                fileName ??= $"image_{Guid.NewGuid()}.png";                

                var bytes = MemoryMarshal.TryGetArray(dataContent.Data, out var segment)
                    ? segment.Array!
                    : dataContent.Data.ToArray();

                var cacheUri = await ImageCache.CacheImageAsync(bytes, dataContent.MediaType, fileName);

                cachedUris[fileName] = Nav.ToAbsoluteUri(cacheUri).ToString();
            }
        }

        // Find each image link in the text content and replace it with the cached URI
        // Use a regex to find image links in the format ![description](url)
        // Replace each in order of appearance
        textContent.Text = Regex.Replace(textContent.Text,
            @"!\[([^\]]*)\]\(([^)]+)\)",
            match =>
            {
                var fileName = match.Groups[1].Value;
                if (!cachedUris.TryGetValue(fileName, out var uri))
                {
                    return match.Value; // no cached URI found, return original
                }
                cachedUris.Remove(fileName);
                return $"![{fileName}]({uri})";
            });

        // if we have any images remaining, just add them to the end
        while (cachedUris.Count > 0)
        {
            var imageUri = cachedUris.First();
            cachedUris.Remove(imageUri.Key);

            textContent.Text += $"![{imageUri.Key}]({imageUri.Value})\n";
        }
    }

    private void CancelAnyCurrentResponse()
    {
        // If a response was cancelled while streaming, include it in the conversation so it's not lost
        if (currentResponseMessage is not null)
        {
            messages.Add(currentResponseMessage);
        }

        currentResponseCancellation?.Cancel();
        currentResponseMessage = null;
    }

    private async Task ResetConversationAsync()
    {
        CancelAnyCurrentResponse();
        messages.Clear();
        messages.Add(new(ChatRole.System, SystemPrompt));
        chatSuggestions?.Clear();
        await chatInput!.FocusAsync();
    }

    [Description("Searches for information using a phrase or keyword")]
    private async Task<IEnumerable<string>> SearchAsync(
        [Description("The phrase to search for.")] string searchPhrase,
        [Description("If possible, specify the filename to search that file only. If not provided or empty, the search includes all files.")] string? filenameFilter = null)
    {
        await InvokeAsync(StateHasChanged);
        var results = await Search.SearchAsync(searchPhrase, filenameFilter, maxResults: 5);
        return results.Select(result =>
            $"<result filename=\"{result.DocumentId}\" page_number=\"{result.PageNumber}\">{result.Text}</result>");
    }

    public void Dispose()
        => currentResponseCancellation?.Cancel();
}