diff --git a/examples/abort/specific-request.ts b/examples/abort/specific-request.ts index 778c514..3bfbe7b 100644 --- a/examples/abort/specific-request.ts +++ b/examples/abort/specific-request.ts @@ -1,5 +1,4 @@ -import ollama from 'ollama' -import { AbortableAsyncIterator } from '../../src/utils' +import ollama, { AbortableAsyncIterator } from '../../src' let stream: AbortableAsyncIterator @@ -7,6 +6,7 @@ let stream: AbortableAsyncIterator 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({ diff --git a/src/browser.ts b/src/browser.ts index 7fa23cd..17d59b2 100644 --- a/src/browser.ts +++ b/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 @@ import type { ShowResponse, StatusResponse, } from './interfaces.js' +import { AbortableAsyncIterator } from './index' export class Ollama { protected readonly config: Config diff --git a/src/index.ts b/src/index.ts index 598304c..4947c30 100644 --- a/src/index.ts +++ b/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 class Ollama extends OllamaBrowser { export default new Ollama() +/** + * An AsyncIterator which can be aborted + */ +export class AbortableAsyncIterator { + private readonly abortController: AbortController + private readonly itr: AsyncGenerator + private readonly doneCallback: () => void + + constructor(abortController: AbortController, itr: AsyncGenerator, 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' diff --git a/src/utils.ts b/src/utils.ts index 7ab32d4..98650e9 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -19,41 +19,6 @@ class ResponseError extends Error { } } -/** - * An AsyncIterator which can be aborted - */ -export class AbortableAsyncIterator { - private readonly abortController: AbortController - private readonly itr: AsyncGenerator - private readonly doneCallback: () => void - - constructor(abortController: AbortController, itr: AsyncGenerator, 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.