AI-Studio/app/MindWork AI Studio/Chat/ChatThread.cs

using System.Globalization;
using System.Text.Json.Serialization;

using AIStudio.Components;
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.ToolCallingSystem;
using AIStudio.Tools.ERIClient.DataModel;

namespace AIStudio.Chat;

/// <summary>
/// Data structure for a chat thread.
/// </summary>
public sealed record ChatThread
{
    private static readonly ILogger<ChatThread> LOGGER = Program.LOGGER_FACTORY.CreateLogger<ChatThread>();
    
    /// <summary>
    /// The unique identifier of the chat thread.
    /// </summary>
    public Guid ChatId { get; init; }
    
    /// <summary>
    /// The unique identifier of the workspace.
    /// </summary>
    public Guid WorkspaceId { get; set; }

    /// <summary>
    /// The monotonically increasing number used for managed media transcript filenames.
    /// </summary>
    public ulong LastMediaTranscriptNumber { get; set; }

    /// <summary>
    /// Managed transcript attachments prepared for the composer but not sent yet.
    /// Empty by default so older serialized threads require no migration.
    /// </summary>
    public List<ManagedTranscriptAttachment> PendingMediaTranscripts { get; set; } = [];

    /// <summary>
    /// Specifies the provider selected for the chat thread.
    /// </summary>
    public string SelectedProvider { get; set; } = string.Empty;

    /// <summary>
    /// Specifies the profile selected for the chat thread.
    /// </summary>
    public string SelectedProfile { get; set; } = string.Empty;
    
    /// <summary>
    /// Specifies the profile selected for the chat thread.
    /// </summary>
    public string SelectedChatTemplate { get; set; } = string.Empty;

    /// <summary>
    /// Specifies the tools selected for the chat thread, as the user chose them.
    /// </summary>
    /// <remarks>
    /// Null means the thread never stored a selection, which is the case for every chat written
    /// before tools existed: those open with the defaults of their component. An empty set is the
    /// opposite statement — the user switched every tool off and wants it to stay that way.<br/><br/>
    /// This is the unfiltered selection. What a provider may actually run is decided per request,
    /// because a provider with too little confidence must not cost the user a tool permanently.
    /// </remarks>
    public HashSet<string>? SelectedToolIds { get; set; }

    /// <summary>
    /// Indicates whether to include the current date and time in the system prompt.
    /// False by default for backward compatibility.
    /// </summary>
    public bool IncludeDateTime { get; set; } = false;

    /// <summary>
    /// The data source options for this chat thread.
    /// </summary>
    public DataSourceOptions DataSourceOptions { get; set; } = new();

    /// <summary>
    /// The AI-selected data sources for this chat thread.
    /// </summary>
    public IReadOnlyList<DataSourceAgentSelected> AISelectedDataSources { get; set; } = [];

    /// <summary>
    /// The augmented data for this chat thread. Will be inserted into the system prompt.
    /// </summary>
    public string AugmentedData { get; set; } = string.Empty;

    /// <summary>
    /// The data security to use, derived from the data sources used so far.
    /// </summary>
    public DataSourceSecurity DataSecurity { get; set; } = DataSourceSecurity.NOT_SPECIFIED;

    /// <summary>
    /// The minimum confidence required for providers that continue this chat. It is raised whenever
    /// a tool returned sensitive data, and whenever a data source was used which demands a higher
    /// level. Both cases share one rule: once such data is in the thread, every provider which
    /// continues it must meet the level.
    /// </summary>
    [JsonInclude]
    public ConfidenceLevel RequiredProviderConfidence { get; private set; } = ConfidenceLevel.NONE;

    public void RequireProviderConfidence(ConfidenceLevel minimumProviderConfidence)
    {
        if (minimumProviderConfidence > this.RequiredProviderConfidence)
            this.RequiredProviderConfidence = minimumProviderConfidence;
    }

    /// <summary>
    /// The name of the chat thread. Usually generated by an AI model or manually edited by the user.
    /// </summary>
    public string Name { get; set; } = string.Empty;
    
    /// <summary>
    /// The current system prompt for the chat thread.
    /// </summary>
    public string SystemPrompt { get; set; } = string.Empty;

    /// <summary>
    /// The content blocks of the chat thread.
    /// </summary>
    public List<ContentBlock> Blocks { get; init; } = [];

    [JsonIgnore]
    public AIStudio.Tools.Components RuntimeComponent { get; set; } = AIStudio.Tools.Components.CHAT;

    [JsonIgnore]
    public HashSet<string> RuntimeSelectedToolIds { get; set; } = [];

