Merge pull request #2949 from danswer-ai/avoid_image_confusion

avoid image generation tool confusion

Author: hagen-danswer

backend/danswer/server/features/persona/api.py

@@ -30,6 +30,7 @@
 from danswer.file_store.models import ChatFileType
 from danswer.llm.answering.prompts.utils import build_dummy_prompt
 from danswer.server.features.persona.models import CreatePersonaRequest
+from danswer.server.features.persona.models import ImageGenerationToolStatus
 from danswer.server.features.persona.models import PersonaSharedNotificationData
 from danswer.server.features.persona.models import PersonaSnapshot
 from danswer.server.features.persona.models import PromptTemplateResponse
@@ -227,6 +228,16 @@ def delete_persona(
     )
 
 
+@basic_router.get("/image-generation-tool")
+def get_image_generation_tool(
+    _: User
+    | None = Depends(current_user),  # User param not used but kept for consistency
+    db_session: Session = Depends(get_session),
+) -> ImageGenerationToolStatus:  # Use bool instead of str for boolean values
+    is_available = is_image_generation_available(db_session=db_session)
+    return ImageGenerationToolStatus(is_available=is_available)
+
+
 @basic_router.get("")
 def list_personas(
     user: User | None = Depends(current_user),

backend/danswer/server/features/persona/models.py

@@ -124,3 +124,7 @@ class PromptTemplateResponse(BaseModel):
 
 class PersonaSharedNotificationData(BaseModel):
     persona_id: int
+
+
+class ImageGenerationToolStatus(BaseModel):
+    is_available: bool

web/src/app/admin/assistants/AssistantEditor.tsx

@@ -45,12 +45,7 @@ import { FullLLMProvider } from "../configuration/llm/interfaces";
 import CollapsibleSection from "./CollapsibleSection";
 import { SuccessfulPersonaUpdateRedirectType } from "./enums";
 import { Persona, StarterMessage } from "./interfaces";
-import {
-  buildFinalPrompt,
-  createPersona,
-  providersContainImageGeneratingSupport,
-  updatePersona,
-} from "./lib";
+import { buildFinalPrompt, createPersona, updatePersona } from "./lib";
 import { Popover } from "@/components/popover/Popover";
 import {
   CameraIcon,
@@ -106,7 +101,7 @@ export function AssistantEditor({
   shouldAddAssistantToUserPreferences?: boolean;
   admin?: boolean;
 }) {
-  const { refreshAssistants } = useAssistants();
+  const { refreshAssistants, isImageGenerationAvailable } = useAssistants();
   const router = useRouter();
 
   const { popup, setPopup } = usePopup();
@@ -138,42 +133,13 @@ export function AssistantEditor({
 
   const [isIconDropdownOpen, setIsIconDropdownOpen] = useState(false);
 
-  const [finalPrompt, setFinalPrompt] = useState<string | null>("");
-  const [finalPromptError, setFinalPromptError] = useState<string>("");
   const [removePersonaImage, setRemovePersonaImage] = useState(false);
 
-  const triggerFinalPromptUpdate = async (
-    systemPrompt: string,
-    taskPrompt: string,
-    retrievalDisabled: boolean
-  ) => {
-    const response = await buildFinalPrompt(
-      systemPrompt,
-      taskPrompt,
-      retrievalDisabled
-    );
-    if (response.ok) {
-      setFinalPrompt((await response.json()).final_prompt_template);
-    }
-  };
-
   const isUpdate = existingPersona !== undefined && existingPersona !== null;
   const existingPrompt = existingPersona?.prompts[0] ?? null;
-
-  useEffect(() => {
-    if (isUpdate && existingPrompt) {
-      triggerFinalPromptUpdate(
-        existingPrompt.system_prompt,
-        existingPrompt.task_prompt,
-        existingPersona.num_chunks === 0
-      );
-    }
-  }, [isUpdate, existingPrompt, existingPersona?.num_chunks]);
-
   const defaultProvider = llmProviders.find(
     (llmProvider) => llmProvider.is_default_provider
   );
-  const defaultProviderName = defaultProvider?.provider;
   const defaultModelName = defaultProvider?.default_model_name;
   const providerDisplayNameToProviderName = new Map<string, string>();
   llmProviders.forEach((llmProvider) => {
@@ -314,14 +280,6 @@ export function AssistantEditor({
             }
           )}
         onSubmit={async (values, formikHelpers) => {
-          if (finalPromptError) {
-            setPopup({
-              type: "error",
-              message: "Cannot submit while there are errors in the form",
-            });
-            return;
-          }
-
           if (
             values.llm_model_provider_override &&
             !values.llm_model_version_override
@@ -642,13 +600,7 @@ export function AssistantEditor({
                 placeholder="e.g. 'You are a professional email writing assistant that always uses a polite enthusiastic tone, emphasizes action items, and leaves blanks for the human to fill in when you have unknowns'"
                 onChange={(e) => {
                   setFieldValue("system_prompt", e.target.value);
-                  triggerFinalPromptUpdate(
-                    e.target.value,
-                    values.task_prompt,
-                    searchToolEnabled()
-                  );
                 }}
-                error={finalPromptError}
               />
 
               <div>
@@ -775,7 +727,8 @@ export function AssistantEditor({
                         <TooltipTrigger asChild>
                           <div
                             className={`w-fit ${
-                              !currentLLMSupportsImageOutput
+                              !currentLLMSupportsImageOutput ||
+                              !isImageGenerationAvailable
                                 ? "opacity-70 cursor-not-allowed"
                                 : ""
                             }`}
@@ -787,18 +740,30 @@ export function AssistantEditor({
                               onChange={() => {
                                 toggleToolInValues(imageGenerationTool.id);
                               }}
-                              disabled={!currentLLMSupportsImageOutput}
+                              disabled={
+                                !currentLLMSupportsImageOutput ||
+                                !isImageGenerationAvailable
+                              }
                             />
                           </div>
                         </TooltipTrigger>
-                        {!currentLLMSupportsImageOutput && (
+                        {!currentLLMSupportsImageOutput ? (
                           <TooltipContent side="top" align="center">
                             <p className="bg-background-900 max-w-[200px] mb-1 text-sm rounded-lg p-1.5 text-white">
                               To use Image Generation, select GPT-4o or another
                               image compatible model as the default model for
                               this Assistant.
                             </p>
                           </TooltipContent>
+                        ) : (
+                          !isImageGenerationAvailable && (
+                            <TooltipContent side="top" align="center">
+                              <p className="bg-background-900 max-w-[200px] mb-1 text-sm rounded-lg p-1.5 text-white">
+                                Image Generation requires an OpenAI or Azure
+                                Dalle configuration.
+                              </p>
+                            </TooltipContent>
+                          )
                         )}
                       </Tooltip>
                     </TooltipProvider>
@@ -1024,11 +989,6 @@ export function AssistantEditor({
                         placeholder="e.g. 'Remember to reference all of the points mentioned in my message to you and focus on identifying action items that can move things forward'"
                         onChange={(e) => {
                           setFieldValue("task_prompt", e.target.value);
-                          triggerFinalPromptUpdate(
-                            values.system_prompt,
-                            e.target.value,
-                            searchToolEnabled()
-                          );
                         }}
                         explanationText="Learn about prompting in our docs!"
                         explanationLink="https://docs.danswer.dev/guides/assistants"

web/src/components/context/AssistantsContext.tsx

@@ -21,6 +21,7 @@ interface AssistantsContextProps {
   finalAssistants: Persona[];
   ownedButHiddenAssistants: Persona[];
   refreshAssistants: () => Promise<void>;
+  isImageGenerationAvailable: boolean;
 
   // Admin only
   editablePersonas: Persona[];
@@ -47,51 +48,54 @@ export const AssistantsProvider: React.FC<{
   );
   const { user, isLoadingUser, isAdmin } = useUser();
   const [editablePersonas, setEditablePersonas] = useState<Persona[]>([]);
+  const [allAssistants, setAllAssistants] = useState<Persona[]>([]);
 
-  useEffect(() => {
-    const fetchEditablePersonas = async () => {
-      if (!isAdmin) {
-        return;
-      }
+  const [isImageGenerationAvailable, setIsImageGenerationAvailable] =
+    useState<boolean>(false);
 
+  useEffect(() => {
+    const checkImageGenerationAvailability = async () => {
       try {
-        const response = await fetch("/api/admin/persona?get_editable=true");
-        if (!response.ok) {
-          console.error("Failed to fetch editable personas");
-          return;
+        const response = await fetch("/api/persona/image-generation-tool");
+        if (response.ok) {
+          const { is_available } = await response.json();
+          setIsImageGenerationAvailable(is_available);
         }
-        const personas = await response.json();
-        setEditablePersonas(personas);
       } catch (error) {
-        console.error("Error fetching editable personas:", error);
+        console.error("Error checking image generation availability:", error);
       }
     };
 
-    fetchEditablePersonas();
-  }, [isAdmin]);
-
-  const [allAssistants, setAllAssistants] = useState<Persona[]>([]);
+    checkImageGenerationAvailability();
+  }, []);
 
   useEffect(() => {
-    const fetchAllAssistants = async () => {
+    const fetchPersonas = async () => {
       if (!isAdmin) {
         return;
       }
 
       try {
-        const response = await fetch("/api/admin/persona");
-        if (!response.ok) {
-          console.error("Failed to fetch all personas");
-          return;
+        const [editableResponse, allResponse] = await Promise.all([
+          fetch("/api/admin/persona?get_editable=true"),
+          fetch("/api/admin/persona"),
+        ]);
+
+        if (editableResponse.ok) {
+          const editablePersonas = await editableResponse.json();
+          setEditablePersonas(editablePersonas);
+        }
+
+        if (allResponse.ok) {
+          const allPersonas = await allResponse.json();
+          setAllAssistants(allPersonas);
         }
-        const personas = await response.json();
-        setAllAssistants(personas);
       } catch (error) {
-        console.error("Error fetching all personas:", error);
+        console.error("Error fetching personas:", error);
       }
     };
 
-    fetchAllAssistants();
+    fetchPersonas();
   }, [isAdmin]);
 
   const refreshAssistants = async () => {
@@ -162,6 +166,7 @@ export const AssistantsProvider: React.FC<{
         refreshAssistants,
         editablePersonas,
         allAssistants,
+        isImageGenerationAvailable,
       }}
     >
       {children}