@ai-sdk/google/src/google-batch.ts

import {
  InvalidArgumentError,
  InvalidResponseDataError,
  UnsupportedFunctionalityError,
  type Experimental_BatchV4 as BatchV4,
  type Experimental_BatchV4CancelResult as BatchV4CancelResult,
  type Experimental_BatchV4Error as BatchV4Error,
  type Experimental_BatchV4ItemResult as BatchV4ItemResult,
  type Experimental_BatchV4ListOptions as BatchV4ListOptions,
  type Experimental_BatchV4ListResult as BatchV4ListResult,
  type Experimental_BatchV4OperationOptions as BatchV4OperationOptions,
  type Experimental_BatchV4StartResult as BatchV4StartResult,
  type Experimental_BatchV4Status as BatchV4Status,
  type Experimental_ImageBatchV4Request as ImageBatchV4Request,
  type Experimental_BatchV4StartOptions as BatchV4StartOptions,
  type LanguageModelV4GenerateResult,
  type ImageModelV4Result,
  type LanguageModelV4Prompt,
  type SharedV4Warning,
} from '@ai-sdk/provider';
import {
  combineHeaders,
  convertAsyncIteratorToReadableStream,
  createJsonLinesResponseHandler,
  createJsonResponseHandler,
  generateId,
  getFromApi,
  lazySchema,
  normalizeBatchRequestCounts,
  parseProviderOptions,
  postJsonToApi,
  postToApi,
  resolve,
  safeValidateTypes,
  zodSchema,
  convertToBase64,
  type InferSchema,
  type ResponseHandler,
} from '@ai-sdk/provider-utils';
import { z } from 'zod/v4';
import { getModelPath } from './get-model-path';
import { googleFailedResponseHandler } from './google-error';
import {
  GoogleLanguageModel,
  responseSchema,
  type GoogleLanguageModelConfig,
} from './google-language-model';
import {
  googleLanguageModelOptions,
  type GoogleModelId,
} from './google-language-model-options';
import type { GoogleImageModelId } from './google-image-settings';
import { googleImageModelOptionsSchema } from './google-image-model-options';

const googleBatchInputFileMaxBytes = 2 * 1024 * 1024 * 1024;
const googleBatchInlineCreationMaxBytes = 20_000_000;
const supportedGoogleBatchContentTypes = new Set<
  LanguageModelV4GenerateResult['content'][number]['type']
>(['text', 'reasoning', 'source', 'tool-call', 'tool-result']);

type GoogleImageBatchRequest = ImageBatchV4Request<GoogleImageModelId>;
type GoogleBatchModelIds = {
  readonly text: GoogleModelId;
  readonly image: GoogleImageModelId;
};

function assertSupportedBatchRequests(
  requests: BatchV4StartOptions<GoogleBatchModelIds>['requests'],
) {
  for (const request of requests) {
    const requestType = request.type;
    if (requestType !== 'text' && requestType !== 'image') {
      throw new UnsupportedFunctionalityError({
        functionality: `batch request type: ${requestType}`,
        message: `The Google Batch API does not support batch requests with type "${requestType}".`,
      });
    }
  }
}

const googleRpcStatusSchema = () =>
  z.object({
    code: z.union([z.number(), z.string()]).nullish(),
    message: z.string().nullish(),
    status: z.string().nullish(),
  });

const googleBatchStatsSchema = () =>
  z.object({
    requestCount: z.union([z.string(), z.number()]).nullish(),
    successfulRequestCount: z.union([z.string(), z.number()]).nullish(),
    failedRequestCount: z.union([z.string(), z.number()]).nullish(),
    pendingRequestCount: z.union([z.string(), z.number()]).nullish(),
  });

const googleBatchOutputSchema = () =>
  z.object({
    responsesFile: z.string().nullish(),
    inlinedResponses: z
      .object({
        inlinedResponses: z.array(
          z.object({
            metadata: z.object({
              key: z.string(),
            }),
            response: z.unknown().nullish(),
            error: googleRpcStatusSchema().nullish(),
          }),
        ),
      })
      .nullish(),
  });

