ai-retry 0.1.0 → 0.2.0

package/README.md

@@ -1,8 +1,15 @@
-<a href="https://www.npmjs.com/package/ai-retry" alt="ai-retry"><img src="https://img.shields.io/npm/dt/ai-retry?label=ai-retry"></a> <a href="https://github.com/zirkelc/ai-retry/actions/workflows/ci.yml" alt="CI"><img src="https://img.shields.io/github/actions/workflow/status/zirkelc/ai-retry/ci.yml?branch=main"></a>
+<div align='center'>
 
-# ai-retry: Retry and fallback mechanisms for AI SDK
+# ai-retry
 
-Automatically handle API failures, content filtering and timeouts by switching between different AI models.
+<p align="center">Retry and fallback mechanisms for AI SDK</p>
+<p align="center">
+  <a href="https://www.npmjs.com/package/ai-retry" alt="ai-retry"><img src="https://img.shields.io/npm/dt/ai-retry?label=ai-retry"></a> <a href="https://github.com/zirkelc/ai-retry/actions/workflows/ci.yml" alt="CI"><img src="https://img.shields.io/github/actions/workflow/status/zirkelc/ai-retry/ci.yml?branch=main"></a>
+</p>
+
+</div>
+
+Automatically handle API failures, content filtering, timeouts and other errors by switching between different AI models and providers.
 
 `ai-retry` wraps the provided base model with a set of retry conditions (retryables). When a request fails with an error or the response is not satisfying, it iterates through the given retryables to find a suitable fallback model. It automatically tracks which models have been tried and how many attempts have been made to prevent infinite loops.
 
@@ -26,16 +33,17 @@ npm install ai-retry
 Create a retryable model by providing a base model and a list of retryables or fallback models.
 When an error occurs, it will evaluate each retryable in order and use the first one that indicates a retry should be attempted with a different model.
 
+> [!NOTE]
+> `ai-retry` supports both language models and embedding models.
+
 ```typescript
-import { azure } from '@ai-sdk/azure';
 import { openai } from '@ai-sdk/openai';
 import { generateText, streamText } from 'ai';
 import { createRetryable } from 'ai-retry';
-import { contentFilterTriggered, requestTimeout } from 'ai-retry/retryables';
 
 // Create a retryable model
 const retryableModel = createRetryable({
-  model: azure('gpt-4-mini'), // Base model
+  model: openai('gpt-4-mini'), // Base model
   retries: [
     // Retry strategies and fallbacks...
   ],
@@ -47,6 +55,8 @@ const result = await generateText({
   prompt: 'Hello world!',
 });
 
+console.log(result.text);
+
 // Or with streaming
 const result = streamText({
   model: retryableModel,
@@ -58,52 +68,47 @@ for await (const chunk of result.textStream) {
 }
 ```
 
-#### Content Filter
-
-Automatically switch to a different model when content filtering blocks your request.
-
-> [!WARNING]  
-> This retryable currently does not work with streaming requests, because the content filter is only indicated in the final response.
+This also works with embedding models:
 
 ```typescript
-import { contentFilterTriggered } from 'ai-retry/retryables';
+import { openai } from '@ai-sdk/openai';
+import { embed } from 'ai';
+import { createRetryable } from 'ai-retry';
 
+// Create a retryable model
 const retryableModel = createRetryable({
-  model: azure('gpt-4-mini'),
+  model: openai.textEmbedding('text-embedding-3-large'), // Base model
   retries: [
-    contentFilterTriggered(openai('gpt-4-mini')), // Try OpenAI if Azure filters
+    // Retry strategies and fallbacks...
   ],
 });
+
+// Use like any other AI SDK model
+const result = await embed({
+  model: retryableModel,
+  value: 'Hello world!',
+});
+
+console.log(result.embedding);
 ```
 
-<!--
-##### Response Schema Mismatch
+#### Content Filter
+
+Automatically switch to a different model when content filtering blocks your request.
 
-Retry with different models when structured output validation fails:
+> [!WARNING]  
+> This retryable currently does not work with streaming requests, because the content filter is only indicated in the final response.
 
 ```typescript
-import { responseSchemaMismatch } from 'ai-retry/retryables';
+import { contentFilterTriggered } from 'ai-retry/retryables';
 
 const retryableModel = createRetryable({
   model: azure('gpt-4-mini'),
   retries: [
-    responseSchemaMismatch(azure('gpt-4')), // Try full model for better structured output
+    contentFilterTriggered(openai('gpt-4-mini')), // Try OpenAI if Azure filters
   ],
 });
-
-const result = await generateObject({
-  model: retryableModel,
-  schema: z.object({
-    recipe: z.object({
-      name: z.string(),
-      ingredients: z.array(z.object({ name: z.string(), amount: z.string() })),
-      steps: z.array(z.string()),
-    }),
-  }),
-  prompt: 'Generate a lasagna recipe.',
-});
 ```
--->
 
 #### Request Timeout
 
@@ -287,16 +292,16 @@ By default, each retryable will only attempt to retry once per model to avoid in
 
 ### API Reference
 
-#### `createRetryable(options: CreateRetryableOptions): LanguageModelV2`
+#### `createRetryable(options: RetryableModelOptions): LanguageModelV2 | EmbeddingModelV2`
 
