From 37762cfbabce82d7ad2c26cf0f9047c5a97e840f Mon Sep 17 00:00:00 2001
From: Jean Schmitz <jean@schmitzonline.name>
Date: Tue, 6 Feb 2024 17:26:15 +0100
Subject: [PATCH] Added Button to cancel run requests

---
 .../controller/AssistantController.java       |  8 +++++++-
 .../assistant/service/AssistantService.java   |  6 +++++-
 frontend/src/components/chat/ChatControl.vue  | 19 +++++++++++++------
 frontend/src/components/chat/ChatMessage.vue  |  9 +++++++--
 frontend/src/service/assistant.service.ts     |  6 +++++-
 frontend/src/store/chat-store.ts              | 18 +++++++++++++++---
 .../assistant/OpenAIAssistantService.java     | 15 +++++++++++++--
 7 files changed, 65 insertions(+), 16 deletions(-)

diff --git a/backend/src/main/java/com/talkforgeai/backend/assistant/controller/AssistantController.java b/backend/src/main/java/com/talkforgeai/backend/assistant/controller/AssistantController.java
index b7a38382..e13758d2 100644
--- a/backend/src/main/java/com/talkforgeai/backend/assistant/controller/AssistantController.java
+++ b/backend/src/main/java/com/talkforgeai/backend/assistant/controller/AssistantController.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023 Jean Schmitz.
+ * Copyright (c) 2023-2024 Jean Schmitz.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -183,6 +183,12 @@ public Run getRun(@PathVariable("threadId") String threadId,
     return assistantService.retrieveRun(threadId, runId);
   }
 
