fix(ai-sdk): use embedding models for embedding requests

packages/openai-adapters/src/apis/AiSdk.ts

@@ -28,7 +28,10 @@ type AiSdkProviderCreator = (options: {
   apiKey?: string;
   baseURL?: string;
   fetch?: typeof fetch;
-}) => (modelId: string) => any;
+}) => {
+  (modelId: string): any;
+  embeddingModel(modelId: string): any;
+};
 
 const PROVIDER_MAP: Record<string, AiSdkProviderCreator> = {
   openai: createOpenAI,
@@ -49,7 +52,7 @@ const PROVIDER_MAP: Record<string, AiSdkProviderCreator> = {
 };
 
 export class AiSdkApi implements BaseLlmApi {
-  private provider?: (modelId: string) => any;
+  private provider?: ReturnType<AiSdkProviderCreator>;
   private config: AiSdkConfig;
   private providerId: string;
   private modelId: string;
@@ -267,8 +270,8 @@ export class AiSdkApi implements BaseLlmApi {
 
     const { embed: aiEmbed, embedMany } = await import("ai");
 
-    const modelId = typeof body.model === "string" ? body.model : body.model;
-    const model = this.provider!(modelId);
+    const modelId = this.modelId;
+    const model = this.provider!.embeddingModel(modelId);
 
     const inputs = Array.isArray(body.input) ? body.input : [body.input];
     const stringInputs = inputs.map((input) =>

packages/openai-adapters/src/test/ai-sdk-embeddings.test.ts

@@ -0,0 +1,63 @@
+import { afterEach, describe, expect, test, vi } from "vitest";
+import { AiSdkApi } from "../apis/AiSdk.js";
+
+afterEach(() => {
+  vi.unstubAllGlobals();
+});
+
+describe("AI SDK embeddings", () => {
+  test.each([{ input: "hello" }, { input: ["hello", "world"] }])(
+    "uses an embedding model for input $input",
+    async ({ input }) => {
+      const requests: { url: string; body: Record<string, unknown> }[] = [];
+      vi.stubGlobal(
+        "fetch",
+        vi.fn(async (url: string | URL | Request, init?: RequestInit) => {
+          const body = JSON.parse(String(init?.body));
+          requests.push({ url: String(url), body });
+          return new Response(
+            JSON.stringify({
+              model: "text-embedding-3-small",
+              data: body.input.map((_: string, index: number) => ({
+                embedding: [index, 0.5],
+                index,
+              })),
+              usage: { prompt_tokens: 2, total_tokens: 2 },
+            }),
+            { status: 200, headers: { "Content-Type": "application/json" } },
+          );
+        }),
+      );
+      const api = new AiSdkApi({
+        provider: "ai-sdk",
+        model: "openai/text-embedding-3-small",
+        apiKey: "test-key",
+      });
+
+      const result = await api.embed({
+        model: "openai/text-embedding-3-small",
+        input,
+      });
+
+      expect(requests).toEqual([
+        {
+          url: "https://api.openai.com/v1/embeddings",
+          body: {
+            model: "text-embedding-3-small",
+            input: Array.isArray(input) ? input : [input],
+            encoding_format: "float",
+          },
+        },
+      ]);
+      expect(result.model).toBe("text-embedding-3-small");
+      expect(result.data).toEqual(
+        (Array.isArray(input) ? input : [input]).map((_, index) => ({
+          object: "embedding",
+          embedding: [index, 0.5],
+          index,
+        })),
+      );
+      expect(result.usage).toEqual({ prompt_tokens: 2, total_tokens: 2 });
+    },
+  );
+});