Add regenerate last response feature (#46)

Co-authored-by: Loïc Carrère <[email protected]>
LM-Kit · Dec 7, 2024 · dbb6b57 · dbb6b57
1 parent 84401b3
commit dbb6b57
Show file tree

Hide file tree

Showing 11 changed files with 312 additions and 207 deletions.
diff --git a/LM-Kit-Maestro/Models/Message.cs b/LM-Kit-Maestro/Models/Message.cs
diff --git a/LM-Kit-Maestro/Services/Enumerations.cs b/LM-Kit-Maestro/Services/Enumerations.cs
@@ -18,5 +18,5 @@ public enum LMKitTextGenerationStatus
 {
     Undefined,
     Cancelled,
-    UnknownError
+    GenericError
 }
diff --git a/LM-Kit-Maestro/Services/LMKitService.cs b/LM-Kit-Maestro/Services/LMKitService.cs
@@ -3,6 +3,7 @@
 using LMKit.TextGeneration.Sampling;
 using LMKit.TextGeneration.Chat;
 using LMKit.Translation;
+using System.Diagnostics;
 
 namespace LMKit.Maestro.Services;
 
@@ -171,10 +172,10 @@ public async Task<LMKitResult> SubmitPrompt(Conversation conversation, string pr
 
     public async Task<LMKitResult> RegenerateResponse(Conversation conversation, ChatHistory.Message message)
     {
-        // Ignoring message parameter, only regenerate the latest response for now.
-        var prompt = conversation.ChatHistory!.Messages[conversation.ChatHistory.Messages.Count - 1];
+        var regenerateResponseRequest = new LMKitRequest(LMKitRequest.LMKitRequestType.RegenerateResponse,
+            new LMKitRequest.RegenerateResponseParameters(conversation, message), LMKitConfig.RequestTimeout);
 
-        var regenerateResponseRequest = new LMKitRequest(LMKitRequest.LMKitRequestType.RegenerateResponse, prompt.Content, LMKitConfig.RequestTimeout);
+        ScheduleRequest(regenerateResponseRequest);
 
         return await HandleLmKitRequest(regenerateResponseRequest);
     }
@@ -184,14 +185,17 @@ public async Task CancelPrompt(Conversation conversation, bool shouldAwaitTermin
         var conversationPrompt = _requestSchedule.Unschedule(conversation);
 
         if (conversationPrompt != null)
-        {
+        {
+            _lmKitServiceSemaphore.Wait();
             conversationPrompt.CancellationTokenSource.Cancel();
             conversationPrompt.ResponseTask.TrySetCanceled();
+            _lmKitServiceSemaphore.Release();
 
             if (shouldAwaitTermination)
             {
                 await conversationPrompt.ResponseTask.Task.WaitAsync(TimeSpan.FromSeconds(10));
             }
+
         }
     }
 
@@ -205,14 +209,14 @@ private void ScheduleRequest(LMKitRequest request)
         }
     }
 
-    private async Task<LMKitResult> HandleLmKitRequest(LMKitRequest promptRequest)
+    private async Task<LMKitResult> HandleLmKitRequest(LMKitRequest request)
     {
-        // Ensuring we don't touch anything until Lm-Kit objects' state has been set to handle this prompt request.
+        // Ensuring we don't touch anything until Lm-Kit objects' state has been set to handle this request.
         _lmKitServiceSemaphore.Wait();
 
         LMKitResult result;
 
-        if (promptRequest.CancellationTokenSource.IsCancellationRequested || ModelLoadingState == LMKitModelLoadingState.Unloaded)
+        if (request.CancellationTokenSource.IsCancellationRequested || ModelLoadingState == LMKitModelLoadingState.Unloaded)
         {
             result = new LMKitResult()
             {
@@ -223,22 +227,27 @@ private async Task<LMKitResult> HandleLmKitRequest(LMKitRequest promptRequest)
         }
         else
         {
-            if (promptRequest.RequestType == LMKitRequest.LMKitRequestType.Prompt)
+            if (request.RequestType == LMKitRequest.LMKitRequestType.Prompt || request.RequestType == LMKitRequest.LMKitRequestType.RegenerateResponse)
             {
-                BeforeSubmittingPrompt(((LMKitRequest.PromptRequestParameters)promptRequest.Parameters!).Conversation);
+                var conversation = request.RequestType == LMKitRequest.LMKitRequestType.Prompt ?
+                    ((LMKitRequest.PromptRequestParameters)request.Parameters!).Conversation :
+                    ((LMKitRequest.RegenerateResponseParameters)request.Parameters!).Conversation;
+
+                BeforeSubmittingPrompt(conversation);
             }
 
             _lmKitServiceSemaphore.Release();
 
-            result = await SubmitRequest(promptRequest);
+
+            result = await SubmitRequest(request);
         }
 
-        if (_requestSchedule.Contains(promptRequest))
+        if (_requestSchedule.Contains(request))
         {
-            _requestSchedule.Remove(promptRequest);
+            _requestSchedule.Remove(request);
         }
 
-        promptRequest.ResponseTask.TrySetResult(result);
+        request.ResponseTask.TrySetResult(result);
 
         return result;
 
@@ -271,7 +280,7 @@ private async Task<LMKitResult> SubmitRequest(LMKitRequest request)
                 }
             }
             catch (Exception exception)
-            {
+            { 
                 result.Exception = exception;
 
                 if (result.Exception is OperationCanceledException)
@@ -280,7 +289,7 @@ private async Task<LMKitResult> SubmitRequest(LMKitRequest request)
                 }
                 else
                 {
-                    result.Status = LMKitTextGenerationStatus.UnknownError;
+                    result.Status = LMKitTextGenerationStatus.GenericError;
                 }
             }
 
@@ -291,7 +300,7 @@ private async Task<LMKitResult> SubmitRequest(LMKitRequest request)
                 parameter.Conversation.ChatHistory = _multiTurnConversation.ChatHistory;
                 parameter.Conversation.LatestChatHistoryData = _multiTurnConversation.ChatHistory.Serialize();
 
-                if (parameter.Conversation.GeneratedTitleSummary == null && result.Status == LMKitTextGenerationStatus.Undefined 
+                if (parameter.Conversation.GeneratedTitleSummary == null && result.Status == LMKitTextGenerationStatus.Undefined
                     && !string.IsNullOrEmpty(((TextGenerationResult)result.Result!).Completion))
                 {
                     GenerateConversationSummaryTitle(parameter.Conversation, parameter.Prompt);
@@ -310,7 +319,7 @@ private async Task<LMKitResult> SubmitRequest(LMKitRequest request)
             return new LMKitResult()
             {
                 Exception = exception,
-                Status = LMKitTextGenerationStatus.UnknownError
+                Status = LMKitTextGenerationStatus.GenericError
             };
         }
         finally

diff --git a/LM-Kit-Maestro/Services/LmKitService.LmKitRequest.cs b/LM-Kit-Maestro/Services/LmKitService.LmKitRequest.cs
@@ -1,4 +1,5 @@
 using LMKit.TextGeneration;
+using LMKit.TextGeneration.Chat;
 using System.ComponentModel;
 
 namespace LMKit.Maestro.Services;
@@ -48,6 +49,19 @@ public PromptRequestParameters(Conversation conversation, string prompt)
             }
         }
 