    /// <summary>
    /// Whether the tools of this run were named by the assistant's own rules instead of chosen by
    /// the user.
    /// </summary>
    /// <remarks>
    /// A user who cannot see a tool selection must not get tools they never picked, which is why
    /// running tools normally requires a visible selection. That rule misses the case where nobody
    /// asked the user in the first place: a document analysis policy or an assistant plugin names
    /// its tools, and hiding the selection is the point rather than an obstacle. This flag tells
    /// the providers which of the two they are looking at.
    /// </remarks>
    [JsonIgnore]
    public bool RuntimeToolsAreAssistantManaged { get; set; }

    /// <summary>
    /// Whether this thread may run tools at all.
    /// </summary>
    public bool MayRunTools(SettingsManager settingsManager) => this.RuntimeToolsAreAssistantManaged || settingsManager.IsToolSelectionVisible(this.RuntimeComponent);
    
    private bool allowProfile = true;

    /// <summary>
    /// Prepares the system prompt for the chat thread.
    /// </summary>
    /// <remarks>
    /// The actual system prompt depends on the selected profile. If no profile is selected,
    /// the system prompt is returned as is. When a profile is selected, the system prompt
    /// is extended with the profile chosen.
    /// </remarks>
    /// <param name="settingsManager">The settings manager instance to use.</param>
    /// <param name="runnableToolDefinitions">The tools which may run in this thread. Their instructions become part of the system prompt. Null when the thread runs without tools.</param>
    /// <returns>The prepared system prompt.</returns>
    public string PrepareSystemPrompt(SettingsManager settingsManager, IEnumerable<ToolDefinition>? runnableToolDefinitions = null)
    {
        this.allowProfile = true;

        //
        // Use the information from the chat template, if provided. Otherwise, use the default system prompt
        //
        string systemPromptTextWithChatTemplate;
        var logMessage = $"Using no chat template for chat thread '{this.Name}'.";
        if (string.IsNullOrWhiteSpace(this.SelectedChatTemplate))
            systemPromptTextWithChatTemplate = this.SystemPrompt;
        else
        {
            if(!Guid.TryParse(this.SelectedChatTemplate, out var chatTemplateId))
                systemPromptTextWithChatTemplate = this.SystemPrompt;
            else
            {
                if(this.SelectedChatTemplate == ChatTemplate.NO_CHAT_TEMPLATE.Id || chatTemplateId == Guid.Empty)
                    systemPromptTextWithChatTemplate = this.SystemPrompt;
                else
                {
                    var chatTemplate = settingsManager.GetChatTemplateById(this.SelectedChatTemplate);
                    if(chatTemplate == ChatTemplate.NO_CHAT_TEMPLATE)
                        systemPromptTextWithChatTemplate = this.SystemPrompt;
                    else
                    {
                        logMessage = $"Using chat template '{chatTemplate.Name}' for chat thread '{this.Name}'.";
                        this.allowProfile = chatTemplate.AllowProfileUsage;
                        systemPromptTextWithChatTemplate = chatTemplate.ToSystemPrompt();
                    }
                }
            }
        }
        
        // We need a way to save the changed system prompt in our chat thread.
        // Otherwise, the chat thread will always tell us that it is using the
        // default system prompt:
        this.SystemPrompt = systemPromptTextWithChatTemplate;
        LOGGER.LogInformation(logMessage);

        //
        // Add augmented data, if available:
        //
        var isAugmentedDataAvailable = !string.IsNullOrWhiteSpace(this.AugmentedData);
        var systemPromptWithAugmentedData = isAugmentedDataAvailable switch
        {
            true => $"""
                     {systemPromptTextWithChatTemplate}
                     
                     {this.AugmentedData}
                     """,

            false => systemPromptTextWithChatTemplate,
        };
        
        if(isAugmentedDataAvailable)
            LOGGER.LogInformation("Augmented data is available for the chat thread.");
        else
            LOGGER.LogInformation("No augmented data is available for the chat thread.");
        
        
        //
        // Add information from the profile if available and allowed:
        //
        string systemPromptText;
        logMessage = $"Using no profile for chat thread '{this.Name}'.";
        if (string.IsNullOrWhiteSpace(this.SelectedProfile) || !this.allowProfile)
            systemPromptText = systemPromptWithAugmentedData;
        else
        {
            if(!Guid.TryParse(this.SelectedProfile, out var profileId))
                systemPromptText = systemPromptWithAugmentedData;
            else
            {
                if(this.SelectedProfile == Profile.NO_PROFILE.Id || profileId == Guid.Empty)
                    systemPromptText = systemPromptWithAugmentedData;
                else
                {
                    var profile = settingsManager.GetProfileById(this.SelectedProfile);
                    if(profile == Profile.NO_PROFILE)
                        systemPromptText = systemPromptWithAugmentedData;
                    else
                    {
                        logMessage = $"Using profile '{profile.Name}' for chat thread '{this.Name}'.";
                        systemPromptText = $"""
                                            {systemPromptWithAugmentedData}

                                            {profile.ToSystemPrompt()}
                                            """;
                    }
                }
            }
        }
        
        LOGGER.LogInformation(logMessage);

        var toolPolicy = ToolSelectionRules.BuildToolPolicyPrompt(runnableToolDefinitions ?? []);
        if (!string.IsNullOrWhiteSpace(toolPolicy))
        {
            systemPromptText = $"""
                                {systemPromptText}

                                {toolPolicy}
                                """;
        }

        if(!this.IncludeDateTime)
            return systemPromptText;
        
        //
        // Prepend the current date and time to the system prompt:
        //
        var nowUtc = DateTime.UtcNow;
        var nowLocal = DateTime.Now;
        var currentDateTime = string.Create(
            new CultureInfo("en-US"),
            $"Today is {nowUtc:dddd, MMMM d, yyyy h:mm tt} (UTC) and {nowLocal:dddd, MMMM d, yyyy h:mm tt} (local time)."
        );

        return $"""
                {currentDateTime}

                {systemPromptText}
                """;
    }

