add chat cancel method

README.md

@@ -147,6 +147,30 @@ try {
 }
 ```
 
+Alternatively the `Chat` instance has a method `cancelStream` which which allows equivalent functionality to the above without having to use the `cancel` property of the `onUpdate` callback. (This is useful if you want to have a way to cancel the stream from something removed for the details of handling the streams.)
+
+```ts
+import { createChat, CancelledCompletionError } from "completions";
+
+const chat = createChat({
+  apiKey: OPENAI_API_KEY,
+  model: "gpt-3.5-turbo",
+});
+
+try {
+  const responsePromise = chat.sendMessage("continue sequence: a b c");
+  // in practice this would be linked to something external
+  chat.cancelStream();
+  await responsePromise;
+} catch (error) {
+  if (error instanceof CancelledCompletion) {
+    console.log("cancelled");
+  }
+
+  throw error;
+}
+```
+
 ### Overriding API
 
 If you want to use `completions` library against another API endpoint that is compatible with the official API, you can do so by passing `apiUrl` parameter:

src/createChat.test.ts

@@ -103,6 +103,14 @@ test("cancel response", async () => {
       },
     })
   );
+
+  const chatPromise = chat.sendMessage("continue sequence: a b c");
+  chat.cancelStream();
+  await assert.rejects(chatPromise);
+
+  // make sure that the chat still functions after cancelling
+  const response = await chat.sendMessage("continue sequence: a b c");
+  assert(response.content);
 });
 
 test("calls user defined function", async () => {

src/createChat.ts

@@ -56,6 +56,7 @@ export type Chat = {
   addMessage: (message: Message) => void;
   getMessages: () => Message[];
   sendMessage: SendMessage;
+  cancelStream: () => void;
 };
 
 /**
@@ -97,6 +98,7 @@ export const createChat = (
   options: Omit<CompletionsOptions, "messages" | "n" | "onUpdate">
 ): Chat => {
   const messages: Message[] = [];
+  const cancel = { cancelled: false };
 
   const userFunctions: Record<string, UserFunction> = {};
 
@@ -106,12 +108,15 @@ export const createChat = (
 
   const complete = async (messageOptions?: MessageOptions) => {
     const response = await retry(() => {
-      return createCompletions({
-        messages,
-        onUpdate: messageOptions?.onUpdate,
-        ...options,
-        ...(messageOptions ? omit(messageOptions, "expect") : {}),
-      });
+      return createCompletions(
+        {
+          messages,
+          onUpdate: messageOptions?.onUpdate,
+          ...options,
+          ...(messageOptions ? omit(messageOptions, "expect") : {}),
+        },
+        cancel
+      );
     });
 
     if (!response) {
@@ -173,6 +178,8 @@ export const createChat = (
       role: "user",
     };
 
+    cancel.cancelled = false;
+
     messages.push(message);
 
     let choice = await complete(messageOptions);
@@ -218,9 +225,14 @@ export const createChat = (
     messages.push(message);
   };
 
+  const cancelStream = () => {
+    cancel.cancelled = true;
+  };
+
   return {
     addMessage,
     getMessages: () => messages,
     sendMessage,
+    cancelStream,
   };
 };

src/createCompletions.ts

@@ -169,7 +169,8 @@ export type CompletionResponse = {
 };
 
 export const createCompletions = async (
-  options: CompletionsOptions
+  options: CompletionsOptions,
+  cancel: { cancelled: boolean }
 ): Promise<CompletionResponse> => {
   CompletionsOptionsZodSchema.parse(options);
 
@@ -208,8 +209,6 @@ export const createCompletions = async (
 
   const choices: Choice[] = [];
 
-  let cancelled = false;
-
   while (true) {
     const { value, done } = await reader.read();
 
@@ -260,7 +259,7 @@ export const createCompletions = async (
 
       options.onUpdate?.({
         cancel: () => {
-          cancelled = true;
+          cancel.cancelled = true;
 
           reader.cancel();
         },
@@ -302,7 +301,7 @@ export const createCompletions = async (
     }
   }
 
-  if (cancelled) {
+  if (cancel.cancelled) {
     throw new CancelledCompletionError(choices);
   }