Improve Chat error handling - refactor error handling, show snackbar upon error

LM-Kit-Maestro/Services/LMKitTextGenerationStatus.cs → LM-Kit-Maestro/Services/LMKitRequestStatus.cs

@@ -1,8 +1,8 @@
 namespace LMKit.Maestro.Services;
 
-public enum LMKitTextGenerationStatus
+public enum LMKitRequestStatus
 {
-    Undefined,
+    OK,
     Cancelled,
     GenericError
 }

LM-Kit-Maestro/Services/LMKitService.LMKitResult.cs

@@ -6,7 +6,7 @@ public sealed class LMKitResult
     {
         public Exception? Exception { get; set; }
 
-        public LMKitTextGenerationStatus Status { get; set; }
+        public LMKitRequestStatus Status { get; set; }
 
         public object? Result { get; set; }
     }

LM-Kit-Maestro/Services/LMKitService.cs

@@ -211,18 +211,16 @@ private async Task<LMKitResult> HandleLMKitRequest(LMKitRequest request)
         // Ensuring we don't touch anything until Lm-Kit objects' state has been set to handle this request.
         _lmKitServiceSemaphore.Wait();
 
+        LMKitResult result;
+
         try
         {
-            LMKitResult result;
-
             if (request.CancellationTokenSource.IsCancellationRequested || ModelLoadingState == LMKitModelLoadingState.Unloaded)
             {
                 result = new LMKitResult()
                 {
-                    Status = LMKitTextGenerationStatus.Cancelled
+                    Status = LMKitRequestStatus.Cancelled
                 };
-
-                _lmKitServiceSemaphore.Release();
             }
             else
             {
@@ -255,16 +253,26 @@ private async Task<LMKitResult> HandleLMKitRequest(LMKitRequest request)
             }
 
             request.ResponseTask.TrySetResult(result);
-
-            return result;
+        }
+        catch (Exception exception)
+        {
+            result = new LMKitResult()
+            {
+                Exception = exception,
+                Status = LMKitRequestStatus.GenericError
+            };
         }
         finally
         {
+            _lmKitServiceSemaphore.Release();
+
             if (_requestSchedule.Contains(request))
             {
                 _requestSchedule.Remove(request);
             }
         }
+
+        return result;
     }
 
     private async Task<LMKitResult> SubmitRequest(LMKitRequest request)
@@ -300,11 +308,11 @@ private async Task<LMKitResult> SubmitRequest(LMKitRequest request)
 
                 if (result.Exception is OperationCanceledException)
                 {
-                    result.Status = LMKitTextGenerationStatus.Cancelled;
+                    result.Status = LMKitRequestStatus.Cancelled;
                 }
                 else
                 {
-                    result.Status = LMKitTextGenerationStatus.GenericError;
+                    result.Status = LMKitRequestStatus.GenericError;
                 }
             }
 
@@ -320,7 +328,7 @@ private async Task<LMKitResult> SubmitRequest(LMKitRequest request)
 
                 if (request.RequestType == LMKitRequest.LMKitRequestType.Prompt &&
                     conversation.GeneratedTitleSummary == null &&
-                    result.Status == LMKitTextGenerationStatus.Undefined &&
+                    result.Status == LMKitRequestStatus.OK &&
                     !string.IsNullOrEmpty(((TextGenerationResult)result.Result!).Completion))
                 {
                     GenerateConversationSummaryTitle(conversation);
@@ -329,7 +337,7 @@ private async Task<LMKitResult> SubmitRequest(LMKitRequest request)
 
             if (result.Exception != null && request.CancellationTokenSource.IsCancellationRequested)
             {
-                result.Status = LMKitTextGenerationStatus.Cancelled;
+                result.Status = LMKitRequestStatus.Cancelled;
             }
 
             return result;
@@ -339,7 +347,7 @@ private async Task<LMKitResult> SubmitRequest(LMKitRequest request)
             return new LMKitResult()
             {
                 Exception = exception,
-                Status = LMKitTextGenerationStatus.GenericError
+                Status = LMKitRequestStatus.GenericError
             };
         }
         finally

