diff --git a/front/lib/api/assistant/configuration.ts b/front/lib/api/assistant/configuration.ts index 349051fbb497..da8332b2bb60 100644 --- a/front/lib/api/assistant/configuration.ts +++ b/front/lib/api/assistant/configuration.ts @@ -34,7 +34,7 @@ import { agentConfigurationWasUpdatedBy } from "@app/lib/api/assistant/recent_au import { agentUserListStatus } from "@app/lib/api/assistant/user_relation"; import { compareAgentsForSort } from "@app/lib/assistant"; import type { Authenticator } from "@app/lib/auth"; -import { getFileContentType } from "@app/lib/dfs"; +import { getPublicUploadBucket } from "@app/lib/dfs"; import { AgentConfiguration, AgentDataSourceConfiguration, @@ -761,7 +761,9 @@ async function isSelfHostedImageWithValidContentType(pictureUrl: string) { return false; } - const contentType = await getFileContentType("PUBLIC_UPLOAD", filename); + const contentType = await getPublicUploadBucket().getFileContentType( + filename + ); if (!contentType) { return false; } diff --git a/front/lib/dfs/config.ts b/front/lib/dfs/config.ts index e24107759337..b63598983c31 100644 --- a/front/lib/dfs/config.ts +++ b/front/lib/dfs/config.ts @@ -4,9 +4,12 @@ const config = { getServiceAccount: (): string => { return EnvironmentConfig.getEnvVariable("SERVICE_ACCOUNT"); }, - getPublicUploadBucket: (): string => { + getGcsPublicUploadBucket: (): string => { return EnvironmentConfig.getEnvVariable("DUST_UPLOAD_BUCKET"); }, + getGcsPrivateUploadsBucket: (): string => { + return EnvironmentConfig.getEnvVariable("DUST_PRIVATE_UPLOADS_BUCKET"); + }, }; export default config; diff --git a/front/lib/dfs/index.ts b/front/lib/dfs/index.ts index f912658cde28..6fd3bf0bec2a 100644 --- a/front/lib/dfs/index.ts +++ b/front/lib/dfs/index.ts @@ -2,51 +2,91 @@ import type { Bucket } from "@google-cloud/storage"; import { Storage } from "@google-cloud/storage"; import type formidable from "formidable"; import fs from "fs"; +import { pipeline } from "stream/promises"; import config from "@app/lib/dfs/config"; -type SupportedBucketKeyType = "PUBLIC_UPLOAD"; +type BucketKeyType = "PRIVATE_UPLOAD" | "PUBLIC_UPLOAD"; const storage = new Storage({ keyFilename: config.getServiceAccount(), }); -const bucketKeysToBucket: Record = { - PUBLIC_UPLOAD: storage.bucket(config.getPublicUploadBucket()), +const bucketKeysToBucket: Record = { + PRIVATE_UPLOAD: storage.bucket(config.getGcsPrivateUploadsBucket()), + PUBLIC_UPLOAD: storage.bucket(config.getGcsPublicUploadBucket()), }; -export async function uploadToBucket( - bucketKey: SupportedBucketKeyType, - file: formidable.File -) { - const bucket = bucketKeysToBucket[bucketKey]; - - const gcsFile = bucket.file(file.newFilename); - const fileStream = fs.createReadStream(file.filepath); - - return new Promise((resolve, reject) => - fileStream - .pipe( - gcsFile.createWriteStream({ - metadata: { - contentType: file.mimetype, - }, - }) - ) - .on("error", reject) - .on("finish", () => resolve(gcsFile)) - ); -} +class DFS { + private readonly bucket: Bucket; + + constructor(bucketKey: BucketKeyType) { + this.bucket = bucketKeysToBucket[bucketKey]; + } + + /** + * Upload functions. + */ + + async uploadFileToBucket(file: formidable.File, destPath: string) { + const gcsFile = this.file(destPath); + const fileStream = fs.createReadStream(file.filepath); + + await pipeline( + fileStream, + gcsFile.createWriteStream({ + metadata: { + contentType: file.mimetype, + }, + }) + ); + } + + async uploadRawContentToBucket({ + content, + contentType, + filePath, + }: { + content: string; + contentType: string; + filePath: string; + }) { + const gcsFile = this.file(filePath); + + await gcsFile.save(content, { + contentType, + }); + } -export async function getFileContentType( - bucketKey: SupportedBucketKeyType, - filename: string -): Promise { - const bucket = bucketKeysToBucket[bucketKey]; + /** + * Download functions. + */ - const gcsFile = bucket.file(filename); + async fetchFileContent(filePath: string) { + const gcsFile = this.file(filePath); - const [metadata] = await gcsFile.getMetadata(); + const [content] = await gcsFile.download(); - return metadata.contentType; + return content.toString(); + } + + async getFileContentType(filename: string): Promise { + const gcsFile = this.file(filename); + + const [metadata] = await gcsFile.getMetadata(); + + return metadata.contentType; + } + + file(filename: string) { + return this.bucket.file(filename); + } + + get name() { + return this.bucket.name; + } } + +export const getPrivateUploadBucket = () => new DFS("PRIVATE_UPLOAD"); + +export const getPublicUploadBucket = () => new DFS("PUBLIC_UPLOAD"); diff --git a/front/lib/resources/content_fragment_resource.ts b/front/lib/resources/content_fragment_resource.ts index b0a633e206c7..eec89e70f31b 100644 --- a/front/lib/resources/content_fragment_resource.ts +++ b/front/lib/resources/content_fragment_resource.ts @@ -1,6 +1,5 @@ import type { ContentFragmentType, ModelId, Result } from "@dust-tt/types"; import { Err, Ok } from "@dust-tt/types"; -import { Storage } from "@google-cloud/storage"; import type { Attributes, CreationAttributes, @@ -9,12 +8,14 @@ import type { } from "sequelize"; import appConfig from "@app/lib/api/config"; +import { getPrivateUploadBucket } from "@app/lib/dfs"; import { Message } from "@app/lib/models"; import { BaseResource } from "@app/lib/resources/base_resource"; -import { gcsConfig } from "@app/lib/resources/storage/config"; import { ContentFragmentModel } from "@app/lib/resources/storage/models/content_fragment"; import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types"; +const privateUploadGcs = getPrivateUploadBucket(); + // Attributes are marked as read-only to reflect the stateless nature of our Resource. // This design will be moved up to BaseResource once we transition away from Sequelize. // eslint-disable-next-line @typescript-eslint/no-empty-interface @@ -143,7 +144,7 @@ export function fileAttachmentLocation({ const filePath = `content_fragments/w/${workspaceId}/assistant/conversations/${conversationId}/content_fragment/${messageId}/${contentFormat}`; return { filePath, - internalUrl: `https://storage.googleapis.com/${gcsConfig.getGcsPrivateUploadsBucket()}/${filePath}`, + internalUrl: `https://storage.googleapis.com/${privateUploadGcs.name}/${filePath}`, downloadUrl: `${appConfig.getAppUrl()}/api/w/${workspaceId}/assistant/conversations/${conversationId}/messages/${messageId}/raw_content_fragment`, }; } @@ -169,15 +170,11 @@ export async function storeContentFragmentText({ messageId, contentFormat: "text", }); - const storage = new Storage({ - keyFilename: appConfig.getServiceAccount(), - }); - const bucket = storage.bucket(gcsConfig.getGcsPrivateUploadsBucket()); - const gcsFile = bucket.file(filePath); - - await gcsFile.save(content, { + await privateUploadGcs.uploadRawContentToBucket({ + content, contentType: "text/plain", + filePath, }); return Buffer.byteLength(content); @@ -192,10 +189,6 @@ export async function getContentFragmentText({ conversationId: string; messageId: string; }): Promise { - const storage = new Storage({ - keyFilename: appConfig.getServiceAccount(), - }); - const { filePath } = fileAttachmentLocation({ workspaceId, conversationId, @@ -203,9 +196,5 @@ export async function getContentFragmentText({ contentFormat: "text", }); - const bucket = storage.bucket(gcsConfig.getGcsPrivateUploadsBucket()); - const gcsFile = bucket.file(filePath); - - const [content] = await gcsFile.download(); - return content.toString(); + return privateUploadGcs.fetchFileContent(filePath); } diff --git a/front/lib/resources/storage/config.ts b/front/lib/resources/storage/config.ts index 1af617feec04..c43eac89c34a 100644 --- a/front/lib/resources/storage/config.ts +++ b/front/lib/resources/storage/config.ts @@ -5,9 +5,3 @@ export const dbConfig = { return EnvironmentConfig.getEnvVariable("FRONT_DATABASE_URI"); }, }; - -export const gcsConfig = { - getGcsPrivateUploadsBucket: (): string => { - return EnvironmentConfig.getEnvVariable("DUST_PRIVATE_UPLOADS_BUCKET"); - }, -}; diff --git a/front/pages/api/w/[wId]/assistant/agent_configurations/avatar.ts b/front/pages/api/w/[wId]/assistant/agent_configurations/avatar.ts index 3392ce138f09..926a110d973f 100644 --- a/front/pages/api/w/[wId]/assistant/agent_configurations/avatar.ts +++ b/front/pages/api/w/[wId]/assistant/agent_configurations/avatar.ts @@ -1,17 +1,17 @@ import { IncomingForm } from "formidable"; import type { NextApiRequest, NextApiResponse } from "next"; -import { uploadToBucket } from "@app/lib/dfs"; +import { getPublicUploadBucket } from "@app/lib/dfs"; import { withLogging } from "@app/logger/withlogging"; -const { DUST_UPLOAD_BUCKET = "" } = process.env; - export const config = { api: { bodyParser: false, // Disabling Next.js's body parser as formidable has its own }, }; +const publicUploadGcs = getPublicUploadBucket(); + async function handler( req: NextApiRequest, res: NextApiResponse @@ -41,9 +41,9 @@ async function handler( const [file] = maybeFiles; - await uploadToBucket("PUBLIC_UPLOAD", file); + await publicUploadGcs.uploadFileToBucket(file, file.newFilename); - const fileUrl = `https://storage.googleapis.com/${DUST_UPLOAD_BUCKET}/${file.newFilename}`; + const fileUrl = `https://storage.googleapis.com/${publicUploadGcs.name}/${file.newFilename}`; res.status(200).json({ fileUrl }); } catch (error) { diff --git a/front/pages/api/w/[wId]/assistant/conversations/[cId]/messages/[mId]/raw_content_fragment/index.ts b/front/pages/api/w/[wId]/assistant/conversations/[cId]/messages/[mId]/raw_content_fragment/index.ts index 67bf51b9c0ec..1256e7cfb4b0 100644 --- a/front/pages/api/w/[wId]/assistant/conversations/[cId]/messages/[mId]/raw_content_fragment/index.ts +++ b/front/pages/api/w/[wId]/assistant/conversations/[cId]/messages/[mId]/raw_content_fragment/index.ts @@ -1,16 +1,12 @@ import type { WithAPIErrorReponse } from "@dust-tt/types"; import { isContentFragmentType } from "@dust-tt/types"; -import { Storage } from "@google-cloud/storage"; import { IncomingForm } from "formidable"; -import fs from "fs"; import type { NextApiRequest, NextApiResponse } from "next"; -import { pipeline } from "stream/promises"; import { getConversation } from "@app/lib/api/assistant/conversation"; -import appConfig from "@app/lib/api/config"; import { Authenticator, getSession } from "@app/lib/auth"; +import { getPrivateUploadBucket } from "@app/lib/dfs"; import { fileAttachmentLocation } from "@app/lib/resources/content_fragment_resource"; -import { gcsConfig } from "@app/lib/resources/storage/config"; import { apiError, withLogging } from "@app/logger/withlogging"; export const config = { @@ -19,6 +15,8 @@ export const config = { }, }; +const privateUploadGcs = getPrivateUploadBucket(); + async function handler( req: NextApiRequest, res: NextApiResponse> @@ -112,16 +110,10 @@ async function handler( contentFormat: "raw", }); - const storage = new Storage({ - keyFilename: appConfig.getServiceAccount(), - }); - const bucket = storage.bucket(gcsConfig.getGcsPrivateUploadsBucket()); - const gcsFile = bucket.file(filePath); - switch (req.method) { case "GET": // redirect to a signed URL - const [url] = await gcsFile.getSignedUrl({ + const [url] = await privateUploadGcs.file(filePath).getSignedUrl({ version: "v4", action: "read", // since we redirect, the use is immediate so expiry can be short @@ -150,16 +142,7 @@ async function handler( const [file] = maybeFiles; - const fileStream = fs.createReadStream(file.filepath); - - await pipeline( - fileStream, - gcsFile.createWriteStream({ - metadata: { - contentType: file.mimetype, - }, - }) - ); + await privateUploadGcs.uploadFileToBucket(file, filePath); res.status(200).json({ sourceUrl: downloadUrl }); return;