const googleBatchOperationZodSchema = () =>
  z.object({
    name: z.string(),
    metadata: z
      .object({
        state: z.string().nullish(),
        createTime: z.string().nullish(),
        batchStats: googleBatchStatsSchema().nullish(),
        output: googleBatchOutputSchema().nullish(),
      })
      .nullish(),
    done: z.boolean().nullish(),
    error: googleRpcStatusSchema().nullish(),
    response: googleBatchOutputSchema().nullish(),
  });

const googleBatchOperationSchema = lazySchema(() =>
  zodSchema(googleBatchOperationZodSchema()),
);

type GoogleBatchOperation = InferSchema<typeof googleBatchOperationSchema>;

const googleBatchListResponseSchema = lazySchema(() =>
  zodSchema(
    z.object({
      operations: z.array(googleBatchOperationZodSchema()).nullish(),
      nextPageToken: z.string().nullish(),
    }),
  ),
);

const googleBatchCancelResponseSchema = lazySchema(() =>
  zodSchema(z.object({})),
);

const googleFileUploadResponseSchema = lazySchema(() =>
  zodSchema(
    z.object({
      file: z.object({
        name: z.string(),
        expirationTime: z.string().nullish(),
      }),
    }),
  ),
);

const googleBatchResultLineSchema = lazySchema(() =>
  zodSchema(
    z.object({
      key: z.string(),
      response: z.unknown().nullish(),
      error: googleRpcStatusSchema().nullish(),
    }),
  ),
);

type GoogleBatchResultLine = InferSchema<typeof googleBatchResultLineSchema>;

const googleBatchResponsePreviewSchema = lazySchema(() =>
  zodSchema(
    z.object({
      candidates: z.array(z.unknown()).nullish(),
      promptFeedback: z
        .object({
          blockReason: z.string().nullish(),
        })
        .nullish(),
    }),
  ),
);

export class GoogleBatch implements BatchV4<GoogleBatchModelIds> {
  readonly specificationVersion = 'v4' as const;
  readonly provider: string;
  readonly supportedUrls: Record<string, RegExp[]>;
  private readonly batchConfig: GoogleLanguageModelConfig;
  private readonly batchGenerateId: () => string;
  private readonly maxLineBytes?: number;

  constructor(options: {
    provider: string;
    config: GoogleLanguageModelConfig;
    supportedUrls: Record<string, RegExp[]>;
    maxLineBytes?: number;
  }) {
    this.provider = options.provider;
    this.batchConfig = options.config;
    this.supportedUrls = options.supportedUrls;
    this.batchGenerateId = options.config.generateId ?? generateId;
    this.maxLineBytes = options.maxLineBytes;
  }

