Enhance external image hosting in agent configuration

front/components/assistant_builder/AssistantBuilderAvatarPicker.tsx

@@ -194,6 +194,7 @@ export function AvatarPicker({
         style={{ display: "none" }}
         onChange={onFileChange}
         ref={fileInputRef}
+        accept=".png,.jpg,.jpeg"
       />
 
       <div className="h-full w-full overflow-visible">

front/lib/api/assistant/configuration.ts

@@ -34,6 +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 {
   AgentConfiguration,
   AgentDataSourceConfiguration,
@@ -754,6 +755,20 @@ export async function getAgentNames(auth: Authenticator): Promise<string[]> {
   return agents.map((a) => a.name);
 }
 
+async function isSelfHostedImageWithValidContentType(pictureUrl: string) {
+  const filename = pictureUrl.split("/").at(-1);
+  if (!filename) {
+    return false;
+  }
+
+  const contentType = await getFileContentType("PUBLIC_UPLOAD", filename);
+  if (!contentType) {
+    return false;
+  }
+
+  return contentType.includes("image");
+}
+
 /**
  * Create Agent Configuration
  */
@@ -789,6 +804,13 @@ export async function createAgentConfiguration(
     throw new Error("Unexpected `auth` without `user`.");
   }
 
+  const isValidPictureUrl = await isSelfHostedImageWithValidContentType(
+    pictureUrl
+  );
+  if (!isValidPictureUrl) {
+    return new Err(new Error("Invalid picture url."));
+  }
+
   let version = 0;
   let listStatusOverride: AgentUserListStatus | null = null;
 

front/lib/dfs/config.ts

@@ -0,0 +1,12 @@
+import { EnvironmentConfig } from "@dust-tt/types";
+
+const config = {
+  getServiceAccount: (): string => {
+    return EnvironmentConfig.getEnvVariable("SERVICE_ACCOUNT");
+  },
+  getPublicUploadBucket: (): string => {
+    return EnvironmentConfig.getEnvVariable("DUST_UPLOAD_BUCKET");
+  },
+};
+
+export default config;

front/lib/dfs/index.ts

@@ -0,0 +1,52 @@
+import type { Bucket } from "@google-cloud/storage";
+import { Storage } from "@google-cloud/storage";
+import type formidable from "formidable";
+import fs from "fs";
+
+import config from "@app/lib/dfs/config";
+
+type SupportedBucketKeyType = "PUBLIC_UPLOAD";
+
+const storage = new Storage({
+  keyFilename: config.getServiceAccount(),
+});
+
+const bucketKeysToBucket: Record<SupportedBucketKeyType, Bucket> = {
+  PUBLIC_UPLOAD: storage.bucket(config.getPublicUploadBucket()),
+};
+
+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))
+  );
+}
+
+export async function getFileContentType(
+  bucketKey: SupportedBucketKeyType,
+  filename: string
+): Promise<string | null> {
+  const bucket = bucketKeysToBucket[bucketKey];
+
+  const gcsFile = bucket.file(filename);
+
+  const [metadata] = await gcsFile.getMetadata();
+
+  return metadata.contentType;
+}

front/pages/api/w/[wId]/assistant/agent_configurations/avatar.ts

@@ -1,11 +1,10 @@
-import { Storage } from "@google-cloud/storage";
 import { IncomingForm } from "formidable";
-import fs from "fs";
 import type { NextApiRequest, NextApiResponse } from "next";
 
+import { uploadToBucket } from "@app/lib/dfs";
 import { withLogging } from "@app/logger/withlogging";
 
-const { DUST_UPLOAD_BUCKET = "", SERVICE_ACCOUNT } = process.env;
+const { DUST_UPLOAD_BUCKET = "" } = process.env;
 
 export const config = {
   api: {
@@ -19,39 +18,30 @@ async function handler(
 ): Promise<void> {
   if (req.method === "POST") {
     try {
-      const form = new IncomingForm();
-      const [_fields, files] = await form.parse(req);
-      void _fields;
+      const form = new IncomingForm({
+        filter: ({ mimetype }) => {
+          if (!mimetype) {
+            return false;
+          }
+
+          // Only allow uploading image.
+          return mimetype.includes("image");
+        },
+        maxFileSize: 3 * 1024 * 1024, // 3 mb.
+      });
+
+      const [, files] = await form.parse(req);
 
-      const maybeFiles = files.file;
+      const { file: maybeFiles } = files;
 
       if (!maybeFiles) {
         res.status(400).send("No file uploaded.");
         return;
       }
 
-      const file = maybeFiles[0];
-
-      const storage = new Storage({
-        keyFilename: SERVICE_ACCOUNT,
-      });
-
-      const bucket = storage.bucket(DUST_UPLOAD_BUCKET);
-      const gcsFile = await bucket.file(file.newFilename);
-      const fileStream = fs.createReadStream(file.filepath);
+      const [file] = maybeFiles;
 
-      await new Promise((resolve, reject) =>
-        fileStream
-          .pipe(
-            gcsFile.createWriteStream({
-              metadata: {
-                contentType: file.mimetype,
-              },
-            })
-          )
-          .on("error", reject)
-          .on("finish", resolve)
-      );
+      await uploadToBucket("PUBLIC_UPLOAD", file);
 
       const fileUrl = `https://storage.googleapis.com/${DUST_UPLOAD_BUCKET}/${file.newFilename}`;