diff --git a/surfsense_web/components/new-chat/chat-header.tsx b/surfsense_web/components/new-chat/chat-header.tsx index 8a8fa11a0..2f1d9d845 100644 --- a/surfsense_web/components/new-chat/chat-header.tsx +++ b/surfsense_web/components/new-chat/chat-header.tsx @@ -8,7 +8,6 @@ import type { NewLLMConfigPublic, } from "@/contracts/types/new-llm-config.types"; import { ImageConfigSidebar } from "./image-config-sidebar"; -import { ImageModelSelector } from "./image-model-selector"; import { ModelConfigSidebar } from "./model-config-sidebar"; import { ModelSelector } from "./model-selector"; @@ -34,7 +33,7 @@ export function ChatHeader({ searchSpaceId }: ChatHeaderProps) { const [imageSidebarMode, setImageSidebarMode] = useState<"create" | "edit" | "view">("view"); // LLM handlers - const handleEditConfig = useCallback( + const handleEditLLMConfig = useCallback( (config: NewLLMConfigPublic | GlobalNewLLMConfig, global: boolean) => { setSelectedConfig(config); setIsGlobal(global); @@ -44,7 +43,7 @@ export function ChatHeader({ searchSpaceId }: ChatHeaderProps) { [] ); - const handleAddNew = useCallback(() => { + const handleAddNewLLM = useCallback(() => { setSelectedConfig(null); setIsGlobal(false); setSidebarMode("create"); @@ -81,8 +80,12 @@ export function ChatHeader({ searchSpaceId }: ChatHeaderProps) { return (
- - + void; - onEdit?: (config: ImageGenerationConfig | GlobalImageGenConfig, isGlobal: boolean) => void; -} - -export function ImageModelSelector({ className, onAddNew, onEdit }: ImageModelSelectorProps) { - const [open, setOpen] = useState(false); - const [searchQuery, setSearchQuery] = useState(""); - - const { data: globalConfigs, isLoading: globalLoading } = useAtomValue(globalImageGenConfigsAtom); - const { data: userConfigs, isLoading: userLoading } = useAtomValue(imageGenConfigsAtom); - const { data: preferences, isLoading: prefsLoading } = useAtomValue(llmPreferencesAtom); - const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom); - const { mutateAsync: updatePreferences } = useAtomValue(updateLLMPreferencesMutationAtom); - - const isLoading = globalLoading || userLoading || prefsLoading; - - const currentConfig = useMemo(() => { - if (!preferences) return null; - const id = preferences.image_generation_config_id; - if (id === null || id === undefined) return null; - const globalMatch = globalConfigs?.find((c) => c.id === id); - if (globalMatch) return globalMatch; - return userConfigs?.find((c) => c.id === id) ?? null; - }, [preferences, globalConfigs, userConfigs]); - - const isCurrentAutoMode = useMemo(() => { - return currentConfig && "is_auto_mode" in currentConfig && currentConfig.is_auto_mode; - }, [currentConfig]); - - const filteredGlobal = useMemo(() => { - if (!globalConfigs) return []; - if (!searchQuery) return globalConfigs; - const q = searchQuery.toLowerCase(); - return globalConfigs.filter( - (c) => - c.name.toLowerCase().includes(q) || - c.model_name.toLowerCase().includes(q) || - c.provider.toLowerCase().includes(q) - ); - }, [globalConfigs, searchQuery]); - - const filteredUser = useMemo(() => { - if (!userConfigs) return []; - if (!searchQuery) return userConfigs; - const q = searchQuery.toLowerCase(); - return userConfigs.filter( - (c) => - c.name.toLowerCase().includes(q) || - c.model_name.toLowerCase().includes(q) || - c.provider.toLowerCase().includes(q) - ); - }, [userConfigs, searchQuery]); - - const totalModels = (globalConfigs?.length ?? 0) + (userConfigs?.length ?? 0); - - const handleSelect = useCallback( - async (configId: number) => { - if (currentConfig?.id === configId) { - setOpen(false); - return; - } - if (!searchSpaceId) { - toast.error("No search space selected"); - return; - } - try { - await updatePreferences({ - search_space_id: Number(searchSpaceId), - data: { image_generation_config_id: configId }, - }); - toast.success("Image model updated"); - setOpen(false); - } catch { - toast.error("Failed to switch image model"); - } - }, - [currentConfig, searchSpaceId, updatePreferences] - ); - - // Don't render if no configs at all - if (!isLoading && totalModels === 0) { - return ( - - ); - } - - return ( - - - - - - - - {totalModels > 3 && ( -
- -
- )} - - -
- -

