diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ba482f6..5f595a8f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - Fix TS2590 when constructing Agent or calling startGeneration with large model catalogs (#414) +- Pass embedding providerMetadata to usageHandler (#393) ## 0.7.8 diff --git a/src/vercel/client/search.test.ts b/src/vercel/client/search.test.ts index 7f74fc26..748d352a 100644 --- a/src/vercel/client/search.test.ts +++ b/src/vercel/client/search.test.ts @@ -7,6 +7,7 @@ import { type MockedFunction, } from "vitest"; import type { ModelMessage } from "ai"; +import { MockEmbeddingModelV4 } from "ai/test"; import { type Auth, type StorageActionWriter, @@ -15,6 +16,7 @@ import { import type { MessageDoc } from "../../validators.js"; import type { ActionCtx, QueryCtx } from "./types.js"; import { + embedMany, fetchContextWithPrompt, fetchContextMessages, filterOutOrphanedToolMessages, @@ -105,6 +107,39 @@ describe("search.ts", () => { process.env.CONVEX_CLOUD_URL = "https://example.convex.cloud"; }); + describe("embedMany", () => { + it("passes the provider metadata to the usage handler", async () => { + const providerMetadata = { testProvider: { cost: 0.000002 } }; + const embeddingModel = new MockEmbeddingModelV4({ + modelId: "test-embedding", + doEmbed: { + embeddings: [[0.1, 0.2]], + usage: { tokens: 3 }, + providerMetadata, + warnings: [], + }, + }); + const usageHandler = vi.fn(); + + await embedMany({} as ActionCtx, { + userId: "user", + threadId: "thread", + values: ["hello"], + embeddingModel, + usageHandler, + }); + + expect(usageHandler).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + model: "test-embedding", + providerMetadata, + usage: expect.objectContaining({ inputTokens: 3 }), + }), + ); + }); + }); + describe("getPromptArray", () => { it("should return empty array for undefined prompt", () => { expect(getPromptArray(undefined)).toEqual([]); diff --git a/src/vercel/client/search.ts b/src/vercel/client/search.ts index 403f8f00..ad7c672f 100644 --- a/src/vercel/client/search.ts +++ b/src/vercel/client/search.ts @@ -499,7 +499,7 @@ export async function embedMany( agentName, model: getModelName(effectiveEmbeddingModel), provider: getProviderName(effectiveEmbeddingModel), - providerMetadata: undefined, + providerMetadata: result.providerMetadata, usage: { inputTokens: result.usage.tokens, outputTokens: 0,