add an abort controller (#40)

examples/abort/abort.ts

@@ -0,0 +1,24 @@
+import ollama from 'ollama'
+
+// Set a timeout to abort the request after 1 second
+setTimeout(() => {
+  console.log('\nAborting request...\n')
+  ollama.abort()
+}, 1000) // 1000 milliseconds = 1 second
+
+try {
+  const stream = await ollama.generate({
+    model: 'llama2',
+    prompt: 'Write a long story',
+    stream: true,
+  })
+  for await (const chunk of stream) {
+    process.stdout.write(chunk.response)
+  }
+} catch (error) {
+  if (error.name === 'AbortError') {
+    console.log('The request has been aborted')
+  } else {
+    console.error('An error occurred:', error)
+  }
+}

src/index.ts

@@ -30,6 +30,7 @@ import type {
 export class Ollama {
   private readonly config: Config
   private readonly fetch: Fetch
+  private abortController: AbortController
 
   constructor(config?: Partial<Config>) {
     this.config = {
@@ -43,16 +44,29 @@ export class Ollama {
     if (config?.fetch != null) {
       this.fetch = config.fetch
     }
+
+    this.abortController = new AbortController()
+  }
+
+  // Abort any ongoing requests to Ollama
+  public abort() {
+    this.abortController.abort()
+    this.abortController = new AbortController()
   }
 
   private async processStreamableRequest<T extends object>(
     endpoint: string,
     request: { stream?: boolean } & Record<string, any>,
   ): Promise<T | AsyncGenerator<T>> {
     request.stream = request.stream ?? false
-    const response = await utils.post(this.fetch, `${this.config.host}/api/${endpoint}`, {
-      ...request,
-    })
+    const response = await utils.post(
+      this.fetch,
+      `${this.config.host}/api/${endpoint}`,
+      {
+        ...request,
+      },
+      { signal: this.abortController.signal },
+    )
 
     if (!response.body) {
       throw new Error('Missing body')

src/utils.ts

@@ -62,6 +62,7 @@ export const post = async (
   fetch: Fetch,
   host: string,
   data?: Record<string, unknown> | BodyInit,
+  options?: { signal: AbortSignal },
 ): Promise<Response> => {
   const isRecord = (input: any): input is Record<string, unknown> => {
     return input !== null && typeof input === 'object' && !Array.isArray(input)
@@ -72,6 +73,7 @@ export const post = async (
   const response = await fetch(host, {
     method: 'POST',
     body: formattedData,
+    signal: options?.signal,
   })
 
   await checkOk(response)