From d2cf3fb3b7eafce01c881e80049346239c84bc6f Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Thu, 2 Apr 2026 09:58:14 +0530 Subject: [PATCH] refactor: update LLM role management logic and enhance UI feedback --- .../new-llm-config-mutation.atoms.ts | 9 +- .../components/settings/llm-role-manager.tsx | 165 +++++------------- 2 files changed, 52 insertions(+), 122 deletions(-) diff --git a/surfsense_web/atoms/new-llm-config/new-llm-config-mutation.atoms.ts b/surfsense_web/atoms/new-llm-config/new-llm-config-mutation.atoms.ts index 861606f80..d6d3aa820 100644 --- a/surfsense_web/atoms/new-llm-config/new-llm-config-mutation.atoms.ts +++ b/surfsense_web/atoms/new-llm-config/new-llm-config-mutation.atoms.ts @@ -109,10 +109,11 @@ export const updateLLMPreferencesMutationAtom = atomWithMutation((get) => { mutationFn: async (request: UpdateLLMPreferencesRequest) => { return newLLMConfigApiService.updateLLMPreferences(request); }, - onSuccess: () => { - queryClient.invalidateQueries({ - queryKey: cacheKeys.newLLMConfigs.preferences(Number(searchSpaceId)), - }); + onSuccess: (_data, request: UpdateLLMPreferencesRequest) => { + queryClient.setQueryData( + cacheKeys.newLLMConfigs.preferences(Number(searchSpaceId)), + (old: Record | undefined) => ({ ...old, ...request.data }) + ); }, onError: (error: Error) => { toast.error(error.message || "Failed to update LLM preferences"); diff --git a/surfsense_web/components/settings/llm-role-manager.tsx b/surfsense_web/components/settings/llm-role-manager.tsx index 07ec492a3..22e17e431 100644 --- a/surfsense_web/components/settings/llm-role-manager.tsx +++ b/surfsense_web/components/settings/llm-role-manager.tsx @@ -4,16 +4,14 @@ import { useAtomValue } from "jotai"; import { AlertCircle, Bot, - CheckCircle, + CircleCheck, CircleDashed, FileText, ImageIcon, RefreshCw, - RotateCcw, - Save, Shuffle, } from "lucide-react"; -import { useEffect, useState } from "react"; +import { useCallback, useEffect, useRef, useState } from "react"; import { toast } from "sonner"; import { globalImageGenConfigsAtom, @@ -40,6 +38,7 @@ import { SelectValue, } from "@/components/ui/select"; import { Skeleton } from "@/components/ui/skeleton"; +import { Spinner } from "@/components/ui/spinner"; import { getProviderIcon } from "@/lib/provider-icons"; import { cn } from "@/lib/utils"; @@ -48,8 +47,8 @@ const ROLE_DESCRIPTIONS = { icon: Bot, title: "Agent LLM", description: "Primary LLM for chat interactions and agent operations", - color: "text-blue-600 dark:text-blue-400", - bgColor: "bg-blue-500/10", + color: "text-muted-foreground", + bgColor: "bg-muted", prefKey: "agent_llm_id" as const, configType: "llm" as const, }, @@ -57,8 +56,8 @@ const ROLE_DESCRIPTIONS = { icon: FileText, title: "Document Summary LLM", description: "Handles document summarization and research synthesis", - color: "text-purple-600 dark:text-purple-400", - bgColor: "bg-purple-500/10", + color: "text-muted-foreground", + bgColor: "bg-muted", prefKey: "document_summary_llm_id" as const, configType: "llm" as const, }, @@ -66,8 +65,8 @@ const ROLE_DESCRIPTIONS = { icon: ImageIcon, title: "Image Generation Model", description: "Model used for AI image generation (DALL-E, GPT Image, etc.)", - color: "text-teal-600 dark:text-teal-400", - bgColor: "bg-teal-500/10", + color: "text-muted-foreground", + bgColor: "bg-muted", prefKey: "image_generation_config_id" as const, configType: "image" as const, }, @@ -118,88 +117,41 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) { image_generation_config_id: preferences.image_generation_config_id ?? "", })); - const [hasChanges, setHasChanges] = useState(false); - const [isSaving, setIsSaving] = useState(false); + const [savingRole, setSavingRole] = useState(null); + const savingRef = useRef(false); useEffect(() => { - const newAssignments = { - agent_llm_id: preferences.agent_llm_id ?? "", - document_summary_llm_id: preferences.document_summary_llm_id ?? "", - image_generation_config_id: preferences.image_generation_config_id ?? "", - }; - setAssignments(newAssignments); - setHasChanges(false); + if (!savingRef.current) { + setAssignments({ + agent_llm_id: preferences.agent_llm_id ?? "", + document_summary_llm_id: preferences.document_summary_llm_id ?? "", + image_generation_config_id: preferences.image_generation_config_id ?? "", + }); + } }, [ preferences?.agent_llm_id, preferences?.document_summary_llm_id, preferences?.image_generation_config_id, ]); - const handleRoleAssignment = (prefKey: string, configId: string) => { - const newAssignments = { - ...assignments, - [prefKey]: configId === "unassigned" ? "" : parseInt(configId), - }; + const handleRoleAssignment = useCallback(async (prefKey: string, configId: string) => { + const value = configId === "unassigned" ? "" : parseInt(configId); - setAssignments(newAssignments); + setAssignments((prev) => ({ ...prev, [prefKey]: value })); + setSavingRole(prefKey); + savingRef.current = true; - const currentPrefs = { - agent_llm_id: preferences.agent_llm_id ?? "", - document_summary_llm_id: preferences.document_summary_llm_id ?? "", - image_generation_config_id: preferences.image_generation_config_id ?? "", - }; - - const hasChangesNow = Object.keys(newAssignments).some( - (key) => - newAssignments[key as keyof typeof newAssignments] !== - currentPrefs[key as keyof typeof currentPrefs] - ); - - setHasChanges(hasChangesNow); - }; - - const handleSave = async () => { - setIsSaving(true); - - const toNumericOrUndefined = (val: string | number) => - typeof val === "string" ? (val ? parseInt(val) : undefined) : val; - - const numericAssignments = { - agent_llm_id: toNumericOrUndefined(assignments.agent_llm_id), - document_summary_llm_id: toNumericOrUndefined(assignments.document_summary_llm_id), - image_generation_config_id: toNumericOrUndefined(assignments.image_generation_config_id), - }; - - await updatePreferences({ - search_space_id: searchSpaceId, - data: numericAssignments, - }); - - setHasChanges(false); - toast.success("Role assignments saved successfully!"); - - setIsSaving(false); - }; - - const handleReset = () => { - setAssignments({ - agent_llm_id: preferences.agent_llm_id ?? "", - document_summary_llm_id: preferences.document_summary_llm_id ?? "", - image_generation_config_id: preferences.image_generation_config_id ?? "", - }); - setHasChanges(false); - }; - - const isAssignmentComplete = - assignments.agent_llm_id !== "" && - assignments.agent_llm_id !== null && - assignments.agent_llm_id !== undefined && - assignments.document_summary_llm_id !== "" && - assignments.document_summary_llm_id !== null && - assignments.document_summary_llm_id !== undefined && - assignments.image_generation_config_id !== "" && - assignments.image_generation_config_id !== null && - assignments.image_generation_config_id !== undefined; + try { + await updatePreferences({ + search_space_id: searchSpaceId, + data: { [prefKey]: value || undefined }, + }); + toast.success("Role assignment updated"); + } finally { + setSavingRole(null); + savingRef.current = false; + } + }, [updatePreferences, searchSpaceId]); // Combine global and custom LLM configs const allLLMConfigs = [ @@ -213,6 +165,11 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) { ...(userImageConfigs ?? []).filter((config) => config.id && config.id.toString().trim() !== ""), ]; + const isAssignmentComplete = + allLLMConfigs.some((c) => c.id === assignments.agent_llm_id) && + allLLMConfigs.some((c) => c.id === assignments.document_summary_llm_id) && + allImageConfigs.some((c) => c.id === assignments.image_generation_config_id); + const isLoading = configsLoading || preferencesLoading || @@ -244,9 +201,9 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) { {isAssignmentComplete && !isLoading && !hasError && ( - + All roles assigned )} @@ -332,10 +289,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) { const roleAllConfigs = isImageRole ? allImageConfigs : allLLMConfigs; const assignedConfig = roleAllConfigs.find((config) => config.id === currentAssignment); - const isAssigned = - currentAssignment !== "" && - currentAssignment !== null && - currentAssignment !== undefined; + const isAssigned = !!assignedConfig; const isAutoMode = assignedConfig && "is_auto_mode" in assignedConfig && assignedConfig.is_auto_mode; @@ -361,8 +315,10 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {

- {isAssigned ? ( - + {savingRole === role.prefKey ? ( + + ) : isAssigned ? ( + ) : ( )} @@ -374,7 +330,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) { Configuration