+        public sealed class RegenerateResponseParameters
+        {
+            public Conversation Conversation { get; set; }
+
+            public ChatHistory.Message Message { get; set; }
+
+            public RegenerateResponseParameters(Conversation conversation, ChatHistory.Message message)
+            {
+                Conversation = conversation;
+                Message = message;
+            }
+        }
+
         public sealed class TranslationRequestParameters
         {
             public string InputText { get; set; }

diff --git a/LM-Kit-Maestro/Services/LmKitService.RequestSchedule.cs b/LM-Kit-Maestro/Services/LmKitService.RequestSchedule.cs
@@ -79,7 +79,12 @@ public void Remove(LMKitRequest scheduledPrompt)
             {
                 foreach (var scheduledPrompt in _scheduledPrompts)
                 {
-                    if (scheduledPrompt.Parameters is LMKitRequest.PromptRequestParameters parameter && parameter.Conversation == conversation)
+                    if (scheduledPrompt.Parameters is LMKitRequest.PromptRequestParameters promptParameters && promptParameters.Conversation == conversation)
+                    {
+                        prompt = scheduledPrompt;
+                        break;
+                    }
+                    else if (scheduledPrompt.Parameters is LMKitRequest.RegenerateResponseParameters regenerateParameters && regenerateParameters.Conversation == conversation)
                     {
                         prompt = scheduledPrompt;
                         break;