Feat/abort

.changeset/thin-waves-breathe.md

@@ -0,0 +1,5 @@
+---
+'outpostkit': patch
+---
+
+added abort feature

src/comet.ts

@@ -53,7 +53,8 @@ export class Comet implements IComet {
             this.cometId,
             this.apiKey,
             payload,
-            handleNewText
+            handleNewText,
+            options?.signal
           );
 
           // @ts-ignore
@@ -68,7 +69,8 @@ export class Comet implements IComet {
             this.cometId,
             this.apiKey,
             payload,
-            handleNewText
+            handleNewText,
+            options?.signal
           );
           // @ts-ignore
           if (resp.error) {
@@ -79,7 +81,12 @@ export class Comet implements IComet {
           }
         }
       } else {
-        const resp = await streamPromptWithAxios(this.cometAPI, payload, handleNewText);
+        const resp = await streamPromptWithAxios(
+          this.cometAPI,
+          payload,
+          handleNewText,
+          options?.signal
+        );
         // @ts-ignore
         if (resp.error) {
           // @ts-ignore
@@ -89,7 +96,7 @@ export class Comet implements IComet {
         }
       }
     } else {
-      const { data } = await this.cometAPI.post(`/prompt`, payload);
+      const { data } = await this.cometAPI.post(`/prompt`, payload, { signal: options?.signal });
       return data as TCometPromptResponse;
     }
   }

src/helpers.ts

@@ -19,7 +19,8 @@ export const streamPromptWithNativeFetch = (
   cometId: string,
   apiKey: string,
   payload: PromptPayload,
-  handleNewText?: (token: string) => void | Promise<void>
+  handleNewText?: (token: string) => void | Promise<void>,
+  signal?: AbortSignal
 ) => {
   return new Promise<TCometPromptResponse | TCometPromptStreamResponseError>((resolve, reject) => {
     (async () => {
@@ -31,6 +32,7 @@ export const streamPromptWithNativeFetch = (
           Authorization: `Bearer ${apiKey}`,
           Accept: 'text/plain, application/json',
         },
+        signal,
       });
       if (response.ok) {
         if (!response.body) return reject('No response body found.');
@@ -77,12 +79,14 @@ export const streamPromptWithNativeFetch = (
 export const streamPromptWithAxios = (
   cometAPI: AxiosInstance,
   payload: PromptPayload,
-  handleNewText?: (token: string) => void | Promise<void>
+  handleNewText?: (token: string) => void | Promise<void>,
+  signal?: AbortSignal
 ) => {
   return new Promise<TCometPromptResponse | TCometPromptStreamResponseError>((resolve, reject) => {
     (async () => {
       const { data: stream } = await cometAPI.post(`/prompt`, payload, {
         responseType: 'stream',
+        signal,
       });
       let responsePrefixReceived = false;
       let responseText: string = '';
@@ -117,7 +121,8 @@ export const streamPromptWithEventStreaming = async (
   cometId: string,
   apiKey: string,
   payload: PromptPayload,
-  handleNewText?: (token: string) => void | Promise<void>
+  handleNewText?: (token: string) => void | Promise<void>,
+  signal?: AbortSignal
 ): Promise<TCometPromptResponse | TCometPromptStreamResponseError> => {
   try {
     let finalResponse: TCometPromptResponse | TCometPromptStreamResponseError;
@@ -128,6 +133,7 @@ export const streamPromptWithEventStreaming = async (
         Accept: 'text/event-stream',
         Authorization: `Bearer ${apiKey}`,
       },
+      signal,
       body: JSON.stringify(payload),
       async onopen(response) {
         const contentType = response.headers.get('content-type');
@@ -156,8 +162,7 @@ export const streamPromptWithEventStreaming = async (
           } catch (e) {
             throw new ClientError('Encountered error while parsing response into JSON.');
           }
-        } else if (msg.event==='error') {
-
+        } else if (msg.event === 'error') {
         }
       },
       // onclose() {

src/types/comet.ts

@@ -32,6 +32,7 @@ export interface IComet {
 
 export interface PromptOptions {
   useNativeFetch?: boolean;
+  signal?: AbortSignal;
 }
 
 export interface PromptPayload {

src/utils/inference/vllm-stream.ts

@@ -1,11 +1,12 @@
 import { EventStreamContentType, fetchEventSource } from '@microsoft/fetch-event-source';
 import { APIError } from 'error';
-import { VLLMOpenAICompletionsOutputType, VLLMPromptParameters } from 'types/inference';
+import { VLLMPromptParameters } from 'types/inference';
 
 export const streamGenericInferenceServer = (
   domain: string,
   payload: { prompt: string } & VLLMPromptParameters,
-  handleNewChunk?: (chunk: string) => void | Promise<void>
+  handleNewChunk?: (chunk: string) => void | Promise<void>,
+  signal?: AbortSignal
 ) => {
   return new Promise<string>((resolve, reject) => {
     (async () => {
@@ -16,6 +17,7 @@ export const streamGenericInferenceServer = (
           'Content-Type': 'application/json',
           Accept: 'text/plain, application/json',
         },
+        signal,
       });
       if (response.ok) {
         if (!response.body) return reject('No response body found.');
@@ -88,12 +90,17 @@ export const streamGenericInferenceServer = (
 class ClientError extends Error {}
 // class FatalError extends Error {}
 
-export const streamOpenAIInferenceServer = async (
-  payload: { prompt: string; model: string } & VLLMPromptParameters,
+export type ChatMessage = { role: 'system' | 'user' | 'assistant'; content: string };
+export async function streamOpenAIInferenceServer<T extends 'chat' | 'text'>(
+  payload: {
+    prompt: T extends 'chat' ? ChatMessage[] : T extends 'text' ? string : never;
+    model: string;
+  } & VLLMPromptParameters,
   domain: string,
-  type: 'chat' | 'text',
-  handleNewChunk?: (chunk: string) => void | Promise<void>
-): Promise<string> => {
+  type: T,
+  handleNewChunk?: (chunk: string) => void | Promise<void>,
+  signal?: AbortSignal
+): Promise<string> {
   try {
     let finalResponse: string;
     await fetchEventSource(`${domain}/v1/${type === 'chat' ? 'chat/completions' : 'completions'}`, {
@@ -102,6 +109,7 @@ export const streamOpenAIInferenceServer = async (
         'Content-Type': 'application/json',
         Accept: 'text/event-stream',
       },
+      signal,
       body: JSON.stringify({ ...payload, stream: true }),
       async onopen(response) {
         const contentType = response.headers.get('content-type');