Add new frontend - fix 2

frontend/src/api/hooks/useTrainingProgress.ts

@@ -0,0 +1,83 @@
+// src/hooks/useTrainingProgress.ts
+import { useState, useEffect } from 'react';
+import { useWebSocket } from '../hooks/useWebSocket';
+
+export interface TrainingProgress {
+  job_id: string;
+  status: 'pending' | 'running' | 'completed' | 'failed';
+  progress: number;
+  current_step: string;
+  total_steps: number;
+  estimated_time_remaining?: number;
+  metrics?: Record<string, any>;
+}
+
+export interface TrainingProgressUpdate {
+  type: 'training_progress' | 'training_completed' | 'training_error';
+  job_id: string;
+  progress?: TrainingProgress;
+  results?: any;
+  error?: string;
+}
+
+export const useTrainingProgress = (jobId: string | null) => {
+  const [progress, setProgress] = useState<TrainingProgress | null>(null);
+  const [error, setError] = useState<string | null>(null);
+  const [isComplete, setIsComplete] = useState(false);
+
+  const handleMessage = (data: TrainingProgressUpdate) => {
+    switch (data.type) {
+      case 'training_progress':
+        setProgress(data.progress!);
+        setError(null);
+        break;
+      
+      case 'training_completed':
+        setProgress(prev => ({
+          ...prev!,
+          status: 'completed',
+          progress: 100
+        }));
+        setIsComplete(true);
+        break;
+      
+      case 'training_error':
+        setError(data.error || 'Training failed');
+        setProgress(prev => prev ? { ...prev, status: 'failed' } : null);
+        break;
+    }
+  };
+
+  const { isConnected } = useWebSocket({
+    endpoint: jobId ? `/training/progress/${jobId}` : '',
+    onMessage: handleMessage,
+    onError: () => setError('Connection lost'),
+    autoConnect: !!jobId
+  });
+
+  // Fetch initial status when job ID changes
+  useEffect(() => {
+    if (jobId) {
+      fetchTrainingStatus(jobId);
+    }
+  }, [jobId]);
+
+  const fetchTrainingStatus = async (id: string) => {
+    try {
+      const response = await fetch(`/api/training/status/${id}`);
+      if (response.ok) {
+        const data = await response.json();
+        setProgress(data);
+      }
+    } catch (err) {
+      console.error('Failed to fetch training status:', err);
+    }
+  };
+
+  return {
+    progress,
+    error,
+    isComplete,
+    isConnected
+  };
+};