Fix file limit (mckaywrigley#1613)

* Add file size limit for uploads and fix error handling in file upload * fix for sidebar upload issue
bolechen · Apr 2, 2024 · 16f8749 · 16f8749
1 parent 7f29d2a
commit 16f8749
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 10 deletions.
diff --git a/.env.local.example b/.env.local.example
@@ -31,3 +31,6 @@ AZURE_EMBEDDINGS_NAME=
 # General Configuration (Optional)
 EMAIL_DOMAIN_WHITELIST=
 EMAIL_WHITELIST=
+
+# File size limit for uploads in bytes
+NEXT_PUBLIC_USER_FILE_SIZE_LIMIT=10485760
diff --git a/.gitignore b/.gitignore
@@ -43,4 +43,4 @@ custom-prompts
 sw.js
 sw.js.map
 workbox-*.js
-workbox-*.js.map
+workbox-*.js.map
diff --git a/app/api/retrieval/process/route.ts b/app/api/retrieval/process/route.ts
@@ -24,13 +24,39 @@ export async function POST(req: Request) {
 
     const formData = await req.formData()
 
-    const file = formData.get("file") as File
     const file_id = formData.get("file_id") as string
     const embeddingsProvider = formData.get("embeddingsProvider") as string
 
+    const { data: fileMetadata, error: metadataError } = await supabaseAdmin
+      .from("files")
+      .select("*")
+      .eq("id", file_id)
+      .single()
+
+    if (metadataError) {
+      throw new Error(
+        `Failed to retrieve file metadata: ${metadataError.message}`
+      )
+    }
+
+    if (!fileMetadata) {
+      throw new Error("File not found")
+    }
+
+    if (fileMetadata.user_id !== profile.user_id) {
+      throw new Error("Unauthorized")
+    }
+
+    const { data: file, error: fileError } = await supabaseAdmin.storage
+      .from("files")
+      .download(fileMetadata.file_path)
+
+    if (fileError)
+      throw new Error(`Failed to retrieve file: ${fileError.message}`)
+
     const fileBuffer = Buffer.from(await file.arrayBuffer())
     const blob = new Blob([fileBuffer])
-    const fileExtension = file.name.split(".").pop()?.toLowerCase()
+    const fileExtension = fileMetadata.name.split(".").pop()?.toLowerCase()
 
     if (embeddingsProvider === "openai") {
       if (profile.use_azure_openai) {

diff --git a/components/chat/chat-hooks/use-select-file-handler.tsx b/components/chat/chat-hooks/use-select-file-handler.tsx
@@ -184,9 +184,10 @@ export const useSelectFileHandler = () => {
               )
             )
           }
-        } catch (error) {
-          toast.error("Failed to upload.")
-
+        } catch (error: any) {
+          toast.error("Failed to upload. " + error?.message, {
+            duration: 10000
+          })
           setNewMessageImages(prev =>
             prev.filter(img => img.messageId !== "temp")
           )

diff --git a/db/files.ts b/db/files.ts
@@ -91,6 +91,15 @@ export const createFile = async (
   workspace_id: string,
   embeddingsProvider: "openai" | "local"
 ) => {
+  let validFilename = fileRecord.name.replace(/[^a-z0-9.]/gi, "_").toLowerCase()
+  const extension = file.name.split(".").pop()
+  const baseName = validFilename.substring(0, validFilename.lastIndexOf("."))
+  const maxBaseNameLength = 100 - (extension?.length || 0) - 1
+  if (baseName.length > maxBaseNameLength) {
+    fileRecord.name = baseName.substring(0, maxBaseNameLength) + "." + extension
+  } else {
+    fileRecord.name = baseName + "." + extension
+  }
   const { data: createdFile, error } = await supabase
     .from("files")
     .insert([fileRecord])
@@ -118,7 +127,6 @@ export const createFile = async (
   })
 
   const formData = new FormData()
-  formData.append("file", file)
   formData.append("file_id", createdFile.id)
   formData.append("embeddingsProvider", embeddingsProvider)
 

diff --git a/db/storage/files.ts b/db/storage/files.ts
@@ -9,10 +9,14 @@ export const uploadFile = async (
     file_id: string
   }
 ) => {
-  const SIZE_LIMIT = 10000000 // 10MB
+  const SIZE_LIMIT = parseInt(
+    process.env.NEXT_PUBLIC_USER_FILE_SIZE_LIMIT || "10000000"
+  )
 
   if (file.size > SIZE_LIMIT) {
-    throw new Error(`File must be less than ${SIZE_LIMIT / 1000000}MB`)
+    throw new Error(
+      `File must be less than ${Math.floor(SIZE_LIMIT / 1000000)}MB`
+    )
   }
 
   const filePath = `${payload.user_id}/${Buffer.from(payload.file_id).toString("base64")}`
@@ -24,7 +28,6 @@ export const uploadFile = async (
     })
 
   if (error) {
-    console.error(`Error uploading file with path: ${filePath}`, error)
     throw new Error("Error uploading file")
   }
 
@@ -46,6 +49,7 @@ export const getFileFromStorage = async (filePath: string) => {
     .createSignedUrl(filePath, 60 * 60 * 24) // 24hrs
 
   if (error) {
+    console.error(`Error uploading file with path: ${filePath}`, error)
     throw new Error("Error downloading file")
   }