@ai-sdk/google 3.0.9 → 3.0.11

package/src/convert-to-google-generative-ai-messages.ts

@@ -0,0 +1,232 @@
+import {
+  LanguageModelV3Prompt,
+  UnsupportedFunctionalityError,
+} from '@ai-sdk/provider';
+import {
+  GoogleGenerativeAIContent,
+  GoogleGenerativeAIContentPart,
+  GoogleGenerativeAIPrompt,
+} from './google-generative-ai-prompt';
+import { convertToBase64 } from '@ai-sdk/provider-utils';
+
+export function convertToGoogleGenerativeAIMessages(
+  prompt: LanguageModelV3Prompt,
+  options?: { isGemmaModel?: boolean; providerOptionsName?: string },
+): GoogleGenerativeAIPrompt {
+  const systemInstructionParts: Array<{ text: string }> = [];
+  const contents: Array<GoogleGenerativeAIContent> = [];
+  let systemMessagesAllowed = true;
+  const isGemmaModel = options?.isGemmaModel ?? false;
+  const providerOptionsName = options?.providerOptionsName ?? 'google';
+
+  for (const { role, content } of prompt) {
+    switch (role) {
+      case 'system': {
+        if (!systemMessagesAllowed) {
+          throw new UnsupportedFunctionalityError({
+            functionality:
+              'system messages are only supported at the beginning of the conversation',
+          });
+        }
+
+        systemInstructionParts.push({ text: content });
+        break;
+      }
+
+      case 'user': {
+        systemMessagesAllowed = false;
+
+        const parts: GoogleGenerativeAIContentPart[] = [];
+
+        for (const part of content) {
+          switch (part.type) {
+            case 'text': {
+              parts.push({ text: part.text });
+              break;
+            }
+
+            case 'file': {
+              // default to image/jpeg for unknown image/* types
+              const mediaType =
+                part.mediaType === 'image/*' ? 'image/jpeg' : part.mediaType;
+
+              parts.push(
+                part.data instanceof URL
+                  ? {
+                      fileData: {
+                        mimeType: mediaType,
+                        fileUri: part.data.toString(),
+                      },
+                    }
+                  : {
+                      inlineData: {
+                        mimeType: mediaType,
+                        data: convertToBase64(part.data),
+                      },
+                    },
+              );
+
+              break;
+            }
+          }
+        }
+
+        contents.push({ role: 'user', parts });
+        break;
+      }
+
+      case 'assistant': {
+        systemMessagesAllowed = false;
+
+        contents.push({
+          role: 'model',
+          parts: content
+            .map(part => {
+              const providerOpts = part.providerOptions?.[providerOptionsName];
+              const thoughtSignature =
+                providerOpts?.thoughtSignature != null
+                  ? String(providerOpts.thoughtSignature)
+                  : undefined;
+
+              switch (part.type) {
+                case 'text': {
+                  return part.text.length === 0
+                    ? undefined
+                    : {
+                        text: part.text,
+                        thoughtSignature,
+                      };
+                }
+
+                case 'reasoning': {
+                  return part.text.length === 0
+                    ? undefined
+                    : {
+                        text: part.text,
+                        thought: true,
+                        thoughtSignature,
+                      };
+                }
+
+                case 'file': {
+                  if (part.data instanceof URL) {
+                    throw new UnsupportedFunctionalityError({
+                      functionality:
+                        'File data URLs in assistant messages are not supported',
+                    });
+                  }
+
+                  return {
+                    inlineData: {
+                      mimeType: part.mediaType,
+                      data: convertToBase64(part.data),
+                    },
+                    thoughtSignature,
+                  };
+                }
+
+                case 'tool-call': {
+                  return {
+                    functionCall: {
+                      name: part.toolName,
+                      args: part.input,
+                    },
+                    thoughtSignature,
+                  };
+                }
+              }
+            })
+            .filter(part => part !== undefined),
+        });
+        break;
+      }
+
+      case 'tool': {
+        systemMessagesAllowed = false;
+
+        const parts: GoogleGenerativeAIContentPart[] = [];
+
+        for (const part of content) {
+          if (part.type === 'tool-approval-response') {
+            continue;
+          }
+          const output = part.output;
+
+          if (output.type === 'content') {
+            for (const contentPart of output.value) {
+              switch (contentPart.type) {
+                case 'text':
+                  parts.push({
+                    functionResponse: {
+                      name: part.toolName,
+                      response: {
+                        name: part.toolName,
+                        content: contentPart.text,
+                      },
+                    },
+                  });
+                  break;
+                case 'image-data':
+                  parts.push(
+                    {
+                      inlineData: {
+                        mimeType: contentPart.mediaType,
+                        data: contentPart.data,
+                      },
+                    },
+                    {
+                      text: 'Tool executed successfully and returned this image as a response',
+                    },
+                  );
+                  break;
+                default:
+                  parts.push({ text: JSON.stringify(contentPart) });
+                  break;
+              }
+            }
+          } else {
+            parts.push({
+              functionResponse: {
+                name: part.toolName,
+                response: {
+                  name: part.toolName,
+                  content:
+                    output.type === 'execution-denied'
+                      ? (output.reason ?? 'Tool execution denied.')
+                      : output.value,
+                },
+              },
+            });
+          }
+        }
+
+        contents.push({
+          role: 'user',
+          parts,
+        });
+        break;
+      }
+    }
+  }
+
+  if (
+    isGemmaModel &&
+    systemInstructionParts.length > 0 &&
+    contents.length > 0 &&
+    contents[0].role === 'user'
+  ) {
+    const systemText = systemInstructionParts
+      .map(part => part.text)
+      .join('\n\n');
+
+    contents[0].parts.unshift({ text: systemText + '\n\n' });
+  }
+
+  return {
+    systemInstruction:
+      systemInstructionParts.length > 0 && !isGemmaModel
+        ? { parts: systemInstructionParts }
+        : undefined,
+    contents,
+  };
+}