  async doStartBatch(
    options: BatchV4StartOptions<GoogleBatchModelIds>,
  ): Promise<BatchV4StartResult> {
    assertSupportedBatchRequests(options.requests);
    const modelId = getGoogleBatchModelId(options.requests);
    const warnings: BatchV4StartResult['warnings'] = [];
    const displayName = `ai-sdk-batch-${this.batchGenerateId()}`;
    const inlinedRequests: Array<{
      request: unknown;
      metadata: { key: string };
    }> = [];
    const inlineBatchBody = {
      batch: {
        displayName,
        ...(options.webhookUrl != null && {
          webhookConfig: { uris: [options.webhookUrl] },
        }),
        inputConfig: {
          requests: { requests: inlinedRequests },
        },
      },
    };
    const textEncoder = new TextEncoder();
    let inlineInputBytes = textEncoder.encode(
      JSON.stringify(inlineBatchBody),
    ).byteLength;
    let fileParts: string[] | undefined;

    for (const request of options.requests) {
      const preparedRequest =
        request.type === 'text'
          ? await GoogleLanguageModel.prepareRequest({
              modelId: request.modelId,
              config: this.batchConfig,
              options: request.options,
            })
          : await this.prepareImageRequest(request);
      const inlinedRequest = {
        request: preparedRequest.args,
        metadata: { key: request.id },
      };

      if (fileParts == null) {
        const requestBytes = textEncoder.encode(
          JSON.stringify(inlinedRequest),
        ).byteLength;
        const nextInlineInputBytes =
          inlineInputBytes +
          requestBytes +
          (inlinedRequests.length > 0 ? 1 : 0);

        if (nextInlineInputBytes < googleBatchInlineCreationMaxBytes) {
          inlinedRequests.push(inlinedRequest);
          inlineInputBytes = nextInlineInputBytes;
        } else {
          fileParts = [];
          for (const previousRequest of inlinedRequests) {
            fileParts.push(
              JSON.stringify({
                key: previousRequest.metadata.key,
                request: previousRequest.request,
              }),
              '\n',
            );
          }
          inlinedRequests.length = 0;
          fileParts.push(
            JSON.stringify({
              key: request.id,
              request: preparedRequest.args,
            }),
            '\n',
          );
        }
      } else {
        fileParts.push(
          JSON.stringify({
            key: request.id,
            request: preparedRequest.args,
          }),
          '\n',
        );
      }

      for (const warning of preparedRequest.warnings) {
        warnings.push({ requestId: request.id, warning });
      }
    }

    const createUrl = `${this.batchConfig.baseURL}/${getModelPath(
      modelId,
    )}:batchGenerateContent`;

    if (fileParts == null) {
      const { value: operation } = await postJsonToApi({
        url: createUrl,
        headers: await this.getHeaders(options.headers),
        body: inlineBatchBody,
        failedResponseHandler: googleFailedResponseHandler,
        successfulResponseHandler: createJsonResponseHandler(
          googleBatchOperationSchema,
        ),
        abortSignal: options.abortSignal,
        fetch: this.batchConfig.fetch,
      });

      return {
        batchId: operation.name,
        ...convertGoogleBatchStatus(operation),
        warnings,
      };
    }

    const inputFile = new Blob(fileParts, { type: 'application/jsonl' });
    // Blob snapshots the strings, so release the potentially large input array.
    fileParts.length = 0;
    if (inputFile.size > googleBatchInputFileMaxBytes) {
      throw new InvalidArgumentError({
        argument: 'requests',
        message: 'Google batch input files must not exceed 2 GB.',
      });
    }

    const { value: uploadUrl } = await postJsonToApi({
      url: `${this.getBaseOrigin()}/upload/v1beta/files`,
      headers: combineHeaders(await this.getHeaders(options.headers), {
        'X-Goog-Upload-Protocol': 'resumable',
        'X-Goog-Upload-Command': 'start',
        'X-Goog-Upload-Header-Content-Length': String(inputFile.size),
        'X-Goog-Upload-Header-Content-Type': 'application/jsonl',
      }),
      body: {
        file: {
          display_name: `${displayName}-input`,
        },
      },
      failedResponseHandler: googleFailedResponseHandler,
      successfulResponseHandler: googleUploadUrlResponseHandler,
      abortSignal: options.abortSignal,
      fetch: this.batchConfig.fetch,
    });

    const { value: uploadedFile } = await postToApi({
      url: uploadUrl,
      headers: {
        'X-Goog-Upload-Offset': '0',
        'X-Goog-Upload-Command': 'upload, finalize',
        'Content-Type': 'application/jsonl',
      },
      body: {
        content: inputFile,
        values: {
          byteLength: inputFile.size,
          mediaType: 'application/jsonl',
        },
      },
      failedResponseHandler: googleFailedResponseHandler,
      successfulResponseHandler: createJsonResponseHandler(
        googleFileUploadResponseSchema,
      ),
      abortSignal: options.abortSignal,
      fetch: this.batchConfig.fetch,
    });

    const { value: operation } = await postJsonToApi({
      url: createUrl,
      headers: await this.getHeaders(options.headers),
      body: {
        batch: {
          displayName,
          ...(options.webhookUrl != null && {
            webhookConfig: { uris: [options.webhookUrl] },
          }),
          inputConfig: {
            fileName: uploadedFile.file.name,
          },
        },
      },
      failedResponseHandler: googleFailedResponseHandler,
      successfulResponseHandler: createJsonResponseHandler(
        googleBatchOperationSchema,
      ),
      abortSignal: options.abortSignal,
      fetch: this.batchConfig.fetch,
    });

    return {
      batchId: operation.name,
      ...convertGoogleBatchStatus(operation),
      providerMetadata: {
        google: {
          inputFileId: uploadedFile.file.name,
          ...(uploadedFile.file.expirationTime != null
            ? { inputFileExpiresAt: uploadedFile.file.expirationTime }
            : {}),
        },
      },
      warnings,
    };
  }