LM-Kit-Maestro/UI/Razor/Components/Chat.razor

@@ -6,28 +6,28 @@
 @inject IScrollHandler ScrollHandler
 @inject IResizeHandler ResizeHandler
 @inject ILogger<Chat> Logger
-
+@inject ISnackbar Snackbar
 @inherits MvvmComponentBase<ChatPageViewModel>
 
 <div id="chat-container">
     <div id="conversation-container">
         <div id="conversation-content"
-             class="chat-element dark @(ViewModel.ConversationListViewModel.CurrentConversation is { IsEmpty: true } ? "centered-container" : "top-align-container")">
+        class="chat-element dark @(ViewModel.ConversationListViewModel.CurrentConversation is { IsEmpty: true } ? "centered-container" : "top-align-container")">
             @if (ViewModel?.ConversationListViewModel.CurrentConversation != null)
             {
                 if (ViewModel.ConversationListViewModel.CurrentConversation.IsEmpty)
                 {
                     <div id="empty-conversation" class="vertical-stack">
                         <div class="vertical-stack spacing-4">
-                                                    <div class="welcome-message"><b>Maestro</b> at your service—let’s orchestrate something amazing!</div>
-                                    <div class="welcome-message">
-                                                Feel free to ask questions, explore ideas, or engage in meaningful conversations. 
+                            <div class="welcome-message"><b>Maestro</b> at your service—let’s orchestrate something amazing!</div>
+                            <div class="welcome-message">
+                                Feel free to ask questions, explore ideas, or engage in meaningful conversations.
 
-                                    </div>
-                                               <div class="welcome-message">
-                                                Whether you need assistance, inspiration, or just some lighthearted chat, I'm here to help.
+                            </div>
+                            <div class="welcome-message">
+                                Whether you need assistance, inspiration, or just some lighthearted chat, I'm here to help.
 
-                                    </div>
+                            </div>
 
                         </div>
 
@@ -79,7 +79,7 @@
                 Tokens: @ViewModel.ConversationListViewModel.CurrentConversation.LMKitConversation.ContextUsedSpace /
                 @ViewModel.ConversationListViewModel.CurrentConversation.LMKitConversation.ContextSize
                 (@CalculateUsagePercentage(ViewModel.ConversationListViewModel.CurrentConversation.LMKitConversation.ContextUsedSpace,
-                        ViewModel.ConversationListViewModel.CurrentConversation.LMKitConversation.ContextSize)%)
+                                        ViewModel.ConversationListViewModel.CurrentConversation.LMKitConversation.ContextSize)%)
             </text>
         }
     </div>
@@ -128,8 +128,7 @@
         }
 
         ViewModel.ConversationListViewModel.PropertyChanged += OnConversationListViewModelPropertyChanged;
-
         ViewModel.ConversationListViewModel.CurrentConversation.LMKitConversation.PropertyChanged += OnLMKitConversationPropertyChanged;
 
         await ResizeHandler.RegisterPageResizeAsync(Resized);
         await JS.InvokeVoidAsync("initializeScrollHandler", DotNetObjectReference.Create(this));
@@ -219,13 +218,16 @@
         if (_previousConversationViewModel != null)
         {
             _previousConversationViewModel.Messages.CollectionChanged -= OnConversationMessagesCollectionChanged;
+            _previousConversationViewModel.TextGenerationCompleted -= OnTextGenerationCompleted;
         }
 
         _previousConversationViewModel = ViewModel.ConversationListViewModel.CurrentConversation;
 
         if (ViewModel.ConversationListViewModel.CurrentConversation != null)
         {
             ViewModel.ConversationListViewModel.CurrentConversation.Messages.CollectionChanged += OnConversationMessagesCollectionChanged;
+            ViewModel.ConversationListViewModel.CurrentConversation.TextGenerationCompleted += OnTextGenerationCompleted;
+
             _previousScrollTop = null;
             _ignoreScrollsUntilNextScrollUp = true;
             IsScrolledToEnd = true;
@@ -251,6 +253,18 @@
         }
     }
 