-Creates a retryable language model.
+Creates a retryable model that works with both language models and embedding models.
 
 ```ts
-interface CreateRetryableOptions {
-  model: LanguageModelV2;
-  retries: Array<Retryable | LanguageModelV2>;
-  onError?: (context: RetryContext) => void;
-  onRetry?: (context: RetryContext) => void; 
+interface RetryableModelOptions<MODEL extends LanguageModelV2 | EmbeddingModelV2> {
+  model: MODEL;
+  retries: Array<Retryable<MODEL> | MODEL>;
+  onError?: (context: RetryContext<MODEL>) => void;
+  onRetry?: (context: RetryContext<MODEL>) => void;
 }
 ```
 
@@ -306,7 +311,9 @@ A `Retryable` is a function that receives a `RetryContext` with the current erro
 It should evaluate the error/result and decide whether to retry by returning a `RetryModel` or to skip by returning `undefined`.
 
 ```ts
-type Retryable = (context: RetryContext) => RetryModel | Promise<RetryModel> | undefined;
+type Retryable = (
+  context: RetryContext
+) => RetryModel | Promise<RetryModel> | undefined;
 ```
 
 #### `RetryModel`
@@ -316,7 +323,7 @@ By default, each retryable will only attempt to retry once per model. This can b
 
 ```typescript
 interface RetryModel {
-  model: LanguageModelV2;
+  model: LanguageModelV2 | EmbeddingModelV2;
   maxAttempts?: number;
 }
 ```
