ai 6.0.32 → 6.0.34

package/src/middleware/simulate-streaming-middleware.ts

@@ -0,0 +1,79 @@
+import type { LanguageModelV3StreamPart } from '@ai-sdk/provider';
+import { LanguageModelMiddleware } from '../types';
+
+/**
+ * Simulates streaming chunks with the response from a generate call.
+ */
+export function simulateStreamingMiddleware(): LanguageModelMiddleware {
+  return {
+    specificationVersion: 'v3',
+    wrapStream: async ({ doGenerate }) => {
+      const result = await doGenerate();
+
+      let id = 0;
+
+      const simulatedStream = new ReadableStream<LanguageModelV3StreamPart>({
+        start(controller) {
+          controller.enqueue({
+            type: 'stream-start',
+            warnings: result.warnings,
+          });
+
+          controller.enqueue({ type: 'response-metadata', ...result.response });
+
+          for (const part of result.content) {
+            switch (part.type) {
+              case 'text': {
+                if (part.text.length > 0) {
+                  controller.enqueue({ type: 'text-start', id: String(id) });
+                  controller.enqueue({
+                    type: 'text-delta',
+                    id: String(id),
+                    delta: part.text,
+                  });
+                  controller.enqueue({ type: 'text-end', id: String(id) });
+                  id++;
+                }
+                break;
+              }
+              case 'reasoning': {
+                controller.enqueue({
+                  type: 'reasoning-start',
+                  id: String(id),
+                  providerMetadata: part.providerMetadata,
+                });
+                controller.enqueue({
+                  type: 'reasoning-delta',
+                  id: String(id),
+                  delta: part.text,
+                });
+                controller.enqueue({ type: 'reasoning-end', id: String(id) });
+                id++;
+                break;
+              }
+              default: {
+                controller.enqueue(part);
+                break;
+              }
+            }
+          }
+
+          controller.enqueue({
+            type: 'finish',
+            finishReason: result.finishReason,
+            usage: result.usage,
+            providerMetadata: result.providerMetadata,
+          });
+
+          controller.close();
+        },
+      });
+
+      return {
+        stream: simulatedStream,
+        request: result.request,
+        response: result.response,
+      };
+    },
+  };
+}

package/src/middleware/wrap-embedding-model.test.ts