No image models found

-
- - - {/* Global Image Gen Configs */} - {filteredGlobal.length > 0 && ( - -
- - Global Image Models -
- {filteredGlobal.map((config) => { - const isSelected = currentConfig?.id === config.id; - const isAuto = "is_auto_mode" in config && config.is_auto_mode; - return ( - handleSelect(config.id)} - className={cn( - "mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50", - isSelected && "bg-accent/80", - isAuto && "border border-violet-200 dark:border-violet-800/50" - )} - > -
-
- {isAuto ? ( - - ) : ( - - )} -
-
-
- {config.name} - {isAuto && ( - - Recommended - - )} - {isSelected && } -
- - {isAuto ? "Auto load balancing" : config.model_name} - -
- {onEdit && ( - { - e.stopPropagation(); - setOpen(false); - onEdit(config, true); - }} - /> - )} -
- - ); - })} - - )} - - {/* User Image Gen Configs */} - {filteredUser.length > 0 && ( - <> - {filteredGlobal.length > 0 && } - -
- - Your Image Models -
- {filteredUser.map((config) => { - const isSelected = currentConfig?.id === config.id; - return ( - handleSelect(config.id)} - className={cn( - "mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50", - isSelected && "bg-accent/80" - )} - > -
-
- -
-
-
- {config.name} - {isSelected && } -
- - {config.model_name} - -
- {onEdit && ( - - )} -
- - ); - })} - - - )} - - {/* Add New */} - {onAddNew && ( -
- -
- )} - - - - - ); -} diff --git a/surfsense_web/components/new-chat/model-selector.tsx b/surfsense_web/components/new-chat/model-selector.tsx index 988dc7209..d27594ee6 100644 --- a/surfsense_web/components/new-chat/model-selector.tsx +++ b/surfsense_web/components/new-chat/model-selector.tsx @@ -6,10 +6,16 @@ import { Check, ChevronDown, Edit3, + ImageIcon, Plus, + Zap, } from "lucide-react"; import { useCallback, useMemo, useState } from "react"; import { toast } from "sonner"; +import { + globalImageGenConfigsAtom, + imageGenConfigsAtom, +} from "@/atoms/image-gen-config/image-gen-config-query.atoms"; import { updateLLMPreferencesMutationAtom } from "@/atoms/new-llm-config/new-llm-config-mutation.atoms"; import { globalNewLLMConfigsAtom, @@ -30,102 +36,150 @@ import { } from "@/components/ui/command"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { Spinner } from "@/components/ui/spinner"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import type { + GlobalImageGenConfig, GlobalNewLLMConfig, + ImageGenerationConfig, NewLLMConfigPublic, } from "@/contracts/types/new-llm-config.types"; -import { cn } from "@/lib/utils"; import { getProviderIcon } from "@/lib/provider-icons"; +import { cn } from "@/lib/utils"; interface ModelSelectorProps { - onEdit: (config: NewLLMConfigPublic | GlobalNewLLMConfig, isGlobal: boolean) => void; - onAddNew: () => void; + onEditLLM: (config: NewLLMConfigPublic | GlobalNewLLMConfig, isGlobal: boolean) => void; + onAddNewLLM: () => void; + onEditImage?: (config: ImageGenerationConfig | GlobalImageGenConfig, isGlobal: boolean) => void; + onAddNewImage?: () => void; className?: string; } -export function ModelSelector({ onEdit, onAddNew, className }: ModelSelectorProps) { +export function ModelSelector({ + onEditLLM, + onAddNewLLM, + onEditImage, + onAddNewImage, + className, +}: ModelSelectorProps) { const [open, setOpen] = useState(false); - const [searchQuery, setSearchQuery] = useState(""); + const [activeTab, setActiveTab] = useState<"llm" | "image">("llm"); + const [llmSearchQuery, setLlmSearchQuery] = useState(""); + const [imageSearchQuery, setImageSearchQuery] = useState(""); - // Fetch configs - const { data: userConfigs, isLoading: userConfigsLoading } = useAtomValue(newLLMConfigsAtom); - const { data: globalConfigs, isLoading: globalConfigsLoading } = + // LLM data + const { data: llmUserConfigs, isLoading: llmUserLoading } = useAtomValue(newLLMConfigsAtom); + const { data: llmGlobalConfigs, isLoading: llmGlobalLoading } = useAtomValue(globalNewLLMConfigsAtom); - const { data: preferences, isLoading: preferencesLoading } = useAtomValue(llmPreferencesAtom); + const { data: preferences, isLoading: prefsLoading } = useAtomValue(llmPreferencesAtom); const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom); const { mutateAsync: updatePreferences } = useAtomValue(updateLLMPreferencesMutationAtom); - const isLoading = userConfigsLoading || globalConfigsLoading || preferencesLoading; + // Image data + const { data: imageGlobalConfigs, isLoading: imageGlobalLoading } = + useAtomValue(globalImageGenConfigsAtom); + const { data: imageUserConfigs, isLoading: imageUserLoading } = + useAtomValue(imageGenConfigsAtom); - // Get current agent LLM config - const currentConfig = useMemo(() => { + const isLoading = llmUserLoading || llmGlobalLoading || prefsLoading || imageGlobalLoading || imageUserLoading; + + // ─── LLM current config ─── + const currentLLMConfig = useMemo(() => { if (!preferences) return null; - const agentLlmId = preferences.agent_llm_id; if (agentLlmId === null || agentLlmId === undefined) return null; - - // Check if it's Auto mode (ID 0) or global config (negative ID) if (agentLlmId <= 0) { - return globalConfigs?.find((c) => c.id === agentLlmId) ?? null; + return llmGlobalConfigs?.find((c) => c.id === agentLlmId) ?? null; } - // Otherwise, check user configs - return userConfigs?.find((c) => c.id === agentLlmId) ?? null; - }, [preferences, globalConfigs, userConfigs]); + return llmUserConfigs?.find((c) => c.id === agentLlmId) ?? null; + }, [preferences, llmGlobalConfigs, llmUserConfigs]); - // Check if current config is Auto mode - const isCurrentAutoMode = useMemo(() => { - return currentConfig && "is_auto_mode" in currentConfig && currentConfig.is_auto_mode; - }, [currentConfig]); + const isLLMAutoMode = useMemo(() => { + return currentLLMConfig && "is_auto_mode" in currentLLMConfig && currentLLMConfig.is_auto_mode; + }, [currentLLMConfig]); - // Filter configs based on search - const filteredGlobalConfigs = useMemo(() => { - if (!globalConfigs) return []; - if (!searchQuery) return globalConfigs; - const query = searchQuery.toLowerCase(); - return globalConfigs.filter( + // ─── Image current config ─── + const currentImageConfig = useMemo(() => { + if (!preferences) return null; + const id = preferences.image_generation_config_id; + if (id === null || id === undefined) return null; + const globalMatch = imageGlobalConfigs?.find((c) => c.id === id); + if (globalMatch) return globalMatch; + return imageUserConfigs?.find((c) => c.id === id) ?? null; + }, [preferences, imageGlobalConfigs, imageUserConfigs]); + + const isImageAutoMode = useMemo(() => { + return currentImageConfig && "is_auto_mode" in currentImageConfig && currentImageConfig.is_auto_mode; + }, [currentImageConfig]); + + // ─── LLM filtering ─── + const filteredLLMGlobal = useMemo(() => { + if (!llmGlobalConfigs) return []; + if (!llmSearchQuery) return llmGlobalConfigs; + const q = llmSearchQuery.toLowerCase(); + return llmGlobalConfigs.filter( (c) => - c.name.toLowerCase().includes(query) || - c.model_name.toLowerCase().includes(query) || - c.provider.toLowerCase().includes(query) + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) ); - }, [globalConfigs, searchQuery]); + }, [llmGlobalConfigs, llmSearchQuery]); - const filteredUserConfigs = useMemo(() => { - if (!userConfigs) return []; - if (!searchQuery) return userConfigs; - const query = searchQuery.toLowerCase(); - return userConfigs.filter( + const filteredLLMUser = useMemo(() => { + if (!llmUserConfigs) return []; + if (!llmSearchQuery) return llmUserConfigs; + const q = llmSearchQuery.toLowerCase(); + return llmUserConfigs.filter( (c) => - c.name.toLowerCase().includes(query) || - c.model_name.toLowerCase().includes(query) || - c.provider.toLowerCase().includes(query) + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) ); - }, [userConfigs, searchQuery]); + }, [llmUserConfigs, llmSearchQuery]); - // Total model count for conditional search display - const totalModels = useMemo(() => { - return (globalConfigs?.length ?? 0) + (userConfigs?.length ?? 0); - }, [globalConfigs, userConfigs]); + const totalLLMModels = (llmGlobalConfigs?.length ?? 0) + (llmUserConfigs?.length ?? 0); - const handleSelectConfig = useCallback( + // ─── Image filtering ─── + const filteredImageGlobal = useMemo(() => { + if (!imageGlobalConfigs) return []; + if (!imageSearchQuery) return imageGlobalConfigs; + const q = imageSearchQuery.toLowerCase(); + return imageGlobalConfigs.filter( + (c) => + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) + ); + }, [imageGlobalConfigs, imageSearchQuery]); + + const filteredImageUser = useMemo(() => { + if (!imageUserConfigs) return []; + if (!imageSearchQuery) return imageUserConfigs; + const q = imageSearchQuery.toLowerCase(); + return imageUserConfigs.filter( + (c) => + c.name.toLowerCase().includes(q) || + c.model_name.toLowerCase().includes(q) || + c.provider.toLowerCase().includes(q) + ); + }, [imageUserConfigs, imageSearchQuery]); + + const totalImageModels = (imageGlobalConfigs?.length ?? 0) + (imageUserConfigs?.length ?? 0); + + // ─── Handlers ─── + const handleSelectLLM = useCallback( async (config: NewLLMConfigPublic | GlobalNewLLMConfig) => { - // If already selected, just close - if (currentConfig?.id === config.id) { + if (currentLLMConfig?.id === config.id) { setOpen(false); return; } - if (!searchSpaceId) { toast.error("No search space selected"); return; } - try { await updatePreferences({ search_space_id: Number(searchSpaceId), - data: { - agent_llm_id: config.id, - }, + data: { agent_llm_id: config.id }, }); toast.success(`Switched to ${config.name}`); setOpen(false); @@ -134,16 +188,40 @@ export function ModelSelector({ onEdit, onAddNew, className }: ModelSelectorProp toast.error("Failed to switch model"); } }, - [currentConfig, searchSpaceId, updatePreferences] + [currentLLMConfig, searchSpaceId, updatePreferences] ); - const handleEditConfig = useCallback( + const handleEditLLMConfig = useCallback( (e: React.MouseEvent, config: NewLLMConfigPublic | GlobalNewLLMConfig, isGlobal: boolean) => { e.stopPropagation(); - onEdit(config, isGlobal); + onEditLLM(config, isGlobal); setOpen(false); }, - [onEdit] + [onEditLLM] + ); + + const handleSelectImage = useCallback( + async (configId: number) => { + if (currentImageConfig?.id === configId) { + setOpen(false); + return; + } + if (!searchSpaceId) { + toast.error("No search space selected"); + return; + } + try { + await updatePreferences({ + search_space_id: Number(searchSpaceId), + data: { image_generation_config_id: configId }, + }); + toast.success("Image model updated"); + setOpen(false); + } catch { + toast.error("Failed to switch image model"); + } + }, + [currentImageConfig, searchSpaceId, updatePreferences] ); return ( @@ -161,30 +239,37 @@ export function ModelSelector({ onEdit, onAddNew, className }: ModelSelectorProp Loading - ) : currentConfig ? ( - <> - {getProviderIcon(currentConfig.provider, { isAutoMode: isCurrentAutoMode ?? false })} - - {currentConfig.name} - - {isCurrentAutoMode ? ( - - Balanced - - ) : ( - - {currentConfig.model_name.split("/").pop()?.slice(0, 10) || - currentConfig.model_name.slice(0, 10)} - - )} - ) : ( <> - - Select Model + {/* LLM section */} + {currentLLMConfig ? ( + <> + {getProviderIcon(currentLLMConfig.provider, { isAutoMode: isLLMAutoMode ?? false })} + + {currentLLMConfig.name} + + + ) : ( + <> + + Select Model + + )} + + {/* Divider */} +
+ + {/* Image section */} + {currentImageConfig ? ( + <> + {getProviderIcon(currentImageConfig.provider, { isAutoMode: isImageAutoMode ?? false })} + + {currentImageConfig.name} + + + ) : ( + + )} )} - setActiveTab(v as "llm" | "image")} + className="w-full" > - {totalModels > 3 && ( -
- -
- )} - - - -
- -

