85 lines
2.5 KiB
TypeScript
85 lines
2.5 KiB
TypeScript
import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'
|
|
import { datasetsApi } from '../api/endpoints'
|
|
import type {
|
|
DatasetCreateRequest,
|
|
DatasetDetailResponse,
|
|
DatasetListResponse,
|
|
DatasetTrainRequest,
|
|
} from '../api/types'
|
|
|
|
export const useDatasets = (params?: {
|
|
status?: string
|
|
limit?: number
|
|
offset?: number
|
|
}) => {
|
|
const queryClient = useQueryClient()
|
|
|
|
const { data, isLoading, error, refetch } = useQuery<DatasetListResponse>({
|
|
queryKey: ['datasets', params],
|
|
queryFn: () => datasetsApi.list(params),
|
|
staleTime: 30000,
|
|
// Poll every 5 seconds when there's an active training task
|
|
refetchInterval: (query) => {
|
|
const datasets = query.state.data?.datasets ?? []
|
|
const hasActiveTraining = datasets.some(
|
|
d => d.training_status === 'running' || d.training_status === 'pending' || d.training_status === 'scheduled'
|
|
)
|
|
return hasActiveTraining ? 5000 : false
|
|
},
|
|
})
|
|
|
|
const createMutation = useMutation({
|
|
mutationFn: (req: DatasetCreateRequest) => datasetsApi.create(req),
|
|
onSuccess: () => {
|
|
queryClient.invalidateQueries({ queryKey: ['datasets'] })
|
|
},
|
|
})
|
|
|
|
const deleteMutation = useMutation({
|
|
mutationFn: (datasetId: string) => datasetsApi.remove(datasetId),
|
|
onSuccess: () => {
|
|
queryClient.invalidateQueries({ queryKey: ['datasets'] })
|
|
},
|
|
})
|
|
|
|
const trainMutation = useMutation({
|
|
mutationFn: ({ datasetId, req }: { datasetId: string; req: DatasetTrainRequest }) =>
|
|
datasetsApi.trainFromDataset(datasetId, req),
|
|
onSuccess: () => {
|
|
queryClient.invalidateQueries({ queryKey: ['datasets'] })
|
|
queryClient.invalidateQueries({ queryKey: ['training', 'models'] })
|
|
},
|
|
})
|
|
|
|
return {
|
|
datasets: data?.datasets ?? [],
|
|
total: data?.total ?? 0,
|
|
isLoading,
|
|
error,
|
|
refetch,
|
|
createDataset: createMutation.mutate,
|
|
createDatasetAsync: createMutation.mutateAsync,
|
|
isCreating: createMutation.isPending,
|
|
deleteDataset: deleteMutation.mutate,
|
|
isDeleting: deleteMutation.isPending,
|
|
trainFromDataset: trainMutation.mutate,
|
|
trainFromDatasetAsync: trainMutation.mutateAsync,
|
|
isTraining: trainMutation.isPending,
|
|
}
|
|
}
|
|
|
|
export const useDatasetDetail = (datasetId: string | null) => {
|
|
const { data, isLoading, error } = useQuery<DatasetDetailResponse>({
|
|
queryKey: ['datasets', datasetId],
|
|
queryFn: () => datasetsApi.getDetail(datasetId!),
|
|
enabled: !!datasetId,
|
|
staleTime: 30000,
|
|
})
|
|
|
|
return {
|
|
dataset: data ?? null,
|
|
isLoading,
|
|
error,
|
|
}
|
|
}
|