feat: update LLM role manager to include image generation model preferences and improve loading/error handling

This commit is contained in:
Anish Sarkar 2026-02-10 18:09:27 +05:30
parent 7557f5d2be
commit 4d7132c16c
2 changed files with 94 additions and 233 deletions

View file

@ -6,18 +6,16 @@ import {
Check,
ChevronsUpDown,
Edit3,
ImageIcon,
Key,
Plus,
RefreshCw,
Shuffle,
Info,
Trash2,
Wand2,
} from "lucide-react";
import { AnimatePresence, motion } from "motion/react";
import Image from "next/image";
import { useCallback, useEffect, useMemo, useState } from "react";
import { useCallback, useMemo, useState } from "react";
import { toast } from "sonner";
import { membersAtom, myAccessAtom } from "@/atoms/members/members-query.atoms";
import {
@ -30,7 +28,6 @@ import {
imageGenConfigsAtom,
} from "@/atoms/image-gen-config/image-gen-config-query.atoms";
import { updateLLMPreferencesMutationAtom } from "@/atoms/new-llm-config/new-llm-config-mutation.atoms";
import { llmPreferencesAtom } from "@/atoms/new-llm-config/new-llm-config-query.atoms";
import { Alert, AlertDescription } from "@/components/ui/alert";
import {
AlertDialog,
@ -42,9 +39,8 @@ import {
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
import { Card, CardContent } from "@/components/ui/card";
import {
Command,
CommandEmpty,
@ -131,7 +127,6 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) {
} = useAtomValue(imageGenConfigsAtom);
const { data: globalConfigs = [], isFetching: globalLoading } =
useAtomValue(globalImageGenConfigsAtom);
const { data: preferences = {}, isFetching: prefsLoading } = useAtomValue(llmPreferencesAtom);
// Members for user resolution
const { data: members } = useAtomValue(membersAtom);
@ -170,20 +165,8 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) {
const [editingConfig, setEditingConfig] = useState<ImageGenerationConfig | null>(null);
const [configToDelete, setConfigToDelete] = useState<ImageGenerationConfig | null>(null);
// Preference state
const [selectedPrefId, setSelectedPrefId] = useState<string | number>(
preferences.image_generation_config_id ?? ""
);
const [hasPrefChanges, setHasPrefChanges] = useState(false);
const [isSavingPref, setIsSavingPref] = useState(false);
useEffect(() => {
setSelectedPrefId(preferences.image_generation_config_id ?? "");
setHasPrefChanges(false);
}, [preferences]);
const isSubmitting = isCreating || isUpdating;
const isLoading = configsLoading || globalLoading || prefsLoading;
const isLoading = configsLoading || globalLoading;
const errors = [createError, updateError, deleteError, fetchError].filter(Boolean) as Error[];
// Form state for create/edit dialog
@ -291,40 +274,6 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) {
setIsDialogOpen(true);
};
const handlePrefChange = (value: string) => {
const newVal = value === "unassigned" ? "" : parseInt(value);
setSelectedPrefId(newVal);
setHasPrefChanges(newVal !== (preferences.image_generation_config_id ?? ""));
};
const handleSavePref = async () => {
setIsSavingPref(true);
try {
await updatePreferences({
search_space_id: searchSpaceId,
data: {
image_generation_config_id:
typeof selectedPrefId === "string"
? selectedPrefId
? parseInt(selectedPrefId)
: undefined
: selectedPrefId,
},
});
setHasPrefChanges(false);
toast.success("Image generation model preference saved!");
} catch {
toast.error("Failed to save preference");
} finally {
setIsSavingPref(false);
}
};
const allConfigs = [
...globalConfigs.map((c) => ({ ...c, _source: "global" as const })),
...(userConfigs ?? []).map((c) => ({ ...c, _source: "user" as const })),
];
const selectedProvider = IMAGE_GEN_PROVIDERS.find((p) => p.value === formData.provider);
const suggestedModels = getImageGenModelsByProvider(formData.provider);
@ -342,6 +291,14 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) {
<RefreshCw className={cn("h-3 w-3 md:h-4 md:w-4", configsLoading && "animate-spin")} />
Refresh
</Button>
{canCreate && (
<Button
onClick={openNewDialog}
className="flex items-center gap-2 text-xs md:text-sm h-8 md:h-9"
>
Add Image Model
</Button>
)}
</div>
{/* Errors */}
@ -403,135 +360,9 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) {
</Alert>
)}
{/* Active Preference Card */}
{!isLoading && allConfigs.length > 0 && (
<motion.div initial={{ opacity: 0, y: 10 }} animate={{ opacity: 1, y: 0 }}>
<Card className="border-l-4 border-l-teal-500">
<CardHeader className="pb-2 px-3 md:px-6 pt-3 md:pt-6">
<div className="flex items-center gap-2 md:gap-3">
<div className="p-1.5 md:p-2 rounded-lg bg-teal-100 text-teal-800">
<ImageIcon className="w-4 h-4 md:w-5 md:h-5" />
</div>
<div>
<CardTitle className="text-base md:text-lg">Active Image Model</CardTitle>
<CardDescription className="text-xs md:text-sm">
Select which model to use for image generation
</CardDescription>
</div>
</div>
</CardHeader>
<CardContent className="space-y-3 px-3 md:px-6 pb-3 md:pb-6">
<Select
value={selectedPrefId?.toString() || "unassigned"}
onValueChange={handlePrefChange}
>
<SelectTrigger className="h-9 md:h-10 text-xs md:text-sm">
<SelectValue placeholder="Select an image model" />
</SelectTrigger>
<SelectContent>
<SelectItem value="unassigned">
<span className="text-muted-foreground">Unassigned</span>
</SelectItem>
{globalConfigs.length > 0 && (
<>
<div className="px-2 py-1.5 text-xs font-semibold text-muted-foreground">
Global
</div>
{globalConfigs.map((c) => {
const isAuto = "is_auto_mode" in c && c.is_auto_mode;
return (
<SelectItem key={`g-${c.id}`} value={c.id.toString()}>
<div className="flex items-center gap-2">
{isAuto ? (
<Badge
variant="outline"
className="text-xs bg-violet-100 text-violet-700 dark:bg-violet-900/30 dark:text-violet-300 border-violet-200"
>
<Shuffle className="size-3 mr-1" />
AUTO
</Badge>
) : (
<Badge
variant="outline"
className="text-xs bg-teal-50 text-teal-700 dark:bg-teal-900/30 dark:text-teal-300 border-teal-200"
>
{c.provider}
</Badge>
)}
<span>{c.name}</span>
</div>
</SelectItem>
);
})}
</>
)}
{(userConfigs?.length ?? 0) > 0 && (
<>
<div className="px-2 py-1.5 text-xs font-semibold text-muted-foreground">
Your Models
</div>
{userConfigs?.map((c) => (
<SelectItem key={`u-${c.id}`} value={c.id.toString()}>
<div className="flex items-center gap-2">
<Badge variant="outline" className="text-xs">
{c.provider}
</Badge>
<span>{c.name}</span>
<span className="text-muted-foreground">({c.model_name})</span>
</div>
</SelectItem>
))}
</>
)}
</SelectContent>
</Select>
{hasPrefChanges && (
<div className="flex gap-2 pt-1">
<Button
size="sm"
onClick={handleSavePref}
disabled={isSavingPref}
className="text-xs h-8"
>
{isSavingPref ? "Saving..." : "Save"}
</Button>
<Button
size="sm"
variant="outline"
onClick={() => {
setSelectedPrefId(preferences.image_generation_config_id ?? "");
setHasPrefChanges(false);
}}
className="text-xs h-8"
>
Reset
</Button>
</div>
)}
</CardContent>
</Card>
</motion.div>
)}
{/* Loading Skeleton */}
{isLoading && (
<div className="space-y-4 md:space-y-6">
{/* Active Preference Skeleton */}
<Card className="border-l-4 border-l-teal-500/30">
<CardHeader className="pb-2 px-3 md:px-6 pt-3 md:pt-6">
<div className="flex items-center gap-2 md:gap-3">
<Skeleton className="h-9 w-9 md:h-11 md:w-11 rounded-lg" />
<div className="space-y-2 flex-1">
<Skeleton className="h-5 md:h-6 w-36 md:w-44" />
<Skeleton className="h-3 md:h-4 w-56 md:w-72" />
</div>
</div>
</CardHeader>
<CardContent className="px-3 md:px-6 pb-3 md:pb-6">
<Skeleton className="h-9 md:h-10 w-full rounded-md" />
</CardContent>
</Card>
{/* Your Image Models Section Skeleton */}
<div className="space-y-4">
<div className="flex items-center justify-between">
@ -573,18 +404,6 @@ export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) {
{/* User Configs */}
{!isLoading && (
<div className="space-y-4 md:space-y-6">
<div className="flex flex-col space-y-4 sm:flex-row sm:items-center sm:justify-between sm:space-y-0">
<h3 className="text-lg md:text-xl font-semibold tracking-tight">Your Image Models</h3>
{canCreate && (
<Button
onClick={openNewDialog}
className="flex items-center gap-2 text-xs md:text-sm h-8 md:h-9"
>
Add Image Model
</Button>
)}
</div>
{(userConfigs?.length ?? 0) === 0 ? (
<Card className="border-dashed border-2 border-muted-foreground/25">
<CardContent className="flex flex-col items-center justify-center py-10 md:py-16 text-center">

View file

@ -7,6 +7,7 @@ import {
CheckCircle,
CircleDashed,
FileText,
ImageIcon,
RefreshCw,
RotateCcw,
Save,
@ -15,6 +16,10 @@ import {
import { AnimatePresence, motion } from "motion/react";
import { useEffect, 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,
@ -46,6 +51,8 @@ const ROLE_DESCRIPTIONS = {
description: "Primary LLM for chat interactions and agent operations",
color: "text-blue-600 dark:text-blue-400",
bgColor: "bg-blue-500/10",
prefKey: "agent_llm_id" as const,
configType: "llm" as const,
},
document_summary: {
icon: FileText,
@ -53,6 +60,17 @@ const ROLE_DESCRIPTIONS = {
description: "Handles document summarization and research synthesis",
color: "text-purple-600 dark:text-purple-400",
bgColor: "bg-purple-500/10",
prefKey: "document_summary_llm_id" as const,
configType: "llm" as const,
},
image_generation: {
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",
prefKey: "image_generation_config_id" as const,
configType: "image" as const,
},
};
@ -61,6 +79,7 @@ interface LLMRoleManagerProps {
}
export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
// LLM configs
const {
data: newLLMConfigs = [],
isFetching: configsLoading,
@ -72,6 +91,20 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
isFetching: globalConfigsLoading,
error: globalConfigsError,
} = useAtomValue(globalNewLLMConfigsAtom);
// Image gen configs
const {
data: userImageConfigs = [],
isFetching: imageConfigsLoading,
error: imageConfigsError,
} = useAtomValue(imageGenConfigsAtom);
const {
data: globalImageConfigs = [],
isFetching: globalImageConfigsLoading,
error: globalImageConfigsError,
} = useAtomValue(globalImageGenConfigsAtom);
// Preferences
const {
data: preferences = {},
isFetching: preferencesLoading,
@ -83,6 +116,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
const [assignments, setAssignments] = useState({
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 [hasChanges, setHasChanges] = useState(false);
@ -92,15 +126,16 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
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);
}, [preferences]);
const handleRoleAssignment = (role: string, configId: string) => {
const handleRoleAssignment = (prefKey: string, configId: string) => {
const newAssignments = {
...assignments,
[role]: configId === "unassigned" ? "" : parseInt(configId),
[prefKey]: configId === "unassigned" ? "" : parseInt(configId),
};
setAssignments(newAssignments);
@ -108,6 +143,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
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(
@ -122,19 +158,13 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
const handleSave = async () => {
setIsSaving(true);
const toNumericOrUndefined = (val: string | number) =>
typeof val === "string" ? (val ? parseInt(val) : undefined) : val;
const numericAssignments = {
agent_llm_id:
typeof assignments.agent_llm_id === "string"
? assignments.agent_llm_id
? parseInt(assignments.agent_llm_id)
: undefined
: assignments.agent_llm_id,
document_summary_llm_id:
typeof assignments.document_summary_llm_id === "string"
? assignments.document_summary_llm_id
? parseInt(assignments.document_summary_llm_id)
: undefined
: assignments.document_summary_llm_id,
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({
@ -143,7 +173,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
});
setHasChanges(false);
toast.success("LLM role assignments saved successfully!");
toast.success("Role assignments saved successfully!");
setIsSaving(false);
};
@ -152,6 +182,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
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);
};
@ -162,16 +193,26 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
assignments.agent_llm_id !== undefined &&
assignments.document_summary_llm_id !== "" &&
assignments.document_summary_llm_id !== null &&
assignments.document_summary_llm_id !== undefined;
assignments.document_summary_llm_id !== undefined &&
assignments.image_generation_config_id !== "" &&
assignments.image_generation_config_id !== null &&
assignments.image_generation_config_id !== undefined;
// Combine global and custom configs
const allConfigs = [
// Combine global and custom LLM configs
const allLLMConfigs = [
...globalConfigs.map((config) => ({ ...config, is_global: true })),
...newLLMConfigs.filter((config) => config.id && config.id.toString().trim() !== ""),
];
const isLoading = configsLoading || preferencesLoading || globalConfigsLoading;
const hasError = configsError || preferencesError || globalConfigsError;
// Combine global and custom image gen configs
const allImageConfigs = [
...globalImageConfigs.map((config) => ({ ...config, is_global: true })),
...(userImageConfigs ?? []).filter((config) => config.id && config.id.toString().trim() !== ""),
];
const isLoading = configsLoading || preferencesLoading || globalConfigsLoading || imageConfigsLoading || globalImageConfigsLoading;
const hasError = configsError || preferencesError || globalConfigsError || imageConfigsError || globalImageConfigsError;
const hasAnyConfigs = allLLMConfigs.length > 0 || allImageConfigs.length > 0;
return (
<div className="space-y-5 md:space-y-6">
@ -223,7 +264,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
{/* Loading Skeleton */}
{isLoading && (
<div className="grid gap-4 grid-cols-1 lg:grid-cols-2">
{["skeleton-a", "skeleton-b"].map((key) => (
{["skeleton-a", "skeleton-b", "skeleton-c"].map((key) => (
<Card key={key} className="border-border/60">
<CardContent className="p-4 md:p-5 space-y-4">
{/* Header: icon + title + status */}
@ -260,18 +301,18 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
)}
{/* No configs warning */}
{!isLoading && !hasError && allConfigs.length === 0 && (
{!isLoading && !hasError && !hasAnyConfigs && (
<Alert variant="destructive" className="py-3 md:py-4">
<AlertCircle className="h-3 w-3 md:h-4 md:w-4 shrink-0" />
<AlertDescription className="text-xs md:text-sm">
No LLM configurations found. Please add at least one LLM provider in the
Agent Configs tab before assigning roles.
No configurations found. Please add at least one LLM provider or image model
in the respective settings tabs before assigning roles.
</AlertDescription>
</Alert>
)}
{/* Role Assignment Cards */}
{!isLoading && !hasError && allConfigs.length > 0 && (
{!isLoading && !hasError && hasAnyConfigs && (
<motion.div
initial={{ opacity: 0 }}
animate={{ opacity: 1 }}
@ -280,9 +321,18 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
>
{Object.entries(ROLE_DESCRIPTIONS).map(([key, role], index) => {
const IconComponent = role.icon;
const isImageRole = role.configType === "image";
const currentAssignment =
assignments[`${key}_llm_id` as keyof typeof assignments];
const assignedConfig = allConfigs.find(
assignments[role.prefKey as keyof typeof assignments];
// Pick the right config lists based on role type
const roleGlobalConfigs = isImageRole ? globalImageConfigs : globalConfigs;
const roleUserConfigs = isImageRole
? (userImageConfigs ?? []).filter((c) => c.id && c.id.toString().trim() !== "")
: newLLMConfigs.filter((c) => c.id && c.id.toString().trim() !== "");
const roleAllConfigs = isImageRole ? allImageConfigs : allLLMConfigs;
const assignedConfig = roleAllConfigs.find(
(config) => config.id === currentAssignment
);
const isAssigned =
@ -340,7 +390,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
<Select
value={currentAssignment?.toString() || "unassigned"}
onValueChange={(value) =>
handleRoleAssignment(`${key}_llm_id`, value)
handleRoleAssignment(role.prefKey, value)
}
>
<SelectTrigger className="w-full h-9 md:h-10 text-xs md:text-sm">
@ -357,12 +407,12 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
</SelectItem>
{/* Global Configurations */}
{globalConfigs.length > 0 && (
{roleGlobalConfigs.length > 0 && (
<SelectGroup>
<SelectLabel className="text-[11px] md:text-xs font-semibold text-muted-foreground px-2 py-1 md:py-1.5">
Global Configurations
</SelectLabel>
{globalConfigs.map((config) => {
{roleGlobalConfigs.map((config) => {
const isAuto =
"is_auto_mode" in config &&
config.is_auto_mode;
@ -412,20 +462,12 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
)}
{/* Custom Configurations */}
{newLLMConfigs.length > 0 && (
{roleUserConfigs.length > 0 && (
<SelectGroup>
<SelectLabel className="text-[11px] md:text-xs font-semibold text-muted-foreground px-2 py-1 md:py-1.5">
Your Configurations
</SelectLabel>
{newLLMConfigs
.filter(
(config) =>
config.id &&
config.id
.toString()
.trim() !== ""
)
.map((config) => (
{roleUserConfigs.map((config) => (
<SelectItem
key={config.id}
value={config.id.toString()}
@ -480,7 +522,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
</div>
) : (
<div className="flex items-start gap-2">
<Bot className="w-3.5 h-3.5 shrink-0 mt-0.5 text-muted-foreground" />
<IconComponent className="w-3.5 h-3.5 shrink-0 mt-0.5 text-muted-foreground" />
<div className="min-w-0 flex-1">
<div className="flex items-center gap-1.5 flex-wrap">
<span className="text-xs font-medium">