From 0471e064404d8956a394ca5f771107ecf48bccd7 Mon Sep 17 00:00:00 2001
From: Ivy <ivy.jeong@gmail.com>
Date: Fri, 20 Dec 2024 15:56:07 -0500
Subject: [PATCH] move sort & suggested label out of ComboBox

---
 ui/admin/app/components/composed/ComboBox.tsx | 57 ++++++-------------
 .../model/DefaultModelAliasForm.tsx           | 55 +++++++++++++-----
 2 files changed, 58 insertions(+), 54 deletions(-)

diff --git a/ui/admin/app/components/composed/ComboBox.tsx b/ui/admin/app/components/composed/ComboBox.tsx
index 5c2b34c1..06d819f6 100644
--- a/ui/admin/app/components/composed/ComboBox.tsx
+++ b/ui/admin/app/components/composed/ComboBox.tsx
@@ -35,7 +35,7 @@ type ComboBoxProps<T extends BaseOption> = {
     onChange: (option: T | null) => void;
     options: (T | GroupedOption<T>)[];
     placeholder?: string;
-    suggested?: string[];
+    renderOption?: (option: T) => ReactNode;
     value?: T | null;
 };
 
@@ -43,7 +43,7 @@ export function ComboBox<T extends BaseOption>({
     disabled,
     placeholder,
     value,
-    suggested,
+    renderOption,
     ...props
 }: {
     disabled?: boolean;
@@ -58,7 +58,7 @@ export function ComboBox<T extends BaseOption>({
                 <PopoverContent className="w-full p-0" align="start">
                     <ComboBoxList
                         setOpen={setOpen}
-                        suggested={suggested}
+                        renderOption={renderOption}
                         value={value}
                         {...props}
                     />
@@ -74,7 +74,7 @@ export function ComboBox<T extends BaseOption>({
                 <div className="mt-4 border-t">
                     <ComboBoxList
                         setOpen={setOpen}
-                        suggested={suggested}
+                        renderOption={renderOption}
                         value={value}
                         {...props}
                     />
@@ -95,12 +95,9 @@ export function ComboBox<T extends BaseOption>({
                 }}
             >
                 <span className="text-ellipsis overflow-hidden">
-                    {value ? value.name : placeholder}{" "}
-                    {value?.name && suggested?.includes(value.name) && (
-                        <span className="text-muted-foreground">
-                            (Suggested)
-                        </span>
-                    )}
+                    {renderOption && value
+                        ? renderOption(value)
+                        : (value?.name ?? placeholder)}
                 </span>
             </Button>
         );
@@ -113,7 +110,7 @@ function ComboBoxList<T extends BaseOption>({
     onChange,
     options,
     setOpen,
-    suggested,
+    renderOption,
     value,
     placeholder = "Filter...",
     emptyLabel = "No results found.",
@@ -160,22 +157,6 @@ function ComboBoxList<T extends BaseOption>({
             [] as (T | GroupedOption<T>)[]
         );
 
-    const sortBySuggested = (
-        a: T | GroupedOption<T>,
-        b: T | GroupedOption<T>
-    ) => {
-        // Handle nested groups - keep original order
-        if ("heading" in a || "heading" in b) return 0;
-
-        const aIsSuggested = a.name && suggested?.includes(a.name);
-        const bIsSuggested = b.name && suggested?.includes(b.name);
-
-        // If both or neither are suggested, maintain original order
-        if (aIsSuggested === bIsSuggested) return 0;
-        // Sort suggested items first
-        return aIsSuggested ? -1 : 1;
-    };
-
     const handleValueChange = (value: string) => {
         setFilteredOptions(filterOptions(options, value));
     };
@@ -215,14 +196,11 @@ function ComboBoxList<T extends BaseOption>({
     function renderGroupedOption(group: GroupedOption<T>) {
         return (
             <CommandGroup key={group.heading} heading={group.heading}>
-                {group.value
-                    .slice() // Create a copy to avoid mutating original array
-                    .sort(sortBySuggested)
-                    .map((option) =>
-                        "heading" in option
-                            ? renderGroupedOption(option)
-                            : renderUngroupedOption(option)
-                    )}
+                {group.value.map((option) =>
+                    "heading" in option
+                        ? renderGroupedOption(option)
+                        : renderUngroupedOption(option)
+                )}
             </CommandGroup>
         );
     }
@@ -239,12 +217,9 @@ function ComboBoxList<T extends BaseOption>({
                 className="justify-between"
             >
                 <span>
-                    {option.name || option.id}{" "}
-                    {option?.name && suggested?.includes(option.name) && (
-                        <span className="text-muted-foreground">
-                            (Suggested)
-                        </span>
-                    )}
+                    {renderOption
+                        ? renderOption(option)
+                        : (option.name ?? option.id)}
                 </span>
                 {value?.id === option.id && <CheckIcon className="w-4 h-4" />}
             </CommandItem>
diff --git a/ui/admin/app/components/model/DefaultModelAliasForm.tsx b/ui/admin/app/components/model/DefaultModelAliasForm.tsx
index e5abfa0b..c4ca5e5c 100644
--- a/ui/admin/app/components/model/DefaultModelAliasForm.tsx
+++ b/ui/admin/app/components/model/DefaultModelAliasForm.tsx
@@ -180,9 +180,15 @@ export function DefaultModelAliasForm({
                                                 }
                                                 options={getOptionsByUsageAndProvider(
                                                     activeModelOptions,
-                                                    usage
+                                                    usage,
+                                                    alias
                                                 )}
-                                                suggested={getSuggested(alias)}
+                                                renderOption={(option) =>
+                                                    renderDisplayOption(
+                                                        option,
+                                                        alias
+                                                    )
+                                                }
                                                 value={
                                                     field.value
                                                         ? models?.find(
@@ -219,21 +225,33 @@ export function DefaultModelAliasForm({
         </Form>
     );
 
-    function getSuggested(alias: ModelAlias) {
-        return alias && SUGGESTED_MODEL_SELECTIONS[alias]
-            ? [SUGGESTED_MODEL_SELECTIONS[alias]]
-            : [];
+    function renderDisplayOption(option: Model, alias: ModelAlias) {
+        const suggestion = alias && SUGGESTED_MODEL_SELECTIONS[alias];
+
+        return (
+            <span>
+                {option.name}{" "}
+                {suggestion === option.name && (
+                    <span className="text-muted-foreground">(Suggested)</span>
+                )}
+            </span>
+        );
     }
 
     function getOptionsByUsageAndProvider(
         modelOptions: Model[] | undefined,
-        usage: ModelUsage
+        usage: ModelUsage,
+        aliasFor: ModelAlias
     ) {
         if (!modelOptions) return [];
 
+        const suggested = aliasFor && SUGGESTED_MODEL_SELECTIONS[aliasFor];
         const usageGroupName = getModelUsageLabel(usage);
-        const usageModelProviderGroups =
-            getModelOptionsByModelProvider(modelOptions);
+        const usageModelProviderGroups = getModelOptionsByModelProvider(
+            modelOptions,
+            suggested ? [suggested] : []
+        );
+
         const otherModelProviderGroups =
             getModelOptionsByModelProvider(otherModels);
         const usageGroup = {
@@ -290,7 +308,10 @@ export function DefaultModelAliasFormDialog({
     );
 }
 
-export function getModelOptionsByModelProvider(models: Model[]) {
+export function getModelOptionsByModelProvider(
+    models: Model[],
+    suggestions?: string[]
+) {
     const byModelProviderGroups = filterModelsByActive(models).reduce(
         (acc, model) => {
             acc[model.modelProvider] = acc[model.modelProvider] || [];
@@ -302,9 +323,17 @@ export function getModelOptionsByModelProvider(models: Model[]) {
 
     return Object.entries(byModelProviderGroups).map(
         ([modelProvider, models]) => {
-            const sorted = models.sort((a, b) =>
-                (a.name ?? "").localeCompare(b.name ?? "")
-            );
+            const sorted = models.sort((a, b) => {
+                const aIsSuggested = a.name && suggestions?.includes(a.name);
+                const bIsSuggested = b.name && suggestions?.includes(b.name);
+
+                if (aIsSuggested === bIsSuggested) return 0;
+                if (aIsSuggested || bIsSuggested) {
+                    return aIsSuggested ? -1 : 1;
+                }
+
+                return (a.name ?? "").localeCompare(b.name ?? "");
+            });
             return {
                 heading: modelProvider,
                 value: sorted,