@@ -0,0 +1,358 @@
+import {
+  EmbeddingModelV3CallOptions,
+  EmbeddingModelV3Middleware,
+} from '@ai-sdk/provider';
+import { wrapEmbeddingModel } from '../middleware/wrap-embedding-model';
+import { describe, it, expect, vi } from 'vitest';
+import { MockEmbeddingModelV3 } from '../test/mock-embedding-model-v3';
+
+describe('wrapEmbeddingModel', () => {
+  describe('model property', () => {
+    it('should pass through by default', () => {
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({
+          modelId: 'test-model',
+        }),
+        middleware: {
+          specificationVersion: 'v3',
+        },
+      });
+
+      expect(wrappedModel.modelId).toBe('test-model');
+    });
+
+    it('should use middleware overrideModelId if provided', () => {
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({
+          modelId: 'test-model',
+        }),
+        middleware: {
+          specificationVersion: 'v3',
+          overrideModelId: ({ model }) => 'override-model',
+        },
+      });
+
+      expect(wrappedModel.modelId).toBe('override-model');
+    });
+
+    it('should use modelId parameter if provided', () => {
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({
+          modelId: 'test-model',
+        }),
+        middleware: {
+          specificationVersion: 'v3',
+        },
+        modelId: 'override-model',
+      });
+
+      expect(wrappedModel.modelId).toBe('override-model');
+    });
+  });
+
+  describe('provider property', () => {
+    it('should pass through by default', () => {
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({
+          provider: 'test-provider',
+        }),
+        middleware: {
+          specificationVersion: 'v3',
+        },
+      });
+
+      expect(wrappedModel.provider).toBe('test-provider');
+    });
+
+    it('should use middleware overrideProvider if provided', () => {
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({
+          provider: 'test-provider',
+        }),
+        middleware: {
+          specificationVersion: 'v3',
+          overrideProvider: ({ model }) => 'override-provider',
+        },
+      });
+
+      expect(wrappedModel.provider).toBe('override-provider');
+    });
+
+    it('should use providerId parameter if provided', () => {
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({
+          provider: 'test-provider',
+        }),
+        middleware: {
+          specificationVersion: 'v3',
+        },
+        providerId: 'override-provider',
+      });
+
+      expect(wrappedModel.provider).toBe('override-provider');
+    });
+  });
+
+  describe('maxEmbeddingsPerCall property', () => {
+    it('should pass through by default', async () => {
+      const maxEmbeddingsPerCall = 2;
+
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({ maxEmbeddingsPerCall }),
+        middleware: {
+          specificationVersion: 'v3',
+        },
+      });
+
+      expect(await wrappedModel.maxEmbeddingsPerCall).toStrictEqual(2);
+    });
+
+    it('should use middleware overrideSupportedUrls if provided', () => {
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({
+          maxEmbeddingsPerCall: 2,
+        }),
+        middleware: {
+          specificationVersion: 'v3',
+          overrideMaxEmbeddingsPerCall: ({ model }) => 3,
+        },
+      });
+
+      expect(wrappedModel.maxEmbeddingsPerCall).toStrictEqual(3);
+    });
+  });
+
+  describe('supportsParallelCalls property', () => {
+    it('should pass through by default', async () => {
+      const supportsParallelCalls = true;
+
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({ supportsParallelCalls }),
+        middleware: {
+          specificationVersion: 'v3',
+        },
+      });
+
+      expect(await wrappedModel.supportsParallelCalls).toStrictEqual(true);
+    });
+
+    it('should use middleware overrideSupportsParallelCalls if provided', () => {
+      const wrappedModel = wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3({
+          supportsParallelCalls: false,
+        }),
+        middleware: {
+          specificationVersion: 'v3',
+          overrideSupportsParallelCalls: ({ model }) => true,
+        },
+      });
+
+      expect(wrappedModel.supportsParallelCalls).toStrictEqual(true);
+    });
+  });
+
+  it('should call transformParams middleware for doEmbed', async () => {
+    const mockModel = new MockEmbeddingModelV3({
+      doEmbed: [],
+    });
+
+    const transformParams = vi.fn().mockImplementation(({ params }) => ({
+      ...params,
+      transformed: true,
+    }));
+
+    const wrappedModel = wrapEmbeddingModel({
+      model: mockModel,
+      middleware: {
+        specificationVersion: 'v3',
+        transformParams,
+      },
+    });
+
+    const params: EmbeddingModelV3CallOptions = {
+      values: [
+        'sunny day at the beach',
+        'rainy afternoon in the city',
+        'snowy night in the mountains',
+      ],
+    };
+
+    await wrappedModel.doEmbed(params);
+
+    expect(transformParams).toHaveBeenCalledWith({
+      params,
+      model: expect.any(Object),
+    });
+
+    expect(mockModel.doEmbedCalls[0]).toStrictEqual({
+      ...params,
+      transformed: true,
+    });
+  });
+
+  it('should call wrapEmbed middleware', async () => {
+    const mockModel = new MockEmbeddingModelV3({
+      doEmbed: vi.fn().mockResolvedValue('mock result'),
+    });
+
+    const wrapEmbed = vi.fn().mockImplementation(({ doEmbed }) => doEmbed());
+
+    const wrappedModel = wrapEmbeddingModel({
+      model: mockModel,
+      middleware: {
+        specificationVersion: 'v3',
+        wrapEmbed,
+      },
+    });
+
+    const params: EmbeddingModelV3CallOptions = {
+      values: [
+        'sunny day at the beach',
+        'rainy afternoon in the city',
+        'snowy night in the mountains',
+      ],
+    };
+
+    await wrappedModel.doEmbed(params);
+
+    expect(wrapEmbed).toHaveBeenCalledWith({
+      doEmbed: expect.any(Function),
+      params,
+      model: mockModel,
+    });
+  });
+
+  describe('multiple middlewares', () => {
+    it('should call multiple transformParams middlewares in sequence for doEmbed', async () => {
+      const mockModel = new MockEmbeddingModelV3({
+        doEmbed: [],
+      });
+
+      const transformParams1 = vi.fn().mockImplementation(({ params }) => ({
+        ...params,
+        transformationStep1: true,
+      }));
+
+      const transformParams2 = vi.fn().mockImplementation(({ params }) => ({
+        ...params,
+        transformationStep2: true,
+      }));
+
+      const wrappedModel = wrapEmbeddingModel({
+        model: mockModel,
+        middleware: [
+          {
+            specificationVersion: 'v3',
+            transformParams: transformParams1,
+          },
+          {
+            specificationVersion: 'v3',
+            transformParams: transformParams2,
+          },
+        ],
+      });
+
+      const params: EmbeddingModelV3CallOptions = {
+        values: [
+          'sunny day at the beach',
+          'rainy afternoon in the city',
+          'snowy night in the mountains',
+        ],
+      };
+
+      await wrappedModel.doEmbed(params);
+
+      expect(transformParams1).toHaveBeenCalledWith({
+        params,
+        model: expect.any(Object),
+      });
+
+      expect(transformParams2).toHaveBeenCalledWith({
+        params: { ...params, transformationStep1: true },
+        model: expect.any(Object),
+      });
+
+      expect(mockModel.doEmbedCalls[0]).toStrictEqual(
+        expect.objectContaining({
+          transformationStep1: true,
+          transformationStep2: true,
+        }),
+      );
+    });
+
+    it('should chain multiple wrapEmbed middlewares in the correct order', async () => {
+      const mockModel = new MockEmbeddingModelV3({
+        doEmbed: vi.fn().mockResolvedValue('final generate result'),
+      });
+
+      const wrapEmbed1 = vi
+        .fn()
+        .mockImplementation(async ({ doEmbed, params, model }) => {
+          const result = await doEmbed();
+          return `wrapEmbed1(${result})`;
+        });
+
+      const wrapEmbed2 = vi
+        .fn()
+        .mockImplementation(async ({ doEmbed, params, model }) => {
+          const result = await doEmbed();
+          return `wrapEmbed2(${result})`;
+        });
+
+      const wrappedModel = wrapEmbeddingModel({
+        model: mockModel,
+        middleware: [
+          {
+            specificationVersion: 'v3',
+            wrapEmbed: wrapEmbed1,
+          },
+          {
+            specificationVersion: 'v3',
+            wrapEmbed: wrapEmbed2,
+          },
+        ],
+      });
+
+      const params: EmbeddingModelV3CallOptions = {
+        values: [
+          'sunny day at the beach',
+          'rainy afternoon in the city',
+          'snowy night in the mountains',
+        ],
+      };
+
+      const result = await wrappedModel.doEmbed(params);
+
+      // The middlewares should wrap in order, applying wrapEmbed2 last
+      expect(result).toBe('wrapEmbed1(wrapEmbed2(final generate result))');
+      expect(wrapEmbed1).toHaveBeenCalled();
+      expect(wrapEmbed2).toHaveBeenCalled();
+    });
+
+    it('should not mutate the middleware array argument', async () => {
+      const middleware1 = {
+        specificationVersion: 'v3',
+        wrapStream: vi.fn(),
+      };
+
+      const middleware2 = {
+        specificationVersion: 'v3',
+        wrapStream: vi.fn(),
+      };
+
+      const middlewares = [
+        middleware1,
+        middleware2,
+      ] as EmbeddingModelV3Middleware[];
+
+      wrapEmbeddingModel({
+        model: new MockEmbeddingModelV3(),
+        middleware: middlewares,
+      });
+
+      expect(middlewares.length).toBe(2);
+      expect(middlewares[0]).toBe(middleware1);
+      expect(middlewares[1]).toBe(middleware2);
+    });
+  });
+});

