Stream file uploads to google generative AI

stream-ai-uploads.ts

import {
  FileMetadata,
  GoogleAIFileManager,
  UploadFileResponse,
} from "@google/generative-ai/server";

// GoogleAIFileManager can't use streams yet, so we need to do it on our own.
// See: https://github.com/google-gemini/generative-ai-js/pull/252

// Note that this is a direct copy of the existing implementation; it would
// likely be better to use a cryptographically secure random.
function generateBoundary(): string {
  let str = "";
  for (let i = 0; i < 2; i++) {
    str = str + Math.random().toString().slice(2);
  }
  return str;
}

export async function uploadFile(
  fileManager: GoogleAIFileManager,
  fileMetadata: FileMetadata,
  fileContentStream: AsyncIterable<Uint8Array>,
): Promise<UploadFileResponse> {
  const boundary = generateBoundary();

  const headers = new Headers();
  headers.append("x-goog-api-key", fileManager.apiKey);
  headers.append("x-goog-upload-protocol", "multipart");
  headers.append("content-type", `multipart/related; boundary=${boundary}`);

  const preBlobPart = new TextEncoder().encode(
    "--" +
      boundary +
      "\r\n" +
      "Content-Type: application/json; charset=utf-8\r\n\r\n" +
      JSON.stringify({
        file: {
          mimeType: fileMetadata.mimeType,
          displayName: `files/${fileMetadata.displayName}`,
        },
      }) +
      "\r\n--" +
      boundary +
      "\r\n" +
      "Content-Type: " +
      fileMetadata.mimeType +
      "\r\n\r\n",
  );
  const postBlobPart = new TextEncoder().encode("\r\n--" + boundary + "--");

  const response = await fetch(
    `https://generativelanguage.googleapis.com/upload/v1beta/files`,
    {
      method: "POST",
      headers,
      body: (async function* () {
        yield preBlobPart;
        yield* fileContentStream;
        yield postBlobPart;
      })(),
      duplex: "half",
    },
  );

  if (!response.ok) {
    throw new Error("Failed to upload file");
  }

  return (await response.json()) as UploadFileResponse;
}