Add cancelStream function.

Author: cjcenizal

src/apiV1/client.ts

@@ -93,7 +93,7 @@ export const streamQueryV1 = async (
 
   const url = config.endpoint ?? `${DEFAULT_DOMAIN}/v1/stream-query`;
 
-  const stream = await generateStream(requestHeaders, requestBody, url);
+  const { stream } = await generateStream(requestHeaders, requestBody, url);
 
   let previousAnswerText = "";
 

src/apiV2/client.ts

@@ -143,17 +143,32 @@ export const streamQueryV2 = async (
   const url = `${endpoint ?? DEFAULT_DOMAIN}${path}`;
 
   try {
-    const stream = await generateStream(headers, JSON.stringify(body), url);
-    const buffer = new EventBuffer(onStreamEvent);
+    const { cancelStream, stream } = await generateStream(
+      headers,
+      JSON.stringify(body),
+      url
+    );
 
-    for await (const chunk of stream) {
+    new Promise(async (resolve, reject) => {
       try {
-        buffer.consumeChunk(chunk);
-        buffer.drainEvents();
+        const buffer = new EventBuffer(onStreamEvent);
+
+        for await (const chunk of stream) {
+          try {
+            buffer.consumeChunk(chunk);
+            buffer.drainEvents();
+          } catch (error) {
+            console.log("error", error);
+          }
+        }
+
+        resolve(undefined);
       } catch (error) {
-        console.log("error", error);
+        reject(error);
       }
-    }
+    });
+
+    return { cancelStream };
   } catch (error) {
     console.log("error", error);
   }

src/common/generateStream.ts

@@ -2,15 +2,23 @@ export const generateStream = async (
   headers: Record<string, string>,
   body: string,
   url: string
-): Promise<AsyncIterable<string>> => {
+) => {
+  let controller = new AbortController();
+
   const response = await fetch(url, {
     method: "POST",
     headers,
     body,
+    signal: controller.signal,
   });
+
   if (response.status !== 200) throw new Error(response.status.toString());
   if (!response.body) throw new Error("Response body does not exist");
-  return getIterableStream(response.body);
+
+  return {
+    stream: getIterableStream(response.body),
+    cancelStream: () => controller.abort(),
+  };
 };
 
 async function* getIterableStream(