+    private void OnTextGenerationCompleted(object? sender, ConversationViewModel.TextGenerationCompletedEventArgs e)
+    {
+        if (e.Exception != null)
+        {
+            Snackbar.Clear();
+            Snackbar.Configuration.PositionClass = Defaults.Classes.Position.BottomRight;
+            Snackbar.Add($"Text generation failed unexpectedly:\n{e.Exception.Message}", Severity.Error);
+        }
+
+        UpdateUIAsync();
+    }
+
     private async void OnLatestAssistantResponseProgressed()
     {
         if (_shouldAutoScrollEnd)
@@ -305,7 +319,7 @@
     }
 
 
     [JSInvokable]
     public async Task OnConversationContainerScrolled(double scrollTop)
     {
         _scrollTop = scrollTop;

LM-Kit-Maestro/UI/Razor/Components/ChatMessage.razor

@@ -279,17 +279,17 @@
         return html;
     }
 
-    private static string GetStatusText(LMKitTextGenerationStatus status)
+    private static string GetStatusText(LMKitRequestStatus status)
     {
         switch (status)
         {
             default:
                 return string.Empty;
 
-            case LMKitTextGenerationStatus.Cancelled:
+            case LMKitRequestStatus.Cancelled:
                 return "Cancelled";
 
-            case LMKitTextGenerationStatus.GenericError:
+            case LMKitRequestStatus.GenericError:
                 return "Unknown error";
         }
     }

LM-Kit-Maestro/ViewModels/ConversationViewModel.cs

@@ -36,7 +36,7 @@ public partial class ConversationViewModel : AssistantViewModelBase
     bool _isInitialized;
 
     [ObservableProperty]
-    public LMKitTextGenerationStatus _latestPromptStatus;
+    public LMKitRequestStatus _latestPromptStatus;
 
     [ObservableProperty]
     bool _isSelected;
@@ -87,8 +87,9 @@ public Uri? LastUsedModel
         }
     }
 
-    public EventHandler? TextGenerationCompleted;
-    public EventHandler? TextGenerationFailed;
+    public delegate void TextGenerationCompletedEventHandler(object sender, TextGenerationCompletedEventArgs e);
+
+    public TextGenerationCompletedEventHandler? TextGenerationCompleted;
     public EventHandler? DatabaseSaveOperationCompleted;
     public EventHandler? DatabaseSaveOperationFailed;
 
@@ -239,7 +240,7 @@ private void OnNewlySubmittedPrompt()
     {
         InputText = string.Empty;
         UsedDifferentModel &= false;
-        LatestPromptStatus = LMKitTextGenerationStatus.Undefined;
+        LatestPromptStatus = LMKitRequestStatus.OK;
         AwaitingResponse = true;
     }
 
@@ -249,26 +250,15 @@ private void OnTextGenerationResult(LMKitService.LMKitResult? result, Exception?
 
         if (Messages.Count >= 2)
         {
-            // An error may occur before messages consecutive from a prompt have been added to the list, add count check.  
-            Messages.Last().Status = result != null ? result.Status : exception is OperationCanceledException ? LMKitTextGenerationStatus.Cancelled : LMKitTextGenerationStatus.GenericError;
+            // Setting error status for the last assistant message if the response generation failed.
+            Messages.Last().Status = result != null ? result.Status : exception is OperationCanceledException ? LMKitRequestStatus.Cancelled : LMKitRequestStatus.GenericError;
         }
 
-        if (exception != null || result?.Exception != null)
-        {
-            // todo: provide more error info with event args.
-            OnTextGenerationFailure();
-        }
-        else if (result != null)
-        {
-            if (result.Status == LMKitTextGenerationStatus.Undefined && result.Result is TextGenerationResult textGenerationResult)
-            {
-                OnTextGenerationSuccess(textGenerationResult);
-            }
-            else
-            {
-                OnTextGenerationFailure();
-            }
-        }
+        var textGenerationResult = result?.Result is TextGenerationResult ? (TextGenerationResult)result.Result : null;
+
+        TextGenerationCompleted?.Invoke(this, 
+            new TextGenerationCompletedEventArgs(result?.Result is TextGenerationResult ? (TextGenerationResult)result.Result : null,
+            exception ?? (result?.Exception), result?.Status));
 
         if (!_isSynchedWithLog)
         {
@@ -299,17 +289,6 @@ private void SaveConversation()
         });
     }
 
