Add: support uploading gif and webp

dust-tt · Dec 11, 2024 · 289e7f8 · 289e7f8
1 parent fcb081b
commit 289e7f8
Show file tree

Hide file tree

Showing 10 changed files with 239 additions and 223 deletions.
diff --git a/front/components/assistant/conversation/input_bar/InputBarContainer.tsx b/front/components/assistant/conversation/input_bar/InputBarContainer.tsx
@@ -10,7 +10,7 @@ import type {
   LightAgentConfigurationType,
   WorkspaceType,
 } from "@dust-tt/types";
-import { supportedFileExtensions } from "@dust-tt/types";
+import { getSupportedFileExtensions } from "@dust-tt/types";
 import { EditorContent } from "@tiptap/react";
 import React, { useContext, useEffect, useRef, useState } from "react";
 
@@ -125,7 +125,7 @@ const InputBarContainer = ({
           {actions.includes("attachment") && (
             <>
               <input
-                accept={supportedFileExtensions.join(",")}
+                accept={getSupportedFileExtensions().join(",")}
                 onChange={async (e) => {
                   await fileUploaderService.handleFileChange(e);
                   if (fileInputRef.current) {
@@ -142,7 +142,7 @@ const InputBarContainer = ({
                 variant="ghost-secondary"
                 icon={AttachmentIcon}
                 size="xs"
-                tooltip={`Add a document to the conversation (${supportedFileExtensions.join(", ")}).`}
+                tooltip={`Add a document to the conversation (${getSupportedFileExtensions().join(", ")}).`}
                 onClick={() => {
                   fileInputRef.current?.click();
                 }}

diff --git a/front/components/data_source/DocumentUploadOrEditModal.tsx b/front/components/data_source/DocumentUploadOrEditModal.tsx
@@ -1,4 +1,3 @@
-import { supportedPlainTextExtensions } from "@dust-tt/client";
 import {
   Button,
   DocumentPlusIcon,
@@ -21,7 +20,7 @@ import type {
   PlanType,
   WorkspaceType,
 } from "@dust-tt/types";
-import { Err } from "@dust-tt/types";
+import { Err, getSupportedNonImageFileExtensions } from "@dust-tt/types";
 import React, { useCallback, useEffect, useRef, useState } from "react";
 
 import { useFileUploaderService } from "@app/hooks/useFileUploaderService";
@@ -349,7 +348,7 @@ export const DocumentUploadOrEditModal = ({
               <div>
                 <Page.SectionHeader
                   title="Text content"
-                  description={`Copy paste content or upload a file (${supportedPlainTextExtensions.join(", ")}). \n Up to ${
+                  description={`Copy paste content or upload a file (${getSupportedNonImageFileExtensions().join(", ")}). \n Up to ${
                     plan.limits.dataSources.documents.sizeMb === -1
                       ? "2"
                       : plan.limits.dataSources.documents.sizeMb
@@ -370,7 +369,7 @@ export const DocumentUploadOrEditModal = ({
                   type="file"
                   ref={fileInputRef}
                   style={{ display: "none" }}
-                  accept={supportedPlainTextExtensions.join(", ")}
+                  accept={getSupportedNonImageFileExtensions().join(", ")}
                   onChange={handleFileChange}
                 />
                 <TextArea

diff --git a/front/components/data_source/MultipleDocumentsUpload.tsx b/front/components/data_source/MultipleDocumentsUpload.tsx
@@ -4,7 +4,7 @@ import type {
   LightWorkspaceType,
   PlanType,
 } from "@dust-tt/types";
-import { Err, supportedPlainTextExtensions } from "@dust-tt/types";
+import { Err, getSupportedNonImageFileExtensions } from "@dust-tt/types";
 import type { ChangeEvent } from "react";
 import { useCallback, useEffect, useRef, useState } from "react";
 
@@ -229,7 +229,7 @@ export const MultipleDocumentsUpload = ({
       <input
         className="hidden"
         type="file"
-        accept={supportedPlainTextExtensions.join(", ")}
+        accept={getSupportedNonImageFileExtensions().join(", ")}
         ref={fileInputRef}
         multiple={true}
         onChange={handleFileChange}

diff --git a/front/components/data_source/TableUploadOrEditModal.tsx b/front/components/data_source/TableUploadOrEditModal.tsx
@@ -17,6 +17,7 @@ import type {
 } from "@dust-tt/types";
 import {
   Err,
+  getSupportedFileExtensions,
   isBigFileSize,
   isSlugified,
   MAX_FILE_SIZES,
@@ -353,7 +354,7 @@ export const TableUploadOrEditModal = ({
               <div>
                 <Page.SectionHeader
                   title="CSV File"
-                  description={`Select the CSV file for data extraction. The maximum file size allowed is ${maxFileSizeToHumanReadable(MAX_FILE_SIZES.plainText)}.`}
+                  description={`Select the CSV file for data extraction. The maximum file size allowed is ${maxFileSizeToHumanReadable(MAX_FILE_SIZES.delimited)}.`}
                   action={{
                     label:
                       fileUploaderService.isProcessingFiles || isContentLoading
@@ -372,7 +373,7 @@ export const TableUploadOrEditModal = ({
                   type="file"
                   ref={fileInputRef}
                   style={{ display: "none" }}
-                  accept=".csv, .tsv"
+                  accept={getSupportedFileExtensions("delimited").join(",")}
                   onChange={handleFileChange}
                 />
                 {isBigFile && (

diff --git a/front/hooks/useFileUploaderService.ts b/front/hooks/useFileUploaderService.ts
@@ -1,5 +1,6 @@
 import { useSendNotification } from "@dust-tt/sparkle";
 import type {
+  FileFormatCategory,
   FileUseCase,
   FileUseCaseMetadata,
   LightWorkspaceType,
@@ -8,6 +9,7 @@ import type {
 } from "@dust-tt/types";
 import {
   Err,
+  getFileFormatCategory,
   isAPIErrorResponse,
   isSupportedFileContentType,
   isSupportedImageContentType,
@@ -46,9 +48,6 @@ class FileBlobUploadError extends Error {
   }
 }
 
-const COMBINED_MAX_TEXT_FILES_SIZE = MAX_FILE_SIZES["plainText"] * 2;
-const COMBINED_MAX_IMAGE_FILES_SIZE = MAX_FILE_SIZES["image"] * 5;
-
 export function useFileUploaderService({
   owner,
   useCase,
@@ -66,30 +65,26 @@ export function useFileUploaderService({
   const handleFilesUpload = async (files: File[]) => {
     setIsProcessingFiles(true);
 
-    const { totalTextualSize, totalImageSize } = [
-      ...fileBlobs,
-      ...files,
-    ].reduce(
-      (acc, content) => {
-        const { size } = content;
-        if (
-          isSupportedImageContentType(
-            content instanceof File ? content.type : content.contentType
-          )
-        ) {
-          acc.totalImageSize += size;
-        } else {
-          acc.totalTextualSize += size;
-        }
-        return acc;
-      },
-      { totalTextualSize: 0, totalImageSize: 0 }
-    );
+    const categoryToSize: Map<FileFormatCategory, number> = new Map();
+
+    for (const f of [...fileBlobs, ...files]) {
+      const contentType = f instanceof File ? f.type : f.contentType;
+
+      const category = getFileFormatCategory(contentType) || "data";
+      // The fallback should never happen but let's avoid a crash.
+
+      categoryToSize.set(
+        category,
+        (categoryToSize.get(category) || 0) + f.size
+      );
+    }
+
+    const isTooBig = [...categoryToSize].some(([cat, size]) => {
+      const multiplier = cat === "image" ? 5 : 2;
+      return size > MAX_FILE_SIZES[cat] * multiplier;
+    });
 
-    if (
-      totalTextualSize > COMBINED_MAX_TEXT_FILES_SIZE ||
-      totalImageSize > COMBINED_MAX_IMAGE_FILES_SIZE
-    ) {
+    if (isTooBig) {
       sendNotification({
         type: "error",
         title: "Files too large.",

diff --git a/sdks/js/src/types.ts b/sdks/js/src/types.ts
@@ -69,7 +69,13 @@ export type ConnectorsAPIErrorType = z.infer<
 >;
 
 // Supported content types that are plain text and can be sent as file-less content fragment.
-export const supportedRawText = {
+const supportedOtherFileFormats = {
+  "application/msword": [".doc", ".docx"],
+  "application/vnd.openxmlformats-officedocument.wordprocessingml.document": [
+    ".doc",
+    ".docx",
+  ],
+  "application/pdf": [".pdf"],
   "text/comma-separated-values": [".csv"],
   "text/csv": [".csv"],
   "text/markdown": [".md", ".markdown"],
@@ -87,75 +93,57 @@ export const supportedRawText = {
   "application/x-sh": [".sh"],
 } as const;
 
-// Supported content types for plain text (after processing).
-export const supportedPlainText = {
-  "application/msword": [".doc", ".docx"],
-  "application/vnd.openxmlformats-officedocument.wordprocessingml.document": [
-    ".doc",
-    ".docx",
-  ],
-  "application/pdf": [".pdf"],
-  ...supportedRawText,
-} as const;
-
 // Supported content types for images.
-export const supportedImage = {
+const supportedImageFileFormats = {
   "image/jpeg": [".jpg", ".jpeg"],
   "image/png": [".png"],
+  "image/gif": [".gif"],
+  "image/webp": [".webp"],
 } as const;
 
 // Legacy content types still retuned by the API when rendering old messages.
-export const supportedLegacy = {
+const supportedLegacy = {
   "dust-application/slack": [],
 } as const;
 
-export type PlainTextContentType = keyof typeof supportedPlainText;
-export type ImageContentType = keyof typeof supportedImage;
-export type LegacyContentType = keyof typeof supportedLegacy;
+type OtherContentType = keyof typeof supportedOtherFileFormats;
+type ImageContentType = keyof typeof supportedImageFileFormats;
 
-export const supportedPlainTextContentTypes = Object.keys(
-  supportedPlainText
-) as PlainTextContentType[];
-export const supportedImageContentTypes = Object.keys(
-  supportedImage
-) as ImageContentType[];
-export const supportedLegacyContentTypes = Object.keys(
-  supportedImage
+const supportedOtherContentTypes = Object.keys(
+  supportedOtherFileFormats
+) as OtherContentType[];
+const supportedImageContentTypes = Object.keys(
+  supportedImageFileFormats
 ) as ImageContentType[];
 
-export type SupportedFileContentType = PlainTextContentType | ImageContentType;
+export type SupportedFileContentType = OtherContentType | ImageContentType;
 const supportedUploadableContentType = [
-  ...supportedPlainTextContentTypes,
+  ...supportedOtherContentTypes,
   ...supportedImageContentTypes,
 ] as SupportedFileContentType[];
 
 const SupportedContentFragmentTypeSchema = FlexibleEnumSchema([
-  ...(Object.keys(supportedPlainText) as [keyof typeof supportedPlainText]),
-  ...(Object.keys(supportedImage) as [keyof typeof supportedImage]),
+  ...(Object.keys(supportedOtherFileFormats) as [
+    keyof typeof supportedOtherFileFormats
+  ]),
+  ...(Object.keys(supportedImageFileFormats) as [
+    keyof typeof supportedImageFileFormats
+  ]),
   ...(Object.keys(supportedLegacy) as [keyof typeof supportedLegacy]),
 ]);
 
 const SupportedInlinedContentFragmentTypeSchema = FlexibleEnumSchema([
-  ...(Object.keys(supportedRawText) as [keyof typeof supportedRawText]),
+  ...(Object.keys(supportedOtherFileFormats) as [
+    keyof typeof supportedOtherFileFormats
+  ]),
 ]);
 const SupportedFileContentFragmentTypeSchema = FlexibleEnumSchema([
-  ...(Object.keys(supportedPlainText) as [keyof typeof supportedPlainText]),
-  ...(Object.keys(supportedImage) as [keyof typeof supportedImage]),
-]);
-
-const uniq = <T>(arr: T[]): T[] => Array.from(new Set(arr));
-
-export const supportedPlainTextExtensions = uniq(
-  Object.values(supportedPlainText).flat()
-);
-
-export const supportedImageExtensions = uniq(
-  Object.values(supportedImage).flat()
-);
-
-export const supportedFileExtensions = uniq([
-  ...supportedPlainTextExtensions,
-  ...supportedImageExtensions,
+  ...(Object.keys(supportedOtherFileFormats) as [
+    keyof typeof supportedOtherFileFormats
+  ]),
+  ...(Object.keys(supportedImageFileFormats) as [
+    keyof typeof supportedImageFileFormats
+  ]),
 ]);
 
 export function isSupportedFileContentType(
@@ -168,10 +156,8 @@ export function isSupportedFileContentType(
 
 export function isSupportedPlainTextContentType(
   contentType: string
-): contentType is PlainTextContentType {
-  return supportedPlainTextContentTypes.includes(
-    contentType as PlainTextContentType
-  );
+): contentType is OtherContentType {
+  return supportedOtherContentTypes.includes(contentType as OtherContentType);
 }
 
 export function isSupportedImageContentType(

diff --git a/types/src/front/api_handlers/internal/assistant.ts b/types/src/front/api_handlers/internal/assistant.ts
@@ -1,6 +1,6 @@
 import * as t from "io-ts";
 
-import { getSupportedInlinedContentTypeCodec } from "../../content_fragment";
+import { getSupportedNonImageMimeTypes } from "../../files";
 
 export const InternalPostMessagesRequestBodySchema = t.type({
   content: t.string,
@@ -21,11 +21,20 @@ const ContentFragmentBaseSchema = t.intersection([
   }),
 ]);
 
+export const getSupportedInlinedContentType = () => {
+  const [first, second, ...rest] = getSupportedNonImageMimeTypes();
+  return t.union([
+    t.literal(first),
+    t.literal(second),
+    ...rest.map((value) => t.literal(value)),
+  ]);
+};
+
 const ContentFragmentInputWithContentSchema = t.intersection([
   ContentFragmentBaseSchema,
   t.type({
     content: t.string,
-    contentType: getSupportedInlinedContentTypeCodec(),
+    contentType: getSupportedInlinedContentType(),
   }),
 ]);
 
@@ -44,7 +53,7 @@ export type ContentFragmentInputWithFileIdType = t.TypeOf<
   typeof ContentFragmentInputWithFileIdSchema
 >;
 
-export type ContentFragmentInputType =
+type ContentFragmentInputType =
   | ContentFragmentInputWithContentType
   | ContentFragmentInputWithFileIdType;
 

diff --git a/types/src/front/assistant/actions/index.ts b/types/src/front/assistant/actions/index.ts
@@ -4,13 +4,13 @@ import {
 } from "../../../front/assistant/generation";
 import { ModelConfigurationType } from "../../../front/lib/assistant";
 import { ModelId } from "../../../shared/model_id";
-import { SupportedContentFragmentType } from "../../content_fragment";
+import { SupportedFileContentType } from "../../files";
 import { ConversationType } from "../conversation";
 
 export type ActionGeneratedFileType = {
   fileId: string;
   title: string;
-  contentType: SupportedContentFragmentType;
+  contentType: SupportedFileContentType;
   snippet: string | null;
 };