package/src/middleware/wrap-embedding-model.ts

@@ -0,0 +1,86 @@
+import {
+  EmbeddingModelV3,
+  EmbeddingModelV3CallOptions,
+} from '@ai-sdk/provider';
+import { EmbeddingModelMiddleware } from '../types';
+import { asArray } from '../util/as-array';
+
+/**
+ * Wraps a EmbeddingModelV3 instance with middleware functionality.
+ * This function allows you to apply middleware to transform parameters,
+ * wrap embed operations of a language model.
+ *
+ * @param options - Configuration options for wrapping the embedding model.
+ * @param options.model - The original EmbeddingModelV3 instance to be wrapped.
+ * @param options.middleware - The middleware to be applied to the embedding model. When multiple middlewares are provided, the first middleware will transform the input first, and the last middleware will be wrapped directly around the model.
+ * @param options.modelId - Optional custom model ID to override the original model's ID.
+ * @param options.providerId - Optional custom provider ID to override the original model's provider ID.
+ * @returns A new EmbeddingModelV3 instance with middleware applied.
+ */
+export const wrapEmbeddingModel = ({
+  model,
+  middleware: middlewareArg,
+  modelId,
+  providerId,
+}: {
+  model: EmbeddingModelV3;
+  middleware: EmbeddingModelMiddleware | EmbeddingModelMiddleware[];
+  modelId?: string;
+  providerId?: string;
+}): EmbeddingModelV3 => {
+  return [...asArray(middlewareArg)]
+    .reverse()
+    .reduce((wrappedModel, middleware) => {
+      return doWrap({ model: wrappedModel, middleware, modelId, providerId });
+    }, model);
+};
+
+const doWrap = ({
+  model,
+  middleware: {
+    transformParams,
+    wrapEmbed,
+    overrideProvider,
+    overrideModelId,
+    overrideMaxEmbeddingsPerCall,
+    overrideSupportsParallelCalls,
+  },
+  modelId,
+  providerId,
+}: {
+  model: EmbeddingModelV3;
+  middleware: EmbeddingModelMiddleware;
+  modelId?: string;
+  providerId?: string;
+}): EmbeddingModelV3 => {
+  async function doTransform({
+    params,
+  }: {
+    params: EmbeddingModelV3CallOptions;
+  }) {
+    return transformParams ? await transformParams({ params, model }) : params;
+  }
+
+  return {
+    specificationVersion: 'v3',
+    provider: providerId ?? overrideProvider?.({ model }) ?? model.provider,
+    modelId: modelId ?? overrideModelId?.({ model }) ?? model.modelId,
+    maxEmbeddingsPerCall:
+      overrideMaxEmbeddingsPerCall?.({ model }) ?? model.maxEmbeddingsPerCall,
+    supportsParallelCalls:
+      overrideSupportsParallelCalls?.({ model }) ?? model.supportsParallelCalls,
+    async doEmbed(
+      params: EmbeddingModelV3CallOptions,
+    ): Promise<Awaited<ReturnType<EmbeddingModelV3['doEmbed']>>> {
+      const transformedParams = await doTransform({ params });
+      const doEmbed = async () => model.doEmbed(transformedParams);
+      return wrapEmbed
+        ? wrapEmbed({
+            doEmbed,
+            params: transformedParams,
+            model,
+          })
+        : doEmbed();
+    },
+  };
+};