package/src/get-model-path.test.ts

@@ -0,0 +1,16 @@
+import { getModelPath } from './get-model-path';
+import { it, expect } from 'vitest';
+
+it('should pass through model path for models/*', async () => {
+  expect(getModelPath('models/some-model')).toEqual('models/some-model');
+});
+
+it('should pass through model path for tunedModels/*', async () => {
+  expect(getModelPath('tunedModels/some-model')).toEqual(
+    'tunedModels/some-model',
+  );
+});
+
+it('should add model path prefix to models without slash', async () => {
+  expect(getModelPath('some-model')).toEqual('models/some-model');
+});

package/src/get-model-path.ts

@@ -0,0 +1,3 @@
+export function getModelPath(modelId: string): string {
+  return modelId.includes('/') ? modelId : `models/${modelId}`;
+}

package/src/google-error.ts

@@ -0,0 +1,26 @@
+import {
+  createJsonErrorResponseHandler,
+  type InferSchema,
+  lazySchema,
+  zodSchema,
+} from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+
+const googleErrorDataSchema = lazySchema(() =>
+  zodSchema(
+    z.object({
+      error: z.object({
+        code: z.number().nullable(),
+        message: z.string(),
+        status: z.string(),
+      }),
+    }),
+  ),
+);
+
+export type GoogleErrorData = InferSchema<typeof googleErrorDataSchema>;
+
+export const googleFailedResponseHandler = createJsonErrorResponseHandler({
+  errorSchema: googleErrorDataSchema,
+  errorToMessage: data => data.error.message,
+});

package/src/google-generative-ai-embedding-model.test.ts

