Merge pull request #13 from anthropics/justin/cancellation

Add cancellation support using AbortController/AbortSignal
anthropics · Apr 4, 2023 · ec96afb · ec96afb
2 parents 493075d + 607afa5
commit ec96afb
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 5 deletions.
diff --git a/examples/cancellation.ts b/examples/cancellation.ts
@@ -0,0 +1,52 @@
+import "dotenv/config";
+import { AI_PROMPT, Client, HUMAN_PROMPT } from "../src";
+
+const apiKey = process.env.ANTHROPIC_API_KEY;
+if (!apiKey) {
+  throw new Error("The ANTHROPIC_API_KEY environment variable must be set");
+}
+
+const client = new Client(apiKey);
+const abortController = new AbortController();
+
+client
+  .complete(
+    {
+      prompt: `${HUMAN_PROMPT} How many toes do dogs have?${AI_PROMPT}`,
+      stop_sequences: [HUMAN_PROMPT],
+      max_tokens_to_sample: 200,
+      model: "claude-v1",
+    },
+    { signal: abortController.signal }
+  )
+  .catch((error) => {
+    if (error.name === "AbortError") {
+      console.log("Cancelled complete()");
+    }
+  });
+
+client
+  .completeStream(
+    {
+      prompt: `${HUMAN_PROMPT} How many toes do dogs have?${AI_PROMPT}`,
+      stop_sequences: [HUMAN_PROMPT],
+      max_tokens_to_sample: 200,
+      model: "claude-v1",
+    },
+    {
+      onOpen: (response) => {
+        console.log("Opened stream, HTTP status code", response.status);
+      },
+      onUpdate: (completion) => {
+        console.log(completion.completion);
+      },
+      signal: abortController.signal,
+    }
+  )
+  .catch((error) => {
+    if (error.name === "AbortError") {
+      console.log("Cancelled completeStream()");
+    }
+  });
+
+abortController.abort();
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@anthropic-ai/sdk",
-  "version": "0.4.2",
+  "version": "0.4.3",
   "description": "Library for accessing the Anthropic API",
   "repository": "https://github.com/anthropics/anthropic-sdk-typescript",
   "license": "MIT",
@@ -15,7 +15,8 @@
     "lint": "eslint src && prettier --check .",
     "format": "prettier --write .",
     "example:basic_stream": "tsc --build && node ./build/examples/basic_stream.js",
-    "example:basic_sync": "tsc --build && node ./build/examples/basic_sync.js"
+    "example:basic_sync": "tsc --build && node ./build/examples/basic_sync.js",
+    "example:cancellation": "tsc --build && node ./build/examples/cancellation.js"
   },
   "dependencies": {
     "@fortaine/fetch-event-source": "^3.0.6",

diff --git a/src/index.ts b/src/index.ts
@@ -18,7 +18,7 @@ export type OnUpdate = (completion: CompletionResponse) => void | Promise<void>;
 export const HUMAN_PROMPT = "\n\nHuman:";
 export const AI_PROMPT = "\n\nAssistant:";
 
-const CLIENT_ID = "anthropic-typescript/0.4.2";
+const CLIENT_ID = "anthropic-typescript/0.4.3";
 const DEFAULT_API_URL = "https://api.anthropic.com";
 
 enum Event {
@@ -43,7 +43,10 @@ export class Client {
     this.apiUrl = options?.apiUrl ?? DEFAULT_API_URL;
   }
 
-  async complete(params: SamplingParameters): Promise<CompletionResponse> {
+  async complete(
+    params: SamplingParameters,
+    options?: { signal?: AbortSignal }
+  ): Promise<CompletionResponse> {
     const response = await fetch(`${this.apiUrl}/v1/complete`, {
       method: "POST",
       headers: {
@@ -53,6 +56,7 @@ export class Client {
         "X-API-Key": this.apiKey,
       },
       body: JSON.stringify({ ...params, stream: false }),
+      signal: options?.signal,
     });
 
     if (!response.ok) {
@@ -69,10 +73,23 @@ export class Client {
 
   completeStream(
     params: SamplingParameters,
-    { onOpen, onUpdate }: { onOpen?: OnOpen; onUpdate?: OnUpdate }
+    {
+      onOpen,
+      onUpdate,
+      signal,
+    }: { onOpen?: OnOpen; onUpdate?: OnUpdate; signal?: AbortSignal }
   ): Promise<CompletionResponse> {
     const abortController = new AbortController();
+
     return new Promise((resolve, reject) => {
+      signal?.addEventListener("abort", (event) => {
+        abortController.abort(event);
+        reject({
+          name: "AbortError",
+          message: "Caller aborted completeStream",
+        });
+      });
+
       fetchEventSource(`${this.apiUrl}/v1/complete`, {
         method: "POST",
         headers: {