error handling

.changeset/sweet-kangaroos-sip.md

@@ -0,0 +1,6 @@
+---
+'aibitat': patch
+---
+
+Gracefully handle API errors. Added `.onError` and `.retry` methods to let devs
+decide what to do with it

README.md

@@ -14,9 +14,10 @@ provider agnostic and can be used with any provider that implements the
 `AIProvider` interface. Also, it is stateless and can be used in a serverless
 environment.
 
-By default, it uses **OpenAI** and **GPT-3.5-TURBO** as the provider but you can
-change it by passing `provider` and `model` to the `AIbitat` constructor or by
-setting them on the node config.
+By default, aibitat uses **OpenAI** and **GPT-3.5-TURBO** as the provider for
+the conversation and **GPT-4** for predicting the next agent to speak but you
+can change it by passing `provider` and `model` to the `AIbitat` constructor or
+by setting them on the specific node config.
 
 ### Features
 
@@ -40,7 +41,7 @@ setting them on the node config.
       to the conversation.
 - [ ] **Cache**. Store conversation history in a cache to improve performance
       and reduce the number of API calls.
-- [ ] **Error handling.** Handle API errors gracefully.
+- [x] **Error handling.** Handle API errors gracefully.
 - [ ] **Code execution.** Agents can execute code and return the result to the
       conversation.
 
@@ -153,11 +154,13 @@ aibitat.onMessage(({from, to, content}) => console.log(`${from}: ${content}`))
 
 The following events are available:
 
-- `start`: Called when the chat starts.
-- `message`: Called when a message is added to the chat history.
-- `terminate`: Called when the conversation is terminated. Generally means there
-  is nothing else to do and a new conversation should be started.
-- `interrupt`: Called when the conversation is interrupted by an agent.
+- `onStart`: Called when the chat starts.
+- `onError`: Called when there's a known error (see `src/error.ts`). To retry,
+  call `.retry()`.
+- `onMessage`: Called when a message is added to the chat history.
+- `onTerminate`: Called when the conversation is terminated. Generally means
+  there is nothing else to do and a new conversation should be started.
+- `onInterrupt`: Called when the conversation is interrupted by an agent.
   Generally means the agent has a question or needs help. The conversation can
   be resumed by calling `.continue(feedback)`.
 

src/aibitat.test.ts

@@ -2,6 +2,7 @@ import {beforeEach, describe, expect, mock, test} from 'bun:test'
 import OpenAI from 'openai'
 
 import {AIbitat, type AIbitatProps} from './aibitat.ts'
+import {RateLimitError} from './error.ts'
 import {AIProvider} from './providers/index.ts'
 import {type Message} from './types.ts'
 
@@ -303,3 +304,85 @@ test.todo('should call a function', async () => {
 })
 
 test.todo('should execute code', async () => {})
+
+describe('when errors happen', () => {
+  test('should escape unknown errors', async () => {
+    const customError = new Error('unknown error')
+
+    ai.create.mockImplementation(() => {
+      throw customError
+    })
+
+    const aibitat = new AIbitat(defaultaibitat)
+
+    try {
+      await aibitat.start(defaultStart)
+    } catch (error) {
+      expect(error).toEqual(customError)
+    }
+  })
+
+  test('should handle known errors', async () => {
+    const error = new RateLimitError('known error!!!')
+    ai.create.mockImplementation(() => {
+      throw error
+    })
+
+    const aibitat = new AIbitat(defaultaibitat)
+    aibitat.onError((_, error) => {
+      expect(error).toEqual(error)
+    })
+
+    await aibitat.start(defaultStart)
+
+    expect(aibitat.chats).toHaveLength(2)
+    expect(aibitat.chats.at(-1)).toEqual({
+      from: '🤖',
+      to: '🧑',
+      content: 'known error!!!',
+      state: 'error',
+    })
+  })
+
+  test('should trigger the error event', async () => {
+    const error = new RateLimitError('401: Rate limit')
+    ai.create.mockImplementation(() => {
+      throw error
+    })
+
+    const aibitat = new AIbitat(defaultaibitat)
+
+    const callback = mock((error: unknown) => {})
+    aibitat.onError(callback)
+
+    await aibitat.start(defaultStart)
+
+    expect(callback).toHaveBeenCalledTimes(1)
+    expect(callback.mock.calls[0][0]).toEqual(error)
+  })
+
+  test('should be able to retry', async () => {
+    let i = 0
+    const error = new RateLimitError('401: Rate limit')
+    ai.create.mockImplementation(() => {
+      if (i++ === 0) {
+        throw error
+      }
+
+      return Promise.resolve('TERMINATE')
+    })
+
+    const aibitat = new AIbitat(defaultaibitat)
+    await aibitat.start(defaultStart)
+    await aibitat.retry()
+
+    expect(ai.create).toHaveBeenCalledTimes(2)
+    expect(aibitat.chats).toHaveLength(2)
+    expect(aibitat.chats.at(-1)).toEqual({
+      from: '🤖',
+      to: '🧑',
+      content: 'TERMINATE',
+      state: 'success',
+    })
+  })
+})

