refactor lmkitservice error handling

LM-Kit · Dec 28, 2024 · 09dcaed · 09dcaed
1 parent d25df8b
commit 09dcaed
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 51 deletions.
diff --git a/...tro/Services/LMKitTextGenerationStatus.cs → ...it-Maestro/Services/LMKitRequestStatus.cs b/...tro/Services/LMKitTextGenerationStatus.cs → ...it-Maestro/Services/LMKitRequestStatus.cs
@@ -1,8 +1,8 @@
 namespace LMKit.Maestro.Services;
 
-public enum LMKitTextGenerationStatus
+public enum LMKitRequestStatus
 {
-    Undefined,
+    OK,
     Cancelled,
     GenericError
 }
diff --git a/LM-Kit-Maestro/Services/LMKitService.LMKitResult.cs b/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; }
     }

diff --git a/LM-Kit-Maestro/Services/LMKitService.cs b/LM-Kit-Maestro/Services/LMKitService.cs
@@ -211,15 +211,15 @@ 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();
@@ -255,8 +255,14 @@ 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
         {
@@ -265,6 +271,8 @@ private async Task<LMKitResult> HandleLMKitRequest(LMKitRequest 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

diff --git a/LM-Kit-Maestro/UI/Razor/Components/ChatMessage.razor b/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";
         }
     }

diff --git a/LM-Kit-Maestro/ViewModels/ConversationViewModel.cs b/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;
@@ -239,7 +239,7 @@ private void OnNewlySubmittedPrompt()
     {
         InputText = string.Empty;
         UsedDifferentModel &= false;
-        LatestPromptStatus = LMKitTextGenerationStatus.Undefined;
+        LatestPromptStatus = LMKitRequestStatus.OK;
         AwaitingResponse = true;
     }
 
@@ -249,26 +249,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 +288,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 +358,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;
         }
     }
 }
diff --git a/LM-Kit-Maestro/ViewModels/MessageViewModel.cs b/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;