+  @PostMapping("/threads/{threadId}/runs/{runId}/cancel")
+  public Run cancelRun(@PathVariable("threadId") String threadId,
+      @PathVariable("runId") String runId) {
+    return assistantService.cancelRun(threadId, runId);
+  }
+
   @PostMapping("/threads/{threadId}/messages/{messageId}/postprocess")
   public ParsedMessageDto postProcessMessage(@PathVariable("threadId") String threadId,
       @PathVariable("messageId") String messageId) {
diff --git a/backend/src/main/java/com/talkforgeai/backend/assistant/service/AssistantService.java b/backend/src/main/java/com/talkforgeai/backend/assistant/service/AssistantService.java
index c4a6af7c..b86dc24d 100644
--- a/backend/src/main/java/com/talkforgeai/backend/assistant/service/AssistantService.java
+++ b/backend/src/main/java/com/talkforgeai/backend/assistant/service/AssistantService.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023 Jean Schmitz.
+ * Copyright (c) 2023-2024 Jean Schmitz.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -239,6 +239,10 @@ public Run retrieveRun(String threadId, String runId) {
     return this.openAIAssistantService.retrieveRun(threadId, runId);
   }
 
+  public Run cancelRun(String threadId, String runId) {
+    return openAIAssistantService.cancelRun(threadId, runId);
+  }
+
   private ThreadDto mapToDto(ThreadEntity threadEntity) {
     return new ThreadDto(threadEntity.getId(), threadEntity.getTitle(),
         threadEntity.getCreatedAt());
diff --git a/frontend/src/components/chat/ChatControl.vue b/frontend/src/components/chat/ChatControl.vue
index 2474300f..fad2c5cd 100644
--- a/frontend/src/components/chat/ChatControl.vue
+++ b/frontend/src/components/chat/ChatControl.vue
@@ -17,9 +17,9 @@
 <template>
   <div class="p-2 m-1 rounded shadow">
 
-    <div class="row">
-      <div class="col-12">
-        <div class="form-check form-switch d-flex switch-panel">
+    <div class="d-flex flex-row">
+      <div class="flex-grow-1">
+        <div class="form-check form-switch d-flex switch-panel mt-3">
           <input id="flexCheckDefault" v-model="localIsAutoSpeak"
                  class="form-check-input" role="switch" type="checkbox">
           <label class="form-check-label mx-2" for="flexCheckDefault">
@@ -27,6 +27,9 @@
           </label>
         </div>
       </div>
+      <button :hidden="isCancelRequestHidden" class="btn btn-outline-warning my-2"
+              @click="onCancelRun()">Cancel request
+      </button>
     </div>
     <div class="row">
       <ChatMessageInput @submit-result-received="submitResultReceived"
@@ -48,9 +51,9 @@ export default {
       isAutoSpeak: Boolean,
     };
   },
-  getters: {
-    isAutoSpeak() {
-      return this.store.autoSpeak;
+  computed: {
+    isCancelRequestHidden() {
+      return this.store.runId === '';
     },
   },
   methods: {
@@ -62,6 +65,10 @@ export default {
       console.log('Chat Control - Chunk Update Received');
       this.$emit('chunkUpdateReceived');
     },
+    onCancelRun() {
+      console.log('Chat Control - Cancel Run');
+      this.store.cancelCurrentRun();
+    },
   },
   setup() {
     const store = useChatStore(); // Call useMyStore() inside the setup function
diff --git a/frontend/src/components/chat/ChatMessage.vue b/frontend/src/components/chat/ChatMessage.vue
index e30abe36..781343cb 100644
--- a/frontend/src/components/chat/ChatMessage.vue
+++ b/frontend/src/components/chat/ChatMessage.vue
@@ -36,6 +36,9 @@ export default {
     messageIndex: Number,
   },
   computed: {
+    errorMessage() {
+      return this.store.currentStatusMessage;
+    },
     personaImage() {
       // Find the assistant in the store
       const assistantId = this.message.assistant_id;
@@ -131,8 +134,10 @@ export default {
           <div v-if="getMessageStatusType() === 'running'" class="spinner-grow text-primary"
                role="status">
           </div>
-          <i v-else-if="getMessageStatusType() === 'error'"
-             class="bi bi-exclamation-lg bg-danger"></i>
+          <span v-else-if="getMessageStatusType() === 'error'">
+            <i class="bi bi-exclamation-lg bg-danger"></i>
+            {{ errorMessage }}
+          </span>
           <div v-else class="card-text text-start" v-html="getContent()"></div>
         </div>
       </div>
diff --git a/frontend/src/service/assistant.service.ts b/frontend/src/service/assistant.service.ts
index 9f1112cb..8dd42d7c 100644
--- a/frontend/src/service/assistant.service.ts
+++ b/frontend/src/service/assistant.service.ts
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023 Jean Schmitz.
+ * Copyright (c) 2023-2024 Jean Schmitz.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -193,6 +193,10 @@ class AssistantService {
     return result.data;
   }
 
+  cancelRun(threadId: string, id: string) {
+    return axios.post(`/api/v1/threads/${threadId}/runs/${id}/cancel`);
+  }
+
   private async retrieveThread(threadId: string) {
     const result = await axios.get(
         `/api/v1/threads/${threadId}`,
diff --git a/frontend/src/store/chat-store.ts b/frontend/src/store/chat-store.ts
index 913dd2b6..3764a8f3 100644
--- a/frontend/src/store/chat-store.ts
+++ b/frontend/src/store/chat-store.ts
@@ -33,6 +33,7 @@ export const useChatStore = defineStore('chat', {
       threadMessages: [] as Array<ThreadMessage>,
       parsedMessages: {} as any,
       threadId: '',
+      runId: '',
       threads: [] as Array<Thread>,
       threadEditMode: false,
       threadDeleteMode: false,
@@ -72,8 +73,7 @@ export const useChatStore = defineStore('chat', {
     },
     async retrieveMessages(threadId: string) {
       const parsedMessageList = await assistantService.retrieveMessages(threadId);
-      console.log('Parsed message list', parsedMessageList);
-      // Make this.threadMessages a list of ThreadMessage objects
+
       const list = parsedMessageList.message_list?.data?.map((message: any) => {
         return new ThreadMessage(message.id, message.role, message.content[0].text.value,
             message.assistant_id);
@@ -115,13 +115,21 @@ export const useChatStore = defineStore('chat', {
       try {
         const run = await assistantService.runConversation(this.threadId,
             this.selectedAssistant.id);
+        this.runId = run.id;
+
         pollingInterval = setInterval(async () => {
           console.log('Polling for run status');
           const runT = await assistantService.retrieveRun(this.threadId, run.id);
           if (runT.status === 'completed') {
             console.log('Run completed');
             clearInterval(pollingInterval);
+            this.runId = '';
             await this.handleResult();
+          } else if (runT.status === 'cancelled') {
+            console.log('Run cancelled');
+            clearInterval(pollingInterval);
+            this.runId = '';
+            this.updateStatus('Cancelled', 'error');
           }
         }, 2000);
       } catch (e) {
@@ -129,7 +137,11 @@ export const useChatStore = defineStore('chat', {
         if (pollingInterval) clearInterval(pollingInterval);
       }
     },
-
+    async cancelCurrentRun() {
+      if (this.runId.length > 0) {
+        await assistantService.cancelRun(this.threadId, this.runId);
+      }
+    },
     // Usage of the refactored functions
     async submitUserMessage(message: string) {
       await this.createThreadIfNeeded();
diff --git a/service/src/main/java/com/talkforgeai/service/openai/assistant/OpenAIAssistantService.java b/service/src/main/java/com/talkforgeai/service/openai/assistant/OpenAIAssistantService.java
index cc53708a..c82e8129 100644
--- a/service/src/main/java/com/talkforgeai/service/openai/assistant/OpenAIAssistantService.java
+++ b/service/src/main/java/com/talkforgeai/service/openai/assistant/OpenAIAssistantService.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023 Jean Schmitz.
+ * Copyright (c) 2023-2024 Jean Schmitz.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -168,6 +168,10 @@ private Request createPostRequest(String body, String path) {
     return createRequest("POST", path, body);
   }
 
+  private Request createPostRequest(String path) {
+    return createRequest("POST", path, null);
+  }
+
   private Request createGetRequest(String path) {
     return createRequest("GET", path, null);
   }
@@ -180,7 +184,9 @@ private Request createRequest(String method, String path, String body) {
     String apiUrl = openAIProperties.apiUrl() + path;
 
     RequestBody requestBody = null;
-    if (body != null) {
+    if (body == null && method.equals("POST")) {
+      requestBody = RequestBody.create("", null);
+    } else if (body != null) {
       requestBody = RequestBody.create(body, JSON);
     }
 
@@ -201,4 +207,9 @@ public void deleteAssistant(String assistantId) {
     Request request = createDeleteRequest("/assistants/" + assistantId);
     executeRequest(request, Void.class);
   }
+
+  public Run cancelRun(String threadId, String runId) {
+    Request request = createPostRequest("/threads/" + threadId + "/runs/" + runId + "/cancel");
+    return executeRequest(request, Run.class);
+  }
 }