From 886b242b66e8981275c9bacece9c5d4a126164aa Mon Sep 17 00:00:00 2001
From: leo <longshuang@msn.cn>
Date: Fri, 13 Sep 2024 14:25:38 +0800
Subject: [PATCH] enhance: added cancellation logic to the OpenAI chat response
 API call.

---
 src/Commands/GenerateCommitMessage.cs |  4 ++--
 src/Models/OpenAI.cs                  | 29 ++++++++++++++++++---------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/Commands/GenerateCommitMessage.cs b/src/Commands/GenerateCommitMessage.cs
index 1bbf6cd8..b5c978d3 100644
--- a/src/Commands/GenerateCommitMessage.cs
+++ b/src/Commands/GenerateCommitMessage.cs
@@ -78,7 +78,7 @@ namespace SourceGit.Commands
             prompt.AppendLine("- Simply describe the MAIN GOAL of the changes.");
             prompt.AppendLine("- Output directly the summary in plain text.`");
 
-            var rsp = Models.OpenAI.Chat(prompt.ToString(), $"Here is the `git diff` output: {diff}");
+            var rsp = Models.OpenAI.Chat(prompt.ToString(), $"Here is the `git diff` output: {diff}", _cancelToken);
             if (rsp != null && rsp.Choices.Count > 0)
                 return rsp.Choices[0].Message.Content;
 
@@ -104,7 +104,7 @@ namespace SourceGit.Commands
             prompt.AppendLine("- Output directly only one commit message in plain text with the next format: {type}: {commit_message}.");
             prompt.AppendLine("- Be as concise as possible, keep the message under 50 characters.");
 
-            var rsp = Models.OpenAI.Chat(prompt.ToString(), $"Here are the summaries changes: {summary}");
+            var rsp = Models.OpenAI.Chat(prompt.ToString(), $"Here are the summaries changes: {summary}", _cancelToken);
             if (rsp != null && rsp.Choices.Count > 0)
                 return rsp.Choices[0].Message.Content;
 
diff --git a/src/Models/OpenAI.cs b/src/Models/OpenAI.cs
index 4acbe371..94a64189 100644
--- a/src/Models/OpenAI.cs
+++ b/src/Models/OpenAI.cs
@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.Net.Http;
 using System.Text.Json;
 using System.Text.Json.Serialization;
+using System.Threading;
 
 namespace SourceGit.Models
 {
@@ -97,7 +98,7 @@ namespace SourceGit.Models
             get => !string.IsNullOrEmpty(Server) && !string.IsNullOrEmpty(ApiKey) && !string.IsNullOrEmpty(Model);
         }
 
-        public static OpenAIChatResponse Chat(string prompt, string question)
+        public static OpenAIChatResponse Chat(string prompt, string question, CancellationToken cancellation)
         {
             var chat = new OpenAIChatRequest() { Model = Model };
             chat.AddMessage("system", prompt);
@@ -107,17 +108,27 @@ namespace SourceGit.Models
             client.DefaultRequestHeaders.Add("Authorization", $"Bearer {ApiKey}");
 
             var req = new StringContent(JsonSerializer.Serialize(chat, JsonCodeGen.Default.OpenAIChatRequest));
-            var task = client.PostAsync(Server, req);
-            task.Wait();
+            try
+            {
+                var task = client.PostAsync(Server, req, cancellation);
+                task.Wait();
 
-            var rsp = task.Result;
-            if (!rsp.IsSuccessStatusCode)
-                throw new Exception($"AI service returns error code {rsp.StatusCode}");
+                var rsp = task.Result;
+                if (!rsp.IsSuccessStatusCode)
+                    throw new Exception($"AI service returns error code {rsp.StatusCode}");
 
-            var reader = rsp.Content.ReadAsStringAsync();
-            reader.Wait();
+                var reader = rsp.Content.ReadAsStringAsync(cancellation);
+                reader.Wait();
 
-            return JsonSerializer.Deserialize(reader.Result, JsonCodeGen.Default.OpenAIChatResponse);
+                return JsonSerializer.Deserialize(reader.Result, JsonCodeGen.Default.OpenAIChatResponse);
+            } 
+            catch
+            {
+                if (cancellation.IsCancellationRequested)
+                    return null;
+
+                throw;
+            }
         }
     }
 }