backend/node_modules/@google/generative-ai/src/server/file-manager.test.ts

/**
 * @license
 * Copyright 2024 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import { expect, use } from "chai";
import { GoogleAIFileManager, getUploadMetadata } from "./file-manager";
import * as sinonChai from "sinon-chai";
import * as chaiAsPromised from "chai-as-promised";
import { restore, stub } from "sinon";
import * as request from "./request";
import { RpcTask } from "./constants";
import { DEFAULT_API_VERSION } from "../requests/request";
import { FileMetadata } from "../../types/server";

use(sinonChai);
use(chaiAsPromised);

const FAKE_URI = "https://yourfile.here/filename";
const fakeUploadJson: () => Promise<{}> = () =>
  Promise.resolve({ file: { uri: FAKE_URI } });

describe("GoogleAIFileManager", () => {
  afterEach(() => {
    restore();
  });

  it("stores api key", () => {
    const fileManager = new GoogleAIFileManager("apiKey");
    expect(fileManager.apiKey).to.equal("apiKey");
  });
  it("passes uploadFile request info", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: fakeUploadJson,
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    const result = await fileManager.uploadFile("./test-utils/cat.png", {
      mimeType: "image/png",
    });
    expect(result.file.uri).to.equal(FAKE_URI);
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.UPLOAD);
    expect(makeRequestStub.args[0][0].toString()).to.include("/upload/");
    expect(makeRequestStub.args[0][1]).to.be.instanceOf(Headers);
    expect(makeRequestStub.args[0][1].get("X-Goog-Upload-Protocol")).to.equal(
      "multipart",
    );
    expect(makeRequestStub.args[0][2]).to.be.instanceOf(Blob);
    const bodyBlob = makeRequestStub.args[0][2];
    const blobText = await (bodyBlob as Blob).text();
    expect(blobText).to.include("Content-Type: image/png");
  });
  it("passes uploadFile request info and metadata", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: fakeUploadJson,
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    const result = await fileManager.uploadFile("./test-utils/cat.png", {
      mimeType: "image/png",
      name: "files/customname",
      displayName: "mydisplayname",
    });
    expect(result.file.uri).to.equal(FAKE_URI);
    expect(makeRequestStub.args[0][2]).to.be.instanceOf(Blob);
    const bodyBlob = makeRequestStub.args[0][2];
    const blobText = await (bodyBlob as Blob).text();
    expect(blobText).to.include("Content-Type: image/png");
    expect(blobText).to.include("files/customname");
    expect(blobText).to.include("mydisplayname");
  });
  it("passes uploadFile metadata and formats file name", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: fakeUploadJson,
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    await fileManager.uploadFile("./test-utils/cat.png", {
      mimeType: "image/png",
      name: "customname",
      displayName: "mydisplayname",
    });
    const bodyBlob = makeRequestStub.args[0][2];
    const blobText = await (bodyBlob as Blob).text();
    expect(blobText).to.include("files/customname");
  });
  it("passes uploadFile request info (with options)", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: fakeUploadJson,
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey", {
      apiVersion: "v3000",
      baseUrl: "http://mysite.com",
    });
    const result = await fileManager.uploadFile("./test-utils/cat.png", {
      mimeType: "image/png",
    });
    expect(result.file.uri).to.equal(FAKE_URI);
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.UPLOAD);
    expect(makeRequestStub.args[0][0].toString()).to.include("/upload/");
    expect(makeRequestStub.args[0][1]).to.be.instanceOf(Headers);
    expect(makeRequestStub.args[0][1].get("X-Goog-Upload-Protocol")).to.equal(
      "multipart",
    );
    expect(makeRequestStub.args[0][2]).to.be.instanceOf(Blob);
    const bodyBlob = makeRequestStub.args[0][2];
    const blobText = await (bodyBlob as Blob).text();
    expect(blobText).to.include("Content-Type: image/png");
    expect(makeRequestStub.args[0][0].toString()).to.include("v3000/files");
    expect(makeRequestStub.args[0][0].toString()).to.match(
      /^http:\/\/mysite\.com/,
    );
  });
  it("passes listFiles request info", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({ files: [{ uri: FAKE_URI }] }),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    const result = await fileManager.listFiles();
    expect(result.files[0].uri).to.equal(FAKE_URI);
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.LIST);
    expect(makeRequestStub.args[0][0].toString()).to.match(/\/files$/);
  });
  it("passes listFiles request info with params", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({ files: [{ uri: FAKE_URI }] }),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    const result = await fileManager.listFiles({
      pageSize: 3,
      pageToken: "abc",
    });
    expect(result.files[0].uri).to.equal(FAKE_URI);
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.LIST);
    expect(makeRequestStub.args[0][0].toString()).to.include("pageSize=3");
    expect(makeRequestStub.args[0][0].toString()).to.include("pageToken=abc");
  });
  it("passes listFiles request info with options", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({ files: [{ uri: FAKE_URI }] }),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey", {
      apiVersion: "v3000",
      baseUrl: "http://mysite.com",
    });
    const result = await fileManager.listFiles();
    expect(result.files[0].uri).to.equal(FAKE_URI);
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.LIST);
    expect(makeRequestStub.args[0][0].toString()).to.match(/\/files$/);
    expect(makeRequestStub.args[0][0].toString()).to.include("v3000/files");
    expect(makeRequestStub.args[0][0].toString()).to.match(
      /^http:\/\/mysite\.com/,
    );
  });
  it("passes getFile request info", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({ uri: FAKE_URI }),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    const result = await fileManager.getFile("nameoffile");
    expect(result.uri).to.equal(FAKE_URI);
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.GET);
    expect(makeRequestStub.args[0][0].toString()).to.include(
      `${DEFAULT_API_VERSION}/files/nameoffile`,
    );
  });
  it("passes getFile request info", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({ uri: FAKE_URI }),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    await fileManager.getFile("files/nameoffile");
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.GET);
    expect(makeRequestStub.args[0][0].toString()).to.include(
      `${DEFAULT_API_VERSION}/files/nameoffile`,
    );
  });
  it("passes getFile request info (with options)", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({ uri: FAKE_URI }),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey", {
      apiVersion: "v3000",
      baseUrl: "http://mysite.com",
    });
    const result = await fileManager.getFile("nameoffile");
    expect(result.uri).to.equal(FAKE_URI);
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.GET);
    expect(makeRequestStub.args[0][0].toString()).to.include("/nameoffile");
    expect(makeRequestStub.args[0][0].toString()).to.include("v3000/files");
    expect(makeRequestStub.args[0][0].toString()).to.match(
      /^http:\/\/mysite\.com/,
    );
  });
  it("getFile throws on bad fileId", async () => {
    stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({ uri: FAKE_URI }),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    await expect(fileManager.getFile("")).to.be.rejectedWith("Invalid fileId");
  });
  it("passes deleteFile request info", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({}),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    await fileManager.deleteFile("nameoffile");
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.DELETE);
    expect(makeRequestStub.args[0][0].toString()).to.include("/nameoffile");
  });
  it("passes deleteFile request info (with options)", async () => {
    const makeRequestStub = stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({}),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey", {
      apiVersion: "v3000",
      baseUrl: "http://mysite.com",
    });
    await fileManager.deleteFile("nameoffile");
    expect(makeRequestStub.args[0][0].task).to.equal(RpcTask.DELETE);
    expect(makeRequestStub.args[0][0].toString()).to.include("/nameoffile");
    expect(makeRequestStub.args[0][0].toString()).to.include("v3000/files");
    expect(makeRequestStub.args[0][0].toString()).to.match(
      /^http:\/\/mysite\.com/,
    );
  });
  it("deleteFile throws on bad fileId", async () => {
    stub(request, "makeServerRequest").resolves({
      ok: true,
      json: () => Promise.resolve({}),
    } as Response);
    const fileManager = new GoogleAIFileManager("apiKey");
    await expect(fileManager.deleteFile("")).to.be.rejectedWith(
      "Invalid fileId",
    );
  });

  describe("getUploadMetadata", () => {
    it("getUploadMetadata with only mimeType", () => {
      const uploadMetadata = getUploadMetadata({ mimeType: "image/jpeg" });
      expect(uploadMetadata.mimeType).to.equal("image/jpeg");
      expect(uploadMetadata.displayName).be.undefined;
      expect(uploadMetadata.name).be.undefined;
    });
    it("getUploadMetadata with no mimeType", () => {
      expect(() => getUploadMetadata({} as FileMetadata)).to.throw(
        "Must provide a mimeType.",
      );
    });
    it("getUploadMetadata with all fields defined", () => {
      const uploadMetadata = getUploadMetadata({
        mimeType: "image/jpeg",
        displayName: "display name",
        name: "filename",
      });
      expect(uploadMetadata.mimeType).to.equal("image/jpeg");
      expect(uploadMetadata.displayName).to.equal("display name");
      expect(uploadMetadata.name).to.equal("files/filename");
    });
    it("getUploadMetadata with full file path", () => {
      const uploadMetadata = getUploadMetadata({
        mimeType: "image/jpeg",
        displayName: "display name",
        name: "custom/path/filename",
      });
      expect(uploadMetadata.mimeType).to.equal("image/jpeg");
      expect(uploadMetadata.displayName).to.equal("display name");
      expect(uploadMetadata.name).to.equal("custom/path/filename");
    });
  });
});