    /// <summary>
    /// Removes a content block from this chat thread.
    /// </summary>
    /// <param name="content">The content block to remove.</param>
    /// <param name="removeForRegenerate">Indicates whether the content block is removed for
    /// regeneration purposes. True, when the content block is removed for regeneration purposes,
    /// which will not remove the previous user block if it is hidden from the user.</param>
    public void Remove(IContent content, bool removeForRegenerate = false)
    {
        var block = this.Blocks.FirstOrDefault(x => x.Content == content);
        if(block is null)
            return;

        //
        // Remove the previous user block if it is hidden from the user. Otherwise,
        // the experience might be confusing for the user.
        //
        // Explanation, using the ERI assistant as an example:
        // - The ERI assistant generates for every file a hidden user prompt.
        // - In the UI, the user can only see the AI's responses, not the hidden user prompts.
        // - Now, the user removes one AI response
        // - The hidden user prompt is still there, but the user can't see it.
        // - Since the user prompt is hidden, neither is it possible to remove nor edit it.
        // - This method solves this issue by removing the hidden user prompt when the AI response is removed.
        //
        if (block.Role is ChatRole.AI && !removeForRegenerate)
        {
            var sortedBlocks = this.Blocks.OrderBy(x => x.Time).ToList();
            var index = sortedBlocks.IndexOf(block);
            if (index > 0)
            {
                var previousBlock = sortedBlocks[index - 1];
                if (previousBlock.Role is ChatRole.USER && previousBlock.HideFromUser)
                {
                    DeleteManagedAttachments(previousBlock);
                    this.Blocks.Remove(previousBlock);
                }
            }
        }

        DeleteManagedAttachments(block);

        // Remove the block from the chat thread:
        this.Blocks.Remove(block);
    }

    private static void DeleteManagedAttachments(ContentBlock block)
    {
        if (block.Content is not ContentText textContent)
            return;
            
        foreach (var attachment in textContent.FileAttachments)
            ManagedTranscriptAttachment.TryDeleteOwnedFile(attachment);
    }

    /// <summary>
    /// Transforms this chat thread to an ERI chat thread.
    /// </summary>
    /// <param name="token">The cancellation token.</param>
    /// <returns>The ERI chat thread.</returns>
    public async Task<Tools.ERIClient.DataModel.ChatThread> ToERIChatThread(CancellationToken token = default)
    {
        //
        // Transform the content blocks:
        //
        var contentBlocks = new List<Tools.ERIClient.DataModel.ContentBlock>(this.Blocks.Count);
        foreach (var block in this.Blocks)
        {
            var (contentData, contentType) = block.Content switch
            {
                ContentImage image => (await image.TryAsBase64(token) is (success: true, { } base64Image) ? base64Image : string.Empty, Tools.ERIClient.DataModel.ContentType.IMAGE),
                ContentText text => (text.Text, Tools.ERIClient.DataModel.ContentType.TEXT),
                
                _ => (string.Empty, Tools.ERIClient.DataModel.ContentType.UNKNOWN),
            };
            
            contentBlocks.Add(new Tools.ERIClient.DataModel.ContentBlock
            {
                Role = block.Role switch
                {
                    ChatRole.AI => Role.AI,
                    ChatRole.USER => Role.USER,
                    ChatRole.AGENT => Role.AGENT,
                    ChatRole.SYSTEM => Role.SYSTEM,
                    ChatRole.NONE => Role.NONE,
                    
                    _ => Role.UNKNOWN,
                },
                
                Content = contentData,
                Type = contentType,
            });
        }
        
        return new Tools.ERIClient.DataModel.ChatThread { ContentBlocks = contentBlocks };
    }
}