import { getOpenAIConfigById } from "@/db/openai" import { getAllOpenAIModels } from "@/libs/openai" import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query" import { useTranslation } from "react-i18next" import { Checkbox, Input, Spin, message, Radio } from "antd" import { useState, useMemo } from "react" import { createManyModels } from "@/db/models" import { Popover } from "antd" import { InfoIcon } from "lucide-react" type Props = { openaiId: string setOpenModelModal: (openModelModal: boolean) => void } export const OpenAIFetchModel = ({ openaiId, setOpenModelModal }: Props) => { const { t } = useTranslation(["openai"]) const [selectedModels, setSelectedModels] = useState([]) const [searchTerm, setSearchTerm] = useState("") const [modelType, setModelType] = useState("chat") const queryClient = useQueryClient() const { data, status } = useQuery({ queryKey: ["openAIConfigs", openaiId], queryFn: async () => { const config = await getOpenAIConfigById(openaiId) const models = await getAllOpenAIModels(config.baseUrl, config.apiKey) return models }, enabled: !!openaiId }) const filteredModels = useMemo(() => { return ( data?.filter((model) => (model.name ?? model.id) .toLowerCase() .includes(searchTerm.toLowerCase()) ) || [] ) }, [data, searchTerm]) const handleSelectAll = (checked: boolean) => { if (checked) { setSelectedModels(filteredModels.map((model) => model.id)) } else { setSelectedModels([]) } } const handleModelSelect = (modelId: string, checked: boolean) => { if (checked) { setSelectedModels((prev) => [...prev, modelId]) } else { setSelectedModels((prev) => prev.filter((id) => id !== modelId)) } } const onSave = async (models: string[]) => { const payload = models.map((id) => ({ model_id: id, name: filteredModels.find((model) => model.id === id)?.name ?? id, provider_id: openaiId, model_type: modelType })) await createManyModels(payload) return true } const { mutate: saveModels, isPending: isSaving } = useMutation({ mutationFn: onSave, onSuccess: () => { setOpenModelModal(false) queryClient.invalidateQueries({ queryKey: ["fetchModel"] }) message.success(t("modal.model.success")) } }) const handleSave = () => { saveModels(selectedModels) } if (status === "pending") { return (
) } if (status === "error" || !data || data.length === 0) { return (

{t("noModelFound")}

) } return (

{t("modal.model.subheading")}

setSearchTerm(e.target.value)} className="w-full" />
0 && selectedModels.length < filteredModels.length } onChange={(e) => handleSelectAll(e.target.checked)}> {t("selectAll")}
{`${selectedModels?.length} / ${data?.length}`}
{filteredModels.map((model) => ( handleModelSelect(model.id, e.target.checked)}>
{`${model?.name || model.id}`.replaceAll( /accounts\/[^\/]+\/models\//g, "" )}
))}
setModelType(e.target.value)} value={modelType}> {t("radio.chat")} {t("radio.embedding")}

{t("radio.chat")} {" "} {t("radio.chatInfo")}

{t("radio.embedding")} {" "} {t("radio.embeddingInfo")}

}>
) }