  async doGetBatchStatus(
    options: BatchV4OperationOptions,
  ): Promise<BatchV4Status> {
    return convertGoogleBatchStatus(await this.retrieveBatch(options));
  }

  async doCancelBatch(
    options: BatchV4OperationOptions,
  ): Promise<BatchV4CancelResult> {
    await postJsonToApi({
      url: `${this.batchConfig.baseURL}/${options.batchId}:cancel`,
      headers: await this.getHeaders(options.headers),
      body: {},
      failedResponseHandler: googleFailedResponseHandler,
      successfulResponseHandler: createJsonResponseHandler(
        googleBatchCancelResponseSchema,
      ),
      abortSignal: options.abortSignal,
      fetch: this.batchConfig.fetch,
    });

    return {};
  }

  async doListBatches(options: BatchV4ListOptions): Promise<BatchV4ListResult> {
    const url = new URL(`${this.batchConfig.baseURL}/batches`);
    if (options.limit != null) {
      url.searchParams.set('pageSize', String(options.limit));
    }
    if (options.cursor != null) {
      url.searchParams.set('pageToken', options.cursor);
    }

    const { value: page } = await getFromApi({
      url: url.toString(),
      headers: await this.getHeaders(options.headers),
      failedResponseHandler: googleFailedResponseHandler,
      successfulResponseHandler: createJsonResponseHandler(
        googleBatchListResponseSchema,
      ),
      abortSignal: options.abortSignal,
      fetch: this.batchConfig.fetch,
      validateUrl: false,
    });

    return {
      batches: (page.operations ?? []).map(operation => ({
        batchId: operation.name,
        ...convertGoogleBatchStatus(operation),
      })),
      ...(page.nextPageToken != null ? { nextCursor: page.nextPageToken } : {}),
    };
  }

  async doGetBatchResults(
    options: BatchV4OperationOptions,
  ): Promise<ReadableStream<BatchV4ItemResult>> {
    const operation = await this.retrieveBatch(options);
    const batchStatus = convertGoogleBatchStatus(operation);

    if (batchStatus.status === 'pending') {
      throw new InvalidArgumentError({
        argument: 'batchId',
        message: `Google batch "${options.batchId}" is not complete.`,
      });
    }

    const inlinedResponses =
      operation.metadata?.output?.inlinedResponses?.inlinedResponses ??
      operation.response?.inlinedResponses?.inlinedResponses;

    if (inlinedResponses != null) {
      return convertAsyncIteratorToReadableStream(
        this.iterateBatchResults(
          inlinedResponses.map(result => ({
            key: result.metadata.key,
            response: result.response,
            error: result.error,
          })),
        ),
      );
    }

    const responsesFile =
      operation.metadata?.output?.responsesFile ??
      operation.response?.responsesFile;

    if (responsesFile == null) {
      if (batchStatus.status === 'completed') {
        throw new InvalidResponseDataError({
          data: operation,
          message: `Google batch "${options.batchId}" completed without batch output.`,
        });
      }
      return new ReadableStream({
        start(controller) {
          controller.close();
        },
      });
    }

    const encodedResponsesFile = responsesFile
      .split('/')
      .map(segment => encodeURIComponent(segment))
      .join('/');

    const { value: lines } = await getFromApi({
      url: `${this.getBaseOrigin()}/download/v1beta/${encodedResponsesFile}:download?alt=media`,
      headers: await this.getHeaders(options.headers),
      failedResponseHandler: googleFailedResponseHandler,
      successfulResponseHandler: createJsonLinesResponseHandler(
        googleBatchResultLineSchema,
        { maxLineBytes: this.maxLineBytes },
      ),
      abortSignal: options.abortSignal,
      fetch: this.batchConfig.fetch,
      validateUrl: false,
    });

    return convertAsyncIteratorToReadableStream(
      this.iterateBatchResults(lines),
    );
  }

