diff --git a/lpm_frontend/src/components/ModelStatus/index.tsx b/lpm_frontend/src/components/ModelStatus/index.tsx index 9a96d54..0cbf51f 100644 --- a/lpm_frontend/src/components/ModelStatus/index.tsx +++ b/lpm_frontend/src/components/ModelStatus/index.tsx @@ -2,7 +2,7 @@ import { Status, statusRankMap, useTrainingStore } from '@/store/useTrainingStor import { startService, stopService } from '@/service/train'; import { StatusBar } from '../StatusBar'; import { useRef, useEffect, useState, useMemo } from 'react'; -import { message } from 'antd'; +import { message, Select, Tooltip, Spin } from 'antd'; import { CloudUploadOutlined, CheckCircleOutlined, @@ -15,6 +15,7 @@ import RegisterUploadModal from '../upload/RegisterUploadModal'; import { useLoadInfoStore } from '@/store/useLoadInfoStore'; import TrainingTipModal from '../upload/TraingTipModal'; import { getMemoryList } from '@/service/memory'; +import { getModelList, ModelInfo } from '@/service/model'; const StatusDot = ({ active }: { active: boolean }) => (
( /> ); +// Helper component for option tooltips +const OptionTooltip = ({ model }: { model: ModelInfo }) => { + const fileName = model.model_path.split('/')[1]; + const timeStamp = fileName?.replace('.gguf', '') || 'Unknown version'; + const modelName = model.model_path.split('/')[0]; + + return ( + +
{modelName}
+
Training Parameters:
+ +
+ } + placement="left" + > +
+ {timeStamp} + {modelName} +
+ + ); +}; + export function ModelStatus() { const status = useTrainingStore((state) => state.status); const setStatus = useTrainingStore((state) => state.setStatus); @@ -43,6 +79,11 @@ export function ModelStatus() { const [showRegisterModal, setShowRegisterModal] = useState(false); const [showtrainingModal, setShowtrainingModal] = useState(false); + // Model selection states + const [modelList, setModelList] = useState([]); + const [selectedModel, setSelectedModel] = useState(''); + const [loadingModels, setLoadingModels] = useState(false); + const handleRegistryClick = () => { if (!serviceStarted) { messageApi.info({ @@ -54,6 +95,16 @@ export function ModelStatus() { } }; + const handleModelChange = (value: string) => { + setSelectedModel(value); + // Update configuration in localStorage + const config = JSON.parse(localStorage.getItem('trainingParams') || '{}'); + + // Store the complete model path + config.model_name = value; + localStorage.setItem('trainingParams', JSON.stringify(config)); + }; + const fetchMemories = async () => { try { const memoryRes = await getMemoryList(); @@ -70,9 +121,51 @@ export function ModelStatus() { } }; + const fetchModels = async () => { + setLoadingModels(true); + + try { + const res = await getModelList(); + + if (res.data.code === 0) { + setModelList(res.data.data); + // Read saved configuration, initialize selected model + const config = JSON.parse(localStorage.getItem('trainingParams') || '{}'); + + if (config.model_name) { + // Check if the model exists in the list + const modelExists = res.data.data.some( + (model: ModelInfo) => model.model_path === config.model_name + ); + + if (modelExists) { + setSelectedModel(config.model_name); + } else if (res.data.data.length > 0) { + // Default to the first model in the list + setSelectedModel(res.data.data[0].model_path); + } + } else if (res.data.data.length > 0) { + // Default to the first model in the list + setSelectedModel(res.data.data[0].model_path); + } + } else { + messageApi.error({ content: res.data.message!, duration: 1 }); + } + } catch (error) { + console.error('Error fetching model list:', error); + messageApi.error({ + content: error.response?.data?.message || error.message, + duration: 1 + }); + } finally { + setLoadingModels(false); + } + }; + useEffect(() => { fetchMemories(); fetchServiceStatus(); + fetchModels(); return () => { clearPolling(); @@ -130,19 +223,18 @@ export function ModelStatus() { console.error('Error checking service status:', error); }); }, 3000); - }; - - const handleStartService = () => { + }; const handleStartService = () => { + // Use selected model or get from local storage const config = JSON.parse(localStorage.getItem('trainingParams') || '{}'); + const modelPath = selectedModel || config.model_name; - if (!config.model_name) { - message.error('Please train a base model first'); - + if (!modelPath) { + message.error('Please select a model to start'); return; } setServiceStarting(true); - startService({ model_name: config.model_name }) + startService({ model_name: modelPath }) .then((res) => { if (res.data.code === 0) { messageApi.success({ content: 'Service starting...', duration: 1 }); @@ -204,6 +296,47 @@ export function ModelStatus() {
+ {/* Model Selection - only show when service is not started */} + {!serviceStarted && ( +
+ +