@@ -0,0 +1,204 @@
+import { EmbeddingModelV3Embedding } from '@ai-sdk/provider';
+import { createTestServer } from '@ai-sdk/test-server/with-vitest';
+import { GoogleGenerativeAIEmbeddingModel } from './google-generative-ai-embedding-model';
+import { createGoogleGenerativeAI } from './google-provider';
+import { describe, it, expect, vi } from 'vitest';
+
+vi.mock('./version', () => ({
+  VERSION: '0.0.0-test',
+}));
+
+const dummyEmbeddings = [
+  [0.1, 0.2, 0.3, 0.4, 0.5],
+  [0.6, 0.7, 0.8, 0.9, 1.0],
+];
+const testValues = ['sunny day at the beach', 'rainy day in the city'];
+
+const provider = createGoogleGenerativeAI({ apiKey: 'test-api-key' });
+const model = provider.embeddingModel('gemini-embedding-001');
+
+const URL =
+  'https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:something';
+
+const server = createTestServer({
+  [URL]: {},
+});
+
+describe('GoogleGenerativeAIEmbeddingModel', () => {
+  function prepareBatchJsonResponse({
+    embeddings = dummyEmbeddings,
+    headers,
+  }: {
+    embeddings?: EmbeddingModelV3Embedding[];
+    headers?: Record<string, string>;
+  } = {}) {
+    server.urls[URL].response = {
+      type: 'json-value',
+      headers,
+      body: {
+        embeddings: embeddings.map(embedding => ({ values: embedding })),
+      },
+    };
+  }
+
+  function prepareSingleJsonResponse({
+    embeddings = dummyEmbeddings,
+    headers,
+  }: {
+    embeddings?: EmbeddingModelV3Embedding[];
+    headers?: Record<string, string>;
+  } = {}) {
+    server.urls[URL].response = {
+      type: 'json-value',
+      headers,
+      body: {
+        embedding: { values: embeddings[0] },
+      },
+    };
+  }
+
+  it('should extract embedding', async () => {
+    prepareBatchJsonResponse();
+
+    const { embeddings } = await model.doEmbed({ values: testValues });
+
+    expect(embeddings).toStrictEqual(dummyEmbeddings);
+  });
+
+  it('should expose the raw response', async () => {
+    prepareBatchJsonResponse({
+      headers: {
+        'test-header': 'test-value',
+      },
+    });
+
+    const { response } = await model.doEmbed({ values: testValues });
+
+    expect(response?.headers).toStrictEqual({
+      // default headers:
+      'content-length': '80',
+      'content-type': 'application/json',
+
+      // custom header
+      'test-header': 'test-value',
+    });
+    expect(response).toMatchSnapshot();
+  });
+
+  it('should pass the model and the values', async () => {
+    prepareBatchJsonResponse();
+
+    await model.doEmbed({ values: testValues });
+
+    expect(await server.calls[0].requestBodyJson).toStrictEqual({
+      requests: testValues.map(value => ({
+        model: 'models/gemini-embedding-001',
+        content: { role: 'user', parts: [{ text: value }] },
+      })),
+    });
+  });
+
+  it('should pass the outputDimensionality setting', async () => {
+    prepareBatchJsonResponse();
+
+    await provider.embedding('gemini-embedding-001').doEmbed({
+      values: testValues,
+      providerOptions: {
+        google: { outputDimensionality: 64 },
+      },
+    });
+
+    expect(await server.calls[0].requestBodyJson).toStrictEqual({
+      requests: testValues.map(value => ({
+        model: 'models/gemini-embedding-001',
+        content: { role: 'user', parts: [{ text: value }] },
+        outputDimensionality: 64,
+      })),
+    });
+  });
+
+  it('should pass the taskType setting', async () => {
+    prepareBatchJsonResponse();
+
+    await provider.embedding('gemini-embedding-001').doEmbed({
+      values: testValues,
+      providerOptions: { google: { taskType: 'SEMANTIC_SIMILARITY' } },
+    });
+
+    expect(await server.calls[0].requestBodyJson).toStrictEqual({
+      requests: testValues.map(value => ({
+        model: 'models/gemini-embedding-001',
+        content: { role: 'user', parts: [{ text: value }] },
+        taskType: 'SEMANTIC_SIMILARITY',
+      })),
+    });
+  });
+
+  it('should pass headers', async () => {
+    prepareBatchJsonResponse();
+
+    const provider = createGoogleGenerativeAI({
+      apiKey: 'test-api-key',
+      headers: {
+        'Custom-Provider-Header': 'provider-header-value',
+      },
+    });
+
+    await provider.embedding('gemini-embedding-001').doEmbed({
+      values: testValues,
+      headers: {
+        'Custom-Request-Header': 'request-header-value',
+      },
+    });
+
+    expect(server.calls[0].requestHeaders).toStrictEqual({
+      'x-goog-api-key': 'test-api-key',
+      'content-type': 'application/json',
+      'custom-provider-header': 'provider-header-value',
+      'custom-request-header': 'request-header-value',
+    });
+    expect(server.calls[0].requestUserAgent).toContain(
+      `ai-sdk/google/0.0.0-test`,
+    );
+  });
+
+  it('should throw an error if too many values are provided', async () => {
+    const model = new GoogleGenerativeAIEmbeddingModel('gemini-embedding-001', {
+      provider: 'google.generative-ai',
+      baseURL: 'https://generativelanguage.googleapis.com/v1beta',
+      headers: () => ({}),
+    });
+
+    const tooManyValues = Array(2049).fill('test');
+
+    await expect(model.doEmbed({ values: tooManyValues })).rejects.toThrow(
+      'Too many values for a single embedding call. The google.generative-ai model "gemini-embedding-001" can only embed up to 2048 values per call, but 2049 values were provided.',
+    );
+  });
+
+  it('should use the batch embeddings endpoint', async () => {
+    prepareBatchJsonResponse();
+    const model = provider.embeddingModel('gemini-embedding-001');
+    await model.doEmbed({
+      values: testValues,
+    });
+
+    expect(server.calls[0].requestUrl).toBe(
+      'https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents',
+    );
+  });
+
+  it('should use the single embeddings endpoint', async () => {
+    prepareSingleJsonResponse();
+
+    const model = provider.embeddingModel('gemini-embedding-001');
+
+    await model.doEmbed({
+      values: [testValues[0]],
+    });
+
+    expect(server.calls[0].requestUrl).toBe(
+      'https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:embedContent',
+    );
+  });
+});

