diff --git a/packages/ai/README.md b/packages/ai/README.md index 0b45d216486..9c576009d11 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -409,6 +409,31 @@ OpenAICompatible.configure({ provider: "nvidia", baseURL: "https://integrate.api Every host here uses OpenAI's `/embeddings` API; Cohere through its compatibility API (`https://api.cohere.ai/compatibility/v1`). Package entry points are `@opencode/ai/providers//embeddings`. +Google embeds through `embedContent`, which also takes images, audio, video and PDFs on `gemini-embedding-2`. Each +value is its own content and gets its own vector (several parts in one content would be averaged into one). `dimensions` +is `outputDimensionality`, `inputType` becomes `taskType` `RETRIEVAL_QUERY` or `RETRIEVAL_DOCUMENT`, and `taskType` or +`title` in `providerOptions` replace or add to it. `gemini-embedding-2` ignores task types; Google recommends writing the +task into the text instead (`task: search result | query: …`, `title: none | text: …`). + +```ts +import { Google, GoogleVertex } from "@opencode/ai/providers" + +const response = await ai.embedding.generate({ + model: Google.configure().embedding("gemini-embedding-2"), + values: [await ai.file("./receipt.png"), "A receipt for two coffees."], + dimensions: 768, +}) + +GoogleVertex.configure({ project, location: "global" }).embedding("gemini-embedding-2") +``` + +The Gemini API sends up to 100 values per `batchEmbedContents` request. Vertex `embedContent` takes one value per +request, so a call sends one request per value, four at a time, and needs project credentials: Vertex Express Mode API +keys can't call `embedContent`. `gemini-embedding-2` runs in the `global`, `us` and `eu` Vertex locations. When a call +is split into several requests, the vectors stay in value order, `usage.inputTokens` is their sum, and each request's +metadata is listed in order under `providerMetadata..responses`. Google has no `/embeddings` entry point yet; +use `embeddingModel` from `@opencode/ai/providers/google` or `@opencode/ai/providers/google-vertex`. + ## Alibaba Cloud Model Studio `Alibaba` provides standard Model Studio inference. Configure a region explicitly, then select diff --git a/packages/ai/src/protocols/google-embed-content.ts b/packages/ai/src/protocols/google-embed-content.ts new file mode 100644 index 00000000000..1ed2d78d3f1 --- /dev/null +++ b/packages/ai/src/protocols/google-embed-content.ts @@ -0,0 +1,165 @@ +import { Effect, Schema } from "effect" +import { EmbeddingResponse, EmbeddingUsage, type EmbeddingRequest } from "../embedding.js" +import { EmbeddingRoute } from "../route/embedding.js" +import { Endpoint } from "../route/endpoint.js" +import { MediaProtocol } from "../route/media-protocol.js" +import { mergeJsonRecords, type AIError, type OpenString } from "../schema/index.js" +import { GeminiGenerateContent } from "./utils/gemini-generate-content.js" + +const ID = "google-embed-content" +const NAME = "Google embedContent" +export const DEFAULT_BASE_URL = "https://generativelanguage.googleapis.com/v1beta" +/** `batchEmbedContents` takes at most 100 requests. */ +export const MAX_BATCH = 100 + +export type TaskType = OpenString< + | "RETRIEVAL_QUERY" + | "RETRIEVAL_DOCUMENT" + | "SEMANTIC_SIMILARITY" + | "CLASSIFICATION" + | "CLUSTERING" + | "QUESTION_ANSWERING" + | "FACT_VERIFICATION" + | "CODE_RETRIEVAL_QUERY" +> + +/** + * Per-value embedding fields, sent with each value. `taskType` replaces the one `inputType` maps to; Gemini Embedding 2 + * ignores task types and documents task prefixes in the text instead (`task: search result | query: …`). + */ +export interface OptionsInput { + readonly [key: string]: unknown + readonly taskType?: TaskType + readonly title?: string +} + +/** Vertex `embedContentConfig` fields; Vertex alone takes `autoTruncate`, `documentOcr` and `audioTrackExtraction`. */ +export interface VertexOptionsInput extends OptionsInput { + readonly autoTruncate?: boolean + readonly documentOcr?: boolean + readonly audioTrackExtraction?: boolean +} + +const Values = Schema.Struct({ values: Schema.Array(Schema.Number) }) + +const UsageMetadata = Schema.StructWithRest(Schema.Struct({ promptTokenCount: Schema.optional(Schema.Number) }), [ + Schema.Record(Schema.String, Schema.Unknown), +]) + +const BatchResponse = Schema.Struct({ + embeddings: Schema.Array(Values), + usageMetadata: Schema.optional(UsageMetadata), +}) + +const SingleResponse = Schema.Struct({ + embedding: Values, + usageMetadata: Schema.optional(UsageMetadata), + truncated: Schema.optional(Schema.Boolean), +}) + +const TASK_TYPES = { query: "RETRIEVAL_QUERY", document: "RETRIEVAL_DOCUMENT" } as const + +// Each value is its own content: several parts in one content embed into one aggregated vector. +const content = (value: EmbeddingRequest["values"][number]): Effect.Effect<{ readonly parts: [unknown] }, AIError> => + typeof value === "string" + ? Effect.succeed({ parts: [{ text: value }] }) + : GeminiGenerateContent.mediaPart(NAME, value).pipe(Effect.map((part) => ({ parts: [part] }))) + +const config = (request: EmbeddingRequest) => ({ + ...(request.inputType === undefined ? {} : { taskType: TASK_TYPES[request.inputType] }), + ...(request.dimensions === undefined ? {} : { outputDimensionality: request.dimensions }), + ...request.providerOptions, +}) + +const identity = MediaProtocol.identity({ id: ID, name: NAME, provider: "google" }) +const decodeBatch = identity.decodeJson(BatchResponse) +const decodeSingle = identity.decodeJson(SingleResponse) + +const usage = (metadata: Schema.Schema.Type | undefined) => + metadata?.promptTokenCount === undefined ? undefined : new EmbeddingUsage({ inputTokens: metadata.promptTokenCount }) + +/** Gemini API `batchEmbedContents`: one request per value, so every value gets its own vector. */ +export const protocol = EmbeddingRoute.protocol({ + id: ID, + name: NAME, + body: { + from: Effect.fn("GoogleEmbedContent.fromRequest")(function* (request: EmbeddingRequest) { + const contents = yield* Effect.forEach(request.values, content) + const fields = config(request) + return ( + mergeJsonRecords( + { requests: contents.map((item) => ({ model: `models/${request.model.id}`, content: item, ...fields })) }, + request.http?.body, + ) ?? {} + ) + }), + }, + response: { + decode: Effect.fn("GoogleEmbedContent.decodeResponse")(function* (response, context) { + const output = yield* decodeBatch(response) + return new EmbeddingResponse({ + embeddings: output.value.embeddings.map((item) => item.values), + usage: usage(output.value.usageMetadata), + providerMetadata: { + [context.providerMetadataKey]: + output.value.usageMetadata === undefined ? {} : { usageMetadata: output.value.usageMetadata }, + }, + }) + }), + }, +}) + +/** + * Vertex `embedContent`: one content per request, with options under `embedContentConfig`. Routes using it set + * `maxValues: 1`. + */ +export const vertexProtocol = EmbeddingRoute.protocol({ + id: "google-vertex-embed-content", + name: "Google Vertex embedContent", + body: { + from: Effect.fn("GoogleEmbedContent.vertexFromRequest")(function* (request: EmbeddingRequest) { + const contents = yield* Effect.forEach(request.values, content) + const fields = config(request) + return ( + mergeJsonRecords( + { + content: contents[0], + ...(Object.keys(fields).length === 0 ? {} : { embedContentConfig: fields }), + }, + request.http?.body, + ) ?? {} + ) + }), + }, + response: { + decode: Effect.fn("GoogleEmbedContent.vertexDecodeResponse")(function* (response, context) { + const output = yield* decodeSingle(response) + return new EmbeddingResponse({ + embeddings: [output.value.embedding.values], + usage: usage(output.value.usageMetadata), + providerMetadata: { + [context.providerMetadataKey]: { + ...(output.value.usageMetadata === undefined ? {} : { usageMetadata: output.value.usageMetadata }), + ...(output.value.truncated === undefined ? {} : { truncated: output.value.truncated }), + }, + }, + }) + }), + }, +}) + +export const route = EmbeddingRoute.make({ + id: ID, + provider: "google", + protocol, + endpoint: Endpoint.path(({ request }) => `/models/${request.model.id}:batchEmbedContents`, { + baseURL: DEFAULT_BASE_URL, + }), + maxValues: MAX_BATCH, +}) + +export const GoogleEmbedContent = { + protocol, + vertexProtocol, + route, +} as const diff --git a/packages/ai/src/protocols/index.ts b/packages/ai/src/protocols/index.ts index 7c12ef59db0..b3dde0cfd32 100644 --- a/packages/ai/src/protocols/index.ts +++ b/packages/ai/src/protocols/index.ts @@ -3,6 +3,7 @@ export * as BedrockConverse from "./bedrock-converse.js" export * as CloudflareRunDecisions from "./cloudflare-run-decisions.js" export * as CohereChat from "./cohere-chat.js" export * as Gemini from "./gemini.js" +export * as GoogleEmbedContent from "./google-embed-content.js" export * as GoogleInteractions from "./google-interactions.js" export * as MistralChat from "./mistral-chat.js" export * as OpenAIChat from "./openai-chat.js" diff --git a/packages/ai/src/providers/google-vertex.ts b/packages/ai/src/providers/google-vertex.ts index 6d60157f03b..aaa0c77979c 100644 --- a/packages/ai/src/providers/google-vertex.ts +++ b/packages/ai/src/providers/google-vertex.ts @@ -1,10 +1,12 @@ import { Effect } from "effect" import type { ProviderPackage } from "../provider-package.js" import { Gemini } from "../protocols/gemini.js" +import { GoogleEmbedContent, type VertexOptionsInput } from "../protocols/google-embed-content.js" import { GoogleImages } from "../protocols/google-images.js" import { ProviderShared } from "../protocols/shared.js" import { Auth } from "../route/auth.js" import { Route } from "../route/client.js" +import { EmbeddingRoute } from "../route/embedding.js" import { Endpoint } from "../route/endpoint.js" import { MediaRoute } from "../route/media.js" import { @@ -82,6 +84,18 @@ const route = Route.make({ export const routes = [route] +export type { VertexOptionsInput as GoogleVertexEmbeddingOptions } from "../protocols/google-embed-content.js" + +const embeddingRoute = EmbeddingRoute.make({ + id: "google-vertex-embed-content", + provider: id, + providerMetadataKey: "vertex", + protocol: GoogleEmbedContent.vertexProtocol, + endpoint: Endpoint.path(({ request }) => `/models/${request.model.id}:embedContent`), + // Vertex embedContent embeds one content per request. + maxValues: 1, +}) + const configuredRoute = (input: Config, modelID: string | ModelID) => { const { accessToken: _accessToken, @@ -129,6 +143,18 @@ export const configure = (input: Config = {}) => { providerOptions, }) }, + /** `gemini-embedding-2` through Vertex `embedContent`; it needs project credentials, not an Express Mode key. */ + embedding: (modelID: string | ModelID, providerOptions?: ProviderOptions) => { + const target = connection(input, modelID) + if (GoogleVertexShared.apiKey(input) !== undefined && input.baseURL === undefined) + throw new ProviderConfigurationError({ + provider: id, + message: "Google Vertex embedContent does not support Express Mode API keys; use project credentials", + }) + return embeddingRoute + .with({ endpoint: target.endpoint, auth: target.auth, headers: input.headers, http: input.http }) + .model({ id: modelID, providerOptions }) + }, configure, } } @@ -143,6 +169,10 @@ export const imageModel = (modelID: string, settings?: Settings) => { const input = config(settings) return configure(input).image(modelID, input.providerOptions) } +export const embeddingModel = (modelID: string, settings?: Settings) => { + const input = config(settings) + return configure(input).embedding(modelID, input.providerOptions) +} function config({ accessToken, apiKey, baseURL, body, headers, location, project, ...providerOptions }: Settings = {}) { if (apiKey !== undefined && accessToken !== undefined) diff --git a/packages/ai/src/providers/google.ts b/packages/ai/src/providers/google.ts index d5fb4cb7966..145e8ce5241 100644 --- a/packages/ai/src/providers/google.ts +++ b/packages/ai/src/providers/google.ts @@ -4,6 +4,10 @@ import { MediaRoute } from "../route/media.js" import type { ProviderPackage } from "../provider-package.js" import { ProviderID, type ModelID, type ProviderOptions, type RequestOptions } from "../schema/index.js" import { Gemini } from "../protocols/gemini.js" +import { + GoogleEmbedContent, + type OptionsInput as GoogleEmbeddingOptionsInput, +} from "../protocols/google-embed-content.js" import { GoogleInteractions } from "../protocols/google-interactions.js" import { GoogleImages } from "../protocols/google-images.js" import { GoogleSpeech } from "../protocols/google-speech.js" @@ -11,6 +15,7 @@ import { GoogleTranscription } from "../protocols/google-transcription.js" import { GoogleVideo } from "../protocols/google-video.js" import { toConfig } from "./settings.js" +export type { OptionsInput as GoogleEmbeddingOptions } from "../protocols/google-embed-content.js" export type { GoogleImageOptions } from "../protocols/google-images.js" export type { GoogleSpeechOptions } from "../protocols/google-speech.js" export type { GoogleTranscriptionOptions } from "../protocols/google-transcription.js" @@ -52,6 +57,13 @@ export const configure = (input: Config = {}) => { const route = configuredRoute(input) const media = MediaRoute.deployment(input, auth(input)) const gemini = (modelID: string | ModelID) => route.model({ id: modelID }) + const embeddingRoute = GoogleEmbedContent.route.with({ + provider: id, + endpoint: { baseURL: input.baseURL }, + auth: auth(input), + headers: input.headers, + http: input.http, + }) return { id, model: gemini, @@ -63,6 +75,8 @@ export const configure = (input: Config = {}) => { video: (modelID: string | ModelID) => GoogleVideo.model({ ...media, id: modelID }), speech: (modelID: string | ModelID) => GoogleSpeech.model({ ...media, id: modelID }), transcription: (modelID: string | ModelID) => GoogleTranscription.model({ ...media, id: modelID }), + embedding: (modelID: string | ModelID, providerOptions?: ProviderOptions) => + embeddingRoute.model({ id: modelID, providerOptions }), configure, } } @@ -74,8 +88,13 @@ export const imageModel = (modelID: string, settings?: Settings) => { const input = toConfig(settings) return configure(input).image(modelID, input.providerOptions) } +export const embeddingModel = (modelID: string, settings?: Settings) => { + const input = toConfig(settings) + return configure(input).embedding(modelID, input.providerOptions) +} export const image = provider.image +export const embedding = provider.embedding export const gemini = provider.gemini export const interactions = provider.interactions export const video = provider.video diff --git a/packages/ai/src/route/embedding.ts b/packages/ai/src/route/embedding.ts index d275a3c175b..457e546aebb 100644 --- a/packages/ai/src/route/embedding.ts +++ b/packages/ai/src/route/embedding.ts @@ -1,6 +1,12 @@ import { Effect } from "effect" import type { HttpClientResponse } from "effect/http" -import { EmbeddingModel, type EmbeddingOptions, type EmbeddingRequest, type EmbeddingResponse } from "../embedding.js" +import { + EmbeddingModel, + EmbeddingResponse, + EmbeddingUsage, + type EmbeddingOptions, + type EmbeddingRequest, +} from "../embedding.js" import { ProviderID, type HttpOptions, type AIError, type ProviderOptions } from "../schema/index.js" import { Auth } from "./auth.js" import { Endpoint } from "./endpoint.js" @@ -31,6 +37,11 @@ export interface Patch { readonly auth?: Auth.Definition readonly headers?: Record readonly http?: HttpOptions + /** + * Most values one HTTP request may carry. Longer requests are split into requests of this size, sent a few at a + * time, and joined back in value order. + */ + readonly maxValues?: number } export interface Route { @@ -42,6 +53,7 @@ export interface Route { readonly auth: Auth.Definition readonly headers?: Record readonly http?: HttpOptions + readonly maxValues?: number readonly with: (patch: Patch) => Route readonly model: (input: { readonly id: string @@ -67,6 +79,7 @@ export const make = (input: MakeInput): Route => { auth: input.auth ?? Auth.none, headers: input.headers, http: input.http, + maxValues: input.maxValues, } return { ...route, @@ -80,6 +93,7 @@ export const make = (input: MakeInput): Route => { endpoint: patch.endpoint === undefined ? route.endpoint : Endpoint.merge(route.endpoint, patch.endpoint), headers: patch.headers === undefined ? route.headers : { ...route.headers, ...patch.headers }, http: patch.http ?? route.http, + maxValues: patch.maxValues ?? route.maxValues, }), model: (model: { readonly id: string @@ -90,11 +104,51 @@ export const make = (input: MakeInput): Route => { provider, http: route.http, providerOptions: model.providerOptions, - route: { id: route.id, generate: transport(route).generate }, + route: { id: route.id, generate: split(route.maxValues, transport(route).generate) }, }), } } +const CONCURRENCY = 4 + +/** + * Splits a request whose values exceed `maxValues` and joins the parts: vectors in value order, `inputTokens` summed + * when every part reports it, and each part's metadata in order under `providerMetadata[key].responses`. + */ +const split = + (maxValues: number | undefined, generate: Generate): Generate => + (request, execute) => { + if (maxValues === undefined || request.values.length <= maxValues) return generate(request, execute) + const chunks = Array.from({ length: Math.ceil(request.values.length / maxValues) }, (_, index) => + request.values.slice(index * maxValues, (index + 1) * maxValues), + ) + return Effect.forEach(chunks, (values) => generate({ ...request, values }, execute), { + concurrency: CONCURRENCY, + }).pipe( + Effect.map((responses) => { + const tokens = responses.map((response) => response.usage?.inputTokens) + const keys = [...new Set(responses.flatMap((response) => Object.keys(response.providerMetadata ?? {})))] + return new EmbeddingResponse({ + embeddings: responses.flatMap((response) => response.embeddings), + usage: tokens.every((count) => count !== undefined) + ? new EmbeddingUsage({ inputTokens: tokens.reduce((sum, count) => sum + count, 0) }) + : undefined, + providerMetadata: + keys.length === 0 + ? undefined + : Object.fromEntries( + keys.map((key) => [ + key, + { responses: responses.map((response) => response.providerMetadata?.[key] ?? {}) }, + ]), + ), + }) + }), + ) + } + +type Generate = EmbeddingModel["route"]["generate"] + // The HTTP plumbing (URL, query, auth, headers, `http` overlays, surrogate sanitizing) is the one media routes use. const transport = (route: Omit) => MediaRoute.inline({ diff --git a/packages/ai/test/embedding.test.ts b/packages/ai/test/embedding.test.ts index 9652bccd43f..9976c6db6d1 100644 --- a/packages/ai/test/embedding.test.ts +++ b/packages/ai/test/embedding.test.ts @@ -1,7 +1,8 @@ -import { describe, expect } from "bun:test" +import { describe, expect, test } from "bun:test" import { Effect, Layer } from "effect" import { ai, EmbeddingClient, EmbeddingUsage, Media } from "../src/effect.js" -import { OpenAI, OpenAICompatible } from "../src/providers.js" +import { Google, GoogleVertex, OpenAI, OpenAICompatible } from "../src/providers.js" +import { embeddingModel as googleEmbeddingModel } from "../src/providers/google.js" import { it } from "./lib/effect.js" import { dynamicResponse, json, type Handler } from "./lib/http.js" @@ -130,4 +131,181 @@ describe("Embedding", () => { expect(error.message).toContain("Expected 2 embeddings, got 1") }).pipe(Effect.provide(client((input) => Effect.succeed(json(input, { data: [{ index: 0, embedding: [0.1] }] }))))), ) + + it.effect("batches Google values as one content each and maps inputType to taskType", () => + ai.embedding + .generate({ + model: Google.configure({ apiKey: "google-key" }).embedding("gemini-embedding-2"), + values: ["a caption", Media.base64("iVBORw0KGgo=", "image/png")], + dimensions: 256, + inputType: "document", + providerOptions: { title: "Restroom" }, + }) + .pipe( + Effect.tap((response) => + Effect.sync(() => { + expect(response.embeddings).toEqual([ + [0.1, 0.2], + [0.3, 0.4], + ]) + expect(response.usage).toEqual(new EmbeddingUsage({ inputTokens: 260 })) + expect(response.providerMetadata).toEqual({ google: { usageMetadata: { promptTokenCount: 260 } } }) + }), + ), + Effect.provide( + client((input) => { + expect(input.request.url).toBe( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2:batchEmbedContents", + ) + expect(input.request.headers["x-goog-api-key"]).toBe("google-key") + const fields = { taskType: "RETRIEVAL_DOCUMENT", outputDimensionality: 256, title: "Restroom" } + expect(JSON.parse(input.text)).toEqual({ + requests: [ + { model: "models/gemini-embedding-2", content: { parts: [{ text: "a caption" }] }, ...fields }, + { + model: "models/gemini-embedding-2", + content: { parts: [{ inlineData: { mimeType: "image/png", data: "iVBORw0KGgo=" } }] }, + ...fields, + }, + ], + }) + return Effect.succeed( + json(input, { + embeddings: [{ values: [0.1, 0.2] }, { values: [0.3, 0.4] }], + usageMetadata: { promptTokenCount: 260 }, + }), + ) + }), + ), + ), + ) + + it.effect("lets a Google taskType option replace the inputType mapping and sends settings", () => + ai.embedding + .generate({ + model: googleEmbeddingModel("gemini-embedding-001", { + apiKey: "google-key", + headers: { "x-test": "yes" }, + body: { extra: true }, + }), + values: ["query"], + inputType: "query", + providerOptions: { taskType: "CODE_RETRIEVAL_QUERY" }, + }) + .pipe( + Effect.tap((response) => + Effect.sync(() => { + expect(response.usage).toBeUndefined() + expect(response.providerMetadata).toEqual({ google: {} }) + }), + ), + Effect.provide( + client((input) => { + expect(input.request.headers["x-test"]).toBe("yes") + expect(JSON.parse(input.text)).toEqual({ + requests: [ + { + model: "models/gemini-embedding-001", + content: { parts: [{ text: "query" }] }, + taskType: "CODE_RETRIEVAL_QUERY", + }, + ], + extra: true, + }) + return Effect.succeed(json(input, { embeddings: [{ values: [1] }] })) + }), + ), + ), + ) + + it.effect("splits more than 100 Google values across batches and joins them in order", () => + Effect.gen(function* () { + const sizes: Array = [] + const response = yield* ai.embedding + .generate({ + model: Google.configure({ apiKey: "google-key" }).embedding("gemini-embedding-2"), + values: Array.from({ length: 105 }, (_, index) => `${index}`), + }) + .pipe( + Effect.provide( + client((input) => { + const requests: Array<{ content: { parts: Array<{ text: string }> } }> = JSON.parse(input.text).requests + sizes.push(requests.length) + return Effect.succeed( + json(input, { + embeddings: requests.map((request) => ({ values: [Number(request.content.parts[0].text)] })), + usageMetadata: { promptTokenCount: requests.length }, + }), + ) + }), + ), + ) + expect(sizes.toSorted((a, b) => a - b)).toEqual([5, 100]) + expect(response.embeddings).toEqual(Array.from({ length: 105 }, (_, index) => [index])) + expect(response.usage).toEqual(new EmbeddingUsage({ inputTokens: 105 })) + expect(response.providerMetadata).toEqual({ + google: { + responses: [{ usageMetadata: { promptTokenCount: 100 } }, { usageMetadata: { promptTokenCount: 5 } }], + }, + }) + }), + ) + + it.effect("sends one Vertex embedContent request per value with embedContentConfig", () => + Effect.gen(function* () { + const bodies: Array = [] + const response = yield* ai.embedding + .generate({ + model: GoogleVertex.configure({ accessToken: "vertex-token", project: "proj", location: "us" }).embedding( + "gemini-embedding-2", + ), + values: ["first", Media.base64("JVBERi0=", "application/pdf")], + dimensions: 128, + providerOptions: { documentOcr: true }, + }) + .pipe( + Effect.provide( + client((input) => { + expect(input.request.url).toBe( + "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/proj/locations/us/publishers/google/models/gemini-embedding-2:embedContent", + ) + expect(input.request.headers.authorization).toBe("Bearer vertex-token") + const body = JSON.parse(input.text) + bodies.push(body) + return Effect.succeed( + json(input, { + embedding: { values: ["text" in body.content.parts[0] ? 1 : 2] }, + usageMetadata: { promptTokenCount: 3 }, + truncated: false, + }), + ) + }), + ), + ) + expect(bodies).toContainEqual({ + content: { parts: [{ text: "first" }] }, + embedContentConfig: { outputDimensionality: 128, documentOcr: true }, + }) + expect(bodies).toContainEqual({ + content: { parts: [{ inlineData: { mimeType: "application/pdf", data: "JVBERi0=" } }] }, + embedContentConfig: { outputDimensionality: 128, documentOcr: true }, + }) + expect(response.embeddings).toEqual([[1], [2]]) + expect(response.usage).toEqual(new EmbeddingUsage({ inputTokens: 6 })) + expect(response.providerMetadata).toEqual({ + vertex: { + responses: [ + { usageMetadata: { promptTokenCount: 3 }, truncated: false }, + { usageMetadata: { promptTokenCount: 3 }, truncated: false }, + ], + }, + }) + }), + ) + + test("rejects Vertex Express Mode keys for embedContent", () => { + expect(() => GoogleVertex.configure({ apiKey: "vertex-key" }).embedding("gemini-embedding-2")).toThrow( + "does not support Express Mode API keys", + ) + }) }) diff --git a/packages/ai/test/fixtures/recordings/google-embedding/embeds-an-image-and-text-into-one-vector-each.json b/packages/ai/test/fixtures/recordings/google-embedding/embeds-an-image-and-text-into-one-vector-each.json new file mode 100644 index 00000000000..a4b548b94fb --- /dev/null +++ b/packages/ai/test/fixtures/recordings/google-embedding/embeds-an-image-and-text-into-one-vector-each.json @@ -0,0 +1,32 @@ +{ + "version": 1, + "metadata": { + "tags": [ + "prefix:google-embedding", + "provider:google", + "protocol:google-embed-content" + ], + "name": "google-embedding/embeds-an-image-and-text-into-one-vector-each", + "recordedAt": "2026-11-09T06:56:09.697Z" + }, + "interactions": [ + { + "transport": "http", + "request": { + "method": "POST", + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2:batchEmbedContents", + "headers": { + "content-type": "application/json" + }, + "body": "{\"requests\":[{\"model\":\"models/gemini-embedding-2\",\"content\":{\"parts\":[{\"inlineData\":{\"mimeType\":\"image/jpeg\",\"data\":\"/9j/4AAQSkZJRgABAgAAAQABAAD//gAQTGF2YzYwLjMxLjEwMgD/2wBDAAgQEBMQExYWFhYWFhoYGhsbGxoaGhobGxsdHR0iIiIdHR0bGx0dICAiIiUmJSMjIiMmJigoKDAwLi44ODpFRVP/xABrAAEBAQEBAQAAAAAAAAAAAAAABwgGBAUBAQAAAAAAAAAAAAAAAAAAAAAQAAICAQEFBwMEAwEAAAAAAAABAwIEEQYFEiFhwXGBUUExE6GRU0LwciLh8bEjEQEAAAAAAAAAAAAAAAAAAAAA/8AAEQgAgACAAwEiAAIRAAMRAP/aAAwDAQACEQMRAD8Av4AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAB458iLFjtLNesdK+9rPRf76e4yMiPFivNLbhpHV2s+i7X6LzMhb33xPvabitrWKrfxxa8qrzfnZ+r8EBTc/bPm64USa/LLrz61omvq/AnMu0W9ZXq8q9elFSi+ldTigB2se0W9Ynqsq9ul1S6+tSiYG2b1Vc2Jafli15dXRt/R+BBgBuvHyIcqOssN6yUt7Wq9V/h9HzPYY03RvefdM3HTW0dn/AOkWvKy815WXo/ua+x8iPKhpNFbipJVWq/36r2a8wPaAAAAAAAAAAIDtnnvWLCq+WnyydfSlX9X9iDHZ7QyuXeuW3+m6ou6lUjjAAAAAAAXbYzeDVpcKz5NfLH0ftdLv5P7kJOw2flcW9cRr1k4H3XTr2gbHAAAAAAAAAAGO9oYnFvXLT/VdXXdeqZxhdts8Bq0WZVcmvik6NaujffzX2ISAAAAAADr9wRuXeuGl6ScfhRO3YcgXHYzAdpZcyy/rRfFH1s+dmu5aLxA0MAAAAAAAAAAPDk40WXDJDLXipJXha7V1Xuupj7eu6p91TuORO1Hr8cmnK67LL1Rs8+fk4sOZFaKeOslLe6f/AFP3T8muYGGAXHP2Nkq3bDlV6/jl5WXRXXJ+KROpdxb0ielsOZ/wXGvvRsDkwdXHuLekr0WHMv5V4F97tFDwNjZrtWzJVHX8cb4rvo7acK8NQJtuvdc+9Z1FGtKrT5JNP60r2vyXqbBxMWLCgjgiXDSNaLzfm31b5sYuJBhRKKCOsdF6L1fm37t9WfRAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/9k=\"}}]},\"outputDimensionality\":128},{\"model\":\"models/gemini-embedding-2\",\"content\":{\"parts\":[{\"text\":\"A small square image.\"}]},\"outputDimensionality\":128}]}" + }, + "response": { + "status": 200, + "headers": { + "content-type": "application/json; charset=UTF-8" + }, + "body": "{\n \"embeddings\": [\n {\n \"values\": [\n -0.060695454,\n -0.16676353,\n -0.014622029,\n -0.0369989,\n -0.07764896,\n -0.024174407,\n 0.042032935,\n -0.0015557865,\n -0.023695432,\n -0.1497187,\n -0.025017679,\n 0.06438567,\n 0.010775195,\n -0.1025198,\n 0.08094253,\n 0.020372627,\n -0.0042680614,\n -0.01302046,\n 0.016640898,\n 0.03898399,\n -0.04134967,\n -0.036376182,\n 0.022557853,\n 0.03913328,\n -0.08211433,\n 0.01890385,\n 0.012346725,\n -0.04081329,\n -0.044884365,\n 0.35187453,\n -0.012423337,\n -0.04902881,\n 0.055406027,\n -0.024738494,\n 0.0910491,\n -0.09677539,\n -0.03170299,\n 0.082624085,\n 0.019814415,\n -0.008206397,\n -0.034400895,\n -0.016858248,\n 0.0113725355,\n 0.047203444,\n -0.011291501,\n -0.08147267,\n 0.024802648,\n -0.0113782985,\n 0.0049352786,\n 0.035522662,\n 0.10245429,\n 0.00234546,\n -0.1333667,\n -0.007869423,\n 0.030382834,\n 0.05243392,\n -0.03965319,\n 0.021363147,\n -0.021653442,\n -0.064418085,\n -0.024809368,\n -0.002167328,\n -0.017991243,\n 0.041653287,\n -0.062733434,\n -0.015845643,\n 0.045335785,\n -0.05263401,\n -0.047420107,\n -0.050459415,\n 0.012136268,\n 0.0106795635,\n 0.018208819,\n 0.047718324,\n 0.050123528,\n -0.057227056,\n -0.05015243,\n -0.026977116,\n 0.01494459,\n -0.021070806,\n -0.0010004406,\n -0.02632088,\n -0.03466405,\n -0.010536535,\n -0.060299378,\n -0.021053744,\n 0.06964341,\n -0.022528982,\n 0.045283075,\n 0.015700605,\n -0.012671969,\n -0.030455485,\n -0.015061047,\n 0.008110022,\n -0.032319747,\n -0.07088189,\n -0.020862978,\n 0.0012270943,\n -0.048924483,\n 0.061084464,\n 0.007996045,\n 0.0019241121,\n -0.069225684,\n -0.019252202,\n 0.0012802874,\n -0.052421395,\n 0.06391537,\n -0.030530954,\n -0.07110764,\n -0.0691781,\n -0.036115203,\n -0.36308432,\n -0.022246694,\n -0.05472309,\n 0.017038023,\n 0.6788818,\n -0.04438193,\n 0.022014968,\n 0.036325,\n 0.02110961,\n -0.033575546,\n -0.0011139157,\n -0.0132337855,\n 0.027654596,\n 0.02986413,\n -0.00597513,\n -0.047889307,\n -0.03105767\n ]\n },\n {\n \"values\": [\n -0.07348734,\n 0.058014095,\n -0.07439354,\n 0.005782239,\n 0.03540112,\n -0.020454172,\n -0.043367438,\n 0.0076436116,\n -0.043582547,\n -0.16619165,\n 0.11180406,\n 0.027603088,\n 0.11537567,\n -0.050174788,\n 0.01823103,\n -0.057214037,\n -0.00063154247,\n 0.08172401,\n 0.010149112,\n -0.013878285,\n 0.004687302,\n -0.007362179,\n -0.038548544,\n -0.012590715,\n 0.032411125,\n -0.031378265,\n -0.03343408,\n 0.080075584,\n 0.020124901,\n 0.36560065,\n -0.005742146,\n 0.12878247,\n -0.07061178,\n -0.06762567,\n 0.04209995,\n -0.05458626,\n 0.04457706,\n 0.011975002,\n 0.028867552,\n -0.099761575,\n -0.023624368,\n 0.000761353,\n -0.03569734,\n 0.08086029,\n -0.1480097,\n 0.061061267,\n -0.03427503,\n 0.05064104,\n 0.024667451,\n 0.059212048,\n 0.04548406,\n 0.0013026758,\n 0.05403232,\n 0.040431008,\n 0.08995319,\n 0.015512121,\n -0.045479815,\n 0.018925536,\n -0.06662972,\n 0.0070882696,\n 0.032891102,\n -0.13106298,\n -0.009988906,\n -0.03928157,\n 0.010292592,\n 0.012359754,\n 0.037858266,\n 0.05525646,\n 0.0076784333,\n 0.037685215,\n -0.039831035,\n -0.052651886,\n 0.055160705,\n 0.039733298,\n -0.07184765,\n -0.18210597,\n -0.049278192,\n 0.05887453,\n 0.012229632,\n -0.0106056845,\n 0.117099665,\n 0.04492413,\n 0.027917886,\n -0.019317398,\n -0.057401214,\n -0.13336012,\n -0.07114207,\n 0.0077259336,\n -0.04206231,\n 0.031035932,\n -0.044222888,\n 0.054862186,\n -0.06381428,\n 0.019193633,\n -0.1454552,\n -0.06991198,\n 0.033321287,\n 0.17329463,\n 0.06438023,\n 0.022623677,\n -0.009573466,\n -0.018114265,\n -0.02589199,\n 0.027245803,\n 0.04145004,\n 0.006055637,\n 0.024168594,\n -0.016137246,\n 0.003217981,\n 0.0276925,\n 0.061841067,\n -0.6379371,\n -0.06398904,\n 0.019271554,\n -0.035416443,\n -0.07965318,\n -0.12509108,\n 0.00091902283,\n -0.027904658,\n 0.014234293,\n -0.07321166,\n -0.052702807,\n -0.06417575,\n 0.042428393,\n -0.008938485,\n -0.048144527,\n 0.029265637,\n 0.03570763\n ]\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 263,\n \"promptTokenDetails\": [\n {\n \"modality\": \"IMAGE\",\n \"tokenCount\": 258\n },\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 5\n }\n ]\n }\n}\n" + } + } + ] +} diff --git a/packages/ai/test/fixtures/recordings/google-embedding/embeds-text-through-batchembedcontents.json b/packages/ai/test/fixtures/recordings/google-embedding/embeds-text-through-batchembedcontents.json new file mode 100644 index 00000000000..f5728c5df24 --- /dev/null +++ b/packages/ai/test/fixtures/recordings/google-embedding/embeds-text-through-batchembedcontents.json @@ -0,0 +1,32 @@ +{ + "version": 1, + "metadata": { + "tags": [ + "prefix:google-embedding", + "provider:google", + "protocol:google-embed-content" + ], + "name": "google-embedding/embeds-text-through-batchembedcontents", + "recordedAt": "2026-11-09T06:56:16.394Z" + }, + "interactions": [ + { + "transport": "http", + "request": { + "method": "POST", + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2:batchEmbedContents", + "headers": { + "content-type": "application/json" + }, + "body": "{\"requests\":[{\"model\":\"models/gemini-embedding-2\",\"content\":{\"parts\":[{\"text\":\"The invoice was charged twice.\"}]},\"taskType\":\"RETRIEVAL_DOCUMENT\",\"outputDimensionality\":256},{\"model\":\"models/gemini-embedding-2\",\"content\":{\"parts\":[{\"text\":\"Refund the duplicate payment.\"}]},\"taskType\":\"RETRIEVAL_DOCUMENT\",\"outputDimensionality\":256}]}" + }, + "response": { + "status": 200, + "headers": { + "content-type": "application/json; charset=UTF-8" + }, + "body": "{\n \"embeddings\": [\n {\n \"values\": [\n 0.08826861,\n -0.010503297,\n -0.003593779,\n -0.04362853,\n -0.036237195,\n -0.055654835,\n -0.083461575,\n 0.033482846,\n 0.036843903,\n -0.17836286,\n -0.012956421,\n 0.021917608,\n 0.05855208,\n 0.016266014,\n 0.05865058,\n 0.009627632,\n 0.053421967,\n -0.014696311,\n -0.06453404,\n -0.018494826,\n -0.042978473,\n -0.058279276,\n -0.029961104,\n -0.030288767,\n 0.0073536662,\n -0.046593722,\n 0.015977599,\n 0.019617883,\n 0.03970064,\n 0.38354987,\n -0.033308115,\n 0.031935018,\n -0.03247434,\n -0.03815824,\n 0.05267559,\n -0.10165431,\n -0.041823577,\n 0.06042345,\n 0.0018351023,\n 0.051456507,\n 0.052291792,\n 2.8428964e-05,\n -0.07374836,\n -0.0061268187,\n -0.035403196,\n 0.0024109676,\n -0.0416654,\n -0.0021237924,\n 0.04037082,\n -0.079502486,\n -0.05369692,\n -0.015340817,\n -0.0644507,\n 0.04906819,\n 0.043795276,\n 0.011756599,\n 0.027179904,\n -0.014487515,\n 0.040554237,\n 0.05260127,\n -0.07062641,\n 0.010476285,\n 0.06995466,\n 0.05079814,\n -0.021293104,\n 0.054295633,\n -0.027834542,\n -0.072753094,\n 0.07105707,\n 0.06042876,\n -0.12779789,\n 0.038879916,\n -0.08831736,\n 0.06513875,\n -0.007510305,\n -0.13622166,\n -0.074078396,\n 0.06498969,\n 0.044894755,\n 0.019542063,\n -0.04276829,\n -0.011232061,\n -0.057544988,\n -0.020430364,\n -0.060440596,\n 0.04143041,\n -0.030176722,\n -0.0069830897,\n -0.07553233,\n 0.030186994,\n -0.041128784,\n -0.13769294,\n 0.08790415,\n -0.0050185462,\n -0.089577004,\n 0.009078706,\n 0.07965475,\n 0.059314027,\n -0.014337807,\n -0.009616872,\n 0.017148038,\n -0.0559987,\n -0.084258616,\n -0.087398924,\n 0.023736937,\n -0.006153544,\n 0.012847555,\n -0.051635522,\n -0.023994954,\n -0.017216988,\n 0.005052825,\n -0.48903733,\n -0.031868614,\n -0.00405748,\n 0.049377847,\n -0.11121355,\n 0.03040298,\n 0.0657372,\n -0.052349366,\n 0.027236087,\n 0.059038788,\n 0.009941072,\n -0.016297461,\n -0.011371718,\n -0.010174601,\n 0.08847744,\n 0.07474075,\n -0.018525193,\n -0.0049489983,\n -0.016156295,\n -0.0008972978,\n 0.084703036,\n -8.523724e-05,\n -0.016584247,\n -0.018745024,\n -0.0027212733,\n 0.089199334,\n -0.0036129493,\n 0.026515242,\n -0.0068128984,\n -0.007686669,\n 0.0028604842,\n 0.040065005,\n -0.05489823,\n -0.0055631036,\n 0.058329903,\n 0.011494569,\n -0.034512483,\n -0.008964485,\n -0.046100236,\n -0.028631702,\n 0.053142093,\n 0.10161717,\n 0.008478608,\n -0.025940085,\n 0.041855946,\n 0.036950104,\n 0.03752138,\n -0.031569928,\n -0.09411969,\n -0.07016315,\n -0.077472076,\n 0.054901585,\n 0.032077942,\n 0.0002365121,\n -0.021486815,\n -0.03782185,\n 0.06342147,\n -0.0400071,\n 0.013592747,\n 0.029044827,\n -0.04941764,\n 0.027012195,\n -0.020215036,\n 0.035024326,\n -0.018768141,\n -0.010166947,\n 0.0063576326,\n -0.035777643,\n -0.03318015,\n -0.0013767413,\n -0.016789285,\n -0.023925936,\n -0.02882159,\n 0.09833145,\n 0.03744801,\n -0.0014872245,\n -0.057880018,\n 0.014366774,\n 0.055709057,\n 0.012810761,\n 0.06668003,\n 0.00017016666,\n 0.0012475014,\n 0.00870707,\n -0.070864856,\n -0.04660088,\n 0.011100339,\n 0.027194036,\n 0.0289681,\n -0.060197454,\n 0.005464363,\n 0.010492512,\n -0.041788977,\n -0.037648804,\n -0.086552024,\n -0.043190144,\n 0.02386924,\n -0.011399005,\n 0.06646394,\n 0.050833337,\n 0.03239632,\n 0.049356397,\n 0.036520176,\n 0.008604959,\n 0.020041626,\n 0.01113018,\n 0.15001084,\n 0.016630875,\n 0.018221907,\n 0.006113152,\n 0.09078522,\n 0.027637133,\n 0.018309548,\n -0.01633641,\n -0.02608749,\n 0.052398287,\n 0.011247799,\n -0.060354326,\n 0.021876536,\n 0.12617895,\n 0.01963028,\n 0.080168225,\n -0.043625426,\n 0.008349975,\n 0.025866184,\n 0.018640447,\n 0.060024187,\n -0.015962755,\n -0.0036795249,\n 0.009187861,\n -0.05795868,\n -0.01515489,\n -0.03210024,\n -0.04497872,\n 0.04977143,\n -0.037073787,\n -0.010589404,\n -0.002064704,\n 0.0034019554,\n -0.048233703,\n 0.103998035,\n 0.020032499,\n 0.005160983,\n 0.042482056,\n 0.020384656\n ]\n },\n {\n \"values\": [\n 0.055573214,\n 0.005882677,\n 0.01826917,\n -0.027935615,\n -0.035022806,\n -0.0056345863,\n -0.059873283,\n 0.08781313,\n 0.031654596,\n -0.21276161,\n -0.03063806,\n -0.015950669,\n 0.022213425,\n -0.019672554,\n 0.014765447,\n 0.018805644,\n 0.088912874,\n -0.052033007,\n -0.04261618,\n 0.040153656,\n 0.0002709653,\n 0.027197575,\n -0.028945133,\n -0.023637824,\n -0.07464421,\n 0.0021192082,\n -0.0066155978,\n 0.0103703495,\n 0.046237614,\n 0.37369522,\n -0.026183352,\n -0.0074142427,\n -0.064248495,\n -0.04772917,\n 0.03777998,\n -0.08207684,\n -0.07380202,\n 0.06946735,\n 0.006707171,\n -0.008991639,\n 0.12729676,\n -0.0020389983,\n -0.049856912,\n 0.025080126,\n 0.04332476,\n -0.009537268,\n -0.02804633,\n -0.033539128,\n 0.06978733,\n -0.061974753,\n -0.044053447,\n -0.019087778,\n -0.05079465,\n 0.06880674,\n 0.021709887,\n 0.0013433662,\n 0.03357279,\n -0.020821936,\n 0.06484281,\n 0.047331396,\n -0.06765697,\n 0.006802375,\n -0.021252241,\n 0.05977433,\n 0.00339868,\n -0.013671367,\n 0.0049018324,\n -0.064193174,\n 0.12076415,\n 0.123991035,\n -0.08617238,\n -0.043577645,\n -0.057860445,\n 0.0712765,\n -0.023838129,\n -0.0927374,\n -0.053347606,\n 0.015069004,\n 0.0048751193,\n 0.035117477,\n -0.039658908,\n -0.0112140775,\n -0.05246642,\n -0.013988502,\n -0.096250586,\n -0.0014971048,\n -0.0013022952,\n 0.015804112,\n -0.057270862,\n 0.016260924,\n -0.100514844,\n -0.04636865,\n 0.010188983,\n -0.043096192,\n -0.048539747,\n 0.029445156,\n 0.049847852,\n 0.025705785,\n -0.01587146,\n 0.0136724375,\n 0.019404871,\n 0.011149615,\n -0.12750958,\n -0.08901808,\n -0.006742344,\n 0.023663288,\n -0.023679776,\n -0.024047747,\n -0.05614705,\n 0.024845188,\n 0.0130731305,\n -0.45807788,\n -0.033288486,\n 0.002387504,\n 0.05757872,\n -0.10141078,\n 0.018331088,\n 0.066774756,\n 0.030569423,\n 0.02577125,\n 0.038817436,\n -0.0115171885,\n -0.03519031,\n 0.008516701,\n -0.025135018,\n 0.06304102,\n 0.092472196,\n -0.077763535,\n 0.053744473,\n 0.014108602,\n -0.02527797,\n 0.10871733,\n 0.0029466082,\n -0.0019035314,\n -0.046041057,\n -0.0016688341,\n 0.06353252,\n -0.016068991,\n 0.035331856,\n 0.017738545,\n 0.019702058,\n 0.0352097,\n 0.04297009,\n -0.060130455,\n 0.0014609741,\n 0.08829011,\n 0.065697975,\n -0.050308824,\n 0.008562518,\n -0.02146485,\n 0.009379907,\n 0.080687515,\n 0.08388223,\n 0.01195018,\n -0.027259015,\n -0.08552416,\n -0.022187022,\n 0.054287627,\n -0.016185338,\n -0.06198485,\n -0.11792418,\n -0.07332939,\n -0.051198054,\n 0.03868018,\n -0.03734517,\n 0.00849145,\n -0.036531594,\n 0.025943676,\n -0.12947097,\n 0.009934093,\n -0.012165758,\n -0.018885039,\n -0.0057292483,\n 0.01055107,\n 0.04069054,\n 0.016645007,\n -0.009185493,\n -0.0055974717,\n 0.0031500375,\n -0.025355576,\n -0.02472688,\n -0.026464252,\n -0.03355686,\n 0.015095607,\n 0.061247155,\n 0.03455652,\n 0.047235593,\n -0.067785725,\n 0.027664077,\n 0.0226711,\n 0.07874564,\n 0.019857094,\n -0.028619839,\n 0.011504168,\n 0.041439142,\n 0.011528595,\n 0.012672718,\n 0.011122911,\n 0.042298287,\n 0.024844902,\n -0.0306643,\n 0.056554127,\n 0.034346133,\n -0.052996196,\n -0.029159157,\n -0.14677109,\n -0.03259311,\n 0.029201353,\n -0.00033933605,\n 0.0044265073,\n 0.06981597,\n 0.040167153,\n -0.023513084,\n 0.013591094,\n 0.0060894717,\n 0.047549773,\n 0.032060564,\n 0.1500461,\n -0.0125145,\n -0.007945927,\n -0.030586757,\n 0.079504035,\n 0.041778732,\n 0.015542834,\n 0.0014314849,\n 0.0018921786,\n 0.030846952,\n 0.03526715,\n -0.028502392,\n 0.033490326,\n 0.079798624,\n 0.0016604621,\n 0.093993016,\n -0.004698624,\n 0.022417571,\n 0.004761023,\n 0.014730654,\n 0.05114591,\n 0.034675676,\n 0.02710423,\n 0.074968666,\n -0.035762616,\n -0.015032275,\n -0.10261388,\n -0.10068578,\n 0.073554195,\n -0.0034417463,\n 0.047156025,\n 0.019664466,\n -0.02528637,\n -0.034683526,\n 0.06354508,\n 0.023256633,\n 0.051614694,\n 0.061911386,\n -0.0050556194\n ]\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 11,\n \"promptTokenDetails\": [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 11\n }\n ]\n }\n}\n" + } + } + ] +} diff --git a/packages/ai/test/fixtures/recordings/google-embedding/maps-inputtype-to-a-task-type-for-gemini-embedding-001.json b/packages/ai/test/fixtures/recordings/google-embedding/maps-inputtype-to-a-task-type-for-gemini-embedding-001.json new file mode 100644 index 00000000000..bf73134dacf --- /dev/null +++ b/packages/ai/test/fixtures/recordings/google-embedding/maps-inputtype-to-a-task-type-for-gemini-embedding-001.json @@ -0,0 +1,50 @@ +{ + "version": 1, + "metadata": { + "tags": [ + "prefix:google-embedding", + "provider:google", + "protocol:google-embed-content" + ], + "name": "google-embedding/maps-inputtype-to-a-task-type-for-gemini-embedding-001", + "recordedAt": "2026-11-09T06:56:10.031Z" + }, + "interactions": [ + { + "transport": "http", + "request": { + "method": "POST", + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents", + "headers": { + "content-type": "application/json" + }, + "body": "{\"requests\":[{\"model\":\"models/gemini-embedding-001\",\"content\":{\"parts\":[{\"text\":\"The invoice was charged twice.\"}]},\"taskType\":\"RETRIEVAL_QUERY\",\"outputDimensionality\":64}]}" + }, + "response": { + "status": 200, + "headers": { + "content-type": "application/json; charset=UTF-8" + }, + "body": "{\n \"embeddings\": [\n {\n \"values\": [\n 0.04013612,\n -0.01537641,\n -0.00894363,\n -0.045427274,\n 0.015747963,\n -0.009066882,\n 0.034106165,\n 0.008299592,\n -0.029671155,\n 0.0073614777,\n 0.011705497,\n -0.008064136,\n -0.012201166,\n -0.023549316,\n 0.12342585,\n 0.011068751,\n -0.003649386,\n -0.0046033487,\n -0.011378856,\n 0.008659822,\n -0.024897752,\n 0.001241405,\n -0.0016047506,\n 0.032963928,\n 0.014408621,\n -0.011485396,\n 0.011547241,\n -0.0087663345,\n 0.029018002,\n 0.01399287,\n -0.0005734975,\n 0.009806568,\n -0.033623748,\n 0.013897123,\n 0.027592098,\n -0.001526048,\n 0.024996622,\n -0.0028200862,\n -0.0016486795,\n -0.019535935,\n -0.0060296683,\n 0.020810796,\n -0.014712436,\n 0.0027340506,\n -0.021106219,\n -0.022350833,\n 0.0039701443,\n 0.010802266,\n -0.0124504585,\n -0.0040887143,\n -0.0062737963,\n -0.018470604,\n 0.021161897,\n -0.16764373,\n 0.018159335,\n -0.015530227,\n -0.025055906,\n -0.026749635,\n 0.002768487,\n -0.03824891,\n -0.010511027,\n 0.014685062,\n -0.015881633,\n -0.042950973\n ]\n }\n ]\n}\n" + } + }, + { + "transport": "http", + "request": { + "method": "POST", + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents", + "headers": { + "content-type": "application/json" + }, + "body": "{\"requests\":[{\"model\":\"models/gemini-embedding-001\",\"content\":{\"parts\":[{\"text\":\"The invoice was charged twice.\"}]},\"taskType\":\"RETRIEVAL_DOCUMENT\",\"outputDimensionality\":64}]}" + }, + "response": { + "status": 200, + "headers": { + "content-type": "application/json; charset=UTF-8" + }, + "body": "{\n \"embeddings\": [\n {\n \"values\": [\n 0.0067660585,\n -0.008564701,\n 0.006918798,\n -0.0653632,\n 0.010584768,\n 0.0019771366,\n 0.024187071,\n 0.00468459,\n -0.022482483,\n -0.0011979828,\n 0.01025464,\n -0.012994005,\n 0.010501545,\n -0.008400968,\n 0.1354296,\n -0.003842612,\n -0.009868842,\n 0.0045481804,\n -0.012495539,\n 0.010945229,\n -0.012774331,\n 0.0018108432,\n 0.004564442,\n 0.024119604,\n 0.0096224435,\n 0.0008170349,\n 0.013087967,\n 0.0028655464,\n 0.038301487,\n 0.013128475,\n 0.008370898,\n 0.01710832,\n -0.02379502,\n 0.006764786,\n 0.028454835,\n 0.01613345,\n 0.013660678,\n 0.0013027737,\n 0.003783348,\n -0.0053537926,\n -0.0019457041,\n 0.022436744,\n -0.021374583,\n 0.0048431614,\n 0.0014370559,\n -0.0071425945,\n -0.0075106425,\n 0.0035781835,\n -0.002726382,\n -0.0017613088,\n 0.00242477,\n -0.012590718,\n 0.0015917314,\n -0.23083314,\n 0.006775061,\n -0.013197291,\n -0.020559328,\n -0.01887775,\n 0.011951808,\n -0.03615789,\n -0.017971825,\n 0.011554768,\n -0.017077072,\n -0.033402666\n ]\n }\n ]\n}\n" + } + } + ] +} diff --git a/packages/ai/test/provider/embedding.recorded.test.ts b/packages/ai/test/provider/embedding.recorded.test.ts index b36675dc3ab..12f7e322a35 100644 --- a/packages/ai/test/provider/embedding.recorded.test.ts +++ b/packages/ai/test/provider/embedding.recorded.test.ts @@ -1,7 +1,7 @@ import { describe, expect } from "bun:test" import { Effect } from "effect" -import { ai, type EmbeddingModel, type EmbeddingOptions } from "../../src/effect.js" -import { Cohere, OpenAI } from "../../src/providers.js" +import { ai, Media, type EmbeddingModel, type EmbeddingOptions } from "../../src/effect.js" +import { Cohere, Google, OpenAI } from "../../src/providers.js" import { recordedTests } from "../recorded-test.js" const values = ["The invoice was charged twice.", "Refund the duplicate payment."] @@ -22,6 +22,15 @@ const cohere = recordedTests({ metadata: { model: "embed-v4.0" }, }) +const google = recordedTests({ + prefix: "google-embedding", + provider: "google", + protocol: "google-embed-content", + requires: ["GOOGLE_GENERATIVE_AI_API_KEY"], +}) + +const googleProvider = Google.configure({ apiKey: process.env.GOOGLE_GENERATIVE_AI_API_KEY ?? "fixture" }) + describe("Embedding recorded", () => { openai.effect("embeds text with requested dimensions", () => assertEmbedding( @@ -32,6 +41,44 @@ describe("Embedding recorded", () => { cohere.effect("embeds text through the compatibility API", () => assertEmbedding(Cohere.configure({ apiKey: process.env.COHERE_API_KEY ?? "fixture" }).embedding("embed-v4.0")), ) + + google.effect("embeds text through batchEmbedContents", () => + assertEmbedding(googleProvider.embedding("gemini-embedding-2")), + ) + + google.effect("embeds an image and text into one vector each", () => + Effect.gen(function* () { + const image = Media.bytes( + new Uint8Array( + yield* Effect.promise(() => Bun.file(`${import.meta.dir}/../fixtures/images/edit-source.jpg`).arrayBuffer()), + ), + "image/jpeg", + ) + const response = yield* ai.embedding.generate({ + model: googleProvider.embedding("gemini-embedding-2"), + values: [image, "A small square image."], + dimensions: 128, + }) + expect(response.embeddings).toHaveLength(2) + for (const vector of response.embeddings) expect(vector).toHaveLength(128) + expect(response.usage?.inputTokens).toBeGreaterThan(258) + }), + ) + + google.effect("maps inputType to a task type for gemini-embedding-001", () => + Effect.gen(function* () { + const model = googleProvider.embedding("gemini-embedding-001") + const query = yield* ai.embedding.generate({ model, values: [values[0]], dimensions: 64, inputType: "query" }) + const document = yield* ai.embedding.generate({ + model, + values: [values[0]], + dimensions: 64, + inputType: "document", + }) + expect(query.embeddings[0]).toHaveLength(64) + expect(query.embeddings[0]).not.toEqual(document.embeddings[0]) + }), + ) }) const assertEmbedding = (model: EmbeddingModel) =>