feat (file-manager) support upload with buffer (#365)

Co-authored-by: abdelkd <abdelkarimdjelalda@gmail.com>

Author: hkt74

.changeset/fuzzy-cows-approve.md

@@ -0,0 +1,5 @@
+---
+"@google/generative-ai": minor
+---
+
+added support for buffers in uploadFile

common/api-review/generative-ai-server.api.md

@@ -31,6 +31,8 @@ export interface BooleanSchema extends BaseSchema {
     type: typeof SchemaType.BOOLEAN;
 }
 
+/// <reference types="node" />
+
 // @public
 export interface CachedContent extends CachedContentBase {
     createTime?: string;
@@ -373,7 +375,7 @@ export class GoogleAIFileManager {
     deleteFile(fileId: string): Promise<void>;
     getFile(fileId: string, requestOptions?: SingleRequestOptions): Promise<FileMetadataResponse>;
     listFiles(listParams?: ListParams, requestOptions?: SingleRequestOptions): Promise<ListFilesResponse>;
-    uploadFile(filePath: string, fileMetadata: FileMetadata): Promise<UploadFileResponse>;
+    uploadFile(fileData: string | Buffer, fileMetadata: FileMetadata): Promise<UploadFileResponse>;
 }
 
 // @public

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) |  | <p>Get metadata for file with given ID.</p><p>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.</p> |
 |  [listFiles(listParams, requestOptions)](./generative-ai.googleaifilemanager.listfiles.md) |  | <p>List all uploaded files.</p><p>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.</p> |
-|  [uploadFile(filePath, fileMetadata)](./generative-ai.googleaifilemanager.uploadfile.md) |  | Upload a file. |
+|  [uploadFile(fileData, fileMetadata)](./generative-ai.googleaifilemanager.uploadfile.md) |  | Upload a file. |
 

docs/reference/server/generative-ai.googleaifilemanager.uploadfile.md

@@ -9,14 +9,14 @@ Upload a file.
 **Signature:**
 
 ```typescript
-uploadFile(filePath: string, fileMetadata: FileMetadata): Promise<UploadFileResponse>;
+uploadFile(fileData: string | Buffer, fileMetadata: FileMetadata): Promise<UploadFileResponse>;
 ```
 
 ## Parameters
 
 |  Parameter | Type | Description |
 |  --- | --- | --- |
-|  filePath | string |  |
+|  fileData | string \| Buffer |  |
 |  fileMetadata | [FileMetadata](./generative-ai.filemetadata.md) |  |
 
 **Returns:**

src/server/file-manager.test.ts

@@ -23,6 +23,7 @@ import * as request from "./request";
 import { RpcTask } from "./constants";
 import { DEFAULT_API_VERSION } from "../requests/request";
 import { FileMetadata } from "../../types/server";
+import { readFile } from "fs/promises";
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -61,6 +62,28 @@ describe("GoogleAIFileManager", () => {
     const blobText = await (bodyBlob as Blob).text();
     expect(blobText).to.include("Content-Type: image/png");
   });
+  it("passes uploadFile request info reading from buffer", async () => {
+    const makeRequestStub = stub(request, "makeServerRequest").resolves({
+      ok: true,
+      json: fakeUploadJson,
+    } as Response);
+    const fileManager = new GoogleAIFileManager("apiKey");
+    const fileBuffer = await readFile("./test-utils/cat.png");
+    const result = await fileManager.uploadFile(fileBuffer, {
+      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,
@@ -80,6 +103,26 @@ describe("GoogleAIFileManager", () => {
     expect(blobText).to.include("files/customname");
     expect(blobText).to.include("mydisplayname");
   });
+  it("passes uploadFile request info and metadata from buffer", async () => {
+    const makeRequestStub = stub(request, "makeServerRequest").resolves({
+      ok: true,
+      json: fakeUploadJson,
+    } as Response);
+    const fileManager = new GoogleAIFileManager("apiKey");
+    const fileBuffer = await readFile("./test-utils/cat.png");
+    const result = await fileManager.uploadFile(fileBuffer, {
+      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,
@@ -95,6 +138,22 @@ describe("GoogleAIFileManager", () => {
     const blobText = await (bodyBlob as Blob).text();
     expect(blobText).to.include("files/customname");
   });
+  it("passes uploadFile metadata and formats file name from buffer", async () => {
+    const makeRequestStub = stub(request, "makeServerRequest").resolves({
+      ok: true,
+      json: fakeUploadJson,
+    } as Response);
+    const fileManager = new GoogleAIFileManager("apiKey");
+    const fileBuffer = await readFile("./test-utils/cat.png");
+    await fileManager.uploadFile(fileBuffer, {
+      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,
@@ -123,6 +182,35 @@ describe("GoogleAIFileManager", () => {
       /^http:\/\/mysite\.com/,
     );
   });
+  it("passes uploadFile request info (with options) from buffer", 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 fileBuffer = await readFile("./test-utils/cat.png");
+    const result = await fileManager.uploadFile(fileBuffer, {
+      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,

src/server/file-manager.ts

@@ -51,10 +51,11 @@ export class GoogleAIFileManager {
    * Upload a file.
    */
   async uploadFile(
-    filePath: string,
+    fileData: string | Buffer,
     fileMetadata: FileMetadata,
   ): Promise<UploadFileResponse> {
-    const file = readFileSync(filePath);
+    const file = fileData instanceof Buffer ? fileData : readFileSync(fileData);
+
     const url = new FilesRequestUrl(
       RpcTask.UPLOAD,
       this.apiKey,