  private async retrieveBatch(
    options: BatchV4OperationOptions,
  ): Promise<GoogleBatchOperation> {
    const { value: operation } = await getFromApi({
      url: `${this.batchConfig.baseURL}/${options.batchId}`,
      headers: await this.getHeaders(options.headers),
      failedResponseHandler: googleFailedResponseHandler,
      successfulResponseHandler: createJsonResponseHandler(
        googleBatchOperationSchema,
      ),
      abortSignal: options.abortSignal,
      fetch: this.batchConfig.fetch,
      validateUrl: false,
    });

    return operation;
  }

  private async *iterateBatchResults(
    results:
      | Iterable<GoogleBatchResultLine>
      | AsyncIterable<GoogleBatchResultLine>,
  ): AsyncGenerator<BatchV4ItemResult> {
    for await (const line of results) {
      if (line.error != null) {
        const error = convertGoogleRpcError(
          line.error,
          'Google batch request failed.',
        );
        const status =
          line.error.status === 'CANCELLED' || String(line.error.code) === '1'
            ? 'cancelled'
            : 'failed';

        yield { type: 'text', id: line.key, status, error };
        continue;
      }

      if (line.response == null) {
        yield {
          type: 'text',
          id: line.key,
          status: 'failed',
          error: {
            message:
              'Google returned a batch result without a response or error.',
            code: 'invalid_batch_result',
          },
        };
        continue;
      }

      const preview = await safeValidateTypes({
        value: line.response,
        schema: googleBatchResponsePreviewSchema,
      });
      if (
        preview.success &&
        (preview.value.candidates == null ||
          preview.value.candidates.length === 0)
      ) {
        const promptFeedback = preview.value.promptFeedback ?? undefined;
        const blockReason = promptFeedback?.blockReason ?? undefined;
        yield {
          type: 'text',
          id: line.key,
          status: 'failed',
          error: {
            message:
              blockReason == null
                ? 'Google returned a batch response without any candidates.'
                : `Google blocked the batch request (${blockReason}).`,
            code: blockReason == null ? 'invalid_response' : 'prompt_blocked',
            ...(blockReason != null ? { type: blockReason } : {}),
          },
          ...(promptFeedback != null
            ? {
                providerMetadata: {
                  google: {
                    promptFeedback: {
                      blockReason: promptFeedback.blockReason ?? null,
                    },
                  },
                },
              }
            : {}),
        };
        continue;
      }

      const response = await safeValidateTypes({
        value: line.response,
        schema: responseSchema,
      });
      if (!response.success) {
        yield {
          type: 'text',
          id: line.key,
          status: 'failed',
          error: {
            message: 'Google returned an invalid GenerateContent batch result.',
            code: 'invalid_response',
          },
        };
        continue;
      }

      const result = GoogleLanguageModel.convertGenerateContentResponse({
        config: this.batchConfig,
        response: response.value,
        warnings: [],
        providerOptionsNames: ['google'],
      });
      const imageResult = convertGoogleImageBatchResult(result);
      if (imageResult != null) {
        yield {
          type: 'image',
          id: line.key,
          status: 'succeeded',
          result: imageResult,
        };
        continue;
      }
      const unsupportedPart = result.content.find(
        part => !supportedGoogleBatchContentTypes.has(part.type),
      );

      if (unsupportedPart != null) {
        yield {
          type: 'text',
          id: line.key,
          status: 'failed',
          error: {
            message:
              `Google returned a "${unsupportedPart.type}" content block, ` +
              'but that content is not supported in AI SDK text batches.',
            code: 'unsupported_content',
          },
        };
        continue;
      }

      yield { type: 'text', id: line.key, status: 'succeeded', result };
    }
  }