No models found

-

Try a different search term

-
- - - {/* Global Configs Section */} - {filteredGlobalConfigs.length > 0 && ( - -
- Global Models -
- {filteredGlobalConfigs.map((config) => { - const isSelected = currentConfig?.id === config.id; - const isAutoMode = "is_auto_mode" in config && config.is_auto_mode; - return ( - handleSelectConfig(config)} - className={cn( - "mx-2 rounded-lg mb-1 cursor-pointer group transition-all", - "hover:bg-accent/50", - isSelected && "bg-accent/80", - isAutoMode && "border border-violet-200 dark:border-violet-800/50" - )} - > -
-
-
- {getProviderIcon(config.provider, { isAutoMode })} -
-
-
- {config.name} - {isAutoMode && ( - - Recommended - - )} - {isSelected && } -
-
- - {isAutoMode ? "Auto load balancing" : config.model_name} - - {!isAutoMode && config.citations_enabled && ( - - Citations - - )} -
-
-
- {!isAutoMode && ( - - )} -
- - ); - })} - - )} - - {filteredGlobalConfigs.length > 0 && filteredUserConfigs.length > 0 && ( - - )} - - {/* User Configs Section */} - {filteredUserConfigs.length > 0 && ( - -
- Your Configurations -
- {filteredUserConfigs.map((config) => { - const isSelected = currentConfig?.id === config.id; - return ( - handleSelectConfig(config)} - className={cn( - "mx-2 rounded-lg mb-1 cursor-pointer group transition-all", - "hover:bg-accent/50", - isSelected && "bg-accent/80" - )} - > -
-
-
{getProviderIcon(config.provider)}
-
-
- {config.name} - {isSelected && } -
-
- - {config.model_name} - - {config.citations_enabled && ( - - Citations - - )} -
-
-
- -
- - ); - })} - - )} - - {/* Add New Config Button */} -
- -
- - + + LLM + + + + Image + + +
+ + {/* ─── LLM Tab ─── */} + + + {totalLLMModels > 3 && ( +
+ +
+ )} + + + +
+ +

