diff --git a/.changeset/tiny-spies-pretend.md b/.changeset/tiny-spies-pretend.md new file mode 100644 index 000000000..6b795bf73 --- /dev/null +++ b/.changeset/tiny-spies-pretend.md @@ -0,0 +1,5 @@ +--- +"@google/generative-ai": minor +--- + +Add GoogleAIFileManager for file uploads. diff --git a/README.md b/README.md index 52c60ad7d..f4a2db9d7 100644 --- a/README.md +++ b/README.md @@ -90,7 +90,7 @@ These quickstarts describe how to add your API key and the SDK to your app, init Find complete documentation for the Google AI SDKs and the Gemini model in the Google documentation:\ https://ai.google.dev/docs -Find reference docs for this SDK [here in the repo](/docs/reference/generative-ai.md). +Find reference docs for this SDK [here in the repo](/docs/reference/main/generative-ai.md). ## Changelog - `@google/generative-ai` - [CHANGELOG.md](/packages/main/CHANGELOG.md) diff --git a/packages/main/README.md b/packages/main/README.md index 0ce26643a..c843a0c7c 100644 --- a/packages/main/README.md +++ b/packages/main/README.md @@ -90,7 +90,7 @@ These quickstarts describe how to add your API key and the SDK to your app, init Find complete documentation for the Google AI SDKs and the Gemini model in the Google documentation:\ https://ai.google.dev/docs -Find reference docs for this SDK [here in the repo](/docs/reference/generative-ai.md). +Find reference docs for this SDK [here in the repo](/docs/reference/main/generative-ai.md). ## Changelog - `@google/generative-ai` - [CHANGELOG.md](/main/packages/main/CHANGELOG.md) diff --git a/packages/main/api-extractor.files.json b/packages/main/api-extractor.files.json new file mode 100644 index 000000000..4819da004 --- /dev/null +++ b/packages/main/api-extractor.files.json @@ -0,0 +1,12 @@ +{ + "extends": "../../config/api-extractor.json", + "mainEntryPointFilePath": "/dist/files/src/files/index.d.ts", + "dtsRollup": { + "enabled": true, + "untrimmedFilePath": "/dist/files/files.d.ts" + }, + "docModel": { + "enabled": true, + "apiJsonFilePath": "/temp/files/-files.api.json" + } +} diff --git a/packages/main/api-extractor.json b/packages/main/api-extractor.json index cda92fb6b..f7546e09d 100644 --- a/packages/main/api-extractor.json +++ b/packages/main/api-extractor.json @@ -4,5 +4,9 @@ "dtsRollup": { "enabled": true, "untrimmedFilePath": "/dist/.d.ts" + }, + "docModel": { + "enabled": true, + "apiJsonFilePath": "/temp/main/.api.json" } } diff --git a/packages/main/files/package.json b/packages/main/files/package.json new file mode 100644 index 000000000..01371633f --- /dev/null +++ b/packages/main/files/package.json @@ -0,0 +1,8 @@ +{ + "name": "@google/generative-ai/files", + "description": "GoogleAI file upload manager", + "main": "./dist/files/index.js", + "browser": "./dist/files/index.mjs", + "module": "./dist/files/index.mjs", + "typings": "./dist/files/files.d.ts" +} diff --git a/packages/main/package.json b/packages/main/package.json index aaf550d81..00f23eb5a 100644 --- a/packages/main/package.json +++ b/packages/main/package.json @@ -12,13 +12,20 @@ "import": "./dist/index.mjs", "default": "./dist/index.js" }, + "./files": { + "types": "./dist/files/files.d.ts", + "require": "./dist/files/index.js", + "import": "./dist/files/index.mjs", + "default": "./dist/files/index.js" + }, "./package.json": "./package.json" }, "engines": { "node": ">=18.0.0" }, "files": [ - "dist" + "dist", + "files/package.json" ], "scripts": { "build": "rollup -c && yarn api-report", @@ -27,8 +34,8 @@ "test:node:unit": "TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' mocha \"src/**/*.test.ts\"", "test:node:integration": "yarn build && TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' mocha \"test-integration/node/**/*.test.ts\"", "lint": "eslint -c .eslintrc.js '**/*.ts' --ignore-path '../../.gitignore'", - "api-report": "api-extractor run --local --verbose", - "docs": "yarn build && yarn api-documenter markdown -i ./temp -o ../../docs/reference/" + "api-report": "api-extractor run --local --verbose && api-extractor run -c api-extractor.files.json --local --verbose", + "docs": "yarn build && yarn api-documenter markdown -i ./temp/main -o ../../docs/reference/main && yarn api-documenter markdown -i ./temp/files -o ../../docs/reference/files" }, "repository": { "type": "git", diff --git a/packages/main/rollup.config.mjs b/packages/main/rollup.config.mjs index 170f68d7b..67b9ad3ff 100644 --- a/packages/main/rollup.config.mjs +++ b/packages/main/rollup.config.mjs @@ -20,10 +20,7 @@ import typescriptPlugin from "rollup-plugin-typescript2"; import typescript from "typescript"; import json from "@rollup/plugin-json"; import pkg from "./package.json" assert { type: "json" }; - -const deps = Object.keys( - Object.assign({}, pkg.peerDependencies, pkg.dependencies), -); +import filePkg from "./files/package.json" assert { type: "json" }; const es2017BuildPlugins = [ typescriptPlugin({ @@ -51,8 +48,7 @@ const esmBuilds = [ format: "es", sourcemap: true, }, - external: (id) => - deps.some((dep) => id === dep || id.startsWith(`${dep}/`)), + external: ["fs"], plugins: [...es2017BuildPlugins], }, ]; @@ -61,11 +57,25 @@ const cjsBuilds = [ { input: "src/index.ts", output: [{ file: pkg.main, format: "cjs", sourcemap: true }], - external: (id) => - deps.some((dep) => id === dep || id.startsWith(`${dep}/`)), + external: ["fs"], + plugins: [...es2017BuildPlugins], + }, +]; + +const filesBuilds = [ + { + input: "src/files/index.ts", + output: [{ file: filePkg.module, format: "es", sourcemap: true }], + external: ["fs"], + plugins: [...es2017BuildPlugins], + }, + { + input: "src/files/index.ts", + output: [{ file: filePkg.main, format: "cjs", sourcemap: true }], + external: ["fs"], plugins: [...es2017BuildPlugins], }, ]; // eslint-disable-next-line import/no-default-export -export default [...esmBuilds, ...cjsBuilds]; +export default [...esmBuilds, ...cjsBuilds, ...filesBuilds]; diff --git a/packages/main/src/files/constants.ts b/packages/main/src/files/constants.ts new file mode 100644 index 000000000..d9f6857ca --- /dev/null +++ b/packages/main/src/files/constants.ts @@ -0,0 +1,22 @@ +/** + * @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. + */ +export enum FilesTask { + UPLOAD = "upload", + LIST = "list", + GET = "get", + DELETE = "delete", +} diff --git a/packages/main/src/files/file-manager.test.ts b/packages/main/src/files/file-manager.test.ts new file mode 100644 index 000000000..f3bf5462d --- /dev/null +++ b/packages/main/src/files/file-manager.test.ts @@ -0,0 +1,257 @@ +/** + * @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 } 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 { FilesTask } from "./constants"; +import { DEFAULT_API_VERSION } from "../requests/request"; + +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, "makeFilesRequest").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(FilesTask.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.text(); + expect(blobText).to.include("Content-Type: image/png"); + }); + it("passes uploadFile request info and metadata", async () => { + const makeRequestStub = stub(request, "makeFilesRequest").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.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, "makeFilesRequest").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.text(); + expect(blobText).to.include("files/customname"); + }); + it("passes uploadFile request info (with options)", async () => { + const makeRequestStub = stub(request, "makeFilesRequest").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(FilesTask.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.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, "makeFilesRequest").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(FilesTask.LIST); + expect(makeRequestStub.args[0][0].toString()).to.match(/\/files$/); + }); + it("passes listFiles request info with params", async () => { + const makeRequestStub = stub(request, "makeFilesRequest").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(FilesTask.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, "makeFilesRequest").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(FilesTask.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, "makeFilesRequest").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(FilesTask.GET); + expect(makeRequestStub.args[0][0].toString()).to.include( + `${DEFAULT_API_VERSION}/files/nameoffile`, + ); + }); + it("passes getFile request info", async () => { + const makeRequestStub = stub(request, "makeFilesRequest").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(FilesTask.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, "makeFilesRequest").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(FilesTask.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, "makeFilesRequest").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, "makeFilesRequest").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(FilesTask.DELETE); + expect(makeRequestStub.args[0][0].toString()).to.include("/nameoffile"); + }); + it("passes deleteFile request info (with options)", async () => { + const makeRequestStub = stub(request, "makeFilesRequest").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(FilesTask.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, "makeFilesRequest").resolves({ + ok: true, + json: () => Promise.resolve({}), + } as Response); + const fileManager = new GoogleAIFileManager("apiKey"); + await expect(fileManager.deleteFile("")).to.be.rejectedWith( + "Invalid fileId", + ); + }); +}); diff --git a/packages/main/src/files/file-manager.ts b/packages/main/src/files/file-manager.ts new file mode 100644 index 000000000..e8577d686 --- /dev/null +++ b/packages/main/src/files/file-manager.ts @@ -0,0 +1,171 @@ +/** + * @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 { RequestOptions } from "../../types"; +import { readFileSync } from "fs"; +import { FilesRequestUrl, getHeaders, makeFilesRequest } from "./request"; +import { + FileMetadata, + FileMetadataResponse, + ListFilesResponse, + ListParams, + UploadFileResponse, +} from "./types"; +import { FilesTask } from "./constants"; +import { GoogleGenerativeAIError } from "../errors"; + +// Internal type, metadata sent in the upload +export interface UploadMetadata { + name?: string; + ["display_name"]?: string; +} + +/** + * Class for managing GoogleAI file uploads. + * @public + */ +export class GoogleAIFileManager { + constructor( + public apiKey: string, + private _requestOptions?: RequestOptions, + ) {} + + /** + * Upload a file + */ + async uploadFile( + filePath: string, + fileMetadata: FileMetadata, + ): Promise { + const file = readFileSync(filePath); + const url = new FilesRequestUrl( + FilesTask.UPLOAD, + this.apiKey, + this._requestOptions, + ); + + const uploadHeaders = getHeaders(url); + const boundary = generateBoundary(); + uploadHeaders.append("X-Goog-Upload-Protocol", "multipart"); + uploadHeaders.append( + "Content-Type", + `multipart/related; boundary=${boundary}`, + ); + + const uploadMetadata: FileMetadata = { + mimeType: fileMetadata.mimeType, + displayName: fileMetadata.displayName, + name: fileMetadata.name?.includes("/") + ? fileMetadata.name + : `files/${fileMetadata.name}`, + }; + + // Multipart formatting code taken from @firebase/storage + const metadataString = JSON.stringify({ file: uploadMetadata }); + const preBlobPart = + "--" + + boundary + + "\r\n" + + "Content-Type: application/json; charset=utf-8\r\n\r\n" + + metadataString + + "\r\n--" + + boundary + + "\r\n" + + "Content-Type: " + + fileMetadata.mimeType + + "\r\n\r\n"; + const postBlobPart = "\r\n--" + boundary + "--"; + const blob = new Blob([preBlobPart, file, postBlobPart]); + + const response = await makeFilesRequest(url, uploadHeaders, blob); + return response.json(); + } + + /** + * List all uploaded files + */ + async listFiles(listParams?: ListParams): Promise { + const url = new FilesRequestUrl( + FilesTask.LIST, + this.apiKey, + this._requestOptions, + ); + if (listParams?.pageSize) { + url.appendParam("pageSize", listParams.pageSize.toString()); + } + if (listParams?.pageToken) { + url.appendParam("pageToken", listParams.pageToken); + } + const uploadHeaders = getHeaders(url); + const response = await makeFilesRequest(url, uploadHeaders); + return response.json(); + } + + /** + * Get metadata for file with given ID + */ + async getFile(fileId: string): Promise { + const url = new FilesRequestUrl( + FilesTask.GET, + this.apiKey, + this._requestOptions, + ); + url.appendPath(parseFileId(fileId)); + const uploadHeaders = getHeaders(url); + const response = await makeFilesRequest(url, uploadHeaders); + return response.json(); + } + + /** + * Delete file with given ID + */ + async deleteFile(fileId: string): Promise { + const url = new FilesRequestUrl( + FilesTask.DELETE, + this.apiKey, + this._requestOptions, + ); + url.appendPath(parseFileId(fileId)); + const uploadHeaders = getHeaders(url); + await makeFilesRequest(url, uploadHeaders); + } +} + +/** + * If fileId is prepended with "files/", remove prefix + */ +function parseFileId(fileId: string): string { + if (fileId.startsWith("files/")) { + return fileId.split("files/")[1]; + } + if (!fileId) { + throw new GoogleGenerativeAIError( + `Invalid fileId ${fileId}. ` + + `Must be in the format "files/filename" or "filename"`, + ); + } + + return fileId; +} + +function generateBoundary(): string { + let str = ""; + for (let i = 0; i < 2; i++) { + str = str + Math.random().toString().slice(2); + } + return str; +} diff --git a/packages/main/src/files/index.ts b/packages/main/src/files/index.ts new file mode 100644 index 000000000..6d89670f3 --- /dev/null +++ b/packages/main/src/files/index.ts @@ -0,0 +1,20 @@ +/** + * @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. + */ + +export { GoogleAIFileManager } from "./file-manager"; +export * from "./types"; +export { RequestOptions } from "../../types"; diff --git a/packages/main/src/files/request.test.ts b/packages/main/src/files/request.test.ts new file mode 100644 index 000000000..225f86767 --- /dev/null +++ b/packages/main/src/files/request.test.ts @@ -0,0 +1,161 @@ +/** + * @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 { match, restore, stub } from "sinon"; +import * as sinonChai from "sinon-chai"; +import * as chaiAsPromised from "chai-as-promised"; +import { DEFAULT_API_VERSION, DEFAULT_BASE_URL } from "../requests/request"; +import { FilesRequestUrl, makeFilesRequest } from "./request"; +import { FilesTask } from "./constants"; + +use(sinonChai); +use(chaiAsPromised); + +describe("Files API - request methods", () => { + afterEach(() => { + restore(); + }); + describe("FilesRequestUrl", () => { + it("includes task, apiVersion, baseURL, upload if upload task", async () => { + const url = new FilesRequestUrl(FilesTask.UPLOAD, "key", {}); + expect(url.toString()).to.include("/upload"); + expect(url.toString()).to.not.include("key"); + expect(url.toString()).to.include(DEFAULT_API_VERSION); + expect(url.toString()).to.include(DEFAULT_BASE_URL); + }); + it("includes task, apiVersion, baseURL, no upload if non-upload task", async () => { + const url = new FilesRequestUrl(FilesTask.GET, "key", {}); + expect(url.toString()).to.not.include("/upload"); + expect(url.toString()).to.not.include("key"); + expect(url.toString()).to.include(DEFAULT_API_VERSION); + expect(url.toString()).to.include(DEFAULT_BASE_URL); + }); + it("gets custom apiVersion", async () => { + const url = new FilesRequestUrl(FilesTask.GET, "key", { + apiVersion: "v2beta", + }); + expect(url.toString()).to.include("/v2beta/files"); + }); + it("custom baseUrl", async () => { + const url = new FilesRequestUrl(FilesTask.GET, "key", { + baseUrl: "http://my.staging.website", + }); + expect(url.toString()).to.include("http://my.staging.website"); + }); + it("adds params", async () => { + const url = new FilesRequestUrl(FilesTask.GET, "key", {}); + url.appendParam("param1", "value1"); + expect(url.toString()).to.include("?param1=value1"); + }); + it("adds path segments", async () => { + const url = new FilesRequestUrl(FilesTask.GET, "key", {}); + url.appendPath("newpath"); + expect(url.toString()).to.match(/\/newpath$/); + }); + }); + describe("makeFilesRequest", () => { + it("upload - ok", async () => { + const fetchStub = stub().resolves({ + ok: true, + } as Response); + const url = new FilesRequestUrl(FilesTask.UPLOAD, "key"); + const headers = new Headers(); + const response = await makeFilesRequest( + url, + headers, + new Blob(), + fetchStub as typeof fetch, + ); + expect(fetchStub).to.be.calledWith(match.string, { + method: "POST", + headers: match.instanceOf(Headers), + body: match.instanceOf(Blob), + }); + expect(response.ok).to.be.true; + }); + it("error with timeout", async () => { + const fetchStub = stub().resolves({ + ok: false, + status: 500, + statusText: "AbortError", + } as Response); + + const url = new FilesRequestUrl(FilesTask.GET, "key", { timeout: 0 }); + const headers = new Headers(); + await expect( + makeFilesRequest(url, headers, new Blob(), fetchStub as typeof fetch), + ).to.be.rejectedWith("500 AbortError"); + expect(fetchStub).to.be.calledOnce; + }); + it("Network error, no response.json()", async () => { + const fetchStub = stub().resolves({ + ok: false, + status: 500, + statusText: "Server Error", + } as Response); + const url = new FilesRequestUrl(FilesTask.GET, "key"); + const headers = new Headers(); + await expect( + makeFilesRequest(url, headers, new Blob(), fetchStub as typeof fetch), + ).to.be.rejectedWith(/500 Server Error/); + expect(fetchStub).to.be.calledOnce; + }); + it("Network error, includes response.json()", async () => { + const fetchStub = stub().resolves({ + ok: false, + status: 500, + statusText: "Server Error", + json: () => Promise.resolve({ error: { message: "extra info" } }), + } as Response); + const url = new FilesRequestUrl(FilesTask.GET, "key"); + const headers = new Headers(); + await expect( + makeFilesRequest(url, headers, new Blob(), fetchStub as typeof fetch), + ).to.be.rejectedWith(/500 Server Error.*extra info/); + expect(fetchStub).to.be.calledOnce; + }); + it("Network error, includes response.json() and details", async () => { + const fetchStub = stub().resolves({ + ok: false, + status: 500, + statusText: "Server Error", + json: () => + Promise.resolve({ + error: { + message: "extra info", + details: [ + { + "@type": "type.googleapis.com/google.rpc.DebugInfo", + detail: + "[ORIGINAL ERROR] generic::invalid_argument: invalid status photos.thumbnailer.Status.Code::5: Source image 0 too short", + }, + ], + }, + }), + } as Response); + const url = new FilesRequestUrl(FilesTask.GET, "key"); + const headers = new Headers(); + await expect( + makeFilesRequest(url, headers, new Blob(), fetchStub as typeof fetch), + ).to.be.rejectedWith( + /500 Server Error.*extra info.*generic::invalid_argument/, + ); + expect(fetchStub).to.be.calledOnce; + }); + }); +}); diff --git a/packages/main/src/files/request.ts b/packages/main/src/files/request.ts new file mode 100644 index 000000000..e7b27ea1b --- /dev/null +++ b/packages/main/src/files/request.ts @@ -0,0 +1,130 @@ +/** + * @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 { GoogleGenerativeAIError } from "../errors"; +import { + DEFAULT_API_VERSION, + DEFAULT_BASE_URL, + getClientHeaders, +} from "../requests/request"; +import { RequestOptions } from "../../types"; +import { FilesTask } from "./constants"; + +const taskToMethod = { + [FilesTask.UPLOAD]: "POST", + [FilesTask.LIST]: "GET", + [FilesTask.GET]: "GET", + [FilesTask.DELETE]: "DELETE", +}; + +export class FilesRequestUrl { + private _url: URL; + + constructor( + public task: FilesTask, + public apiKey: string, + public requestOptions?: RequestOptions, + ) { + const apiVersion = this.requestOptions?.apiVersion || DEFAULT_API_VERSION; + const baseUrl = this.requestOptions?.baseUrl || DEFAULT_BASE_URL; + let initialUrl = baseUrl; + if (this.task === FilesTask.UPLOAD) { + initialUrl += `/upload`; + } + initialUrl += `/${apiVersion}/files`; + this._url = new URL(initialUrl); + } + + appendPath(path: string): void { + this._url.pathname = this._url.pathname + `/${path}`; + } + + appendParam(key: string, value: string): void { + this._url.searchParams.append(key, value); + } + + toString(): string { + return this._url.toString(); + } +} + +export function getHeaders(url: FilesRequestUrl): Headers { + const headers = new Headers(); + headers.append("x-goog-api-client", getClientHeaders(url.requestOptions)); + headers.append("x-goog-api-key", url.apiKey); + return headers; +} + +export async function makeFilesRequest( + url: FilesRequestUrl, + headers: Headers, + body?: Blob, + fetchFn: typeof fetch = fetch, +): Promise { + const requestInit: RequestInit = { + method: taskToMethod[url.task], + headers, + }; + + if (body) { + requestInit.body = body; + } + + const signal = getSignal(url.requestOptions); + if (signal) { + requestInit.signal = signal; + } + + try { + const response = await fetchFn(url.toString(), requestInit); + if (!response.ok) { + let message = ""; + try { + const json = await response.json(); + message = json.error.message; + if (json.error.details) { + message += ` ${JSON.stringify(json.error.details)}`; + } + } catch (e) { + // ignored + } + throw new Error(`[${response.status} ${response.statusText}] ${message}`); + } else { + return response; + } + } catch (e) { + const err = new GoogleGenerativeAIError( + `Error on task type: ${url.task} fetching from ${url.toString()}: ${ + e.message + }`, + ); + err.stack = e.stack; + throw err; + } +} + +/** + * Get AbortSignal if timeout is specified + */ +function getSignal(requestOptions?: RequestOptions): AbortSignal | null { + if (requestOptions?.timeout >= 0) { + const abortController = new AbortController(); + const signal = abortController.signal; + setTimeout(() => abortController.abort(), requestOptions.timeout); + return signal; + } +} diff --git a/packages/main/src/files/types.ts b/packages/main/src/files/types.ts new file mode 100644 index 000000000..5378c11bc --- /dev/null +++ b/packages/main/src/files/types.ts @@ -0,0 +1,68 @@ +/** + * @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. + */ + +/** + * Params to pass to {@link GoogleAIFileManager.listFiles} + * @public + */ +export interface ListParams { + pageSize?: number; + pageToken?: string; +} + +/** + * Metadata to provide alongside a file upload + * @public + */ +export interface FileMetadata { + name?: string; + displayName?: string; + mimeType: string; +} + +/** + * File metadata response from server. + * @public + */ +export interface FileMetadataResponse { + name: string; + displayName?: string; + mimeType: string; + sizeBytes: string; + createTime: string; + updateTime: string; + expirationTime: string; + sha256Hash: string; + uri: string; +} + +/** + * Response from calling {@link GoogleAIFileManager.listFiles} + * @public + */ +export interface ListFilesResponse { + files: FileMetadataResponse[]; + nextPageToken?: string; +} + +/** + * Response from calling {@link GoogleAIFileManager.uploadFile} + * @public + */ +export interface UploadFileResponse { + file: FileMetadataResponse; +} diff --git a/packages/main/src/methods/chat-session-helpers.ts b/packages/main/src/methods/chat-session-helpers.ts index bb1544442..665540bff 100644 --- a/packages/main/src/methods/chat-session-helpers.ts +++ b/packages/main/src/methods/chat-session-helpers.ts @@ -79,6 +79,7 @@ export function validateChatHistory(history: Content[]): void { inlineData: 0, functionCall: 0, functionResponse: 0, + fileData: 0, }; for (const part of parts) { diff --git a/packages/main/types/content.ts b/packages/main/types/content.ts index c5e5b60c9..69c50a920 100644 --- a/packages/main/types/content.ts +++ b/packages/main/types/content.ts @@ -32,7 +32,8 @@ export type Part = | TextPart | InlineDataPart | FunctionCallPart - | FunctionResponsePart; + | FunctionResponsePart + | FileDataPart; /** * Content part interface if the part represents a text string. @@ -43,6 +44,7 @@ export interface TextPart { inlineData?: never; functionCall?: never; functionResponse?: never; + fileData?: never; } /** @@ -54,6 +56,7 @@ export interface InlineDataPart { inlineData: GenerativeContentBlob; functionCall?: never; functionResponse?: never; + fileData?: never; } /** @@ -65,6 +68,7 @@ export interface FunctionCallPart { inlineData?: never; functionCall: FunctionCall; functionResponse?: never; + fileData?: never; } /** @@ -76,6 +80,7 @@ export interface FunctionResponsePart { inlineData?: never; functionCall?: never; functionResponse: FunctionResponse; + fileData?: never; } /** @@ -114,3 +119,24 @@ export interface GenerativeContentBlob { */ data: string; } + +/** + * Content part interface if the part represents FunctionResponse. + * @public + */ +export interface FileDataPart { + text?: never; + inlineData?: never; + functionCall?: never; + functionResponse?: never; + fileData: FileData; +} + +/** + * Data pointing to a file uploaded with the Files API. + * @public + */ +export interface FileData { + mimeType: string; + fileUri: string; +} diff --git a/packages/main/types/requests.ts b/packages/main/types/requests.ts index ba3bd3ea5..5f70afdf1 100644 --- a/packages/main/types/requests.ts +++ b/packages/main/types/requests.ts @@ -114,7 +114,7 @@ export interface BatchEmbedContentsRequest { } /** - * Params passed to {@link GoogleGenerativeAI.getGenerativeModel}. + * Params passed to getGenerativeModel() or GoogleAIFileManager(). * @public */ export interface RequestOptions { diff --git a/packages/main/types/responses.ts b/packages/main/types/responses.ts index 1929ad4e1..cb06b2145 100644 --- a/packages/main/types/responses.ts +++ b/packages/main/types/responses.ts @@ -60,7 +60,7 @@ export interface EnhancedGenerateContentResponse text: () => string; /** * Deprecated: use `functionCalls()` instead. - * @deprecated + * @deprecated - use `functionCalls()` instead */ functionCall: () => FunctionCall | undefined; /** diff --git a/samples/node/file-upload.js b/samples/node/file-upload.js new file mode 100644 index 000000000..240502a6f --- /dev/null +++ b/samples/node/file-upload.js @@ -0,0 +1,66 @@ +/** + * @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. + */ + +/** + * Example of uploading a file and referencing it in a call to + * generateContent(). + * + * NOTE: The Files API is only available for use in Node. + * Importing from `@google/generative-ai/files` will crash in the + * browser. + */ + +import { GoogleAIFileManager } from "@google/generative-ai/files"; +import { genAI } from "./utils/common.js"; + +async function run() { + // For text-only inputs, use the gemini-pro model + const model = genAI.getGenerativeModel({ + model: "gemini-1.5-pro-latest", + }); + const fileManager = new GoogleAIFileManager(process.env.API_KEY); + + const fileResult = await fileManager.uploadFile("./utils/cat.jpg", { + mimeType: "image/jpeg", + // It will also add the necessary "files/" prefix if not provided + name: "files/catname", + displayName: "mrcat", + }); + + const result = await model.generateContent({ + contents: [ + { + role: "user", + parts: [ + { text: "What is this?" }, + { + fileData: { + mimeType: fileResult.file.mimeType, + fileUri: fileResult.file.uri + } + }, + ], + }, + ], + }); + + const response = result.response; + const text = response.text(); + console.log(text); +} + +run();