WIP

frontend/src/hooks/useDatasets.ts

@@ -0,0 +1,84 @@
+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,
+  }
+}