Skip to content

Commit

Permalink
Fix bug when no file name is provided for upload (#135)
Browse files Browse the repository at this point in the history
  • Loading branch information
hsubox76 authored May 9, 2024
1 parent 6227ee7 commit 819501f
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 9 deletions.
5 changes: 5 additions & 0 deletions .changeset/dirty-pillows-vanish.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@google/generative-ai": patch
---

Fix a bug that caused file uploads to be named "undefined" if no file name is provided.
37 changes: 36 additions & 1 deletion packages/main/src/files/file-manager.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,14 @@
* limitations under the License.
*/
import { expect, use } from "chai";
import { GoogleAIFileManager } from "./file-manager";
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 { FilesTask } from "./constants";
import { DEFAULT_API_VERSION } from "../requests/request";
import { FileMetadata } from "./types";

use(sinonChai);
use(chaiAsPromised);
Expand Down Expand Up @@ -254,4 +255,38 @@ describe("GoogleAIFileManager", () => {
"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");
});
});
});
30 changes: 22 additions & 8 deletions packages/main/src/files/file-manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,10 @@ import {
UploadFileResponse,
} from "./types";
import { FilesTask } from "./constants";
import { GoogleGenerativeAIError } from "../errors";
import {
GoogleGenerativeAIError,
GoogleGenerativeAIRequestInputError,
} from "../errors";

// Internal type, metadata sent in the upload
export interface UploadMetadata {
Expand Down Expand Up @@ -66,13 +69,7 @@ export class GoogleAIFileManager {
`multipart/related; boundary=${boundary}`,
);

const uploadMetadata: FileMetadata = {
mimeType: fileMetadata.mimeType,
displayName: fileMetadata.displayName,
name: fileMetadata.name?.includes("/")
? fileMetadata.name
: `files/${fileMetadata.name}`,
};
const uploadMetadata = getUploadMetadata(fileMetadata);

// Multipart formatting code taken from @firebase/storage
const metadataString = JSON.stringify({ file: uploadMetadata });
Expand Down Expand Up @@ -169,3 +166,20 @@ function generateBoundary(): string {
}
return str;
}

export function getUploadMetadata(inputMetadata: FileMetadata): FileMetadata {
if (!inputMetadata.mimeType) {
throw new GoogleGenerativeAIRequestInputError("Must provide a mimeType.");
}
const uploadMetadata: FileMetadata = {
mimeType: inputMetadata.mimeType,
displayName: inputMetadata.displayName,
};

if (inputMetadata.name) {
uploadMetadata.name = inputMetadata.name.includes("/")
? inputMetadata.name
: `files/${inputMetadata.name}`;
}
return uploadMetadata;
}

0 comments on commit 819501f

Please sign in to comment.