fix(drift): derive fitColumns from stored metadata when omitted

Apply the same fix as PR #180 to the streaming KS test endpoints:
when fitColumns is not provided, derive it from the model's stored
input schema metadata instead of rejecting with HTTP 400.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Sudip Sinha <Sudip.Sinha@RedHat.com>

Author: SudipSinha

src/endpoints/metrics/drift/kolmogorov_smirnov_streaming.py

@@ -97,9 +97,13 @@ async def compute_ksteststreaming(
         )
 
     if not request.fit_columns:
-        raise HTTPException(
-            status_code=HTTPStatus.BAD_REQUEST,
-            detail="fitColumns is required - specify which features to test for drift",
+        data_source = get_data_source()
+        metadata = await data_source.get_metadata(request.model_id)
+        request.fit_columns = list(metadata.input_schema.items.keys())
+        logger.info(
+            "fitColumns not specified, using all input columns for model %s: %s",
+            request.model_id,
+            request.fit_columns,
         )
 
     try:
@@ -215,6 +219,16 @@ async def schedule_ksteststreaming(
     request: ApproxKSTestMetricRequest,
 ) -> dict[str, str]:
     """Schedule a recurring computation of KS Test Streaming metric."""
+    if not request.fit_columns:
+        data_source = get_data_source()
+        metadata = await data_source.get_metadata(request.model_id)
+        request.fit_columns = list(metadata.input_schema.items.keys())
+        logger.info(
+            "fitColumns not specified, using all input columns for model %s: %s",
+            request.model_id,
+            request.fit_columns,
+        )
+
     # Get the scheduler and validate availability
     scheduler = get_prometheus_scheduler()
     if not scheduler:

tests/endpoints/metrics/drift/factory.py

@@ -73,16 +73,17 @@ def make_compute_endpoint_test(
     def test_impl(_: object, mock_ds: MagicMock) -> None:
         """Test compute endpoint returns valid response structure."""
         # Create sample dataframe (Pandas or Polars based on df_type)
-        sample_df = _create_sample_dataframe(
-            request_payload.get("fitColumns", ["feature1"]),
-            df_type=df_type,
-        )
+        columns = request_payload.get("fitColumns", ["feature1"])
+        sample_df = _create_sample_dataframe(columns, df_type=df_type)
 
         # Mock data source
         mock_data_source = MagicMock()
         mock_data_source.get_dataframe_by_tag = AsyncMock(return_value=sample_df)
         mock_data_source.get_organic_dataframe = AsyncMock(return_value=sample_df)
         mock_data_source.get_dataframe = AsyncMock(return_value=sample_df)
+        mock_metadata = MagicMock()
+        mock_metadata.input_schema.items.keys.return_value = columns
+        mock_data_source.get_metadata = AsyncMock(return_value=mock_metadata)
         mock_ds.return_value = mock_data_source
 
         # Send request
@@ -184,7 +185,9 @@ def test_impl(_: object, mock_ds: MagicMock, mock_sched_fn: MagicMock) -> None:
 
         # Mock data source (needed for KS test registration)
         mock_data_source = MagicMock()
-        mock_data_source.get_metadata = AsyncMock(return_value={"feature1": "type1"})
+        mock_sched_metadata = MagicMock()
+        mock_sched_metadata.input_schema.items.keys.return_value = ["feature1"]
+        mock_data_source.get_metadata = AsyncMock(return_value=mock_sched_metadata)
         mock_ds.return_value = mock_data_source
 
         # Send request
@@ -338,17 +341,17 @@ def make_compute_endpoint_error_test(
     def test_impl(_: object, mock_ds: MagicMock) -> None:
         """Test compute endpoint error handling."""
         if setup_mocks:
-            # Create sample dataframe
-            sample_df = _create_sample_dataframe(
-                ["feature1", "feature2"],
-                df_type=df_type,
-            )
+            columns = ["feature1", "feature2"]
+            sample_df = _create_sample_dataframe(columns, df_type=df_type)
 
             # Mock data source
             mock_data_source = MagicMock()
             mock_data_source.get_dataframe_by_tag = AsyncMock(return_value=sample_df)
             mock_data_source.get_organic_dataframe = AsyncMock(return_value=sample_df)
             mock_data_source.get_dataframe = AsyncMock(return_value=sample_df)
+            mock_metadata = MagicMock()
+            mock_metadata.input_schema.items.keys.return_value = columns
+            mock_data_source.get_metadata = AsyncMock(return_value=mock_metadata)
             mock_ds.return_value = mock_data_source
 
         # Send request
@@ -430,7 +433,9 @@ def test_impl(_: object, mock_ds: MagicMock, mock_sched_fn: MagicMock) -> None:
 
         # Mock data source
         mock_data_source = MagicMock()
-        mock_data_source.get_metadata = AsyncMock(return_value={"feature1": "type1"})
+        mock_sched_metadata = MagicMock()
+        mock_sched_metadata.input_schema.items.keys.return_value = ["feature1"]
+        mock_data_source.get_metadata = AsyncMock(return_value=mock_sched_metadata)
         mock_ds.return_value = mock_data_source
 
         # Send request

tests/endpoints/metrics/drift/test_kolmogorov_smirnov_streaming.py

@@ -126,18 +126,18 @@ class TestKSTestStreamingEndpoints:
         expected_error_substring="referenceTag is required",
     )
 
-    test_compute_missing_fit_columns = factory.make_compute_endpoint_error_test(
-        metric_name="KSTestStreaming",
-        module_path="src.endpoints.metrics.drift.kolmogorov_smirnov_streaming",
-        endpoint_path="/metrics/drift/ksteststreaming",
-        client=client,
-        request_payload={
-            "modelId": "test-model",
-            "referenceTag": "baseline",
-            # Missing fitColumns
-        },
-        expected_status_code=HTTPStatus.BAD_REQUEST,
-        expected_error_substring="fitColumns is required",
+    test_compute_missing_fit_columns_derives_from_metadata = (
+        factory.make_compute_endpoint_test(
+            metric_name="KSTestStreaming",
+            module_path="src.endpoints.metrics.drift.kolmogorov_smirnov_streaming",
+            endpoint_path="/metrics/drift/ksteststreaming",
+            client=client,
+            request_payload={
+                "modelId": "test-model",
+                "referenceTag": "baseline",
+            },
+            expected_response_keys=["status", "value", "drift_detected"],
+        )
     )
 
     test_compute_invalid_feature = factory.make_compute_endpoint_error_test(