Add model check for image input

memfreeme · Dec 19, 2024 · b6681e2 · b6681e2
1 parent 03fe599
commit b6681e2
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 5 deletions.
diff --git a/frontend/components/search/search-bar.tsx b/frontend/components/search/search-bar.tsx
@@ -7,20 +7,21 @@ import { Tooltip, TooltipContent, TooltipTrigger } from '@/components/ui/tooltip
 import { ModelSelection } from '@/components/search/model-selection';
 import { SourceSelection } from '@/components/search/source-selection';
 import TextareaAutosize from 'react-textarea-autosize';
-import { useUIStore, useUserStore } from '@/lib/store';
+import { configStore, useUIStore, useUserStore } from '@/lib/store';
 import { toast } from 'sonner';
 import { Icons } from '@/components/shared/icons';
 import { type FileRejection, useDropzone } from 'react-dropzone';
 import { useUploadFile } from '@/hooks/use-upload-file';
 import { useUpgradeModal } from '@/hooks/use-upgrade-modal';
-import { getFileSizeLimit, processImageFiles } from '@/lib/utils';
+import { getFileSizeLimit, hasImageInput, processImageFiles } from '@/lib/utils';
 import { useTranslations } from 'next-intl';
 import { isProUser } from '@/lib/shared-utils';
 import dynamic from 'next/dynamic';
 import { Switch } from '@/components/ui/switch';
 import { Label } from '@/components/ui/label';
 import { SearchType } from '@/lib/types';
 import WebImageModal, { WebImageFile } from '@/components/modal/web-images-model';
+import { isImageInputModel } from '@/lib/model';
 
 interface Props {
     handleSearch: (key: string, attachments?: string[]) => void;
@@ -83,6 +84,14 @@ const SearchBar: React.FC<Props> = ({
         return false;
     };
 
+    const checkImageInput = (attachments: string[]) => {
+        if (hasImageInput(attachments) && !isImageInputModel(configStore.getState().model)) {
+            toast.error('Image input is not supported for this AI model, please switch to GPT-4o mini, GPT-4o, Claude 3.5 Sonnet AI models');
+            return true;
+        }
+        return false;
+    };
+
     const handleClick = () => {
         if (!user) {
             signInModal.onOpen();
@@ -93,6 +102,9 @@ const SearchBar: React.FC<Props> = ({
         }
         if (uploadedFiles && uploadedFiles.length > 0) {
             const fileUrls = uploadedFiles.map((file) => file.url);
+            if (checkImageInput(fileUrls)) {
+                return;
+            }
             handleSearch(content, fileUrls);
             setFiles([]);
             setUploadedFiles([]);

diff --git a/frontend/components/search/search-window.tsx b/frontend/components/search/search-window.tsx
@@ -109,7 +109,6 @@ export default function SearchWindow({ id, initialMessages, user, isReadOnly = f
             if (!messageValue && attachments && searchType === 'ui') {
                 messageValue = 'Please generate the same UI as the image';
             }
-
             // const imageUrls = extractAllImageUrls(messageValue);
             // if (imageUrls.length > 1 && user && !isProUser(user)) {
             //     toast.error(t('multi-image-free-limit'));

diff --git a/frontend/lib/llm/llm.ts b/frontend/lib/llm/llm.ts
@@ -5,7 +5,6 @@ import { createOpenAI } from '@ai-sdk/openai';
 import { createAnthropic } from '@ai-sdk/anthropic';
 import { Claude_35_Sonnet, GPT_4o, GPT_4o_MIMI, O1_MIMI, O1_PREVIEW } from '@/lib/model';
 import { google } from '@ai-sdk/google';
-import { extractAllImageUrls, replaceImageUrl } from '@/lib/shared-utils';
 import { Message } from '@/lib/types';
 import { OPENAI_BASE_URL } from '@/lib/env';
 

diff --git a/frontend/lib/model.ts b/frontend/lib/model.ts
@@ -9,5 +9,8 @@ const validModels = [GPT_4o_MIMI, GPT_4o, O1_MIMI, O1_PREVIEW, Claude_35_Sonnet,
 
 const proModels = [GPT_4o, Claude_35_Sonnet, Claude_35_Haiku, O1_MIMI, O1_PREVIEW];
 
+export const ImageInputModels = [GPT_4o, GPT_4o_MIMI, Claude_35_Sonnet];
+
 export const validModel = (model: string) => validModels.includes(model);
 export const isProModel = (model: string) => proModels.includes(model);
+export const isImageInputModel = (model: string) => ImageInputModels.includes(model);
diff --git a/frontend/lib/utils.ts b/frontend/lib/utils.ts
@@ -1,4 +1,4 @@
-import { isProUser } from '@/lib/shared-utils';
+import { isProUser, isValidImageUrl } from '@/lib/shared-utils';
 import { type ClassValue, clsx } from 'clsx';
 import { twMerge } from 'tailwind-merge';
 import { toast } from 'sonner';
@@ -82,6 +82,11 @@ export function getFileSizeLimit(user: any) {
     return 4 * 1024 * 1024;
 }
 
+export function hasImageInput(attachments?: string[]): boolean {
+    if (!attachments || attachments.length === 0) return false;
+    return attachments.some((attachment) => isValidImageUrl(attachment));
+}
+
 export async function processImageFiles(imageFiles: File[]): Promise<File[]> {
     if (typeof window === 'undefined') {
         return imageFiles;