SqlDatabaseVectorSearch/SqlDatabaseVectorSearch/Components/Pages/Ask.razor

@page "/ask"
@using Microsoft.Extensions.AI

@inject IServiceScopeFactory ServiceScopeFactory
@inject IJSRuntime JSRuntime

<PageTitle>Chat with your data</PageTitle>

<div class="card mx-auto mt-2">
    <div class="card-body">

        @if (messages.Count == 0)
        {
            <div class="h-100 d-flex flex-column justify-content-center align-items-center text-body-secondary">
                <Icon Name="IconName.ChatSquareQuoteFill" Color="IconColor.Primary" Size="IconSize.x4" />
                <p class="mt-3 mb-1 fw-semibold">Chat with your documents</p>
                <p class="small mb-0">Ask a question about the documents you have uploaded. Press the up arrow key to recall your previous question.</p>
            </div>
        }

        @foreach (var message in messages)
        {
            if (message.Role == "user")
            {
                <div class="d-flex align-items-baseline text-end justify-content-end">
                    <div class="pe-2">
                        <div>
                            <div class="card card-text d-inline-block p-2 px-3 m-1">
                                <Markdown style="overflow-y:auto;">@message.Text</Markdown>
                            </div>
                        </div>
                    </div>
                    <div class="position-relative avatar">
                        <Image src="/images/user.png" class="img-fluid rounded-circle" alt="" />
                    </div>
                </div>
            }
            else if (message.Role == "assistant")
            {
                <div class="d-flex align-items-baseline">
                    <div class="position-relative avatar">
                        <Image src="/images/assistant.png" class="img-fluid rounded-circle" alt="" />
                    </div>
                    <div class="pe-2">
                        <div>
                            @if (message.Text is null)
                            {
                                <div class="card card-text d-inline-block p-3 px-3 m-1">
                                    <div class="typing-indicator" role="status" aria-label="The assistant is typing">
                                        <span class="typing-dot"></span>
                                        <span class="typing-dot"></span>
                                        <span class="typing-dot"></span>
                                    </div>
                                </div>
                            }
                            else
                            {
                                <div class="card card-text d-inline-block p-2 px-3 m-1">
                                    <div class="message-content">
                                        <div class="streaming-content">
                                            <div class="streaming-text @(message.Status == MessageStatus.Streaming ? "streaming-text-with-spinner" : "")">
                                                <Markdown style="overflow-y:auto;">@message.Text</Markdown>
                                            </div>
                                            @if (message.Status == MessageStatus.Streaming)
                                            {
                                                <div class="streaming-spinner-bottom-left">
                                                    <Spinner Size="SpinnerSize.Small" Color="SpinnerColor.Primary" />
                                                </div>
                                            }
                                        </div>
                                    </div>
                                    @if (message.Status == MessageStatus.Completed)
                                    {
                                        <div class="d-flex justify-content-between">
                                            <div class="text-start bg-transparent mt-3">
                                                <Tooltip Title="@message.TokenUsage" IsHtml="true" Color="TooltipColor.Primary" Placement="TooltipPlacement.Bottom">
                                                    <Icon Class="d-flex" Name="IconName.CashCoin"></Icon>
                                                </Tooltip>
                                            </div>
                                            <div class="text-end bg-transparent copy-button">
                                                <Tooltip Title="@(copiedMessage == message ? "Copied!" : "Copy to clipboard")" Color="TooltipColor.Dark" Placement="TooltipPlacement.Bottom">
                                                    <Button Type="ButtonType.Button" Outline="false" @onclick="@(async () => await CopyToClipboardAsync(message))">
                                                        @if (copiedMessage == message)
                                                        {
                                                            <Icon Name="IconName.Check" Class="text-success" />
                                                        }
                                                        else
                                                        {
                                                            <Icon Name="IconName.Clipboard" />
                                                        }
                                                    </Button>
                                                </Tooltip>
                                            </div>
                                        </div>
                                    }
                                </div>
                            }
                        </div>
                    </div>
                </div>
            }
        }

        <div @ref="chat"></div>
    </div>

    <div class="card-footer bg-white w-100 bottom-0 m-0 p-1">
        <div class="input-group">
            <span class="input-group-text bg-transparent border-0">
                <Tooltip Title="Messages aren't stored in any way, either on the client or on the server." Color="TooltipColor.Primary" Placement="TooltipPlacement.Bottom">
                    <Icon Class="d-flex" Color="IconColor.Success" Name="IconName.ShieldLockFill"></Icon>
                </Tooltip>
            </span>
            <input @ref="askInput" type="text" @bind="@question" @bind:event="oninput" placeholder="Ask a question about your documents..." class="form-control border-0" maxlength="2000" @onkeydown="HandleKeyDown" />
            <div class="input-group-text bg-transparent border-0">
                <Button Type="ButtonType.Submit" Color="ButtonColor.Primary" Disabled="@(isAsking || string.IsNullOrWhiteSpace(question))" @onclick="AskQuestion">
                    <Icon Name="IconName.Send" />
                </Button>
                <Tooltip Title="Start a new conversation" Color="TooltipColor.Secondary" Placement="TooltipPlacement.Bottom">
                    <Button Type="ButtonType.Reset" Class="ms-2" Color="ButtonColor.Secondary" Disabled="@isAsking" @onclick="Reset">
                        <Icon CustomIconName="bi bi-x-lg" />
                    </Button>
                </Tooltip>
            </div>
        </div>
    </div>