No models found

+

Try a different search term

+
+ + + {/* Global LLM Configs */} + {filteredLLMGlobal.length > 0 && ( + +
+ Global Models +
+ {filteredLLMGlobal.map((config) => { + const isSelected = currentLLMConfig?.id === config.id; + const isAutoMode = "is_auto_mode" in config && config.is_auto_mode; + return ( + handleSelectLLM(config)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all", + "hover:bg-accent/50", + isSelected && "bg-accent/80", + isAutoMode && "border border-violet-800" + )} + > +
+
+
+ {getProviderIcon(config.provider, { isAutoMode })} +
+
+
+ {config.name} + {isAutoMode && ( + + Recommended + + )} + {isSelected && } +
+
+ + {isAutoMode ? "Auto load balancing" : config.model_name} + + {!isAutoMode && config.citations_enabled && ( + + Citations + + )} +
+
+
+ {!isAutoMode && ( + + )} +
+ + ); + })} + + )} + + {filteredLLMGlobal.length > 0 && filteredLLMUser.length > 0 && ( + + )} + + {/* User LLM Configs */} + {filteredLLMUser.length > 0 && ( + +
+ Your Configurations +
+ {filteredLLMUser.map((config) => { + const isSelected = currentLLMConfig?.id === config.id; + return ( + handleSelectLLM(config)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all", + "hover:bg-accent/50", + isSelected && "bg-accent/80" + )} + > +
+
+
{getProviderIcon(config.provider)}
+
+
+ {config.name} + {isSelected && } +
+
+ + {config.model_name} + + {config.citations_enabled && ( + + Citations + + )} +
+
+
+ +
+ + ); + })} + + )} + + {/* Add New LLM Config */} +
+ +
+ + + + + {/* ─── Image Tab ─── */} + + + {totalImageModels > 3 && ( +
+ +
+ )} + + +
+ +

