Move AbortableAsyncIterator to index.ts

examples/abort/specific-request.ts

@@ -1,12 +1,12 @@
-import ollama from 'ollama'
-import { AbortableAsyncIterator } from '../../src/utils'
+import ollama, { AbortableAsyncIterator } from '../../src'
 
 let stream: AbortableAsyncIterator<object>
 
 // Set a timeout to abort the request after 1 second
 setTimeout(() => {
   console.log('\nAborting request...\n')
   stream.abort()
+  // Note: This will error if Ollama doesn't start responding within 1 second!
 }, 1000) // 1000 milliseconds = 1 second
 
 ollama.generate({

src/browser.ts

@@ -1,5 +1,5 @@
 import * as utils from './utils.js'
-import { AbortableAsyncIterator, parseJSON, post } from './utils.js'
+import { parseJSON, post } from './utils.js'
 import 'whatwg-fetch'
 
 import type {
@@ -25,6 +25,7 @@
   ShowResponse,
   StatusResponse,
 } from './interfaces.js'
+import { AbortableAsyncIterator } from './index'
 
 export class Ollama {
   protected readonly config: Config
@@ -66,7 +67,7 @@
    */
   protected async processStreamableRequest<T extends object>(
     endpoint: string,
     request: { stream?: boolean } & Record<string, any>,
   ): Promise<T | AbortableAsyncIterator<T>> {
     request.stream = request.stream ?? false
     const host = `${this.config.host}/api/${endpoint}`

src/index.ts

@@ -1,11 +1,10 @@
 import * as utils from './utils.js'
-import { AbortableAsyncIterator } from './utils.js'
 
 import fs, { createReadStream, promises } from 'fs'
 import { dirname, join, resolve } from 'path'
 import { createHash } from 'crypto'
 import { homedir } from 'os'
-import { Ollama as OllamaBrowser } from './browser.js'
+import { type ErrorResponse, Ollama as OllamaBrowser } from './browser.js'
 
 import type { CreateRequest, ProgressResponse } from './interfaces.js'
 
@@ -171,5 +170,40 @@
 
 export default new Ollama()
 
+/**
+ * An AsyncIterator which can be aborted
+ */
+export class AbortableAsyncIterator<T extends object> {
+  private readonly abortController: AbortController
+  private readonly itr: AsyncGenerator<T | ErrorResponse>
+  private readonly doneCallback: () => void
+
+  constructor(abortController: AbortController, itr: AsyncGenerator<T | ErrorResponse>, doneCallback: () => void) {
+    this.abortController = abortController
+    this.itr = itr
+    this.doneCallback = doneCallback
+  }
+
+  abort() {
+    this.abortController.abort()
+  }
+
+  async *[Symbol.asyncIterator]() {
+    for await (const message of this.itr) {
+      if ('error' in message) {
+        throw new Error(message.error)
+      }
+      yield message
+      // message will be done in the case of chat and generate
+      // message will be success in the case of a progress response (pull, push, create)
+      if ((message as any).done || (message as any).status === 'success') {
+        this.doneCallback()
+        return
+      }
+    }
+    throw new Error('Did not receive done or success response in stream.')
+  }
+}
+
 // export all types from the main entry point so that packages importing types dont need to specify paths
 export * from './interfaces.js'

src/utils.ts

@@ -19,41 +19,6 @@
   }
 }
 
-/**
- * An AsyncIterator which can be aborted
- */
-export class AbortableAsyncIterator<T extends object> {
-  private readonly abortController: AbortController
-  private readonly itr: AsyncGenerator<T | ErrorResponse>
-  private readonly doneCallback: () => void
-
-  constructor(abortController: AbortController, itr: AsyncGenerator<T | ErrorResponse>, doneCallback: () => void) {
-    this.abortController = abortController
-    this.itr = itr
-    this.doneCallback = doneCallback
-  }
-
-  abort() {
-    this.abortController.abort()
-  }
-
-  async *[Symbol.asyncIterator]() {
-    for await (const message of this.itr) {
-      if ('error' in message) {
-        throw new Error(message.error)
-      }
-      yield message
-      // message will be done in the case of chat and generate
-      // message will be success in the case of a progress response (pull, push, create)
-      if ((message as any).done || (message as any).status === 'success') {
-        this.doneCallback()
-        return
-      }
-    }
-    throw new Error('Did not receive done or success response in stream.')
-  }
-}
-
 /**
  * Checks if the response is ok, if not throws an error.
  * If the response is not ok, it will try to parse the response as JSON and use the error field as the error message.
@@ -171,7 +136,7 @@
   data?: Record<string, unknown> | BodyInit,
   options?: { signal?: AbortSignal, headers?: Headers },
 ): Promise<Response> => {
   const isRecord = (input: any): input is Record<string, unknown> => {
     return input !== null && typeof input === 'object' && !Array.isArray(input)
   }