-
-    private void OnTextGenerationSuccess(TextGenerationResult result)
-    {
-        TextGenerationCompleted?.Invoke(this, new TextGenerationCompletedEventArgs(result.TerminationReason));
-    }
-
-    private void OnTextGenerationFailure()
-    {
-        TextGenerationFailed?.Invoke(this, EventArgs.Empty);
-    }
-
     private void OnLMKitChatHistoryChanged(object? sender, NotifyCollectionChangedEventArgs e)
     {
         _isSynchedWithLog &= false;
@@ -380,11 +359,17 @@ private void OnLMKitConversationPropertyChanged(object? sender, System.Component
 
     public sealed class TextGenerationCompletedEventArgs : EventArgs
     {
-        public TextGenerationResult.StopReason StopReason { get; }
+        public Exception? Exception { get; }
+
+        public LMKitRequestStatus? Status { get; }
+
+        public TextGenerationResult? Result { get; }
 
-        public TextGenerationCompletedEventArgs(TextGenerationResult.StopReason stopReason)
+        public TextGenerationCompletedEventArgs(TextGenerationResult? result = null, Exception? exception = null, LMKitRequestStatus? status = null)
         {
-            StopReason = stopReason;
+            Result = result;
+            Exception = exception;
+            Status = status;
         }
     }
 }

LM-Kit-Maestro/ViewModels/MessageViewModel.cs

@@ -23,7 +23,7 @@ public partial class MessageViewModel : ViewModelBase
     private bool _messageInProgress;
 
     [ObservableProperty]
-    private LMKitTextGenerationStatus _status;
+    private LMKitRequestStatus _status;
 
     [ObservableProperty]
     private bool _isHovered;

tests/LmKitServiceTests.cs

@@ -89,7 +89,7 @@ public async Task HonorsTimeout()
         var response = await testService.LMKitService.SubmitPrompt(conversation, "tell me a story");
 
         Assert.NotNull(response);
-        Assert.Equal(LMKitTextGenerationStatus.Cancelled, response.Status);
+        Assert.Equal(LMKitRequestStatus.Cancelled, response.Status);
         Assert.True(response.Exception is OperationCanceledException operationCancelled);
     }
 
@@ -167,7 +167,7 @@ private async Task SubmitOnePromptChangeModelSubmitAnother2()
 
         var firstPromptResult = await firstResponseTask;
 
-        Assert.True(firstPromptResult.Status == LMKitTextGenerationStatus.Cancelled || firstPromptResult.Status == LMKitTextGenerationStatus.GenericError);
+        Assert.True(firstPromptResult.Status == LMKitRequestStatus.Cancelled || firstPromptResult.Status == LMKitRequestStatus.GenericError);
 
         loadingSuccess = await secondModelLoadingTask;
         Assert.True(loadingSuccess);
@@ -193,10 +193,10 @@ private async Task Submit2PromptsFromDistinctConversationsThenUnloadModel()
         Assert.True(unloadingSuccess);
 
         var result = await conversation1.PromptResultTask.Task;
-        Assert.True(result != null && result.Status == LMKitTextGenerationStatus.Cancelled);
+        Assert.True(result != null && result.Status == LMKitRequestStatus.Cancelled);
 
         result = await conversation2.PromptResultTask.Task;
-        Assert.True(result != null && result.Status == LMKitTextGenerationStatus.Cancelled);
+        Assert.True(result != null && result.Status == LMKitRequestStatus.Cancelled);
     }
 
     [Fact]
@@ -214,7 +214,7 @@ private async Task SubmitOnePromptThenUnloadModel()
         Assert.True(unloadingSuccess);
 
         var result = await conversation1.PromptResultTask.Task;
-        Assert.True(result != null && result.Status == LMKitTextGenerationStatus.Cancelled);
+        Assert.True(result != null && result.Status == LMKitRequestStatus.Cancelled);
     }
 
     [Fact]