No image models found

+
+ + + {/* Global Image Configs */} + {filteredImageGlobal.length > 0 && ( + +
+ Global Image Models +
+ {filteredImageGlobal.map((config) => { + const isSelected = currentImageConfig?.id === config.id; + const isAuto = "is_auto_mode" in config && config.is_auto_mode; + return ( + handleSelectImage(config.id)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50", + isSelected && "bg-accent/80", + isAuto && "border border-violet-800" + )} + > +
+
+ {getProviderIcon(config.provider, { isAutoMode: isAuto })} +
+
+
+ {config.name} + {isAuto && ( + + Recommended + + )} + {isSelected && } +
+ + {isAuto ? "Auto load balancing" : config.model_name} + +
+ {onEditImage && !isAuto && ( + + )} +
+ + ); + })} + + )} + + {/* User Image Configs */} + {filteredImageUser.length > 0 && ( + <> + {filteredImageGlobal.length > 0 && } + +
+ Your Image Models +
+ {filteredImageUser.map((config) => { + const isSelected = currentImageConfig?.id === config.id; + return ( + handleSelectImage(config.id)} + className={cn( + "mx-2 rounded-lg mb-1 cursor-pointer group transition-all hover:bg-accent/50", + isSelected && "bg-accent/80" + )} + > +
+
+ {getProviderIcon(config.provider)} +
+
+
+ {config.name} + {isSelected && } +
+ + {config.model_name} + +
+ {onEditImage && ( + + )} +
+ + ); + })} + + + )} + + {/* Add New Image Config */} + {onAddNewImage && ( +
+ +
+ )} + + + + ); diff --git a/surfsense_web/lib/provider-icons.tsx b/surfsense_web/lib/provider-icons.tsx index ce474c5a9..4a32f0df0 100644 --- a/surfsense_web/lib/provider-icons.tsx +++ b/surfsense_web/lib/provider-icons.tsx @@ -47,7 +47,7 @@ export function getProviderIcon( }: { isAutoMode?: boolean; className?: string } = {} ) { if (isAutoMode || provider?.toUpperCase() === "AUTO") { - return ; + return ; } switch (provider?.toUpperCase()) {