src/aibitat.ts

@@ -2,6 +2,13 @@ import {EventEmitter} from 'events'
 import chalk from 'chalk'
 import debug from 'debug'
 
+import {
+  APIError,
+  AuthorizationError,
+  RateLimitError,
+  ServerError,
+  UnknownError,
+} from './error.ts'
 import {
   AIProvider,
   OpenAIProvider,
@@ -108,8 +115,7 @@ type Chat = {
  */
 type ChatState = Omit<Chat, 'content'> & {
   content?: string
-  // state: 'success' | 'loading' | 'error' | 'interrupt'
-  state: 'success' | 'interrupt'
+  state: 'success' | 'interrupt' | 'error'
 }
 
 /**
@@ -266,7 +272,7 @@ export class AIbitat {
    * @param listener
    * @returns
    */
-  public onTerminate(listener: (node: string, aibitat: AIbitat) => void) {
+  public onTerminate(listener: (node: string) => void) {
     this.emitter.on('terminate', listener)
     return this
   }
@@ -286,9 +292,7 @@ export class AIbitat {
    * @param listener
    * @returns
    */
-  public onInterrupt(
-    listener: (chat: {from: string; to: string}, aibitat: AIbitat) => void,
-  ) {
+  public onInterrupt(listener: (chat: {from: string; to: string}) => void) {
     this.emitter.on('interrupt', listener)
     return this
   }
@@ -314,7 +318,7 @@ export class AIbitat {
    * @param listener
    * @returns
    */
-  public onMessage(listener: (chat: ChatState, aibitat: AIbitat) => void) {
+  public onMessage(listener: (chat: ChatState) => void) {
     this.emitter.on('message', listener)
     return this
   }
@@ -333,9 +337,52 @@ export class AIbitat {
 
     this._chats.push(chat)
     this.emitter.emit('message', chat, this)
+  }
+
+  /**
+   * Triggered when an error occurs during the chat.
+   *
+   * @param listener
+   * @returns
+   */
+  public onError(
+    listener: (
+      /**
+       * The error that occurred.
+       *
+       * Native errors are:
+       * - `APIError`
+       * - `AuthorizationError`
+       * - `UnknownError`
+       * - `RateLimitError`
+       * - `ServerError`
+       */
+      error: unknown,
+      /**
+       * The message when the error occurred.
+       */
+      {}: {from: string; to: string},
+    ) => void,
+  ) {
+    this.emitter.on('replyError', listener)
     return this
   }
 
+  /**
+   * Register an error in the chat history.
+   * This will trigger the `onError` event.
+   *
+   * @param message
+   */
+  private newError(message: {from: string; to: string}, error: unknown) {
+    const chat = {
+      ...message,
+      content: error instanceof Error ? error.message : String(error),
+      state: 'error' as const,
+    }
+    this._chats.push(chat)
+    this.emitter.emit('replyError', error, chat)
+  }
   /**
    * Triggered when a chat is interrupted by a node.
    *
@@ -429,7 +476,15 @@ export class AIbitat {
     }
 
     // If it's a direct message, reply to the message
-    const reply = await this.reply(message)
+    let reply: string
+    try {
+      reply = await this.reply(message)
+    } catch (error: unknown) {
+      if (error instanceof APIError) {
+        return this.newError({from: message.from, to: message.to}, error)
+      }
+      throw error
+    }
 
     if (
       reply === 'TERMINATE' ||
@@ -611,7 +666,6 @@ ${this.getHistory({to})
 
     // get the chat completion
     const content = await provider.create(messages)
-    // TODO: add error handling
     this.newMessage({from, to, content})
 
     return content
@@ -662,6 +716,23 @@ ${this.getHistory({to})
     return this
   }
 
+  /**
+   * Retry the last chat that threw an error.
+   * If the last chat was not an error, it will throw an error.
+   */
+  public async retry() {
+    const lastChat = this._chats.at(-1)
+    if (!lastChat || lastChat.state !== 'error') {
+      throw new Error('No chat to retry')
+    }
+
+    // remove the last chat's that threw an error
+    const {from, to} = this._chats.pop()!
+
+    await this.chat({from, to})
+    return this
+  }
+
   /**
    * Get the chat history between two nodes or all chats to/from a node.
    */

src/error.ts

@@ -0,0 +1,24 @@
+export class AIbitatError extends Error {}
+
+/**
+ * The error for the AIbitat class when the AI provider returns a rate limit error.
+ */
+export class APIError extends AIbitatError {
+  constructor(message?: string | undefined) {
+    super(message)
+  }
+}
+
+export class AuthorizationError extends APIError {}
+export class UnknownError extends APIError {}
+export class RateLimitError extends APIError {}
+export class ServerError extends APIError {}
+
+// // ANTHROPIC
+// 400 - Invalid request: there was an issue with the format or content of your request.
+// 401 - Unauthorized: there's an issue with your API key.
+// 403 - Forbidden: your API key does not have permission to use the specified resource.
+// 404 - Not found: the requested resource was not found.
+// 429 - Your account has hit a rate limit.
+// 500 - An unexpected error has occurred internal to Anthropic's systems.
+// 529 - Anthropic's API is temporarily overloaded.

src/plugins/terminal.ts

@@ -2,6 +2,7 @@ import {input} from '@inquirer/prompts'
 import chalk from 'chalk'
 
 import {AIbitatPlugin} from '..'
+import {RateLimitError, ServerError} from '../error'
 
 /**
  * Print a message on the terminal
@@ -92,6 +93,16 @@ export function terminal({
     setup(aibitat) {
       let printing: Promise<void> | null = null
 
+      aibitat.onError(error => {
+        console.error(chalk.red(`   error: ${(error as Error).message}`))
+
+        if (error instanceof RateLimitError || error instanceof ServerError) {
+          console.error(chalk.red(`   retrying in 60 seconds...`))
+          setTimeout(aibitat.retry, 60000)
+          return
+        }
+      })
+
       aibitat.onStart(() => {
         console.log()
         console.log('🚀 starting chat ...\n')
@@ -106,7 +117,7 @@ export function terminal({
 
       aibitat.onTerminate(() => console.timeEnd('🚀 chat finished'))
 
-      aibitat.onInterrupt(async (node, current) => {
+      aibitat.onInterrupt(async node => {
         await printing
         const feedback = await askForFeedback(node)
         // Add an extra line after the message
@@ -117,7 +128,7 @@ export function terminal({
           return process.exit(0)
         }
 
-        await current.continue(feedback)
+        await aibitat.continue(feedback)
       })
     },
   } as AIbitatPlugin

src/providers/ai-provider.ts

@@ -1,5 +1,8 @@
 import {Message} from '../types.ts'
 
+/**
+ * A service that provides an AI client to create a completion.
+ */
 export abstract class AIProvider<T> {
   private _client: T
 
@@ -17,6 +20,8 @@ export abstract class AIProvider<T> {
 
   /**
    * (Abstract async method) Create a completion based on the received messages.
+   *
+   * @throws It should thrown known treated errors from `src/error.ts`.
    * @param messages A list of messages to send.
    */
   abstract create(messages: Message[]): Promise<string>

src/providers/index.ts

@@ -1,2 +1,2 @@
 export * from './ai-provider.ts'
-export * from './openai-provider.ts'
+export * from './openai.ts'

src/providers/openai.test.ts

@@ -0,0 +1,46 @@
+import {beforeEach, describe, expect, mock, test} from 'bun:test'
+
+import {AuthorizationError} from '../error.ts'
+import {Message} from '../types.ts'
+import {OpenAIProvider} from './openai.ts'
+
+// NOTE: some tests are skipped because it requires a way to mock the http requests.
+
+// // OPENAI
+// 401 - Invalid Authentication	Cause: Invalid Authentication
+// Solution: Ensure the correct API key and requesting organization are being used.
+// 401 - Incorrect API key provided	Cause: The requesting API key is not correct.
+// Solution: Ensure the API key used is correct, clear your browser cache, or generate a new one.
+// 401 - You must be a member of an organization to use the API	Cause: Your account is not part of an organization.
+// Solution: Contact us to get added to a new organization or ask your organization manager to invite you to an organization.
+// 429 - Rate limit reached for requests	Cause: You are sending requests too quickly.
+// Solution: Pace your requests. Read the Rate limit guide.
+// 429 - You exceeded your current quota, please check your plan and billing details	Cause: You have hit your maximum monthly spend (hard limit) which you can view in the account billing section.
+// Solution: Apply for a quota increase.
+// 500 - The server had an error while processing your request	Cause: Issue on our servers.
+// Solution: Retry your request after a brief wait and contact us if the issue persists. Check the status page.
+// 503 - The engine is currently overloaded, please try again later	Cause: Our servers are experiencing high traffic.
+// Solution: Please retry your requests after a brief wait.
+
+const message: Message[] = [
+  {
+    content: 'Hello',
+    role: 'user',
+  },
+]
+
+test('should throw an error when there`s an authorization error', async () => {
+  const provider = new OpenAIProvider({
+    options: {
+      apiKey: 'invalid',
+    },
+  })
+
+  await expect(provider.create(message)).rejects.toBeInstanceOf(
+    AuthorizationError,
+  )
+})
+
+test.todo('should throw a generic error when something else happens', () => {})
+test.todo('should throw a RateLimitError', () => {})
+test.todo('should throw a ServerError', () => {})

src/providers/openai.ts

@@ -1,7 +1,27 @@
-import {OpenAIStream, StreamingTextResponse} from 'ai'
 import debug from 'debug'
-import OpenAI, {ClientOptions} from 'openai'
+import OpenAI, {
+  ClientOptions,
+  APIConnectionError as OpenAIAPIConnectionError,
+  APIConnectionTimeoutError as OpenAIAPIConnectionTimeoutError,
+  APIError as OpenAIAPIError,
+  APIUserAbortError as OpenAIAPIUserAbortError,
+  AuthenticationError as OpenAIAuthenticationError,
+  BadRequestError as OpenAIBadRequestError,
+  ConflictError as OpenAIConflictError,
+  InternalServerError as OpenAIInternalServerError,
+  NotFoundError as OpenAINotFoundError,
+  PermissionDeniedError as OpenAIPermissionDeniedError,
+  RateLimitError as OpenAIRateLimitError,
+  UnprocessableEntityError as OpenAIUnprocessableEntityError,
+} from 'openai'
 
+import {
+  APIError,
+  AuthorizationError,
+  RateLimitError,
+  ServerError,
+  UnknownError,
+} from '../error.ts'
 import {Message} from '../types.ts'
 import {AIProvider} from './ai-provider.ts'
 
@@ -58,6 +78,7 @@ export class OpenAIProvider extends AIProvider<OpenAI> {
     const {
       options = {
         apiKey: process.env.OPENAI_API_KEY,
+        maxRetries: 3,
       },
       model = 'gpt-3.5-turbo',
     } = config
@@ -78,19 +99,58 @@ export class OpenAIProvider extends AIProvider<OpenAI> {
   async create(messages: Message[]) {
     log(`calling 'openai.chat.completions.create' with model '${this.model}'`)
 
-    const response = await this.client.chat.completions.create({
-      model: this.model,
-      // stream: true,
-      messages,
-    })
+    try {
+      const response = await this.client.chat.completions.create({
+        model: this.model,
+        // stream: true,
+        messages,
+      })
 
-    log('cost: ', this.getCost(response.usage))
+      log('cost: ', this.getCost(response.usage))
 
-    return response.choices[0].message.content!
+      return response.choices[0].message.content!
 
-    // const stream = OpenAIStream(response)
-    // const result = new StreamingTextResponse(stream)
-    // return await result.text()
+      // const stream = OpenAIStream(response)
+      // const result = new StreamingTextResponse(stream)
+      // return await result.text()
+    } catch (error) {
+      // if (error instanceof OpenAIBadRequestError) {
+      //   throw new Error(error.message)
+      // }
+
+      if (
+        error instanceof OpenAIAuthenticationError ||
+        error instanceof OpenAIPermissionDeniedError
+      ) {
+        throw new AuthorizationError(error.message)
+      }
+
+      // if (error instanceof OpenAINotFoundError) {
+      //   throw new Error(error.message)
+      // }
+
+      // if (error instanceof OpenAIConflictError) {
+      //   throw new Error(error.message)
+      // }
+
+      // if (error instanceof OpenAIUnprocessableEntityError) {
+      //   throw new Error(error.message)
+      // }
+
+      if (error instanceof OpenAIRateLimitError) {
+        throw new RateLimitError(error.message)
+      }
+
+      if (error instanceof OpenAIInternalServerError) {
+        throw new ServerError(error.message)
+      }
+
+      if (error instanceof OpenAIAPIError) {
+        throw new UnknownError(error.message)
+      }
+
+      throw error
+    }
   }
 
   /**