package/src/google-generative-ai-embedding-model.ts

@@ -0,0 +1,159 @@
+import {
+  EmbeddingModelV3,
+  TooManyEmbeddingValuesForCallError,
+} from '@ai-sdk/provider';
+import {
+  combineHeaders,
+  createJsonResponseHandler,
+  FetchFunction,
+  lazySchema,
+  parseProviderOptions,
+  postJsonToApi,
+  resolve,
+  zodSchema,
+} from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+import { googleFailedResponseHandler } from './google-error';
+import {
+  GoogleGenerativeAIEmbeddingModelId,
+  googleGenerativeAIEmbeddingProviderOptions,
+} from './google-generative-ai-embedding-options';
+
+type GoogleGenerativeAIEmbeddingConfig = {
+  provider: string;
+  baseURL: string;
+  headers: () => Record<string, string | undefined>;
+  fetch?: FetchFunction;
+};
+
+export class GoogleGenerativeAIEmbeddingModel implements EmbeddingModelV3 {
+  readonly specificationVersion = 'v3';
+  readonly modelId: GoogleGenerativeAIEmbeddingModelId;
+  readonly maxEmbeddingsPerCall = 2048;
+  readonly supportsParallelCalls = true;
+
+  private readonly config: GoogleGenerativeAIEmbeddingConfig;
+
+  get provider(): string {
+    return this.config.provider;
+  }
+  constructor(
+    modelId: GoogleGenerativeAIEmbeddingModelId,
+    config: GoogleGenerativeAIEmbeddingConfig,
+  ) {
+    this.modelId = modelId;
+    this.config = config;
+  }
+
+  async doEmbed({
+    values,
+    headers,
+    abortSignal,
+    providerOptions,
+  }: Parameters<EmbeddingModelV3['doEmbed']>[0]): Promise<
+    Awaited<ReturnType<EmbeddingModelV3['doEmbed']>>
+  > {
+    // Parse provider options
+    const googleOptions = await parseProviderOptions({
+      provider: 'google',
+      providerOptions,
+      schema: googleGenerativeAIEmbeddingProviderOptions,
+    });
+
+    if (values.length > this.maxEmbeddingsPerCall) {
+      throw new TooManyEmbeddingValuesForCallError({
+        provider: this.provider,
+        modelId: this.modelId,
+        maxEmbeddingsPerCall: this.maxEmbeddingsPerCall,
+        values,
+      });
+    }
+
+    const mergedHeaders = combineHeaders(
+      await resolve(this.config.headers),
+      headers,
+    );
+
+    // For single embeddings, use the single endpoint (ratelimits, etc.)
+    if (values.length === 1) {
+      const {
+        responseHeaders,
+        value: response,
+        rawValue,
+      } = await postJsonToApi({
+        url: `${this.config.baseURL}/models/${this.modelId}:embedContent`,
+        headers: mergedHeaders,
+        body: {
+          model: `models/${this.modelId}`,
+          content: {
+            parts: [{ text: values[0] }],
+          },
+          outputDimensionality: googleOptions?.outputDimensionality,
+          taskType: googleOptions?.taskType,
+        },
+        failedResponseHandler: googleFailedResponseHandler,
+        successfulResponseHandler: createJsonResponseHandler(
+          googleGenerativeAISingleEmbeddingResponseSchema,
+        ),
+        abortSignal,
+        fetch: this.config.fetch,
+      });
+
+      return {
+        warnings: [],
+        embeddings: [response.embedding.values],
+        usage: undefined,
+        response: { headers: responseHeaders, body: rawValue },
+      };
+    }
+
+    const {
+      responseHeaders,
+      value: response,
+      rawValue,
+    } = await postJsonToApi({
+      url: `${this.config.baseURL}/models/${this.modelId}:batchEmbedContents`,
+      headers: mergedHeaders,
+      body: {
+        requests: values.map(value => ({
+          model: `models/${this.modelId}`,
+          content: { role: 'user', parts: [{ text: value }] },
+          outputDimensionality: googleOptions?.outputDimensionality,
+          taskType: googleOptions?.taskType,
+        })),
+      },
+      failedResponseHandler: googleFailedResponseHandler,
+      successfulResponseHandler: createJsonResponseHandler(
+        googleGenerativeAITextEmbeddingResponseSchema,
+      ),
+      abortSignal,
+      fetch: this.config.fetch,
+    });
+
+    return {
+      warnings: [],
+      embeddings: response.embeddings.map(item => item.values),
+      usage: undefined,
+      response: { headers: responseHeaders, body: rawValue },
+    };
+  }
+}
+
+// minimal version of the schema, focussed on what is needed for the implementation
+// this approach limits breakages when the API changes and increases efficiency
+const googleGenerativeAITextEmbeddingResponseSchema = lazySchema(() =>
+  zodSchema(
+    z.object({
+      embeddings: z.array(z.object({ values: z.array(z.number()) })),
+    }),
+  ),
+);
+
+// Schema for single embedding response
+const googleGenerativeAISingleEmbeddingResponseSchema = lazySchema(() =>
+  zodSchema(
+    z.object({
+      embedding: z.object({ values: z.array(z.number()) }),
+    }),
+  ),
+);

