refactor(backend): add type hints for improved clarity and validation

Modified files (6):
- backend/api/main.py: Updated root and health check endpoints
- backend/api/routers/models.py: Enhanced delete job endpoint
- backend/api/routers/predict.py: Improved prediction info endpoint
- backend/api/services/batch_service.py: Refined job status method
- backend/training/preprocessor.py: Updated feature metadata method
- backend/training/train.py: Enhanced job status update method

Author: cristofima

backend/api/main.py

@@ -2,6 +2,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from mangum import Mangum
+from typing import Dict
 from .routers import upload, training, models, datasets, predict
 from .utils.helpers import get_settings
 
@@ -34,7 +35,7 @@
 
 
 @app.get("/")
-async def root() -> dict[str, str]:
+async def root() -> Dict[str, str]:
     """Health check endpoint"""
     return {
         "message": "AWS AutoML Lite API",
@@ -44,7 +45,7 @@ async def root() -> dict[str, str]:
 
 
 @app.get("/health")
-async def health_check() -> dict[str, str]:
+async def health_check() -> Dict[str, str]:
     """Detailed health check"""
     return {
         "status": "healthy",

backend/api/routers/models.py

@@ -1,5 +1,5 @@
 from fastapi import APIRouter, HTTPException, status, Query
-from typing import Optional
+from typing import Dict, Optional, Any
 from ..models.schemas import (
     JobListResponse, JobResponse, JobStatus, ProblemType, JobUpdateRequest,
     DeployRequest, DeployResponse, PreprocessingInfo
@@ -120,7 +120,7 @@ async def get_job_status(job_id: str) -> JobResponse:
 
 
 @router.delete("/{job_id}")
-async def delete_job(job_id: str, delete_data: bool = True) -> dict:
+async def delete_job(job_id: str, delete_data: bool = True) -> Dict[str, Any]:
     """
     Delete a training job and optionally all associated data (model, report, dataset)
     """

backend/api/routers/predict.py

@@ -13,7 +13,7 @@
 from fastapi import APIRouter, HTTPException, status
 import logging
 import time
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 import numpy as np
 import tempfile
 import os
@@ -128,7 +128,7 @@ def _prepare_input(
     features: Dict[str, Any],
     feature_columns: list,
     session: Any,
-    preprocessing_info: Dict[str, Any] | None = None
+    preprocessing_info: Optional[Dict[str, Any]]
 ) -> np.ndarray:
     """
     Prepare input features for ONNX model inference.
@@ -306,7 +306,7 @@ async def make_prediction(job_id: str, request: PredictionInput) -> PredictionRe
 
 
 @router.get("/{job_id}/info")
-async def get_prediction_info(job_id: str) -> dict:
+async def get_prediction_info(job_id: str) -> Dict[str, Any]:
     """
     Get information about a deployed model for making predictions.
     

backend/api/services/batch_service.py

@@ -1,5 +1,5 @@
 import boto3
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 from botocore.exceptions import ClientError
 from ..utils.helpers import get_settings
 
@@ -49,7 +49,7 @@ def submit_training_job(
         except ClientError as e:
             raise Exception(f"Error submitting batch job: {str(e)}")
 
-    def get_job_status(self, batch_job_id: str) -> Dict[str, Any] | None:
+    def get_job_status(self, batch_job_id: str) -> Optional[Dict[str, Any]]:
         """Get the status of a batch job"""
         try:
             response = self.batch_client.describe_jobs(jobs=[batch_job_id])

backend/tests/api/test_endpoints.py

@@ -5,7 +5,7 @@
 without requiring actual AWS services.
 """
 import pytest
-from unittest.mock import patch, MagicMock
+from unittest.mock import patch
 import sys
 from pathlib import Path
 

backend/tests/api/test_s3_service.py

@@ -287,7 +287,7 @@ def test_object_exists(self, s3_bucket):
         try:
             s3.head_object(Bucket=bucket, Key="not-exists.csv")
             not_exists = True
-        except:
+        except s3.exceptions.ClientError:
             not_exists = False
 
         assert not_exists == False

backend/tests/api/test_schemas.py

@@ -6,7 +6,6 @@
 """
 import pytest
 from pydantic import ValidationError
-from datetime import datetime
 import sys
 from pathlib import Path
 
@@ -20,14 +19,10 @@
     DatasetMetadata,
     ColumnStats,
     JobStatus,
-    JobDetails,
-    JobResponse,
     PredictionInput,
     PredictionResponse,
-    UploadRequest,
     UploadResponse,
     ProblemType,
-    TrainingMetrics,
 )
 
 

backend/training/preprocessor.py

@@ -3,7 +3,7 @@
 from decimal import Decimal
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import LabelEncoder, StandardScaler
-from typing import Tuple, List
+from typing import Dict, Any, Optional, Tuple, List
 
 # Feature-engine for robust feature selection
 from feature_engine.selection import DropConstantFeatures, DropDuplicateFeatures
@@ -154,7 +154,7 @@ def encode_categorical(self, df: pd.DataFrame, fit: bool = True) -> pd.DataFrame
 
         return df
 
-    def get_feature_metadata(self, df: pd.DataFrame | None = None) -> dict:
+    def get_feature_metadata(self, df: Optional[pd.DataFrame]) -> Dict[str, Any]:
         """
         Get metadata about features for inference.
         

backend/training/train.py

@@ -4,7 +4,7 @@
 import pandas as pd
 from datetime import datetime, timezone
 import traceback
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from preprocessor import AutoPreprocessor
 from eda import generate_eda_report
@@ -202,7 +202,7 @@ def update_job_status(
     table: Any,
     job_id: str,
     status: str,
-    error_message: str | None = None
+    error_message: Optional[str]
 ) -> None:
     """Update job status in DynamoDB"""
     now = datetime.now(timezone.utc).isoformat()
@@ -241,14 +241,14 @@ def update_job_completion(
     target_column: str,
     problem_type: str,
     model_path: str,
-    onnx_model_path: str | None,
+    onnx_model_path: Optional[str],
     eda_report_s3_path: str,
     training_report_s3_path: str,
     metrics: Dict[str, Any],
     feature_importance: Dict[str, float],
-    dropped_columns: List[str] | None = None,
-    feature_columns: List[str] | None = None,
-    feature_metadata: Dict[str, Any] | None = None
+    dropped_columns: Optional[List[str]],
+    feature_columns: Optional[List[str]],
+    feature_metadata: Optional[Dict[str, Any]]
 ) -> None:
     """Update job with completion details"""
     from decimal import Decimal