From 664806e7f162b8f97531e3c764fa3f682ab8bf83 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Tue, 21 Nov 2023 20:49:17 +0700
Subject: [PATCH] chore: abort inference stream implementation

---
 plugins/inference-plugin/src/helpers/sse.ts | 3 ++-
 plugins/inference-plugin/src/index.ts       | 7 +++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/plugins/inference-plugin/src/helpers/sse.ts b/plugins/inference-plugin/src/helpers/sse.ts
index f63cc027b..987751221 100644
--- a/plugins/inference-plugin/src/helpers/sse.ts
+++ b/plugins/inference-plugin/src/helpers/sse.ts
@@ -4,7 +4,7 @@ import { Observable } from "rxjs";
  * @param recentMessages - An array of recent messages to use as context for the inference.
  * @returns An Observable that emits the generated response as a string.
  */
-export function requestInference(recentMessages: any[]): Observable<string> {
+export function requestInference(recentMessages: any[], controller?: AbortController): Observable<string> {
   return new Observable((subscriber) => {
     const requestBody = JSON.stringify({
       messages: recentMessages,
@@ -20,6 +20,7 @@ export function requestInference(recentMessages: any[]): Observable<string> {
         "Access-Control-Allow-Origin": "*",
       },
       body: requestBody,
+      signal: controller?.signal
     })
       .then(async (response) => {
         const stream = response.body;
diff --git a/plugins/inference-plugin/src/index.ts b/plugins/inference-plugin/src/index.ts
index 0a02d4de0..ef9409169 100644
--- a/plugins/inference-plugin/src/index.ts
+++ b/plugins/inference-plugin/src/index.ts
@@ -28,6 +28,7 @@ import { fs } from "@janhq/core";
  * It also subscribes to events emitted by the @janhq/core package and handles new message requests.
  */
 export default class JanInferencePlugin implements InferencePlugin {
+  controller = new AbortController();
   /**
    * Returns the type of the plugin.
    * @returns {PluginType} The type of the plugin.
@@ -75,7 +76,7 @@ export default class JanInferencePlugin implements InferencePlugin {
    * @returns {Promise<void>} A promise that resolves when the streaming is stopped.
    */
   async stopInference(): Promise<void> {
-    // TODO: Implementation
+    this.controller.abort();
   }
 
   /**
@@ -121,7 +122,9 @@ export default class JanInferencePlugin implements InferencePlugin {
     };
     events.emit(EventName.OnNewMessageResponse, message);
 
-    requestInference(data.messages).subscribe({
+    this.controller = new AbortController();
+
+    requestInference(data.messages, this.controller).subscribe({
       next: (content) => {
         message.content = content;
         events.emit(EventName.OnMessageResponseUpdate, message);