</div>

@code
{
    private ElementReference askInput = default!;
    private ElementReference chat = default!;

    private IList<Message> messages = [];
    private string? question;

    private Guid conversationId = Guid.NewGuid();
    private bool isAsking;

    private Message? copiedMessage;

    protected override async Task OnAfterRenderAsync(bool firstRender)
    {
        if (!firstRender)
        {
            return;
        }

        await JSRuntime.InvokeVoidAsync("setFocus", askInput);
    }

    private async Task HandleKeyDown(KeyboardEventArgs e)
    {
        if (isAsking)
        {
            return;
        }

        if (e.Key == "Enter" && !string.IsNullOrWhiteSpace(question))
        {
            await AskQuestion();
        }
        else if (e.Key == "ArrowUp" && messages.Count >= 2)
        {
            question = messages[^2].Text;
        }
    }

    private async Task AskQuestion()
    {
        isAsking = true;

        var userQuestion = new Question(conversationId, question!);
        var userMessage = new Message { Text = userQuestion.Text, Role = "user", Status = MessageStatus.Completed };
        messages.Add(userMessage);

        var assistantMessage = new Message { Role = "assistant", Status = MessageStatus.New };
        messages.Add(assistantMessage);

        question = null;
        await InvokeAsync(StateHasChanged);

        await EnsureMessageIsVisibleAsync();

        try
        {
            await using var scope = ServiceScopeFactory.CreateAsyncScope();
            var vectorSearchService = scope.ServiceProvider.GetRequiredService<VectorSearchService>();

            var response = vectorSearchService.AskStreamingAsync(userQuestion);
            await foreach (var update in response)
            {
                if (update.StreamState == StreamState.Start)
                {
                    userMessage.Text = update.ReformulatedQuestion;
                    assistantMessage.TokenUsage = FormatTokenUsage(update.TokenUsage);
                    assistantMessage.Status = MessageStatus.Streaming;
                }
                else if (update.StreamState == StreamState.Delta)
                {
                    // Adds tokens to the assistant message as they are received.
                    assistantMessage.Text += update.Answer;
                }
                else if (update.StreamState == StreamState.End)
                {
                    assistantMessage.Status = MessageStatus.Completed;
                    assistantMessage.TokenUsage += FormatTokenUsage(update.TokenUsage);
                }

                await InvokeAsync(StateHasChanged);

                await EnsureMessageIsVisibleAsync();
            }
        }
        catch (Exception ex)
        {
            assistantMessage.Text = $"There was an error while processing your question: {ex.Message}";
            assistantMessage.Status = MessageStatus.Completed;
        }
        finally
        {
            isAsking = false;
            await InvokeAsync(StateHasChanged);

            await EnsureMessageIsVisibleAsync();
        }
    }

    private void Reset()
    {
        question = null;
        conversationId = Guid.NewGuid();
        copiedMessage = null;
        messages.Clear();
    }

    private async Task CopyToClipboardAsync(Message message)
    {
        if (string.IsNullOrEmpty(message.Text))
        {
            return;
        }

        await JSRuntime.InvokeVoidAsync("navigator.clipboard.writeText", message.Text);

        copiedMessage = message;
        await InvokeAsync(StateHasChanged);

        await Task.Delay(3000); // Shows the checkmark for 3 seconds

        if (copiedMessage == message)
        {
            copiedMessage = null;
            await InvokeAsync(StateHasChanged);
        }
    }

    private static string FormatTokenUsage(TokenUsageResponse? tokenUsageResponse)
    {
        if (tokenUsageResponse is null)
        {
            return string.Empty;
        }

        var reformulation = tokenUsageResponse.Reformulation is not null
            ? $"<p><strong>Reformulation:</strong><br />{FormatTokenUsageDetails(tokenUsageResponse.Reformulation)}</p>"
            : string.Empty;

        var question = tokenUsageResponse.Question is not null
            ? $"<p><strong>Question:</strong><br />{FormatTokenUsageDetails(tokenUsageResponse.Question)}</p>"
            : string.Empty;

        return $"{reformulation}{question}";

        static string FormatTokenUsageDetails(UsageDetails? tokenUsage)
        {
            if (tokenUsage is null)
            {
                return string.Empty;
            }

            return $"Input tokens: {tokenUsage.InputTokenCount}<br />" +
                   $"Output tokens: {tokenUsage.OutputTokenCount}<br />" +
                   $"Total tokens: {tokenUsage.TotalTokenCount}";
        }
    }

    private async Task EnsureMessageIsVisibleAsync()
    {
        await JSRuntime.InvokeVoidAsync("scrollTo", chat);
    }

    public enum MessageStatus
    {
        New,
        Streaming,
        Completed
    }

    public class Message
    {
        public string? Text { get; set; }

        public required string Role { get; set; }

        public MessageStatus Status { get; set; } = MessageStatus.New;

        public string? TokenUsage { get; set; }
    }
}