Skip to content

Commit

Permalink
Moved AbortableAsyncIterator to index.ts
Browse files Browse the repository at this point in the history
Closes #135
  • Loading branch information
hopperelec committed Oct 18, 2024
1 parent 6b6221d commit 98fb3e4
Show file tree
Hide file tree
Showing 4 changed files with 40 additions and 40 deletions.
4 changes: 2 additions & 2 deletions examples/abort/specific-request.ts
Original file line number Diff line number Diff line change
@@ -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({
Expand Down
3 changes: 2 additions & 1 deletion src/browser.ts
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -25,6 +25,7 @@ import type {
ShowResponse,
StatusResponse,
} from './interfaces.js'
import { AbortableAsyncIterator } from './index'

export class Ollama {
protected readonly config: Config
Expand Down
38 changes: 36 additions & 2 deletions src/index.ts
Original file line number Diff line number Diff line change
@@ -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'

Expand Down Expand Up @@ -171,5 +170,40 @@ export class Ollama extends OllamaBrowser {

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') {

Check warning on line 199 in src/index.ts

View workflow job for this annotation

GitHub Actions / test (16)

Unexpected any. Specify a different type

Check warning on line 199 in src/index.ts

View workflow job for this annotation

GitHub Actions / test (16)

Unexpected any. Specify a different type

Check warning on line 199 in src/index.ts

View workflow job for this annotation

GitHub Actions / test (18)

Unexpected any. Specify a different type

Check warning on line 199 in src/index.ts

View workflow job for this annotation

GitHub Actions / test (18)

Unexpected any. Specify a different type

Check warning on line 199 in src/index.ts

View workflow job for this annotation

GitHub Actions / test (20)

Unexpected any. Specify a different type

Check warning on line 199 in src/index.ts

View workflow job for this annotation

GitHub Actions / test (20)

Unexpected any. Specify a different type
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'
35 changes: 0 additions & 35 deletions src/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,41 +19,6 @@ class ResponseError extends Error {
}
}

/**
* 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.
Expand Down

0 comments on commit 98fb3e4

Please sign in to comment.