package/src/google-generative-ai-embedding-options.ts

@@ -0,0 +1,52 @@
+import {
+  type InferSchema,
+  lazySchema,
+  zodSchema,
+} from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+
+export type GoogleGenerativeAIEmbeddingModelId =
+  | 'gemini-embedding-001'
+  | 'text-embedding-004'
+  | (string & {});
+
+export const googleGenerativeAIEmbeddingProviderOptions = lazySchema(() =>
+  zodSchema(
+    z.object({
+      /**
+       * Optional. Optional reduced dimension for the output embedding.
+       * If set, excessive values in the output embedding are truncated from the end.
+       */
+      outputDimensionality: z.number().optional(),
+
+      /**
+       * Optional. Specifies the task type for generating embeddings.
+       * Supported task types:
+       * - SEMANTIC_SIMILARITY: Optimized for text similarity.
+       * - CLASSIFICATION: Optimized for text classification.
+       * - CLUSTERING: Optimized for clustering texts based on similarity.
+       * - RETRIEVAL_DOCUMENT: Optimized for document retrieval.
+       * - RETRIEVAL_QUERY: Optimized for query-based retrieval.
+       * - QUESTION_ANSWERING: Optimized for answering questions.
+       * - FACT_VERIFICATION: Optimized for verifying factual information.
+       * - CODE_RETRIEVAL_QUERY: Optimized for retrieving code blocks based on natural language queries.
+       */
+      taskType: z
+        .enum([
+          'SEMANTIC_SIMILARITY',
+          'CLASSIFICATION',
+          'CLUSTERING',
+          'RETRIEVAL_DOCUMENT',
+          'RETRIEVAL_QUERY',
+          'QUESTION_ANSWERING',
+          'FACT_VERIFICATION',
+          'CODE_RETRIEVAL_QUERY',
+        ])
+        .optional(),
+    }),
+  ),
+);
+
+export type GoogleGenerativeAIEmbeddingProviderOptions = InferSchema<
+  typeof googleGenerativeAIEmbeddingProviderOptions
+>;