fix(ui): Adds related model support to embeddings

invokeai/frontend/web/src/features/prompt/PromptTriggerSelect.tsx

@@ -1,5 +1,5 @@
 import type { ChakraProps, ComboboxOnChange, ComboboxOption } from '@invoke-ai/ui-library';
-import { Combobox, FormControl } from '@invoke-ai/ui-library';
+import { Combobox, Flex, FormControl, Icon } from '@invoke-ai/ui-library';
 import { skipToken } from '@reduxjs/toolkit/query';
 import { useAppSelector } from 'app/store/storeHooks';
 import type { GroupBase } from 'chakra-react-select';
@@ -10,12 +10,16 @@ import type { PromptTriggerSelectProps } from 'features/prompt/types';
 import { t } from 'i18next';
 import { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
+import { PiLinkSimple } from 'react-icons/pi';
+import { useGetRelatedModelIdsBatchQuery } from 'services/api/endpoints/modelRelationships';
 import { useGetModelConfigQuery } from 'services/api/endpoints/models';
 import { useEmbeddingModels, useLoRAModels } from 'services/api/hooks/modelsByType';
 import { isNonRefinerMainModelConfig } from 'services/api/types';
 
 const noOptionsMessage = () => t('prompt.noMatchingTriggers');
 
+type ComboboxOptionWithStarred = ComboboxOption & { starred?: boolean };
+
 export const PromptTriggerSelect = memo(({ onSelect, onClose }: PromptTriggerSelectProps) => {
   const { t } = useTranslation();
 
@@ -27,6 +31,27 @@ export const PromptTriggerSelect = memo(({ onSelect, onClose }: PromptTriggerSel
   const [loraModels, { isLoading: isLoadingLoRAs }] = useLoRAModels();
   const [tiModels, { isLoading: isLoadingTIs }] = useEmbeddingModels();
 
+  // Get related model keys for current selected models
+  const selectedModelKeys = useMemo(() => {
+    const keys: string[] = [];
+    if (mainModel) {
+      keys.push(mainModel.key);
+    }
+    for (const { model } of addedLoRAs) {
+      keys.push(model.key);
+    }
+    return keys;
+  }, [mainModel, addedLoRAs]);
+
+  const { relatedModelKeys } = useGetRelatedModelIdsBatchQuery(selectedModelKeys, {
+    selectFromResult: ({ data }) => {
+      if (!data) {
+        return { relatedModelKeys: [] };
+      }
+      return { relatedModelKeys: data };
+    },
+  });
+
   const _onChange = useCallback<ComboboxOnChange>(
     (v) => {
       if (!v) {
@@ -40,20 +65,7 @@ export const PromptTriggerSelect = memo(({ onSelect, onClose }: PromptTriggerSel
   );
 
   const options = useMemo(() => {
-    const _options: GroupBase<ComboboxOption>[] = [];
-
-    if (tiModels) {
-      const embeddingOptions = tiModels
-        .filter((ti) => ti.base === mainModelConfig?.base)
-        .map((model) => ({ label: model.name, value: `<${model.name}>` }));
-
-      if (embeddingOptions.length > 0) {
-        _options.push({
-          label: t('prompt.compatibleEmbeddings'),
-          options: embeddingOptions,
-        });
-      }
-    }
+    const _options: GroupBase<ComboboxOptionWithStarred>[] = [];
 
     if (loraModels) {
       const triggerPhraseOptions = loraModels
@@ -74,6 +86,35 @@ export const PromptTriggerSelect = memo(({ onSelect, onClose }: PromptTriggerSel
       }
     }
 
+    if (tiModels) {
+      // Create embedding options with starred property for related models
+      const embeddingOptions: ComboboxOptionWithStarred[] = tiModels
+        .filter((ti) => ti.base === mainModelConfig?.base)
+        .map((model) => ({
+          label: model.name,
+          value: `<${model.name}>`,
+          starred: relatedModelKeys.includes(model.key),
+        }));
+
+      // Sort so related embeddings come first
+      embeddingOptions.sort((a, b) => {
+        if (a.starred && !b.starred) {
+          return -1;
+        }
+        if (!a.starred && b.starred) {
+          return 1;
+        }
+        return 0;
+      });
+
+      if (embeddingOptions.length > 0) {
+        _options.push({
+          label: t('prompt.compatibleEmbeddings'),
+          options: embeddingOptions,
+        });
+      }
+    }
+
     if (mainModelConfig && isNonRefinerMainModelConfig(mainModelConfig) && mainModelConfig.trigger_phrases?.length) {
       _options.push({
         label: t('modelManager.mainModelTriggerPhrases'),
@@ -85,16 +126,28 @@ export const PromptTriggerSelect = memo(({ onSelect, onClose }: PromptTriggerSel
     }
 
     return _options;
-  }, [tiModels, loraModels, mainModelConfig, t, addedLoRAs]);
+  }, [tiModels, loraModels, mainModelConfig, t, addedLoRAs, relatedModelKeys]);
+
+  const formatOptionLabel = useCallback((option: ComboboxOptionWithStarred) => {
+    return (
+      <Flex alignItems="center" gap={2}>
+        {option.starred && <Icon as={PiLinkSimple} color="invokeYellow.500" boxSize={4} />}
+        {option.label}
+      </Flex>
+    );
+  }, []);
+
+  const placeholder = useMemo(() => {
+    if (isLoadingLoRAs || isLoadingTIs || isLoadingMainModelConfig) {
+      return t('common.loading');
+    }
+    return t('prompt.addPromptTrigger');
+  }, [t, isLoadingLoRAs, isLoadingTIs, isLoadingMainModelConfig]);
 
   return (
     <FormControl>
       <Combobox
-        placeholder={
-          isLoadingLoRAs || isLoadingTIs || isLoadingMainModelConfig
-            ? t('common.loading')
-            : t('prompt.addPromptTrigger')
-        }
+        placeholder={placeholder}
         defaultMenuIsOpen
         autoFocus
         value={null}
@@ -104,6 +157,7 @@ export const PromptTriggerSelect = memo(({ onSelect, onClose }: PromptTriggerSel
         onMenuClose={onClose}
         data-testid="add-prompt-trigger"
         sx={selectStyles}
+        formatOptionLabel={formatOptionLabel}
       />
     </FormControl>
   );