Add image generation for gpt-4o

mhsabu · Jun 3, 2024 · 0c642f2 · 0c642f2
1 parent 3788041
commit 0c642f2
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/web/src/app/admin/assistants/AssistantEditor.tsx b/web/src/app/admin/assistants/AssistantEditor.tsx
@@ -36,6 +36,7 @@ import { Option } from "@/components/Dropdown";
 import { ToolSnapshot } from "@/lib/tools/interfaces";
 import { checkUserIsNoAuthUser } from "@/lib/user";
 import { addAssistantToList } from "@/lib/assistants/updateAssistantPreferences";
+import { checkLLMSupportsImageInput } from "@/lib/llm/utils";
 
 function findSearchTool(tools: ToolSnapshot[]) {
   return tools.find((tool) => tool.in_code_tool_id === "SearchTool");
@@ -45,10 +46,6 @@ function findImageGenerationTool(tools: ToolSnapshot[]) {
   return tools.find((tool) => tool.in_code_tool_id === "ImageGenerationTool");
 }
 
-function checkLLMSupportsImageGeneration(provider: string, model: string) {
-  return provider === "openai" && model === "gpt-4-turbo";
-}
-
 function Label({ children }: { children: string | JSX.Element }) {
   return (
     <div className="block font-medium text-base text-emphasis">{children}</div>
@@ -261,7 +258,7 @@ export function AssistantEditor({
           if (
             values.image_generation_tool_enabled &&
             imageGenerationTool &&
-            checkLLMSupportsImageGeneration(
+            checkLLMSupportsImageInput(
               providerDisplayNameToProviderName.get(
                 values.llm_model_provider_override || ""
               ) ||
@@ -566,7 +563,7 @@ export function AssistantEditor({
                   )}
 
                   {imageGenerationTool &&
-                    checkLLMSupportsImageGeneration(
+                    checkLLMSupportsImageInput(
                       providerDisplayNameToProviderName.get(
                         values.llm_model_provider_override || ""
                       ) ||