@@ -337,10 +344,13 @@ interface RetryContext {
 A `RetryAttempt` represents a single attempt with a specific model, which can be either an error or a successful result that triggered a retry.
 
 ```typescript
-type RetryAttempt = 
-  | { type: 'error'; error: unknown; model: LanguageModelV2 }
+// For both language and embedding models
+type RetryAttempt =
+  | { type: 'error'; error: unknown; model: LanguageModelV2 | EmbeddingModelV2 }
   | { type: 'result'; result: LanguageModelV2Generate; model: LanguageModelV2 };
 
+// Note: Result-based retries only apply to language models, not embedding models
+
 // Type guards for discriminating attempts
 function isErrorAttempt(attempt: RetryAttempt): attempt is RetryErrorAttempt;
 function isResultAttempt(attempt: RetryAttempt): attempt is RetryResultAttempt;

package/dist/index.d.ts

@@ -1,10 +1,13 @@
-import { CreateRetryableOptions, RetryAttempt, RetryContext, RetryModel, Retryable, createRetryable, isErrorAttempt, isResultAttempt } from "./create-retryable-model-DtspEawi.js";
-import { LanguageModelV2 } from "@ai-sdk/provider";
+import { EmbeddingModelV2, LanguageModelV2, RetryableModelOptions } from "./types-CHhEGL5x.js";
 
+//#region src/create-retryable-model.d.ts
+declare function createRetryable<MODEL extends LanguageModelV2>(options: RetryableModelOptions<MODEL>): LanguageModelV2;
+declare function createRetryable<MODEL extends EmbeddingModelV2>(options: RetryableModelOptions<MODEL>): EmbeddingModelV2;
+//#endregion
 //#region src/get-model-key.d.ts
 /**
  * Generate a unique key for a LanguageModelV2 instance.
  */
-declare const getModelKey: (model: LanguageModelV2) => string;
+declare const getModelKey: (model: LanguageModelV2 | EmbeddingModelV2) => string;
 //#endregion
-export { CreateRetryableOptions, RetryAttempt, RetryContext, RetryModel, Retryable, createRetryable, getModelKey, isErrorAttempt, isResultAttempt };
+export { createRetryable, getModelKey };

package/dist/index.js

@@ -1,3 +1,404 @@
-import { createRetryable, getModelKey, isErrorAttempt, isResultAttempt } from "./create-retryable-model-C3mm1hSN.js";
+import { isErrorAttempt, isGenerateResult, isResultAttempt, isStreamContentPart } from "./utils-lRsC105f.js";
+import "@ai-sdk/provider-utils";
+import { RetryError } from "ai";
+import { getErrorMessage } from "@ai-sdk/provider";
 
-export { createRetryable, getModelKey, isErrorAttempt, isResultAttempt };
+//#region src/get-model-key.ts
+/**
+* Generate a unique key for a LanguageModelV2 instance.
+*/
+const getModelKey = (model) => {
+	return `${model.provider}/${model.modelId}`;
+};
+
+//#endregion
+//#region src/find-retry-model.ts
+/**
+* Find the next model to retry with based on the retry context
+*/
+async function findRetryModel(retries, context) {
+	/**
+	* Filter retryables based on attempt type:
+	* - Result-based attempts: Only consider function retryables (skip plain models)
+	* - Error-based attempts: Consider all retryables (functions + plain models)
+	*/
+	const applicableRetries = isResultAttempt(context.current) ? retries.filter((retry) => typeof retry === "function") : retries;
+	/**
+	* Iterate through the applicable retryables to find a model to retry with
+	*/
+	for (const retry of applicableRetries) {
+		const retryModel = typeof retry === "function" ? await retry(context) : {
+			model: retry,
+			maxAttempts: 1
+		};
+		if (retryModel) {
+			/**
+			* The model key uniquely identifies a model instance (provider + modelId)
+			*/
+			const retryModelKey = getModelKey(retryModel.model);
+			/**
+			* Find all attempts with the same model
+			*/
+			const retryAttempts = context.attempts.filter((a) => getModelKey(a.model) === retryModelKey);
+			const maxAttempts = retryModel.maxAttempts ?? 1;
+			/**
+			* Check if the model can still be retried based on maxAttempts
+			*/
+			if (retryAttempts.length < maxAttempts) return retryModel.model;
+		}
+	}
+}
+
+//#endregion
+//#region src/prepare-retry-error.ts
+/**
+* Prepare a RetryError that includes all errors from previous attempts.
+*/
+function prepareRetryError(error, attempts) {
+	const errorMessage = getErrorMessage(error);
+	const errors = attempts.flatMap((a) => isErrorAttempt(a) ? a.error : `Result with finishReason: ${a.result.finishReason}`);
+	return new RetryError({
+		message: `Failed after ${attempts.length} attempts. Last error: ${errorMessage}`,
+		reason: "maxRetriesExceeded",
+		errors
+	});
+}
+
+//#endregion
+//#region src/retryable-embedding-model.ts
+var RetryableEmbeddingModel = class {
+	specificationVersion = "v2";
+	baseModel;
+	currentModel;
+	options;
+	get modelId() {
+		return this.currentModel.modelId;
+	}
+	get provider() {
+		return this.currentModel.provider;
+	}
+	get maxEmbeddingsPerCall() {
+		return this.currentModel.maxEmbeddingsPerCall;
+	}
+	get supportsParallelCalls() {
+		return this.currentModel.supportsParallelCalls;
+	}
+	constructor(options) {
+		this.options = options;
+		this.baseModel = options.model;
+		this.currentModel = options.model;
+	}
+	/**
+	* Execute a function with retry logic for handling errors
+	*/
+	async withRetry(input) {
+		/**
+		* Track all attempts.
+		*/
+		const attempts = input.attempts ?? [];
+		while (true) {
+			/**
+			* The previous attempt that triggered a retry, or undefined if this is the first attempt
+			*/
+			const previousAttempt = attempts.at(-1);
+			/**
+			* Call the onRetry handler if provided.
+			* Skip on the first attempt since no previous attempt exists yet.
+			*/
+			if (previousAttempt) {
+				const currentAttempt = {
+					...previousAttempt,
+					model: this.currentModel
+				};
+				/**
+				* Create a shallow copy of the attempts for testing purposes
+				*/
+				const updatedAttempts = [...attempts];
+				const context = {
+					current: currentAttempt,
+					attempts: updatedAttempts
+				};
+				this.options.onRetry?.(context);
+			}
+			try {
+				return {
+					result: await input.fn(),
+					attempts
+				};
+			} catch (error) {
+				const { nextModel, attempt } = await this.handleError(error, attempts);
+				attempts.push(attempt);
+				this.currentModel = nextModel;
+			}
+		}
+	}
+	/**
+	* Handle an error and determine if a retry is needed
+	*/
+	async handleError(error, attempts) {
+		const errorAttempt = {
+			type: "error",
+			error,
+			model: this.currentModel
+		};
+		/**
+		* Save the current attempt
+		*/
+		const updatedAttempts = [...attempts, errorAttempt];
+		const context = {
+			current: errorAttempt,
+			attempts: updatedAttempts
+		};
+		this.options.onError?.(context);
+		const nextModel = await findRetryModel(this.options.retries, context);
+		/**
+		* Handler didn't return any models to try next, rethrow the error.
+		* If we retried the request, wrap the error into a `RetryError` for better visibility.
+		*/
+		if (!nextModel) {
+			if (updatedAttempts.length > 1) throw prepareRetryError(error, updatedAttempts);
+			throw error;
+		}
+		return {
+			nextModel,
+			attempt: errorAttempt
+		};
+	}
+	async doEmbed(options) {
+		/**
+		* Always start with the original model
+		*/
+		this.currentModel = this.baseModel;
+		const { result } = await this.withRetry({ fn: async () => await this.currentModel.doEmbed(options) });
+		return result;
+	}
+};
+
+//#endregion
+//#region src/retryable-language-model.ts
+var RetryableLanguageModel = class {
+	specificationVersion = "v2";
+	baseModel;
+	currentModel;
+	options;
+	get modelId() {
+		return this.currentModel.modelId;
+	}
+	get provider() {
+		return this.currentModel.provider;
+	}
+	get supportedUrls() {
+		return this.currentModel.supportedUrls;
+	}
+	constructor(options) {
+		this.options = options;
+		this.baseModel = options.model;
+		this.currentModel = options.model;
+	}
+	/**
+	* Execute a function with retry logic for handling errors
+	*/
+	async withRetry(input) {
+		/**
+		* Track all attempts.
+		*/
+		const attempts = input.attempts ?? [];
+		while (true) {
+			/**
+			* The previous attempt that triggered a retry, or undefined if this is the first attempt
+			*/
+			const previousAttempt = attempts.at(-1);
+			/**
+			* Call the onRetry handler if provided.
+			* Skip on the first attempt since no previous attempt exists yet.
+			*/
+			if (previousAttempt) {
+				const currentAttempt = {
+					...previousAttempt,
+					model: this.currentModel
+				};
+				/**
+				* Create a shallow copy of the attempts for testing purposes
+				*/
+				const updatedAttempts = [...attempts];
+				const context = {
+					current: currentAttempt,
+					attempts: updatedAttempts
+				};
+				this.options.onRetry?.(context);
+			}
+			try {
+				/**
+				* Call the function that may need to be retried
+				*/
+				const result = await input.fn();
+				/**
+				* Check if the result should trigger a retry (only for generate results, not streams)
+				*/
+				if (isGenerateResult(result)) {
+					const { nextModel, attempt } = await this.handleResult(result, attempts);
+					attempts.push(attempt);
+					if (nextModel) {
+						this.currentModel = nextModel;
+						/**
+						* Continue to the next iteration to retry
+						*/
+						continue;
+					}
+				}
+				return {
+					result,
+					attempts
+				};
+			} catch (error) {
+				const { nextModel, attempt } = await this.handleError(error, attempts);
+				attempts.push(attempt);
+				this.currentModel = nextModel;
+			}
+		}
+	}
+	/**
+	* Handle a successful result and determine if a retry is needed
+	*/
+	async handleResult(result, attempts) {
+		const resultAttempt = {
+			type: "result",
+			result,
+			model: this.currentModel
+		};
+		/**
+		* Save the current attempt
+		*/
+		const updatedAttempts = [...attempts, resultAttempt];
+		const context = {
+			current: resultAttempt,
+			attempts: updatedAttempts
+		};
+		return {
+			nextModel: await findRetryModel(this.options.retries, context),
+			attempt: resultAttempt
+		};
+	}
+	/**
+	* Handle an error and determine if a retry is needed
+	*/
+	async handleError(error, attempts) {
+		const errorAttempt = {
+			type: "error",
+			error,
+			model: this.currentModel
+		};
+		/**
+		* Save the current attempt
+		*/
+		const updatedAttempts = [...attempts, errorAttempt];
+		const context = {
+			current: errorAttempt,
+			attempts: updatedAttempts
+		};
+		this.options.onError?.(context);
+		const nextModel = await findRetryModel(this.options.retries, context);
+		/**
+		* Handler didn't return any models to try next, rethrow the error.
+		* If we retried the request, wrap the error into a `RetryError` for better visibility.
+		*/
+		if (!nextModel) {
+			if (updatedAttempts.length > 1) throw prepareRetryError(error, updatedAttempts);
+			throw error;
+		}
+		return {
+			nextModel,
+			attempt: errorAttempt
+		};
+	}
+	async doGenerate(options) {
+		/**
+		* Always start with the original model
+		*/
+		this.currentModel = this.baseModel;
+		const { result } = await this.withRetry({ fn: async () => await this.currentModel.doGenerate(options) });
+		return result;
+	}
+	async doStream(options) {
+		/**
+		* Always start with the original model
+		*/
+		this.currentModel = this.baseModel;
+		/**
+		* Perform the initial call to doStream with retry logic to handle errors before any data is streamed.
+		*/
+		let { result, attempts } = await this.withRetry({ fn: async () => await this.currentModel.doStream(options) });
+		/**
+		* Wrap the original stream to handle retries if an error occurs during streaming.
+		*/
+		const retryableStream = new ReadableStream({ start: async (controller) => {
+			let reader;
+			let isStreaming = false;
+			while (true) try {
+				reader = result.stream.getReader();
+				while (true) {
+					const { done, value } = await reader.read();
+					if (done) break;
+					/**
+					* If the stream part is an error and no data has been streamed yet, we can retry
+					* Throw the error to trigger the retry logic in withRetry
+					*/
+					if (value.type === "error") {
+						if (!isStreaming) throw value.error;
+					}
+					/**
+					* Mark that streaming has started once we receive actual content
+					*/
+					if (isStreamContentPart(value)) isStreaming = true;
+					/**
+					* Enqueue the chunk to the consumer of the stream
+					*/
+					controller.enqueue(value);
+				}
+				controller.close();
+				break;
+			} catch (error) {
+				/**
+				* Check if the error from the stream can be retried.
+				* Otherwise it will rethrow the error.
+				*/
+				const { nextModel, attempt } = await this.handleError(error, attempts);
+				this.currentModel = nextModel;
+				/**
+				* Save the attempt
+				*/
+				attempts.push(attempt);
+				/**
+				* Retry the request by calling doStream again.
+				* This will create a new stream.
+				*/
+				const retriedResult = await this.withRetry({
+					fn: async () => await this.currentModel.doStream(options),
+					attempts
+				});
+				/**
+				* Cancel the previous reader and stream if we are retrying
+				*/
+				await reader?.cancel();
+				result = retriedResult.result;
+				attempts = retriedResult.attempts;
+			} finally {
+				reader?.releaseLock();
+			}
+		} });
+		return {
+			...result,
+			stream: retryableStream
+		};
+	}
+};
+
+//#endregion
+//#region src/create-retryable-model.ts
+function createRetryable(options) {
+	if ("doEmbed" in options.model) return new RetryableEmbeddingModel(options);
+	return new RetryableLanguageModel(options);
+}
+
+//#endregion
+export { createRetryable, getModelKey };

package/dist/retryables/index.d.ts

@@ -1,50 +1,25 @@
-import { RetryModel, Retryable } from "../create-retryable-model-DtspEawi.js";
-import { LanguageModelV2 } from "@ai-sdk/provider";
+import { EmbeddingModelV2, LanguageModelV2, RetryModel, Retryable } from "../types-CHhEGL5x.js";
 
-//#region src/retryables/anthropic-service-overloaded.d.ts
-
-/**
- * Type for Anthropic error responses.
- *
- * @see https://docs.claude.com/en/api/errors#error-shapes
- */
-type AnthropicErrorResponse = {
-  type: 'error';
-  error: {
-    type: string;
-    message: string;
-  };
-};
-/**
- * Fallback if Anthropic returns an "overloaded" error with HTTP 200.
- *
- * ```
- * HTTP 200 OK
- * {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}
- * ```
- *
- * @deprecated Use `serviceOverloaded` instead
- */
-declare function anthropicServiceOverloaded(model: LanguageModelV2, options?: Omit<RetryModel, 'model'>): Retryable;
-//#endregion
 //#region src/retryables/content-filter-triggered.d.ts
+
 /**
  * Fallback to a different model if the content filter was triggered.
  */
-declare function contentFilterTriggered(model: LanguageModelV2, options?: Omit<RetryModel, 'model'>): Retryable;
+declare function contentFilterTriggered<MODEL extends LanguageModelV2>(model: MODEL, options?: Omit<RetryModel<MODEL>, 'model'>): Retryable<MODEL>;
 //#endregion
 //#region src/retryables/request-not-retryable.d.ts
 /**
  * Fallback to a different model if the error is non-retryable.
  */
-declare function requestNotRetryable(model: LanguageModelV2, options?: Omit<RetryModel, 'model'>): Retryable;
+declare function requestNotRetryable<MODEL extends LanguageModelV2 | EmbeddingModelV2>(model: MODEL, options?: Omit<RetryModel<MODEL>, 'model'>): Retryable<MODEL>;
 //#endregion
 //#region src/retryables/request-timeout.d.ts
 /**
  * Fallback to a different model after a timeout/abort error.
- * Use in combination with the `abortSignal` option in `generateText`.
+ * Use in combination with the `abortSignal` option.
+ * Works with both `LanguageModelV2` and `EmbeddingModelV2`.
  */
-declare function requestTimeout(model: LanguageModelV2, options?: Omit<RetryModel, 'model'>): Retryable;
+declare function requestTimeout<MODEL extends LanguageModelV2 | EmbeddingModelV2>(model: MODEL, options?: Omit<RetryModel<MODEL>, 'model'>): Retryable<MODEL>;
 //#endregion
 //#region src/retryables/service-overloaded.d.ts
 /**
@@ -54,6 +29,6 @@ declare function requestTimeout(model: LanguageModelV2, options?: Omit<RetryMode
  * - Response with `type: "overloaded_error"`
  * - Response with a `message` containing "overloaded"
  */
-declare function serviceOverloaded(model: LanguageModelV2, options?: Omit<RetryModel, 'model'>): Retryable;
+declare function serviceOverloaded<MODEL extends LanguageModelV2 | EmbeddingModelV2>(model: MODEL, options?: Omit<RetryModel<MODEL>, 'model'>): Retryable<MODEL>;
 //#endregion
-export { AnthropicErrorResponse, anthropicServiceOverloaded, contentFilterTriggered, requestNotRetryable, requestTimeout, serviceOverloaded };
+export { contentFilterTriggered, requestNotRetryable, requestTimeout, serviceOverloaded };

package/dist/retryables/index.js

@@ -1,38 +1,7 @@
-import { isErrorAttempt, isObject, isResultAttempt, isString } from "../create-retryable-model-C3mm1hSN.js";
+import { isErrorAttempt, isObject, isResultAttempt, isString } from "../utils-lRsC105f.js";
 import { isAbortError } from "@ai-sdk/provider-utils";
 import { APICallError } from "ai";
 
-//#region src/retryables/anthropic-service-overloaded.ts
-/**
-* Fallback if Anthropic returns an "overloaded" error with HTTP 200.
-*
-* ```
-* HTTP 200 OK
-* {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}
-* ```
-*
-* @deprecated Use `serviceOverloaded` instead
-*/
-function anthropicServiceOverloaded(model, options) {
-	return (context) => {
-		const { current } = context;
-		if (isErrorAttempt(current)) {
-			const { error } = current;
-			if (APICallError.isInstance(error) && error.statusCode === 529) return {
-				model,
-				maxAttempts: 1,
-				...options
-			};
-			if (isObject(error) && isString(error.type) && error.type === "overloaded_error") return {
-				model,
-				maxAttempts: 1,
-				...options
-			};
-		}
-	};
-}
-
-//#endregion
 //#region src/retryables/content-filter-triggered.ts
 /**
 * Fallback to a different model if the content filter was triggered.
@@ -81,7 +50,8 @@ function requestNotRetryable(model, options) {
 //#region src/retryables/request-timeout.ts
 /**
 * Fallback to a different model after a timeout/abort error.
-* Use in combination with the `abortSignal` option in `generateText`.
+* Use in combination with the `abortSignal` option.
+* Works with both `LanguageModelV2` and `EmbeddingModelV2`.
 */
 function requestTimeout(model, options) {
 	return (context) => {
@@ -130,4 +100,4 @@ function serviceOverloaded(model, options) {
 }
 
 //#endregion
-export { anthropicServiceOverloaded, contentFilterTriggered, requestNotRetryable, requestTimeout, serviceOverloaded };
+export { contentFilterTriggered, requestNotRetryable, requestTimeout, serviceOverloaded };

package/dist/types-CHhEGL5x.d.ts

@@ -0,0 +1,63 @@
+import { EmbeddingModelV2, LanguageModelV2 } from "@ai-sdk/provider";
+
+//#region src/types.d.ts
+type EmbeddingModelV2$1<VALUE = any> = EmbeddingModelV2<VALUE>;
+/**
+ * Options for creating a retryable model.
+ */
+interface RetryableModelOptions<MODEL extends LanguageModelV2 | EmbeddingModelV2$1> {
+  model: MODEL;
+  retries: Retries<MODEL>;
+  onError?: (context: RetryContext<MODEL>) => void;
+  onRetry?: (context: RetryContext<MODEL>) => void;
+}
+/**
+ * The context provided to Retryables with the current attempt and all previous attempts.
+ */
+type RetryContext<MODEL extends LanguageModelV2 | EmbeddingModelV2$1> = {
+  /**
+   * Current attempt that caused the retry
+   */
+  current: RetryAttempt<MODEL>;
+  /**
+   * All attempts made so far, including the current one
+   */
+  attempts: Array<RetryAttempt<MODEL>>;
+};
+/**
+ * A retry attempt with an error
+ */
+type RetryErrorAttempt<MODEL extends LanguageModelV2 | EmbeddingModelV2$1> = {
+  type: 'error';
+  error: unknown;
+  result?: undefined;
+  model: MODEL;
+};
+/**
+ * A retry attempt with a successful result
+ */
+type RetryResultAttempt = {
+  type: 'result';
+  result: LanguageModelV2Generate;
+  error?: undefined;
+  model: LanguageModelV2;
+};
+/**
+ * A retry attempt with either an error or a result and the model used
+ */
+type RetryAttempt<MODEL extends LanguageModelV2 | EmbeddingModelV2$1> = RetryErrorAttempt<MODEL> | RetryResultAttempt;
+/**
+ * A model to retry with and the maximum number of attempts for that model.
+ */
+type RetryModel<MODEL extends LanguageModelV2 | EmbeddingModelV2$1> = {
+  model: MODEL;
+  maxAttempts?: number;
+};
+/**
+ * A function that determines whether to retry with a different model based on the current attempt and all previous attempts.
+ */
+type Retryable<MODEL extends LanguageModelV2 | EmbeddingModelV2$1> = (context: RetryContext<MODEL>) => RetryModel<MODEL> | Promise<RetryModel<MODEL>> | undefined;
+type Retries<MODEL extends LanguageModelV2 | EmbeddingModelV2$1> = Array<Retryable<MODEL> | MODEL>;
+type LanguageModelV2Generate = Awaited<ReturnType<LanguageModelV2['doGenerate']>>;
+//#endregion
+export { EmbeddingModelV2$1 as EmbeddingModelV2, type LanguageModelV2, RetryModel, Retryable, RetryableModelOptions };

package/dist/utils-lRsC105f.js

@@ -0,0 +1,27 @@
+//#region src/utils.ts
+const isObject = (value) => typeof value === "object" && value !== null;
+const isString = (value) => typeof value === "string";
+const isGenerateResult = (result) => "content" in result;
+/**
+* Type guard to check if a retry attempt is an error attempt
+*/
+function isErrorAttempt(attempt) {
+	return attempt.type === "error";
+}
+/**
+* Type guard to check if a retry attempt is a result attempt
+*/
+function isResultAttempt(attempt) {
+	return attempt.type === "result";
+}
+/**
+* Check if a stream part is a content part (e.g., text delta, reasoning delta, source, tool call, tool result).
+* These types are also emitted by `onChunk` callbacks.
+* @see https://github.com/vercel/ai/blob/1fe4bd4144bff927f5319d9d206e782a73979ccb/packages/ai/src/generate-text/stream-text.ts#L686-L697
+*/
+const isStreamContentPart = (part) => {
+	return part.type === "text-delta" || part.type === "reasoning-delta" || part.type === "source" || part.type === "tool-call" || part.type === "tool-result" || part.type === "tool-input-start" || part.type === "tool-input-delta" || part.type === "raw";
+};
+
+//#endregion
+export { isErrorAttempt, isGenerateResult, isObject, isResultAttempt, isStreamContentPart, isString };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai-retry",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "description": "AI SDK Retry",
   "packageManager": "pnpm@9.0.0",
   "main": "./dist/index.js",

package/dist/create-retryable-model-C3mm1hSN.js

@@ -1,308 +0,0 @@
-import { getErrorMessage } from "@ai-sdk/provider-utils";
-import { RetryError } from "ai";
-
-//#region src/get-model-key.ts
-/**
-* Generate a unique key for a LanguageModelV2 instance.
-*/
-const getModelKey = (model) => {
-	return `${model.provider}/${model.modelId}`;
-};
-
-//#endregion
-//#region src/utils.ts
-const isObject = (value) => typeof value === "object" && value !== null;
-const isString = (value) => typeof value === "string";
-const isGenerateResult = (result) => "content" in result;
-/**
-* Check if a stream part is a content part (e.g., text delta, reasoning delta, source, tool call, tool result).
-* These types are also emitted by `onChunk` callbacks.
-* @see https://github.com/vercel/ai/blob/1fe4bd4144bff927f5319d9d206e782a73979ccb/packages/ai/src/generate-text/stream-text.ts#L686-L697
-*/
-const isStreamContentPart = (part) => {
-	return part.type === "text-delta" || part.type === "reasoning-delta" || part.type === "source" || part.type === "tool-call" || part.type === "tool-result" || part.type === "tool-input-start" || part.type === "tool-input-delta" || part.type === "raw";
-};
-
-//#endregion
-//#region src/create-retryable-model.ts
-/**
-* Type guard to check if a retry attempt is an error attempt
-*/
-function isErrorAttempt(attempt) {
-	return attempt.type === "error";
-}
-/**
-* Type guard to check if a retry attempt is a result attempt
-*/
-function isResultAttempt(attempt) {
-	return attempt.type === "result";
-}
-var RetryableModel = class {
-	specificationVersion = "v2";
-	baseModel;
-	currentModel;
-	options;
-	get modelId() {
-		return this.currentModel.modelId;
-	}
-	get provider() {
-		return this.currentModel.provider;
-	}
-	get supportedUrls() {
-		return this.currentModel.supportedUrls;
-	}
-	constructor(options) {
-		this.options = options;
-		this.baseModel = options.model;
-		this.currentModel = options.model;
-	}
-	/**
-	* Find the next model to retry with based on the retry context
-	*/
-	async findNextModel(context) {
-		/**
-		* Filter retryables based on attempt type:
-		* - Result-based attempts: Only consider function retryables (skip plain models)
-		* - Error-based attempts: Consider all retryables (functions + plain models)
-		*/
-		const applicableRetries = isResultAttempt(context.current) ? this.options.retries.filter((retry) => typeof retry === "function") : this.options.retries;
-		/**
-		* Iterate through the applicable retryables to find a model to retry with
-		*/
-		for (const retry of applicableRetries) {
-			const retryModel = typeof retry === "function" ? await retry(context) : {
-				model: retry,
-				maxAttempts: 1
-			};
-			if (retryModel) {
-				/**
-				* The model key uniquely identifies a model instance (provider + modelId)
-				*/
-				const retryModelKey = getModelKey(retryModel.model);
-				/**
-				* Find all attempts with the same model
-				*/
-				const retryAttempts = context.attempts.filter((a) => getModelKey(a.model) === retryModelKey);
-				const maxAttempts = retryModel.maxAttempts ?? 1;
-				/**
-				* Check if the model can still be retried based on maxAttempts
-				*/
-				if (retryAttempts.length < maxAttempts) return retryModel.model;
-			}
-		}
-	}
-	/**
-	* Execute a function with retry logic for handling errors
-	*/
-	async withRetry(input) {
-		/**
-		* Track all attempts.
-		*/
-		const attempts = input.attempts ?? [];
-		while (true) {
-			/**
-			* The previous attempt that triggered a retry, or undefined if this is the first attempt
-			*/
-			const previousAttempt = attempts.at(-1);
-			/**
-			* Call the onRetry handler if provided.
-			* Skip on the first attempt since no previous attempt exists yet.
-			*/
-			if (previousAttempt) {
-				const currentAttempt = {
-					...previousAttempt,
-					model: this.currentModel
-				};
-				/**
-				* Create a shallow copy of the attempts for testing purposes
-				*/
-				const updatedAttempts = [...attempts];
-				const context = {
-					current: currentAttempt,
-					attempts: updatedAttempts,
-					totalAttempts: updatedAttempts.length
-				};
-				this.options.onRetry?.(context);
-			}
-			try {
-				/**
-				* Call the function that may need to be retried
-				*/
-				const result = await input.fn();
-				/**
-				* Check if the result should trigger a retry (only for generate results, not streams)
-				*/
-				if (isGenerateResult(result)) {
-					const { nextModel, attempt } = await this.handleResult(result, attempts);
-					attempts.push(attempt);
-					if (nextModel) {
-						this.currentModel = nextModel;
-						/**
-						* Continue to the next iteration to retry
-						*/
-						continue;
-					}
-				}
-				return {
-					result,
-					attempts
-				};
-			} catch (error) {
-				const { nextModel, attempt } = await this.handleError(error, attempts);
-				attempts.push(attempt);
-				this.currentModel = nextModel;
-			}
-		}
-	}
-	/**
-	* Handle a successful result and determine if a retry is needed
-	*/
-	async handleResult(result, attempts) {
-		const resultAttempt = {
-			type: "result",
-			result,
-			model: this.currentModel
-		};
-		/**
-		* Save the current attempt
-		*/
-		const updatedAttempts = [...attempts, resultAttempt];
-		const resultContext = {
-			current: resultAttempt,
-			attempts: updatedAttempts,
-			totalAttempts: updatedAttempts.length
-		};
-		return {
-			nextModel: await this.findNextModel(resultContext),
-			attempt: resultAttempt
-		};
-	}
-	/**
-	* Handle an error and determine if a retry is needed
-	*/
-	async handleError(error, attempts) {
-		const errorAttempt = {
-			type: "error",
-			error,
-			model: this.currentModel
-		};
-		/**
-		* Save the current attempt
-		*/
-		const updatedAttempts = [...attempts, errorAttempt];
-		const context = {
-			current: errorAttempt,
-			attempts: updatedAttempts,
-			totalAttempts: updatedAttempts.length
-		};
-		this.options.onError?.(context);
-		const nextModel = await this.findNextModel(context);
-		/**
-		* Handler didn't return any models to try next, rethrow the error.
-		* If we retried the request, wrap the error into a `RetryError` for better visibility.
-		*/
-		if (!nextModel) {
-			if (updatedAttempts.length > 1) throw this.prepareRetryError(error, updatedAttempts);
-			throw error;
-		}
-		return {
-			nextModel,
-			attempt: errorAttempt
-		};
-	}
-	async doGenerate(options) {
-		/**
-		* Always start with the original model
-		*/
-		this.currentModel = this.baseModel;
-		const { result } = await this.withRetry({ fn: async () => await this.currentModel.doGenerate(options) });
-		return result;
-	}
-	async doStream(options) {
-		/**
-		* Always start with the original model
-		*/
-		this.currentModel = this.baseModel;
-		/**
-		* Perform the initial call to doStream with retry logic to handle errors before any data is streamed.
-		*/
-		let { result, attempts } = await this.withRetry({ fn: async () => await this.currentModel.doStream(options) });
-		/**
-		* Wrap the original stream to handle retries if an error occurs during streaming.
-		*/
-		const retryableStream = new ReadableStream({ start: async (controller) => {
-			let reader;
-			let isStreaming = false;
-			while (true) try {
-				reader = result.stream.getReader();
-				while (true) {
-					const { done, value } = await reader.read();
-					if (done) break;
-					/**
-					* If the stream part is an error and no data has been streamed yet, we can retry
-					* Throw the error to trigger the retry logic in withRetry
-					*/
-					if (value.type === "error") {
-						if (!isStreaming) throw value.error;
-					}
-					/**
-					* Mark that streaming has started once we receive actual content
-					*/
-					if (isStreamContentPart(value)) isStreaming = true;
-					/**
-					* Enqueue the chunk to the consumer of the stream
-					*/
-					controller.enqueue(value);
-				}
-				controller.close();
-				break;
-			} catch (error) {
-				/**
-				* Check if the error from the stream can be retried.
-				* Otherwise it will rethrow the error.
-				*/
-				const { nextModel, attempt } = await this.handleError(error, attempts);
-				this.currentModel = nextModel;
-				/**
-				* Save the attempt
-				*/
-				attempts.push(attempt);
-				/**
-				* Retry the request by calling doStream again.
-				* This will create a new stream.
-				*/
-				const retriedResult = await this.withRetry({
-					fn: async () => await this.currentModel.doStream(options),
-					attempts
-				});
-				/**
-				* Cancel the previous reader and stream if we are retrying
-				*/
-				await reader?.cancel();
-				result = retriedResult.result;
-				attempts = retriedResult.attempts;
-			} finally {
-				reader?.releaseLock();
-			}
-		} });
-		return {
-			...result,
-			stream: retryableStream
-		};
-	}
-	prepareRetryError(error, attempts) {
-		const errorMessage = getErrorMessage(error);
-		const errors = attempts.flatMap((a) => isErrorAttempt(a) ? a.error : `Result with finishReason: ${a.result.finishReason}`);
-		return new RetryError(new RetryError({
-			message: `Failed after ${attempts.length} attempts. Last error: ${errorMessage}`,
-			reason: "maxRetriesExceeded",
-			errors
-		}));
-	}
-};
-function createRetryable(config) {
-	return new RetryableModel(config);
-}
-
-//#endregion
-export { createRetryable, getModelKey, isErrorAttempt, isObject, isResultAttempt, isString };

package/dist/create-retryable-model-DtspEawi.d.ts

@@ -1,74 +0,0 @@
-import { LanguageModelV2 } from "@ai-sdk/provider";
-
-//#region src/types.d.ts
-type LanguageModelV2Generate = Awaited<ReturnType<LanguageModelV2['doGenerate']>>;
-//#endregion
-//#region src/create-retryable-model.d.ts
-/**
- * The context provided to Retryables with the current attempt and all previous attempts.
- */
-interface RetryContext<CURRENT extends RetryAttempt = RetryAttempt> {
-  /**
-   * Current attempt that caused the retry
-   */
-  current: CURRENT;
-  /**
-   * All attempts made so far, including the current one
-   */
-  attempts: Array<RetryAttempt>;
-  /**
-   * @deprecated Use `attempts.length` instead
-   */
-  totalAttempts: number;
-}
-/**
- * A retry attempt with an error
- */
-type RetryErrorAttempt = {
-  type: 'error';
-  error: unknown;
-  model: LanguageModelV2;
-};
-/**
- * A retry attempt with a successful result
- */
-type RetryResultAttempt = {
-  type: 'result';
-  result: LanguageModelV2Generate;
-  model: LanguageModelV2;
-};
-/**
- * A retry attempt with either an error or a result and the model used
- */
-type RetryAttempt = RetryErrorAttempt | RetryResultAttempt;
-/**
- * Type guard to check if a retry attempt is an error attempt
- */
-declare function isErrorAttempt(attempt: RetryAttempt): attempt is RetryErrorAttempt;
-/**
- * Type guard to check if a retry attempt is a result attempt
- */
-declare function isResultAttempt(attempt: RetryAttempt): attempt is RetryResultAttempt;
-/**
- * A model to retry with and the maximum number of attempts for that model.
- */
-type RetryModel = {
-  model: LanguageModelV2;
-  maxAttempts?: number;
-};
-/**
- * A function that determines whether to retry with a different model based on the current attempt and all previous attempts.
- */
-type Retryable = (context: RetryContext) => RetryModel | Promise<RetryModel> | undefined;
-/**
- * Options for creating a retryable model.
- */
-interface CreateRetryableOptions {
-  model: LanguageModelV2;
-  retries: Array<Retryable | LanguageModelV2>;
-  onError?: (context: RetryContext<RetryErrorAttempt>) => void;
-  onRetry?: (context: RetryContext<RetryErrorAttempt | RetryResultAttempt>) => void;
-}
-declare function createRetryable(config: CreateRetryableOptions): LanguageModelV2;
-//#endregion
-export { CreateRetryableOptions, RetryAttempt, RetryContext, RetryModel, Retryable, createRetryable, isErrorAttempt, isResultAttempt };