"use client"; import { useAtomValue } from "jotai"; import { AlertCircle, Check, ChevronsUpDown, Clock, Edit3, ImageIcon, Key, Plus, RefreshCw, Shuffle, Sparkles, Trash2, Wand2, } from "lucide-react"; import { AnimatePresence, motion } from "motion/react"; import { useCallback, useEffect, useState } from "react"; import { toast } from "sonner"; import { createImageGenConfigMutationAtom, deleteImageGenConfigMutationAtom, updateImageGenConfigMutationAtom, } from "@/atoms/image-gen-config/image-gen-config-mutation.atoms"; 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 { llmPreferencesAtom } from "@/atoms/new-llm-config/new-llm-config-query.atoms"; import { Alert, AlertDescription } from "@/components/ui/alert"; import { AlertDialog, AlertDialogAction, AlertDialogCancel, AlertDialogContent, AlertDialogDescription, AlertDialogFooter, 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 { Command, CommandEmpty, CommandGroup, CommandInput, CommandItem, CommandList, } from "@/components/ui/command"; import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle, } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue, } from "@/components/ui/select"; import { Separator } from "@/components/ui/separator"; import { Spinner } from "@/components/ui/spinner"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { IMAGE_GEN_PROVIDERS, getImageGenModelsByProvider, } from "@/contracts/enums/image-gen-providers"; import type { ImageGenerationConfig } from "@/contracts/types/new-llm-config.types"; import { cn } from "@/lib/utils"; interface ImageModelManagerProps { searchSpaceId: number; } const container = { hidden: { opacity: 0 }, show: { opacity: 1, transition: { staggerChildren: 0.05 } }, }; const item = { hidden: { opacity: 0, y: 20 }, show: { opacity: 1, y: 0 }, }; export function ImageModelManager({ searchSpaceId }: ImageModelManagerProps) { // Image gen config atoms const { mutateAsync: createConfig, isPending: isCreating, error: createError } = useAtomValue(createImageGenConfigMutationAtom); const { mutateAsync: updateConfig, isPending: isUpdating, error: updateError } = useAtomValue(updateImageGenConfigMutationAtom); const { mutateAsync: deleteConfig, isPending: isDeleting, error: deleteError } = useAtomValue(deleteImageGenConfigMutationAtom); const { mutateAsync: updatePreferences } = useAtomValue(updateLLMPreferencesMutationAtom); const { data: userConfigs, isFetching: configsLoading, error: fetchError, refetch: refreshConfigs } = useAtomValue(imageGenConfigsAtom); const { data: globalConfigs = [], isFetching: globalLoading } = useAtomValue(globalImageGenConfigsAtom); const { data: preferences = {}, isFetching: prefsLoading } = useAtomValue(llmPreferencesAtom); // Local state const [isDialogOpen, setIsDialogOpen] = useState(false); const [editingConfig, setEditingConfig] = useState(null); const [configToDelete, setConfigToDelete] = useState(null); // Preference state const [selectedPrefId, setSelectedPrefId] = useState( 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 errors = [createError, updateError, deleteError, fetchError].filter(Boolean) as Error[]; // Form state for create/edit dialog const [formData, setFormData] = useState({ name: "", description: "", provider: "", custom_provider: "", model_name: "", api_key: "", api_base: "", api_version: "", }); const [modelComboboxOpen, setModelComboboxOpen] = useState(false); const resetForm = () => { setFormData({ name: "", description: "", provider: "", custom_provider: "", model_name: "", api_key: "", api_base: "", api_version: "", }); }; const handleFormSubmit = useCallback(async () => { if (!formData.name || !formData.provider || !formData.model_name || !formData.api_key) { toast.error("Please fill in all required fields"); return; } try { if (editingConfig) { await updateConfig({ id: editingConfig.id, data: { name: formData.name, description: formData.description || undefined, provider: formData.provider as any, custom_provider: formData.custom_provider || undefined, model_name: formData.model_name, api_key: formData.api_key, api_base: formData.api_base || undefined, api_version: formData.api_version || undefined, }, }); } else { const result = await createConfig({ name: formData.name, description: formData.description || undefined, provider: formData.provider as any, custom_provider: formData.custom_provider || undefined, model_name: formData.model_name, api_key: formData.api_key, api_base: formData.api_base || undefined, api_version: formData.api_version || undefined, search_space_id: searchSpaceId, }); // Auto-assign newly created config if (result?.id) { await updatePreferences({ search_space_id: searchSpaceId, data: { image_generation_config_id: result.id }, }); } } setIsDialogOpen(false); setEditingConfig(null); resetForm(); } catch { // Error handled by mutation } }, [editingConfig, formData, searchSpaceId, createConfig, updateConfig, updatePreferences]); const handleDelete = async () => { if (!configToDelete) return; try { await deleteConfig(configToDelete.id); setConfigToDelete(null); } catch { // Error handled by mutation } }; const openEditDialog = (config: ImageGenerationConfig) => { setEditingConfig(config); setFormData({ name: config.name, description: config.description || "", provider: config.provider, custom_provider: config.custom_provider || "", model_name: config.model_name, api_key: config.api_key, api_base: config.api_base || "", api_version: config.api_version || "", }); setIsDialogOpen(true); }; const openNewDialog = () => { setEditingConfig(null); resetForm(); 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); return (
{/* Header */}
{/* Errors */} {errors.map((err) => ( {err?.message} ))} {/* Global info */} {globalConfigs.filter((g) => !("is_auto_mode" in g && g.is_auto_mode)).length > 0 && ( {globalConfigs.filter((g) => !("is_auto_mode" in g && g.is_auto_mode)).length} global image model(s) {" "} available from your administrator. )} {/* Active Preference Card */} {!isLoading && allConfigs.length > 0 && (
Active Image Model Select which model to use for image generation
{hasPrefChanges && (
)}
)} {/* Loading */} {isLoading && ( )} {/* User Configs */} {!isLoading && (

Your Image Models

{(userConfigs?.length ?? 0) === 0 ? (

No Image Models Yet

Add your own image generation model (DALL-E 3, GPT Image 1, etc.)

) : ( {userConfigs?.map((config) => (

{config.name}

{config.provider}
{config.model_name} {config.description && (

{config.description}

)}
{new Date(config.created_at).toLocaleDateString()}
Edit Delete
))} )}
)} {/* Create/Edit Dialog */} { if (!open) { setIsDialogOpen(false); setEditingConfig(null); resetForm(); } }}> {editingConfig ? : } {editingConfig ? "Edit Image Model" : "Add Image Model"} {editingConfig ? "Update your image generation model" : "Configure a new image generation model (DALL-E 3, GPT Image 1, etc.)"}
{/* Name */}
setFormData((p) => ({ ...p, name: e.target.value }))} />
{/* Description */}
setFormData((p) => ({ ...p, description: e.target.value }))} />
{/* Provider */}
{/* Model Name */}
{suggestedModels.length > 0 ? ( setFormData((p) => ({ ...p, model_name: val }))} /> Type a custom model name {suggestedModels.map((m) => ( { setFormData((p) => ({ ...p, model_name: m.value })); setModelComboboxOpen(false); }} > {m.value} {m.label} ))} ) : ( setFormData((p) => ({ ...p, model_name: e.target.value }))} /> )}
{/* API Key */}
setFormData((p) => ({ ...p, api_key: e.target.value }))} />
{/* API Base (optional) */}
setFormData((p) => ({ ...p, api_base: e.target.value }))} />
{/* API Version (Azure) */} {formData.provider === "AZURE_OPENAI" && (
setFormData((p) => ({ ...p, api_version: e.target.value }))} />
)} {/* Actions */}
{/* Delete Confirmation */} !open && setConfigToDelete(null)}> Delete Image Model Are you sure you want to delete {configToDelete?.name}? Cancel {isDeleting ? <>Deleting : <>Delete}
); }