From 3210f1b5d011c9bc8050a1d01633e9b6b9dc1ab3 Mon Sep 17 00:00:00 2001 From: Mike Marcacci Date: Tue, 10 Sep 2024 15:38:40 +0000 Subject: [PATCH 1/2] Support file uploads using a stream The current implementation requires a file to exist in the local filesystem, which is an uncommon scenario for production use cases. This adds the ability to supply the upload contents as a node stream, and improves the filesystem case by streaming instead of loading into memory. --- .changeset/strange-chairs-search.md | 5 ++ common/api-review/generative-ai-server.api.md | 6 ++- .../generative-ai.googleaifilemanager.md | 2 +- ...ative-ai.googleaifilemanager.uploadfile.md | 4 +- src/server/file-manager.test.ts | 24 ++++++---- src/server/file-manager.ts | 47 ++++++++++++------- src/server/request.test.ts | 5 +- src/server/request.ts | 26 ++++++++-- 8 files changed, 84 insertions(+), 35 deletions(-) create mode 100644 .changeset/strange-chairs-search.md diff --git a/.changeset/strange-chairs-search.md b/.changeset/strange-chairs-search.md new file mode 100644 index 00000000..9856d360 --- /dev/null +++ b/.changeset/strange-chairs-search.md @@ -0,0 +1,5 @@ +--- +"@google/generative-ai": minor +--- + +Support file uploads using a stream diff --git a/common/api-review/generative-ai-server.api.md b/common/api-review/generative-ai-server.api.md index b00933b4..f7a02757 100644 --- a/common/api-review/generative-ai-server.api.md +++ b/common/api-review/generative-ai-server.api.md @@ -4,6 +4,10 @@ ```ts +/// + +import { Readable } from 'node:stream'; + // @public export interface CachedContent extends CachedContentBase { createTime?: string; @@ -347,7 +351,7 @@ export class GoogleAIFileManager { deleteFile(fileId: string): Promise; getFile(fileId: string, requestOptions?: SingleRequestOptions): Promise; listFiles(listParams?: ListParams, requestOptions?: SingleRequestOptions): Promise; - uploadFile(filePath: string, fileMetadata: FileMetadata): Promise; + uploadFile(filePathOrStream: string | Readable, fileMetadata: FileMetadata): Promise; } // @public diff --git a/docs/reference/server/generative-ai.googleaifilemanager.md b/docs/reference/server/generative-ai.googleaifilemanager.md index 5bb5823e..10434a60 100644 --- a/docs/reference/server/generative-ai.googleaifilemanager.md +++ b/docs/reference/server/generative-ai.googleaifilemanager.md @@ -31,5 +31,5 @@ export declare class GoogleAIFileManager | [deleteFile(fileId)](./generative-ai.googleaifilemanager.deletefile.md) | | Delete file with given ID. | | [getFile(fileId, requestOptions)](./generative-ai.googleaifilemanager.getfile.md) | |

Get metadata for file with given ID.

Any fields set in the optional [SingleRequestOptions](./generative-ai.singlerequestoptions.md) parameter will take precedence over the [RequestOptions](./generative-ai.requestoptions.md) values provided at the time of the [GoogleAIFileManager](./generative-ai.googleaifilemanager.md) initialization.

| | [listFiles(listParams, requestOptions)](./generative-ai.googleaifilemanager.listfiles.md) | |

List all uploaded files.

Any fields set in the optional [SingleRequestOptions](./generative-ai.singlerequestoptions.md) parameter will take precedence over the [RequestOptions](./generative-ai.requestoptions.md) values provided at the time of the [GoogleAIFileManager](./generative-ai.googleaifilemanager.md) initialization.

| -| [uploadFile(filePath, fileMetadata)](./generative-ai.googleaifilemanager.uploadfile.md) | | Upload a file. | +| [uploadFile(filePathOrStream, fileMetadata)](./generative-ai.googleaifilemanager.uploadfile.md) | | Upload a file. | diff --git a/docs/reference/server/generative-ai.googleaifilemanager.uploadfile.md b/docs/reference/server/generative-ai.googleaifilemanager.uploadfile.md index 71e4f76b..fc0fc9fb 100644 --- a/docs/reference/server/generative-ai.googleaifilemanager.uploadfile.md +++ b/docs/reference/server/generative-ai.googleaifilemanager.uploadfile.md @@ -9,14 +9,14 @@ Upload a file. **Signature:** ```typescript -uploadFile(filePath: string, fileMetadata: FileMetadata): Promise; +uploadFile(filePathOrStream: string | Readable, fileMetadata: FileMetadata): Promise; ``` ## Parameters | Parameter | Type | Description | | --- | --- | --- | -| filePath | string | | +| filePathOrStream | string \| Readable | | | fileMetadata | [FileMetadata](./generative-ai.filemetadata.md) | | **Returns:** diff --git a/src/server/file-manager.test.ts b/src/server/file-manager.test.ts index 93fed994..77bcdd8c 100644 --- a/src/server/file-manager.test.ts +++ b/src/server/file-manager.test.ts @@ -23,7 +23,7 @@ import * as request from "./request"; import { RpcTask } from "./constants"; import { DEFAULT_API_VERSION } from "../requests/request"; import { FileMetadata } from "../../types/server"; - +import { blob } from "node:stream/consumers"; use(sinonChai); use(chaiAsPromised); @@ -56,8 +56,10 @@ describe("GoogleAIFileManager", () => { 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]; + expect(makeRequestStub.args[0][2]).to.have.property("next"); + const bodyBlob = await blob( + makeRequestStub.args[0][2] as any as NodeJS.ReadableStream, + ); const blobText = await (bodyBlob as Blob).text(); expect(blobText).to.include("Content-Type: image/png"); }); @@ -73,8 +75,10 @@ describe("GoogleAIFileManager", () => { 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]; + expect(makeRequestStub.args[0][2]).to.have.property("next"); + const bodyBlob = await blob( + makeRequestStub.args[0][2] as any as NodeJS.ReadableStream, + ); const blobText = await (bodyBlob as Blob).text(); expect(blobText).to.include("Content-Type: image/png"); expect(blobText).to.include("files/customname"); @@ -91,7 +95,9 @@ describe("GoogleAIFileManager", () => { name: "customname", displayName: "mydisplayname", }); - const bodyBlob = makeRequestStub.args[0][2]; + const bodyBlob = await blob( + makeRequestStub.args[0][2] as any as NodeJS.ReadableStream, + ); const blobText = await (bodyBlob as Blob).text(); expect(blobText).to.include("files/customname"); }); @@ -114,8 +120,10 @@ describe("GoogleAIFileManager", () => { 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]; + expect(makeRequestStub.args[0][2]).to.have.property("next"); + const bodyBlob = await blob( + makeRequestStub.args[0][2] as any as NodeJS.ReadableStream, + ); 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"); diff --git a/src/server/file-manager.ts b/src/server/file-manager.ts index c34abb2b..27682baf 100644 --- a/src/server/file-manager.ts +++ b/src/server/file-manager.ts @@ -16,7 +16,7 @@ */ import { RequestOptions, SingleRequestOptions } from "../../types"; -import { readFileSync } from "fs"; +import { createReadStream } from "node:fs"; import { FilesRequestUrl, getHeaders, makeServerRequest } from "./request"; import { FileMetadata, @@ -30,6 +30,7 @@ import { GoogleGenerativeAIError, GoogleGenerativeAIRequestInputError, } from "../errors"; +import { Readable } from "node:stream"; // Internal type, metadata sent in the upload export interface UploadMetadata { @@ -51,10 +52,14 @@ export class GoogleAIFileManager { * Upload a file. */ async uploadFile( - filePath: string, + filePathOrStream: string | Readable, fileMetadata: FileMetadata, ): Promise { - const file = readFileSync(filePath); + const file = + typeof filePathOrStream === "string" + ? createReadStream(filePathOrStream) + : filePathOrStream; + const url = new FilesRequestUrl( RpcTask.UPLOAD, this.apiKey, @@ -73,22 +78,28 @@ export class GoogleAIFileManager { // Multipart formatting code taken from @firebase/storage const metadataString = JSON.stringify({ file: uploadMetadata }); - const preBlobPart = + const preBlobPart = new TextEncoder().encode( "--" + - 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 makeServerRequest(url, uploadHeaders, blob); + 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 = new TextEncoder().encode("\r\n--" + boundary + "--"); + + const stream = (async function* () { + yield preBlobPart; + yield* file; + yield postBlobPart; + })(); + + const response = await makeServerRequest(url, uploadHeaders, stream); return response.json(); } diff --git a/src/server/request.test.ts b/src/server/request.test.ts index dd2c15d2..8dfc5c91 100644 --- a/src/server/request.test.ts +++ b/src/server/request.test.ts @@ -79,13 +79,14 @@ describe("Files API - request methods", () => { const response = await makeServerRequest( url, headers, - new Blob(), + (async function* () {})(), fetchStub as typeof fetch, ); expect(fetchStub).to.be.calledWith(match.string, { method: "POST", headers: match.instanceOf(Headers), - body: match.instanceOf(Blob), + body: match.instanceOf(ReadableStream), + duplex: "half", }); expect(response.ok).to.be.true; }); diff --git a/src/server/request.ts b/src/server/request.ts index 464c1d2e..21bc547f 100644 --- a/src/server/request.ts +++ b/src/server/request.ts @@ -97,16 +97,36 @@ export function getHeaders(url: ServerRequestUrl): Headers { export async function makeServerRequest( url: FilesRequestUrl, headers: Headers, - body?: Blob | string, + body?: Blob | string | AsyncIterable, fetchFn: typeof fetch = fetch, ): Promise { - const requestInit: RequestInit = { + // Add the duplex option, which is required when streaming in newer versions of node. + // See: https://github.com/nodejs/node/issues/46221 + const requestInit: RequestInit & { duplex?: "half" } = { method: taskToMethod[url.task], headers, + duplex: "half", }; - if (body) { + if (typeof body === "string" || body instanceof Blob) { requestInit.body = body; + } else if (body[Symbol.asyncIterator]) { + // Note that in later versions, the signature `fetch` is updated to accept any AsyncIterator, + // and ReadableStream implements AsyncIterator. In this case, `body` can be passed exactly + // as supplied, and the following can be removed: + const iterator = body[Symbol.asyncIterator](); + requestInit.body = new ReadableStream({ + type: "bytes", + async pull(controller) { + const { value, done } = await iterator.next(); + if (done) { + controller.close(); + return; + } + + controller.enqueue(value); + }, + }); } const signal = getSignal(url.requestOptions); From c8c82d0f38088f309b1845f685b6ddf0a659dfe5 Mon Sep 17 00:00:00 2001 From: Mike Marcacci Date: Thu, 12 Sep 2024 12:11:46 -0700 Subject: [PATCH 2/2] Fix scenario where body === undefined --- src/server/request.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/server/request.ts b/src/server/request.ts index 21bc547f..0abd553d 100644 --- a/src/server/request.ts +++ b/src/server/request.ts @@ -110,7 +110,7 @@ export async function makeServerRequest( if (typeof body === "string" || body instanceof Blob) { requestInit.body = body; - } else if (body[Symbol.asyncIterator]) { + } else if (body?.[Symbol.asyncIterator]) { // Note that in later versions, the signature `fetch` is updated to accept any AsyncIterator, // and ReadableStream implements AsyncIterator. In this case, `body` can be passed exactly // as supplied, and the following can be removed: