diff --git a/examples/abort/abort.ts b/examples/abort/abort.ts new file mode 100644 index 0000000..346d906 --- /dev/null +++ b/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) + } +} diff --git a/src/index.ts b/src/index.ts index 8374904..dde83c9 100644 --- a/src/index.ts +++ b/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) { this.config = { @@ -43,6 +44,14 @@ 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( @@ -50,9 +59,14 @@ export class Ollama { request: { stream?: boolean } & Record, ): Promise> { 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') diff --git a/src/utils.ts b/src/utils.ts index 6b33739..f67bf41 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -62,6 +62,7 @@ export const post = async ( fetch: Fetch, host: string, data?: Record | BodyInit, + options?: { signal: AbortSignal }, ): Promise => { const isRecord = (input: any): input is Record => { 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)