using System.Text; using System.Text.Json; using AIStudio.Chat; using AIStudio.Dialogs; using AIStudio.Settings; using Microsoft.AspNetCore.Components; using DialogOptions = AIStudio.Dialogs.DialogOptions; namespace AIStudio.Components; public partial class Workspaces : MSGComponentBase { [Inject] private IDialogService DialogService { get; init; } = null!; [Inject] private ILogger Logger { get; init; } = null!; [Parameter] public ChatThread? CurrentChatThread { get; set; } [Parameter] public EventCallback CurrentChatThreadChanged { get; set; } [Parameter] public bool ExpandRootNodes { get; set; } = true; private const Placement WORKSPACE_ITEM_TOOLTIP_PLACEMENT = Placement.Bottom; private readonly SemaphoreSlim treeLoadingSemaphore = new(1, 1); private readonly List> treeItems = []; private readonly HashSet loadingWorkspaceChatLists = []; private CancellationTokenSource? prefetchCancellationTokenSource; private bool isInitialLoading = true; private bool isDisposed; #region Overrides of ComponentBase protected override async Task OnInitializedAsync() { await base.OnInitializedAsync(); _ = this.LoadTreeItemsAsync(startPrefetch: true); } #endregion private async Task LoadTreeItemsAsync(bool startPrefetch = true, bool forceReload = false) { await this.treeLoadingSemaphore.WaitAsync(); try { if (this.isDisposed) return; if (forceReload) await WorkspaceBehaviour.ForceReloadWorkspaceTreeAsync(); var snapshot = await WorkspaceBehaviour.GetOrLoadWorkspaceTreeShellAsync(); this.BuildTreeItems(snapshot); this.isInitialLoading = false; } finally { this.treeLoadingSemaphore.Release(); } await this.SafeStateHasChanged(); if (startPrefetch) await this.StartPrefetchAsync(); } private void BuildTreeItems(WorkspaceTreeCacheSnapshot snapshot) { this.treeItems.Clear(); var workspaceChildren = new List>(); foreach (var workspace in snapshot.Workspaces) workspaceChildren.Add(this.CreateWorkspaceTreeItem(workspace)); workspaceChildren.Add(new TreeItemData { Expandable = false, Value = new TreeButton(WorkspaceBranch.WORKSPACES, 1, T("Add workspace"), Icons.Material.Filled.LibraryAdd, this.AddWorkspaceAsync), }); this.treeItems.Add(new TreeItemData { Expanded = this.ExpandRootNodes, Expandable = true, Value = new TreeItemData { Depth = 0, Branch = WorkspaceBranch.WORKSPACES, Text = T("Workspaces"), Icon = Icons.Material.Filled.Folder, Expandable = true, Path = "root", Children = workspaceChildren, }, }); this.treeItems.Add(new TreeItemData { Expandable = false, Value = new TreeDivider(), }); var temporaryChatsChildren = new List>(); foreach (var temporaryChat in snapshot.TemporaryChats.OrderByDescending(x => x.LastEditTime)) temporaryChatsChildren.Add(CreateChatTreeItem(temporaryChat, WorkspaceBranch.TEMPORARY_CHATS, depth: 1, icon: Icons.Material.Filled.Timer)); this.treeItems.Add(new TreeItemData { Expanded = this.ExpandRootNodes, Expandable = true, Value = new TreeItemData { Depth = 0, Branch = WorkspaceBranch.TEMPORARY_CHATS, Text = T("Disappearing Chats"), Icon = Icons.Material.Filled.Timer, Expandable = true, Path = "temp", Children = temporaryChatsChildren, }, }); } private TreeItemData CreateWorkspaceTreeItem(WorkspaceTreeWorkspace workspace) { var children = new List>(); if (workspace.ChatsLoaded) { foreach (var workspaceChat in workspace.Chats.OrderByDescending(x => x.LastEditTime)) children.Add(CreateChatTreeItem(workspaceChat, WorkspaceBranch.WORKSPACES, depth: 2, icon: Icons.Material.Filled.Chat)); } else if (this.loadingWorkspaceChatLists.Contains(workspace.WorkspaceId)) children.AddRange(this.CreateLoadingRows(workspace.WorkspacePath)); children.Add(new TreeItemData { Expandable = false, Value = new TreeButton(WorkspaceBranch.WORKSPACES, 2, T("Add chat"), Icons.Material.Filled.AddComment, () => this.AddChatAsync(workspace.WorkspacePath)), }); return new TreeItemData { Expandable = true, Value = new TreeItemData { Type = TreeItemType.WORKSPACE, Depth = 1, Branch = WorkspaceBranch.WORKSPACES, Text = workspace.Name, Icon = Icons.Material.Filled.Description, Expandable = true, Path = workspace.WorkspacePath, Children = children, }, }; } private IReadOnlyCollection> CreateLoadingRows(string workspacePath) { return [ this.CreateLoadingTreeItem(workspacePath, "loading_1"), this.CreateLoadingTreeItem(workspacePath, "loading_2"), this.CreateLoadingTreeItem(workspacePath, "loading_3"), ]; } private TreeItemData CreateLoadingTreeItem(string workspacePath, string suffix) { return new TreeItemData { Expandable = false, Value = new TreeItemData { Type = TreeItemType.LOADING, Depth = 2, Branch = WorkspaceBranch.WORKSPACES, Text = T("Loading chats..."), Icon = Icons.Material.Filled.HourglassTop, Expandable = false, Path = Path.Join(workspacePath, suffix), }, }; } private static TreeItemData CreateChatTreeItem(WorkspaceTreeChat chat, WorkspaceBranch branch, int depth, string icon) { return new TreeItemData { Expandable = false, Value = new TreeItemData { Type = TreeItemType.CHAT, Depth = depth, Branch = branch, Text = chat.Name, Icon = icon, Expandable = false, Path = chat.ChatPath, LastEditTime = chat.LastEditTime, }, }; } private async Task SafeStateHasChanged() { if (this.isDisposed) return; await this.InvokeAsync(this.StateHasChanged); } private async Task StartPrefetchAsync() { if (this.prefetchCancellationTokenSource is not null) { await this.prefetchCancellationTokenSource.CancelAsync(); this.prefetchCancellationTokenSource.Dispose(); } this.prefetchCancellationTokenSource = new CancellationTokenSource(); await this.PrefetchWorkspaceChatsAsync(this.prefetchCancellationTokenSource.Token); } private async Task PrefetchWorkspaceChatsAsync(CancellationToken cancellationToken) { try { await WorkspaceBehaviour.TryPrefetchRemainingChatsAsync(async _ => { if (this.isDisposed || cancellationToken.IsCancellationRequested) return; await this.LoadTreeItemsAsync(startPrefetch: false); }, cancellationToken); } catch (OperationCanceledException) { // Expected when the component is hidden or disposed. } catch (Exception ex) { this.Logger.LogWarning(ex, "Failed while prefetching workspace chats."); } } private async Task OnWorkspaceClicked(TreeItemData treeItem) { if (treeItem.Type is not TreeItemType.WORKSPACE) return; if (!Guid.TryParse(Path.GetFileName(treeItem.Path), out var workspaceId)) return; await this.EnsureWorkspaceChatsLoadedAsync(workspaceId); } private async Task EnsureWorkspaceChatsLoadedAsync(Guid workspaceId) { var snapshot = await WorkspaceBehaviour.GetOrLoadWorkspaceTreeShellAsync(); var hasWorkspace = false; var chatsLoaded = false; foreach (var workspace in snapshot.Workspaces) { if (workspace.WorkspaceId != workspaceId) continue; hasWorkspace = true; chatsLoaded = workspace.ChatsLoaded; break; } if (!hasWorkspace || chatsLoaded || !this.loadingWorkspaceChatLists.Add(workspaceId)) return; await this.LoadTreeItemsAsync(startPrefetch: false); try { await WorkspaceBehaviour.GetWorkspaceChatsAsync(workspaceId); } finally { this.loadingWorkspaceChatLists.Remove(workspaceId); } await this.LoadTreeItemsAsync(startPrefetch: false); } public async Task ForceRefreshFromDiskAsync() { if (this.prefetchCancellationTokenSource is not null) { await this.prefetchCancellationTokenSource.CancelAsync(); this.prefetchCancellationTokenSource.Dispose(); this.prefetchCancellationTokenSource = null; } this.loadingWorkspaceChatLists.Clear(); this.isInitialLoading = true; await this.SafeStateHasChanged(); await this.LoadTreeItemsAsync(startPrefetch: true, forceReload: true); } public async Task StoreChatAsync(ChatThread chat, bool reloadTreeItems = false) { await WorkspaceBehaviour.StoreChatAsync(chat); if (reloadTreeItems) this.loadingWorkspaceChatLists.Clear(); await this.LoadTreeItemsAsync(startPrefetch: false); } private async Task LoadChatAsync(string? chatPath, bool switchToChat) { if (string.IsNullOrWhiteSpace(chatPath)) return null; if (!Directory.Exists(chatPath)) return null; if (switchToChat && await MessageBus.INSTANCE.SendMessageUseFirstResult(this, Event.HAS_CHAT_UNSAVED_CHANGES)) { var dialogParameters = new DialogParameters { { x => x.Message, T("Are you sure you want to load another chat? All unsaved changes will be lost.") }, }; var dialogReference = await this.DialogService.ShowAsync(T("Load Chat"), dialogParameters, DialogOptions.FULLSCREEN); var dialogResult = await dialogReference.Result; if (dialogResult is null || dialogResult.Canceled) return null; } try { var chatData = await File.ReadAllTextAsync(Path.Join(chatPath, "thread.json"), Encoding.UTF8); var chat = JsonSerializer.Deserialize(chatData, WorkspaceBehaviour.JSON_OPTIONS); if (switchToChat) { this.CurrentChatThread = chat; await this.CurrentChatThreadChanged.InvokeAsync(this.CurrentChatThread); await MessageBus.INSTANCE.SendMessage(this, Event.WORKSPACE_LOADED_CHAT_CHANGED); } return chat; } catch (Exception e) { this.Logger.LogError($"Failed to load chat from '{chatPath}': {e.Message}"); } return null; } public async Task DeleteChatAsync(string? chatPath, bool askForConfirmation = true, bool unloadChat = true) { var chat = await this.LoadChatAsync(chatPath, false); if (chat is null) return; if (askForConfirmation) { var workspaceName = await WorkspaceBehaviour.LoadWorkspaceNameAsync(chat.WorkspaceId); var dialogParameters = new DialogParameters { { x => x.Message, (chat.WorkspaceId == Guid.Empty) switch { true => string.Format(T("Are you sure you want to delete the temporary chat '{0}'?"), chat.Name), false => string.Format(T("Are you sure you want to delete the chat '{0}' in the workspace '{1}'?"), chat.Name, workspaceName), } }, }; var dialogReference = await this.DialogService.ShowAsync(T("Delete Chat"), dialogParameters, DialogOptions.FULLSCREEN); var dialogResult = await dialogReference.Result; if (dialogResult is null || dialogResult.Canceled) return; } await WorkspaceBehaviour.DeleteChatAsync(this.DialogService, chat.WorkspaceId, chat.ChatId, askForConfirmation: false); await this.LoadTreeItemsAsync(startPrefetch: false); if (unloadChat && this.CurrentChatThread?.ChatId == chat.ChatId) { this.CurrentChatThread = null; await this.CurrentChatThreadChanged.InvokeAsync(this.CurrentChatThread); await MessageBus.INSTANCE.SendMessage(this, Event.WORKSPACE_LOADED_CHAT_CHANGED); } } private async Task RenameChatAsync(string? chatPath) { var chat = await this.LoadChatAsync(chatPath, false); if (chat is null) return; var dialogParameters = new DialogParameters { { x => x.Message, string.Format(T("Please enter a new or edit the name for your chat '{0}':"), chat.Name) }, { x => x.InputHeaderText, T("Chat Name") }, { x => x.UserInput, chat.Name }, { x => x.ConfirmText, T("Rename") }, { x => x.ConfirmColor, Color.Info }, { x => x.AllowEmptyInput, false }, { x => x.EmptyInputErrorMessage, T("Please enter a chat name.") }, }; var dialogReference = await this.DialogService.ShowAsync(T("Rename Chat"), dialogParameters, DialogOptions.FULLSCREEN); var dialogResult = await dialogReference.Result; if (dialogResult is null || dialogResult.Canceled) return; chat.Name = (dialogResult.Data as string)!; if (this.CurrentChatThread?.ChatId == chat.ChatId) { this.CurrentChatThread.Name = chat.Name; await this.CurrentChatThreadChanged.InvokeAsync(this.CurrentChatThread); await MessageBus.INSTANCE.SendMessage(this, Event.WORKSPACE_LOADED_CHAT_CHANGED); } await WorkspaceBehaviour.StoreChatAsync(chat); await this.LoadTreeItemsAsync(startPrefetch: false); } private async Task RenameWorkspaceAsync(string? workspacePath) { if (workspacePath is null) return; var workspaceId = Guid.Parse(Path.GetFileName(workspacePath)); var workspaceName = await WorkspaceBehaviour.LoadWorkspaceNameAsync(workspaceId); var dialogParameters = new DialogParameters { { x => x.Message, string.Format(T("Please enter a new or edit the name for your workspace '{0}':"), workspaceName) }, { x => x.InputHeaderText, T("Workspace Name") }, { x => x.UserInput, workspaceName }, { x => x.ConfirmText, T("Rename") }, { x => x.ConfirmColor, Color.Info }, { x => x.AllowEmptyInput, false }, { x => x.EmptyInputErrorMessage, T("Please enter a workspace name.") }, }; var dialogReference = await this.DialogService.ShowAsync(T("Rename Workspace"), dialogParameters, DialogOptions.FULLSCREEN); var dialogResult = await dialogReference.Result; if (dialogResult is null || dialogResult.Canceled) return; var alteredWorkspaceName = (dialogResult.Data as string)!; var workspaceNamePath = Path.Join(workspacePath, "name"); await File.WriteAllTextAsync(workspaceNamePath, alteredWorkspaceName, Encoding.UTF8); await WorkspaceBehaviour.UpdateWorkspaceNameInCacheAsync(workspaceId, alteredWorkspaceName); await this.LoadTreeItemsAsync(startPrefetch: false); } private async Task AddWorkspaceAsync() { var dialogParameters = new DialogParameters { { x => x.Message, T("Please name your workspace:") }, { x => x.InputHeaderText, T("Workspace Name") }, { x => x.UserInput, string.Empty }, { x => x.ConfirmText, T("Add workspace") }, { x => x.ConfirmColor, Color.Info }, { x => x.AllowEmptyInput, false }, { x => x.EmptyInputErrorMessage, T("Please enter a workspace name.") }, }; var dialogReference = await this.DialogService.ShowAsync(T("Add Workspace"), dialogParameters, DialogOptions.FULLSCREEN); var dialogResult = await dialogReference.Result; if (dialogResult is null || dialogResult.Canceled) return; var workspaceId = Guid.NewGuid(); var workspacePath = Path.Join(SettingsManager.DataDirectory, "workspaces", workspaceId.ToString()); Directory.CreateDirectory(workspacePath); var workspaceName = (dialogResult.Data as string)!; var workspaceNamePath = Path.Join(workspacePath, "name"); await File.WriteAllTextAsync(workspaceNamePath, workspaceName, Encoding.UTF8); await WorkspaceBehaviour.AddWorkspaceToCacheAsync(workspaceId, workspacePath, workspaceName); await this.LoadTreeItemsAsync(startPrefetch: false); } private async Task DeleteWorkspaceAsync(string? workspacePath) { if (workspacePath is null) return; var workspaceId = Guid.Parse(Path.GetFileName(workspacePath)); var workspaceName = await WorkspaceBehaviour.LoadWorkspaceNameAsync(workspaceId); var chatCount = Directory.EnumerateDirectories(workspacePath).Count(); var dialogParameters = new DialogParameters { { x => x.Message, string.Format(T("Are you sure you want to delete the workspace '{0}'? This will also delete {1} chat(s) in this workspace."), workspaceName, chatCount) }, }; var dialogReference = await this.DialogService.ShowAsync(T("Delete Workspace"), dialogParameters, DialogOptions.FULLSCREEN); var dialogResult = await dialogReference.Result; if (dialogResult is null || dialogResult.Canceled) return; Directory.Delete(workspacePath, true); await WorkspaceBehaviour.RemoveWorkspaceFromCacheAsync(workspaceId); await this.LoadTreeItemsAsync(startPrefetch: false); } private async Task MoveChatAsync(string? chatPath) { var chat = await this.LoadChatAsync(chatPath, false); if (chat is null) return; var dialogParameters = new DialogParameters { { x => x.Message, T("Please select the workspace where you want to move the chat to.") }, { x => x.SelectedWorkspace, chat.WorkspaceId }, { x => x.ConfirmText, T("Move chat") }, }; var dialogReference = await this.DialogService.ShowAsync(T("Move Chat to Workspace"), dialogParameters, DialogOptions.FULLSCREEN); var dialogResult = await dialogReference.Result; if (dialogResult is null || dialogResult.Canceled) return; var workspaceId = dialogResult.Data is Guid id ? id : Guid.Empty; if (workspaceId == Guid.Empty) return; await WorkspaceBehaviour.DeleteChatAsync(this.DialogService, chat.WorkspaceId, chat.ChatId, askForConfirmation: false); chat.WorkspaceId = workspaceId; if (this.CurrentChatThread?.ChatId == chat.ChatId) { this.CurrentChatThread = chat; await this.CurrentChatThreadChanged.InvokeAsync(this.CurrentChatThread); await MessageBus.INSTANCE.SendMessage(this, Event.WORKSPACE_LOADED_CHAT_CHANGED); } await WorkspaceBehaviour.StoreChatAsync(chat); await this.LoadTreeItemsAsync(startPrefetch: false); } private async Task AddChatAsync(string workspacePath) { if (await MessageBus.INSTANCE.SendMessageUseFirstResult(this, Event.HAS_CHAT_UNSAVED_CHANGES)) { var dialogParameters = new DialogParameters { { x => x.Message, T("Are you sure you want to create a another chat? All unsaved changes will be lost.") }, }; var dialogReference = await this.DialogService.ShowAsync(T("Create Chat"), dialogParameters, DialogOptions.FULLSCREEN); var dialogResult = await dialogReference.Result; if (dialogResult is null || dialogResult.Canceled) return; } var workspaceId = Guid.Parse(Path.GetFileName(workspacePath)); var chat = new ChatThread { WorkspaceId = workspaceId, ChatId = Guid.NewGuid(), Name = string.Empty, SystemPrompt = SystemPrompts.DEFAULT, Blocks = [], }; var chatPath = Path.Join(workspacePath, chat.ChatId.ToString()); await WorkspaceBehaviour.StoreChatAsync(chat); await this.LoadChatAsync(chatPath, switchToChat: true); await this.LoadTreeItemsAsync(startPrefetch: false); } #region Overrides of MSGComponentBase protected override async Task ProcessIncomingMessage(ComponentBase? sendingComponent, Event triggeredEvent, T? data) where T : default { switch (triggeredEvent) { case Event.PLUGINS_RELOADED: await this.ForceRefreshFromDiskAsync(); break; } } protected override void DisposeResources() { this.isDisposed = true; this.prefetchCancellationTokenSource?.Cancel(); this.prefetchCancellationTokenSource?.Dispose(); this.prefetchCancellationTokenSource = null; base.DisposeResources(); } #endregion }