Apply new DFS logic to content fragment + refactor

dust-tt · Apr 4, 2024 · dfd21f3 · dfd21f3
1 parent a86e7f9
commit dfd21f3
Show file tree

Hide file tree

Showing 7 changed files with 99 additions and 88 deletions.
diff --git 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
@@ -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
@@ -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<SupportedBucketKeyType, Bucket> = {
-  PUBLIC_UPLOAD: storage.bucket(config.getPublicUploadBucket()),
+const bucketKeysToBucket: Record<BucketKeyType, Bucket> = {
+  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<string | null> {
-  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<string | null> {
+    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
@@ -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,20 +189,12 @@ export async function getContentFragmentText({
   conversationId: string;
   messageId: string;
 }): Promise<string> {
-  const storage = new Storage({
-    keyFilename: appConfig.getServiceAccount(),
-  });
-
   const { filePath } = fileAttachmentLocation({
     workspaceId,
     conversationId,
     messageId,
     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
@@ -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
@@ -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/...es/api/w/[wId]/assistant/conversations/[cId]/messages/[mId]/raw_content_fragment/index.ts b/...es/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<WithAPIErrorReponse<{ sourceUrl: string }>>
@@ -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;