  private async prepareImageRequest(request: GoogleImageBatchRequest) {
    const { prompt, n, size, aspectRatio, seed, files, mask, providerOptions } =
      request.options;
    const warnings: SharedV4Warning[] = [];

    if (mask != null) {
      throw new UnsupportedFunctionalityError({
        functionality: 'mask-based image editing in Google batches',
      });
    }
    if (n > 1) {
      throw new UnsupportedFunctionalityError({
        functionality: 'multiple images per Google batch request',
      });
    }
    if (size != null) {
      warnings.push({
        type: 'unsupported',
        feature: 'size',
        details:
          'This model does not support the `size` option. Use `aspectRatio` instead.',
      });
    }

    const userContent: Extract<
      LanguageModelV4Prompt[number],
      { role: 'user' }
    >['content'] = [];
    if (prompt != null) userContent.push({ type: 'text', text: prompt });
    for (const file of files ?? []) {
      userContent.push(
        file.type === 'url'
          ? {
              type: 'file',
              data: { type: 'url', url: new URL(file.url) },
              mediaType: 'image/*',
            }
          : {
              type: 'file',
              data: { type: 'data', data: file.data },
              mediaType: file.mediaType,
            },
      );
    }

    const googleImageOptions = await parseProviderOptions({
      provider: 'google',
      providerOptions,
      schema: googleImageModelOptionsSchema,
    });
    const {
      responseModalities: _responseModalities,
      imageConfig: userImageConfig,
      ...passthroughGoogleOptions
    } = (await parseProviderOptions({
      provider: 'google',
      providerOptions,
      schema: googleLanguageModelOptions,
    })) ?? {};
    const preparedGoogleOptions = await parseProviderOptions({
      provider: 'google',
      providerOptions: {
        google: {
          ...passthroughGoogleOptions,
          responseModalities: ['IMAGE'],
          imageConfig:
            aspectRatio != null || userImageConfig != null
              ? {
                  ...userImageConfig,
                  ...(aspectRatio != null ? { aspectRatio } : {}),
                }
              : undefined,
        },
      },
      schema: googleLanguageModelOptions,
    });

    const prepared = await GoogleLanguageModel.prepareRequest({
      modelId: request.modelId,
      config: this.batchConfig,
      options: {
        prompt: [{ role: 'user', content: userContent }],
        seed,
        providerOptions: {
          google: preparedGoogleOptions ?? { responseModalities: ['IMAGE'] },
        },
        tools:
          googleImageOptions?.googleSearch != null
            ? [
                {
                  type: 'provider',
                  id: 'google.google_search',
                  name: 'google_search',
                  args: googleImageOptions.googleSearch,
                },
              ]
            : undefined,
      },
    });
    return { ...prepared, warnings: [...warnings, ...prepared.warnings] };
  }

  private async getHeaders(headers?: Record<string, string | undefined>) {
    return combineHeaders(
      this.batchConfig.headers
        ? await resolve(this.batchConfig.headers)
        : undefined,
      headers,
    );
  }

  private getBaseOrigin() {
    return this.batchConfig.baseURL.replace(/\/v1beta$/, '');
  }
}

function convertGoogleBatchStatus(
  operation: GoogleBatchOperation,
): BatchV4Status {
  const rawStatus = operation.metadata?.state ?? undefined;
  const requestCounts = convertGoogleRequestCounts(
    operation.metadata?.batchStats,
  );
  const createdAt = operation.metadata?.createTime ?? undefined;
  const error =
    operation.error != null
      ? convertGoogleRpcError(operation.error, 'Google batch failed.')
      : undefined;

  return {
    status: mapGoogleBatchStatus({
      rawStatus,
      done: operation.done ?? undefined,
      hasError: error != null,
    }),
    ...(rawStatus != null ? { rawStatus } : {}),
    ...(requestCounts != null ? { requestCounts } : {}),
    ...(error != null ? { error } : {}),
    ...(createdAt != null ? { createdAt } : {}),
  };
}

