WIP

frontend/src/hooks/index.ts

@@ -0,0 +1,4 @@
+export { useDocuments } from './useDocuments'
+export { useDocumentDetail } from './useDocumentDetail'
+export { useAnnotations } from './useAnnotations'
+export { useTraining, useTrainingDocuments } from './useTraining'

frontend/src/hooks/useAnnotations.ts

@@ -0,0 +1,70 @@
+import { useMutation, useQueryClient } from '@tanstack/react-query'
+import { annotationsApi } from '../api/endpoints'
+import type { CreateAnnotationRequest, AnnotationOverrideRequest } from '../api/types'
+
+export const useAnnotations = (documentId: string) => {
+  const queryClient = useQueryClient()
+
+  const createMutation = useMutation({
+    mutationFn: (annotation: CreateAnnotationRequest) =>
+      annotationsApi.create(documentId, annotation),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['document', documentId] })
+    },
+  })
+
+  const updateMutation = useMutation({
+    mutationFn: ({
+      annotationId,
+      updates,
+    }: {
+      annotationId: string
+      updates: Partial<CreateAnnotationRequest>
+    }) => annotationsApi.update(documentId, annotationId, updates),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['document', documentId] })
+    },
+  })
+
+  const deleteMutation = useMutation({
+    mutationFn: (annotationId: string) =>
+      annotationsApi.delete(documentId, annotationId),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['document', documentId] })
+    },
+  })
+
+  const verifyMutation = useMutation({
+    mutationFn: (annotationId: string) =>
+      annotationsApi.verify(documentId, annotationId),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['document', documentId] })
+    },
+  })
+
+  const overrideMutation = useMutation({
+    mutationFn: ({
+      annotationId,
+      overrideData,
+    }: {
+      annotationId: string
+      overrideData: AnnotationOverrideRequest
+    }) => annotationsApi.override(documentId, annotationId, overrideData),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['document', documentId] })
+    },
+  })
+
+  return {
+    createAnnotation: createMutation.mutate,
+    isCreating: createMutation.isPending,
+    updateAnnotation: updateMutation.mutate,
+    isUpdating: updateMutation.isPending,
+    deleteAnnotation: deleteMutation.mutate,
+    isDeleting: deleteMutation.isPending,
+    verifyAnnotation: verifyMutation.mutate,
+    isVerifying: verifyMutation.isPending,
+    overrideAnnotation: overrideMutation.mutate,
+    isOverriding: overrideMutation.isPending,
+  }
+}

frontend/src/hooks/useDocumentDetail.ts

@@ -0,0 +1,25 @@
+import { useQuery } from '@tanstack/react-query'
+import { documentsApi } from '../api/endpoints'
+import type { DocumentDetailResponse } from '../api/types'
+
+export const useDocumentDetail = (documentId: string | null) => {
+  const { data, isLoading, error, refetch } = useQuery<DocumentDetailResponse>({
+    queryKey: ['document', documentId],
+    queryFn: () => {
+      if (!documentId) {
+        throw new Error('Document ID is required')
+      }
+      return documentsApi.getDetail(documentId)
+    },
+    enabled: !!documentId,
+    staleTime: 10000,
+  })
+
+  return {
+    document: data || null,
+    annotations: data?.annotations || [],
+    isLoading,
+    error,
+    refetch,
+  }
+}

frontend/src/hooks/useDocuments.ts

@@ -0,0 +1,78 @@
+import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'
+import { documentsApi } from '../api/endpoints'
+import type { DocumentListResponse, UploadDocumentResponse } from '../api/types'
+
+interface UseDocumentsParams {
+  status?: string
+  limit?: number
+  offset?: number
+}
+
+export const useDocuments = (params: UseDocumentsParams = {}) => {
+  const queryClient = useQueryClient()
+
+  const { data, isLoading, error, refetch } = useQuery<DocumentListResponse>({
+    queryKey: ['documents', params],
+    queryFn: () => documentsApi.list(params),
+    staleTime: 30000,
+  })
+
+  const uploadMutation = useMutation({
+    mutationFn: (file: File) => documentsApi.upload(file),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['documents'] })
+    },
+  })
+
+  const batchUploadMutation = useMutation({
+    mutationFn: ({ files, csvFile }: { files: File[]; csvFile?: File }) =>
+      documentsApi.batchUpload(files, csvFile),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['documents'] })
+    },
+  })
+
+  const deleteMutation = useMutation({
+    mutationFn: (documentId: string) => documentsApi.delete(documentId),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['documents'] })
+    },
+  })
+
+  const updateStatusMutation = useMutation({
+    mutationFn: ({ documentId, status }: { documentId: string; status: string }) =>
+      documentsApi.updateStatus(documentId, status),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['documents'] })
+    },
+  })
+
+  const triggerAutoLabelMutation = useMutation({
+    mutationFn: (documentId: string) => documentsApi.triggerAutoLabel(documentId),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['documents'] })
+    },
+  })
+
+  return {
+    documents: data?.documents || [],
+    total: data?.total || 0,
+    limit: data?.limit || params.limit || 20,
+    offset: data?.offset || params.offset || 0,
+    isLoading,
+    error,
+    refetch,
+    uploadDocument: uploadMutation.mutate,
+    uploadDocumentAsync: uploadMutation.mutateAsync,
+    isUploading: uploadMutation.isPending,
+    batchUpload: batchUploadMutation.mutate,
+    batchUploadAsync: batchUploadMutation.mutateAsync,
+    isBatchUploading: batchUploadMutation.isPending,
+    deleteDocument: deleteMutation.mutate,
+    isDeleting: deleteMutation.isPending,
+    updateStatus: updateStatusMutation.mutate,
+    isUpdatingStatus: updateStatusMutation.isPending,
+    triggerAutoLabel: triggerAutoLabelMutation.mutate,
+    isTriggeringAutoLabel: triggerAutoLabelMutation.isPending,
+  }
+}

frontend/src/hooks/useTraining.ts

@@ -0,0 +1,83 @@
+import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'
+import { trainingApi } from '../api/endpoints'
+import type { TrainingModelsResponse } from '../api/types'
+
+export const useTraining = () => {
+  const queryClient = useQueryClient()
+
+  const { data: modelsData, isLoading: isLoadingModels } =
+    useQuery<TrainingModelsResponse>({
+      queryKey: ['training', 'models'],
+      queryFn: () => trainingApi.getModels(),
+      staleTime: 30000,
+    })
+
+  const startTrainingMutation = useMutation({
+    mutationFn: (config: {
+      name: string
+      description?: string
+      document_ids: string[]
+      epochs?: number
+      batch_size?: number
+      model_base?: string
+    }) => trainingApi.startTraining(config),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['training', 'models'] })
+    },
+  })
+
+  const cancelTaskMutation = useMutation({
+    mutationFn: (taskId: string) => trainingApi.cancelTask(taskId),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['training', 'models'] })
+    },
+  })
+
+  const downloadModelMutation = useMutation({
+    mutationFn: (taskId: string) => trainingApi.downloadModel(taskId),
+    onSuccess: (blob, taskId) => {
+      const url = window.URL.createObjectURL(blob)
+      const a = document.createElement('a')
+      a.href = url
+      a.download = `model-${taskId}.pt`
+      document.body.appendChild(a)
+      a.click()
+      window.URL.revokeObjectURL(url)
+      document.body.removeChild(a)
+    },
+  })
+
+  return {
+    models: modelsData?.models || [],
+    total: modelsData?.total || 0,
+    isLoadingModels,
+    startTraining: startTrainingMutation.mutate,
+    startTrainingAsync: startTrainingMutation.mutateAsync,
+    isStartingTraining: startTrainingMutation.isPending,
+    cancelTask: cancelTaskMutation.mutate,
+    isCancelling: cancelTaskMutation.isPending,
+    downloadModel: downloadModelMutation.mutate,
+    isDownloading: downloadModelMutation.isPending,
+  }
+}
+
+export const useTrainingDocuments = (params?: {
+  has_annotations?: boolean
+  min_annotation_count?: number
+  exclude_used_in_training?: boolean
+  limit?: number
+  offset?: number
+}) => {
+  const { data, isLoading, error } = useQuery({
+    queryKey: ['training', 'documents', params],
+    queryFn: () => trainingApi.getDocumentsForTraining(params),
+    staleTime: 30000,
+  })
+
+  return {
+    documents: data?.documents || [],
+    total: data?.total || 0,
+    isLoading,
+    error,
+  }
+}