function mapGoogleBatchStatus({
  rawStatus,
  done,
  hasError,
}: {
  rawStatus?: string;
  done?: boolean;
  hasError: boolean;
}): BatchV4Status['status'] {
  if (hasError) {
    return 'failed';
  }

  if (rawStatus == null) {
    return done ? 'completed' : 'pending';
  }

  const normalizedStatus = rawStatus.replace(/^(?:BATCH|JOB)_STATE_/, '');
  switch (normalizedStatus) {
    case 'SUCCEEDED':
      return 'completed';
    case 'FAILED':
    case 'CANCELLED':
    case 'EXPIRED':
      return 'failed';
    case 'UNSPECIFIED':
    case 'PENDING':
    case 'RUNNING':
    default:
      return 'pending';
  }
}

function convertGoogleRequestCounts(
  counts: NonNullable<GoogleBatchOperation['metadata']>['batchStats'],
): BatchV4Status['requestCounts'] | undefined {
  const total = parseCount(counts?.requestCount);
  const completed = parseCount(counts?.successfulRequestCount ?? 0);
  const failed = parseCount(counts?.failedRequestCount ?? 0);
  const pending = parseCount(counts?.pendingRequestCount ?? 0);

  return normalizeBatchRequestCounts({
    total,
    pending,
    completed,
    failed,
  });
}

function parseCount(value: string | number | null | undefined) {
  const count =
    typeof value === 'string' && /^\d+$/.test(value) ? Number(value) : value;
  return typeof count === 'number' && Number.isSafeInteger(count) && count >= 0
    ? count
    : undefined;
}

function convertGoogleRpcError(
  error: InferSchema<ReturnType<typeof googleRpcStatusSchema>>,
  fallbackMessage: string,
): BatchV4Error {
  return {
    message: error.message ?? fallbackMessage,
    ...(error.status != null ? { type: error.status } : {}),
    ...(error.code != null ? { code: String(error.code) } : {}),
  };
}

const googleUploadUrlResponseHandler: ResponseHandler<string> = async ({
  response,
}) => {
  const uploadUrl = response.headers.get('x-goog-upload-url');
  if (uploadUrl == null) {
    throw new InvalidResponseDataError({
      data: response.headers,
      message: 'Google did not return a resumable upload URL.',
    });
  }

  return { value: uploadUrl };
};

function getGoogleBatchModelId(
  requests: BatchV4StartOptions<GoogleBatchModelIds>['requests'],
): GoogleModelId | GoogleImageModelId {
  const modelId = requests[0]?.modelId;

  if (modelId == null) {
    throw new InvalidArgumentError({
      argument: 'requests',
      message: 'Google batches require at least one request.',
    });
  }

  for (const request of requests) {
    if (request.modelId !== modelId) {
      throw new InvalidArgumentError({
        argument: 'requests',
        message:
          'Google batches require every request to use the same model because the model is part of the batch endpoint.',
      });
    }
  }

  return modelId;
}

function convertGoogleImageBatchResult(
  result: LanguageModelV4GenerateResult,
): ImageModelV4Result | undefined {
  const images = result.content.flatMap(part =>
    part.type === 'file' &&
    part.mediaType.startsWith('image/') &&
    part.data.type === 'data'
      ? [convertToBase64(part.data.data)]
      : [],
  );
  if (images.length === 0) return undefined;

  const googleMetadata =
    (result.providerMetadata?.google as Record<string, unknown> | undefined) ??
    {};
  return {
    images,
    warnings: result.warnings,
    providerMetadata: {
      google: { ...googleMetadata, images: images.map(() => ({})) },
    },
    response: {
      timestamp: new Date(),
      modelId: result.response?.modelId ?? '',
      headers: result.response?.headers,
    },
    usage: {
      inputTokens: result.usage.inputTokens.total,
      outputTokens: result.usage.outputTokens.total,
      totalTokens:
        (result.usage.inputTokens.total ?? 0) +
        (result.usage.outputTokens.total ?? 0),
    },
  };
}