AI column modifications

superset-frontend/src/dashboard/components/WhatIfDrawer/index.tsx

@@ -18,9 +18,14 @@
  */
 import { useCallback, useMemo, useState } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
-import { t } from '@superset-ui/core';
+import { t, logging } from '@superset-ui/core';
 import { css, styled, Alert, useTheme } from '@apache-superset/core/ui';
-import { Button, Select } from '@superset-ui/core/components';
+import {
+  Button,
+  Select,
+  Checkbox,
+  Tooltip,
+} from '@superset-ui/core/components';
 import Slider from '@superset-ui/core/components/Slider';
 import { Icons } from '@superset-ui/core/components/Icons';
 import { setWhatIfModifications } from 'src/dashboard/actions/dashboardState';
@@ -31,6 +36,8 @@ import {
 import { getNumericColumnsForDashboard } from 'src/dashboard/util/whatIf';
 import { RootState, Slice, WhatIfColumn } from 'src/dashboard/types';
 import WhatIfAIInsights from './WhatIfAIInsights';
+import { fetchRelatedColumnSuggestions } from './whatIfApi';
+import { ExtendedWhatIfModification } from './types';
 
 export const WHAT_IF_PANEL_WIDTH = 300;
 
@@ -115,6 +122,69 @@ const SliderContainer = styled.div`
 
 const ApplyButton = styled(Button)`
   width: 100%;
+  min-height: 32px;
+`;
+
+const CheckboxContainer = styled.div`
+  display: flex;
+  align-items: center;
+  gap: ${({ theme }) => theme.sizeUnit}px;
+`;
+
+const ModificationsSection = styled.div`
+  display: flex;
+  flex-direction: column;
+  gap: ${({ theme }) => theme.sizeUnit * 2}px;
+`;
+
+const ModificationsSectionTitle = styled.div`
+  font-weight: ${({ theme }) => theme.fontWeightStrong};
+  color: ${({ theme }) => theme.colorText};
+  font-size: ${({ theme }) => theme.fontSizeSM}px;
+`;
+
+const ModificationCard = styled.div<{ isAISuggested?: boolean }>`
+  padding: ${({ theme }) => theme.sizeUnit * 2}px;
+  background-color: ${({ theme, isAISuggested }) =>
+    isAISuggested ? theme.colorInfoBg : theme.colorBgLayout};
+  border: 1px solid
+    ${({ theme, isAISuggested }) =>
+      isAISuggested ? theme.colorInfoBorder : theme.colorBorderSecondary};
+  border-radius: ${({ theme }) => theme.borderRadius}px;
+`;
+
+const ModificationHeader = styled.div`
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: ${({ theme }) => theme.sizeUnit}px;
+`;
+
+const ModificationColumn = styled.span`
+  font-weight: ${({ theme }) => theme.fontWeightStrong};
+  color: ${({ theme }) => theme.colorText};
+`;
+
+const ModificationValue = styled.span<{ isPositive: boolean }>`
+  font-weight: ${({ theme }) => theme.fontWeightStrong};
+  color: ${({ theme, isPositive }) =>
+    isPositive ? theme.colorSuccess : theme.colorError};
+`;
+
+const AIBadge = styled.span`
+  font-size: ${({ theme }) => theme.fontSizeXS}px;
+  padding: 2px 6px;
+  background-color: ${({ theme }) => theme.colorInfo};
+  color: ${({ theme }) => theme.colorWhite};
+  border-radius: ${({ theme }) => theme.borderRadius}px;
+  font-weight: ${({ theme }) => theme.fontWeightStrong};
+`;
+
+const ModificationReasoning = styled.div`
+  font-size: ${({ theme }) => theme.fontSizeSM}px;
+  color: ${({ theme }) => theme.colorTextSecondary};
+  margin-top: ${({ theme }) => theme.sizeUnit}px;
+  font-style: italic;
 `;
 
 interface WhatIfPanelProps {
@@ -129,6 +199,11 @@ const WhatIfPanel = ({ onClose, topOffset }: WhatIfPanelProps) => {
   const [selectedColumn, setSelectedColumn] = useState<string | null>(null);
   const [sliderValue, setSliderValue] = useState<number>(SLIDER_DEFAULT);
   const [affectedChartIds, setAffectedChartIds] = useState<number[]>([]);
+  const [enableCascadingEffects, setEnableCascadingEffects] = useState(false);
+  const [isLoadingSuggestions, setIsLoadingSuggestions] = useState(false);
+  const [appliedModifications, setAppliedModifications] = useState<
+    ExtendedWhatIfModification[]
+  >([]);
 
   const slices = useSelector(
     (state: RootState) => state.sliceEntities.slices as { [id: number]: Slice },
@@ -166,41 +241,112 @@ const WhatIfPanel = ({ onClose, topOffset }: WhatIfPanelProps) => {
     setSliderValue(value);
   }, []);
 
-  const handleApply = useCallback(() => {
+  const dashboardInfo = useSelector((state: RootState) => state.dashboardInfo);
+
+  const handleApply = useCallback(async () => {
     if (!selectedColumn) return;
 
     const multiplier = 1 + sliderValue / 100;
 
-    // Get affected chart IDs
-    const chartIds = columnToChartIds.get(selectedColumn) || [];
+    // Base user modification
+    const userModification: ExtendedWhatIfModification = {
+      column: selectedColumn,
+      multiplier,
+      isAISuggested: false,
+    };
+
+    let allModifications: ExtendedWhatIfModification[] = [userModification];
+
+    // If cascading effects enabled, fetch AI suggestions
+    if (enableCascadingEffects) {
+      setIsLoadingSuggestions(true);
+      try {
+        const suggestions = await fetchRelatedColumnSuggestions({
+          selectedColumn,
+          userMultiplier: multiplier,
+          availableColumns: numericColumns.map(col => ({
+            columnName: col.columnName,
+            description: col.description,
+            verboseName: col.verboseName,
+            datasourceId: col.datasourceId,
+          })),
+          dashboardName: dashboardInfo?.dash_edit_perm
+            ? dashboardInfo?.dashboard_title
+            : undefined,
+        });
+
+        // Add AI suggestions to modifications
+        const aiModifications: ExtendedWhatIfModification[] =
+          suggestions.suggestedModifications.map(mod => ({
+            column: mod.column,
+            multiplier: mod.multiplier,
+            isAISuggested: true,
+            reasoning: mod.reasoning,
+            confidence: mod.confidence,
+          }));
+
+        allModifications = [...allModifications, ...aiModifications];
+      } catch (error) {
+        logging.error('Failed to get AI suggestions:', error);
+        // Continue with just user modification
+      }
+      setIsLoadingSuggestions(false);
+    }
+
+    setAppliedModifications(allModifications);
+
+    // Collect all affected chart IDs from all modifications
+    const allAffectedChartIds = new Set<number>();
+    allModifications.forEach(mod => {
+      const chartIds = columnToChartIds.get(mod.column) || [];
+      chartIds.forEach(id => allAffectedChartIds.add(id));
+    });
+    const chartIdsArray = Array.from(allAffectedChartIds);
 
     // Save affected chart IDs for AI insights
-    setAffectedChartIds(chartIds);
+    setAffectedChartIds(chartIdsArray);
 
     // Save original chart data before applying what-if modifications
-    chartIds.forEach(chartId => {
+    chartIdsArray.forEach(chartId => {
       dispatch(saveOriginalChartData(chartId));
     });
 
-    // Set the what-if modifications in Redux state
+    // Set the what-if modifications in Redux state (all modifications)
     dispatch(
-      setWhatIfModifications([
-        {
-          column: selectedColumn,
-          multiplier,
-        },
-      ]),
+      setWhatIfModifications(
+        allModifications.map(mod => ({
+          column: mod.column,
+          multiplier: mod.multiplier,
+          filters: mod.filters,
+        })),
+      ),
     );
 
-    // Trigger queries for all charts that use the selected column
-    chartIds.forEach(chartId => {
+    // Trigger queries for all affected charts
+    chartIdsArray.forEach(chartId => {
       dispatch(triggerQuery(true, chartId));
     });
-  }, [dispatch, selectedColumn, sliderValue, columnToChartIds]);
+  }, [
+    dispatch,
+    selectedColumn,
+    sliderValue,
+    columnToChartIds,
+    enableCascadingEffects,
+    numericColumns,
+    dashboardInfo,
+  ]);
 
-  const isApplyDisabled = !selectedColumn || sliderValue === SLIDER_DEFAULT;
+  const isApplyDisabled =
+    !selectedColumn || sliderValue === SLIDER_DEFAULT || isLoadingSuggestions;
   const isSliderDisabled = !selectedColumn;
 
+  // Helper to format percentage change
+  const formatPercentage = (multiplier: number): string => {
+    const pct = (multiplier - 1) * 100;
+    const sign = pct >= 0 ? '+' : '';
+    return `${sign}${pct.toFixed(1)}%`;
+  };
+
   const sliderMarks = {
     [SLIDER_MIN]: `${SLIDER_MIN}%`,
     0: '0%',
@@ -255,22 +401,79 @@ const WhatIfPanel = ({ onClose, topOffset }: WhatIfPanelProps) => {
           </SliderContainer>
         </FormSection>
 
+        <CheckboxContainer>
+          <Checkbox
+            checked={enableCascadingEffects}
+            onChange={e => setEnableCascadingEffects(e.target.checked)}
+          >
+            {t('AI-powered cascading effects')}
+          </Checkbox>
+          <Tooltip
+            title={t(
+              'When enabled, AI will analyze column relationships and automatically suggest related columns that should also be modified.',
+            )}
+          >
+            <Icons.InfoCircleOutlined
+              iconSize="s"
+              css={css`
+                color: ${theme.colorTextSecondary};
+                cursor: help;
+              `}
+            />
+          </Tooltip>
+        </CheckboxContainer>
+
         <ApplyButton
           buttonStyle="primary"
           onClick={handleApply}
           disabled={isApplyDisabled}
+          loading={isLoadingSuggestions}
         >
           <Icons.StarFilled iconSize="s" />
-          {t('See what if')}
+          {isLoadingSuggestions
+            ? t('Analyzing relationships...')
+            : t('See what if')}
         </ApplyButton>
 
-        <Alert
-          type="info"
-          message={t(
-            'Select a column above to simulate changes and preview how it would impact your dashboard in real-time.',
-          )}
-          showIcon
-        />
+        {appliedModifications.length === 0 && (
+          <Alert
+            type="info"
+            message={t(
+              'Select a column above to simulate changes and preview how it would impact your dashboard in real-time.',
+            )}
+            showIcon
+          />
+        )}
+
+        {appliedModifications.length > 0 && (
+          <ModificationsSection>
+            <ModificationsSectionTitle>
+              {t('Applied Modifications')}
+            </ModificationsSectionTitle>
+            {appliedModifications.map((mod, idx) => (
+              <ModificationCard key={idx} isAISuggested={mod.isAISuggested}>
+                <ModificationHeader>
+                  <ModificationColumn>{mod.column}</ModificationColumn>
+                  <div
+                    css={css`
+                      display: flex;
+                      align-items: center;
+                      gap: ${theme.sizeUnit}px;
+                    `}
+                  >
+                    <ModificationValue isPositive={mod.multiplier >= 1}>
+                      {formatPercentage(mod.multiplier)}
+                    </ModificationValue>
+                    {mod.isAISuggested && <AIBadge>{t('AI')}</AIBadge>}
+                  </div>
+                </ModificationHeader>
+                {mod.reasoning && (
+                  <ModificationReasoning>{mod.reasoning}</ModificationReasoning>
+                )}
+              </ModificationCard>
+            ))}
+          </ModificationsSection>
+        )}
 
         {affectedChartIds.length > 0 && (
           <WhatIfAIInsights affectedChartIds={affectedChartIds} />

superset-frontend/src/dashboard/components/WhatIfDrawer/types.ts

@@ -71,3 +71,41 @@ export interface WhatIfInterpretResponse {
 }
 
 export type WhatIfAIStatus = 'idle' | 'loading' | 'success' | 'error';
+
+// Types for suggest_related endpoint
+
+export interface AvailableColumn {
+  columnName: string;
+  description?: string | null;
+  verboseName?: string | null;
+  datasourceId: number;
+}
+
+export interface SuggestedModification {
+  column: string;
+  multiplier: number;
+  reasoning: string;
+  confidence: 'high' | 'medium' | 'low';
+}
+
+export interface WhatIfSuggestRelatedRequest {
+  selectedColumn: string;
+  userMultiplier: number;
+  availableColumns: AvailableColumn[];
+  dashboardName?: string;
+}
+
+export interface WhatIfSuggestRelatedResponse {
+  suggestedModifications: SuggestedModification[];
+  explanation?: string;
+}
+
+// Extended modification type that tracks whether it came from AI
+export interface ExtendedWhatIfModification {
+  column: string;
+  multiplier: number;
+  filters?: WhatIfFilter[];
+  isAISuggested?: boolean;
+  reasoning?: string;
+  confidence?: 'high' | 'medium' | 'low';
+}

superset-frontend/src/dashboard/components/WhatIfDrawer/whatIfApi.ts

@@ -23,6 +23,9 @@ import {
   WhatIfInterpretResponse,
   ChartComparison,
   WhatIfFilter,
+  WhatIfSuggestRelatedRequest,
+  WhatIfSuggestRelatedResponse,
+  SuggestedModification,
 } from './types';
 
 interface ApiResponse {
@@ -84,3 +87,49 @@ export async function fetchWhatIfInterpretation(
     rawResponse: result.raw_response,
   };
 }
+
+interface ApiSuggestRelatedResponse {
+  result: {
+    suggested_modifications: Array<{
+      column: string;
+      multiplier: number;
+      reasoning: string;
+      confidence: string;
+    }>;
+    explanation?: string;
+  };
+}
+
+export async function fetchRelatedColumnSuggestions(
+  request: WhatIfSuggestRelatedRequest,
+): Promise<WhatIfSuggestRelatedResponse> {
+  const response = await SupersetClient.post({
+    endpoint: '/api/v1/what_if/suggest_related',
+    jsonPayload: {
+      selected_column: request.selectedColumn,
+      user_multiplier: request.userMultiplier,
+      available_columns: request.availableColumns.map(col => ({
+        column_name: col.columnName,
+        description: col.description,
+        verbose_name: col.verboseName,
+        datasource_id: col.datasourceId,
+      })),
+      dashboard_name: request.dashboardName,
+    },
+  });
+
+  const data = response.json as ApiSuggestRelatedResponse;
+  const { result } = data;
+
+  return {
+    suggestedModifications: result.suggested_modifications.map(
+      (mod): SuggestedModification => ({
+        column: mod.column,
+        multiplier: mod.multiplier,
+        reasoning: mod.reasoning,
+        confidence: mod.confidence as 'high' | 'medium' | 'low',
+      }),
+    ),
+    explanation: result.explanation,
+  };
+}

superset-frontend/src/dashboard/types.ts

@@ -312,6 +312,8 @@ export interface WhatIfColumn {
   columnName: string;
   datasourceId: number;
   usedByChartIds: number[];
+  description?: string | null;
+  verboseName?: string | null;
 }
 
 export enum MenuKeys {

superset-frontend/src/dashboard/util/whatIf.ts

@@ -154,6 +154,8 @@ export function getNumericColumnsForDashboard(
             columnName: colName,
             datasourceId: datasource.id,
             usedByChartIds: [chartId],
+            description: colMetadata.description,
+            verboseName: colMetadata.verbose_name,
           });
         } else {
           const existing = columnMap.get(key)!;

superset/what_if/api.py

@@ -28,10 +28,13 @@ from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP
 from superset.extensions import event_logger
 from superset.views.base_api import BaseSupersetApi, statsd_metrics
 from superset.what_if.commands.interpret import WhatIfInterpretCommand
+from superset.what_if.commands.suggest_related import WhatIfSuggestRelatedCommand
 from superset.what_if.exceptions import OpenRouterAPIError, OpenRouterConfigError
 from superset.what_if.schemas import (
     WhatIfInterpretRequestSchema,
     WhatIfInterpretResponseSchema,
+    WhatIfSuggestRelatedRequestSchema,
+    WhatIfSuggestRelatedResponseSchema,
 )
 
 logger = logging.getLogger(__name__)
@@ -116,3 +119,73 @@ class WhatIfRestApi(BaseSupersetApi):
         except ValueError as ex:
             logger.warning("Invalid request: %s", ex)
             return self.response_400(message=str(ex))
+
+    @expose("/suggest_related", methods=("POST",))
+    @event_logger.log_this
+    @protect()
+    @safe
+    @statsd_metrics
+    def suggest_related(self) -> Response:
+        """Get AI suggestions for related column modifications.
+        ---
+        post:
+          summary: Get AI-suggested cascading column modifications
+          description: >-
+            Analyzes column relationships and suggests related columns
+            that should be modified when a user modifies a specific column.
+            Uses AI to infer causal, mathematical, and domain-specific relationships.
+          requestBody:
+            required: true
+            content:
+              application/json:
+                schema:
+                  $ref: '#/components/schemas/WhatIfSuggestRelatedRequestSchema'
+          responses:
+            200:
+              description: Related column suggestions generated successfully
+              content:
+                application/json:
+                  schema:
+                    type: object
+                    properties:
+                      result:
+                        $ref: '#/components/schemas/WhatIfSuggestRelatedResponseSchema'
+            400:
+              $ref: '#/components/responses/400'
+            401:
+              $ref: '#/components/responses/401'
+            500:
+              $ref: '#/components/responses/500'
+            502:
+              description: Error communicating with AI service
+              content:
+                application/json:
+                  schema:
+                    type: object
+                    properties:
+                      message:
+                        type: string
+          security:
+            - jwt: []
+        """
+        try:
+            request_data = WhatIfSuggestRelatedRequestSchema().load(request.json)
+        except ValidationError as ex:
+            logger.warning("Invalid request data: %s", ex.messages)
+            return self.response_400(message=str(ex.messages))
+
+        try:
+            command = WhatIfSuggestRelatedCommand(request_data)
+            result = command.run()
+            return self.response(
+                200, result=WhatIfSuggestRelatedResponseSchema().dump(result)
+            )
+        except OpenRouterConfigError as ex:
+            logger.error("OpenRouter configuration error: %s", ex)
+            return self.response(500, message="AI suggestions are not configured")
+        except OpenRouterAPIError as ex:
+            logger.error("OpenRouter API error: %s", ex)
+            return self.response(502, message=str(ex))
+        except ValueError as ex:
+            logger.warning("Invalid request: %s", ex)
+            return self.response_400(message=str(ex))

superset/what_if/commands/suggest_related.py

@@ -0,0 +1,220 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""What-If Analysis suggest related columns command using OpenRouter."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any
+
+import httpx
+from flask import current_app
+
+from superset.commands.base import BaseCommand
+from superset.what_if.exceptions import OpenRouterAPIError, OpenRouterConfigError
+
+logger = logging.getLogger(__name__)
+
+
+class WhatIfSuggestRelatedCommand(BaseCommand):
+    """Command to get AI suggestions for related column modifications."""
+
+    def __init__(self, data: dict[str, Any]) -> None:
+        self._data = data
+
+    def run(self) -> dict[str, Any]:
+        self.validate()
+        return self._get_ai_suggestions()
+
+    def validate(self) -> None:
+        api_key = current_app.config.get("OPENROUTER_API_KEY")
+        if not api_key:
+            raise OpenRouterConfigError("OPENROUTER_API_KEY not configured")
+
+        if not self._data.get("selected_column"):
+            raise ValueError("selected_column is required")
+
+        if not self._data.get("available_columns"):
+            raise ValueError("available_columns list is required")
+
+        if self._data.get("user_multiplier") is None:
+            raise ValueError("user_multiplier is required")
+
+    def _build_prompt(self) -> str:
+        selected_column = self._data["selected_column"]
+        user_multiplier = self._data["user_multiplier"]
+        available_columns = self._data["available_columns"]
+        dashboard_name = self._data.get("dashboard_name") or "Dashboard"
+
+        pct_change = (user_multiplier - 1) * 100
+        sign = "+" if pct_change >= 0 else ""
+
+        # Build column list with descriptions
+        columns_text = []
+        for col in available_columns:
+            # Skip the selected column - we don't want to suggest modifying it again
+            if col["column_name"] == selected_column:
+                continue
+
+            col_desc = f"- **{col['column_name']}**"
+            if col.get("verbose_name"):
+                col_desc += f" ({col['verbose_name']})"
+            if col.get("description"):
+                col_desc += f": {col['description']}"
+            columns_text.append(col_desc)
+
+        if not columns_text:
+            columns_text = ["No other columns available"]
+
+        return f"""You are a business intelligence analyst helping with what-if scenario analysis.
+
+## Context
+A user is working on a "{dashboard_name}" dashboard and wants to simulate the cascading effects of changing a metric.
+
+## User's Modification
+The user is modifying **{selected_column}** by {sign}{pct_change:.1f}%
+
+## Other Available Columns
+These are the other numeric columns available in the dashboard:
+{chr(10).join(columns_text)}
+
+## Your Task
+Analyze the relationships between these columns and suggest which OTHER columns should also be modified as a cascading effect of the user's change to {selected_column}.
+
+Consider:
+1. **Causal relationships**: If column A affects column B in real business scenarios
+2. **Mathematical relationships**: Derived metrics, ratios, calculated fields
+3. **Domain knowledge**: Industry-standard relationships (e.g., increasing customers often increases orders and revenue)
+
+For each suggested column, provide:
+- The appropriate multiplier (proportional, dampened, amplified, or inverse based on the relationship)
+- A brief reasoning explaining the relationship (1 sentence)
+- Your confidence level (high/medium/low)
+
+Guidelines:
+- Only suggest columns that have a clear logical relationship to {selected_column}
+- Be conservative - don't suggest modifications without good reasoning
+- The multiplier should be realistic (e.g., if {selected_column} increases 10%, a related column might increase 5-15%, not 100%)
+- If no clear relationships exist, return an empty suggestions array
+
+Respond in JSON format:
+{{
+  "suggested_modifications": [
+    {{
+      "column": "column_name",
+      "multiplier": 1.08,
+      "reasoning": "Brief explanation of the relationship",
+      "confidence": "high"
+    }}
+  ],
+  "explanation": "Overall summary of the analysis (1-2 sentences)"
+}}"""
+
+    def _get_ai_suggestions(self) -> dict[str, Any]:
+        api_key = current_app.config.get("OPENROUTER_API_KEY")
+        model = current_app.config.get("OPENROUTER_MODEL", "x-ai/grok-4.1-fast")
+        api_base = current_app.config.get(
+            "OPENROUTER_API_BASE", "https://openrouter.ai/api/v1"
+        )
+        timeout = current_app.config.get("OPENROUTER_TIMEOUT", 30)
+
+        prompt = self._build_prompt()
+
+        headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+            "HTTP-Referer": current_app.config.get("WEBDRIVER_BASEURL", ""),
+            "X-Title": "Apache Superset What-If Analysis",
+        }
+
+        payload = {
+            "model": model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": (
+                        "You are a business intelligence analyst specializing in "
+                        "data relationships and cascading effects analysis. "
+                        "Respond only with valid JSON."
+                    ),
+                },
+                {"role": "user", "content": prompt},
+            ],
+            "temperature": 0.3,
+            "max_tokens": 1000,
+            "response_format": {"type": "json_object"},
+        }
+
+        try:
+            with httpx.Client(timeout=timeout) as client:
+                response = client.post(
+                    f"{api_base}/chat/completions",
+                    headers=headers,
+                    json=payload,
+                )
+                response.raise_for_status()
+
+                result = response.json()
+                content = result["choices"][0]["message"]["content"]
+
+                # Parse the JSON response
+                parsed = json.loads(content)
+
+                # Validate and normalize the response
+                suggestions = parsed.get("suggested_modifications", [])
+                validated_suggestions = []
+
+                for suggestion in suggestions:
+                    # Ensure required fields exist
+                    if all(
+                        k in suggestion
+                        for k in ["column", "multiplier", "reasoning", "confidence"]
+                    ):
+                        # Normalize confidence to lowercase
+                        confidence = suggestion["confidence"].lower()
+                        if confidence not in ("high", "medium", "low"):
+                            confidence = "medium"
+
+                        validated_suggestions.append(
+                            {
+                                "column": suggestion["column"],
+                                "multiplier": float(suggestion["multiplier"]),
+                                "reasoning": suggestion["reasoning"],
+                                "confidence": confidence,
+                            }
+                        )
+
+                return {
+                    "suggested_modifications": validated_suggestions,
+                    "explanation": parsed.get("explanation"),
+                }
+
+        except httpx.HTTPStatusError as ex:
+            logger.error("OpenRouter API error: %s", ex.response.status_code)
+            raise OpenRouterAPIError(
+                f"OpenRouter API error: {ex.response.status_code}"
+            ) from ex
+        except json.JSONDecodeError as ex:
+            logger.error("Failed to parse AI response: %s", ex)
+            raise OpenRouterAPIError("Failed to parse AI response") from ex
+        except httpx.TimeoutException as ex:
+            logger.error("OpenRouter API timeout")
+            raise OpenRouterAPIError("AI service timed out") from ex
+        except Exception as ex:
+            logger.exception("Unexpected error calling OpenRouter")
+            raise OpenRouterAPIError(f"Unexpected error: {ex!s}") from ex

superset/what_if/schemas.py

@@ -161,3 +161,90 @@ class WhatIfInterpretResponseSchema(Schema):
         required=False,
         metadata={"description": "Raw AI response (only in debug mode)"},
     )
+
+
+# Schemas for suggest_related endpoint
+
+
+class AvailableColumnSchema(Schema):
+    """Schema for an available column with metadata."""
+
+    column_name = fields.String(
+        required=True,
+        metadata={"description": "Name of the column"},
+    )
+    description = fields.String(
+        required=False,
+        load_default=None,
+        metadata={"description": "Column description/documentation"},
+    )
+    verbose_name = fields.String(
+        required=False,
+        load_default=None,
+        metadata={"description": "Human-readable column name"},
+    )
+    datasource_id = fields.Integer(
+        required=True,
+        metadata={"description": "ID of the datasource containing this column"},
+    )
+
+
+class WhatIfSuggestRelatedRequestSchema(Schema):
+    """Schema for suggest_related request."""
+
+    selected_column = fields.String(
+        required=True,
+        metadata={"description": "The column the user selected to modify"},
+    )
+    user_multiplier = fields.Float(
+        required=True,
+        metadata={
+            "description": "The multiplier the user applied (e.g., 1.1 for +10%)"
+        },
+    )
+    available_columns = fields.List(
+        fields.Nested(AvailableColumnSchema),
+        required=True,
+        metadata={"description": "All numeric columns available in the dashboard"},
+    )
+    dashboard_name = fields.String(
+        required=False,
+        load_default=None,
+        metadata={"description": "Name of the dashboard for context"},
+    )
+
+
+class SuggestedModificationSchema(Schema):
+    """Schema for a single AI-suggested modification."""
+
+    column = fields.String(
+        required=True,
+        metadata={"description": "Column name to modify"},
+    )
+    multiplier = fields.Float(
+        required=True,
+        metadata={"description": "Suggested multiplier for this column"},
+    )
+    reasoning = fields.String(
+        required=True,
+        metadata={"description": "Brief explanation of why this column is related"},
+    )
+    confidence = fields.String(
+        required=True,
+        metadata={"description": "Confidence level: high, medium, or low"},
+    )
+
+
+class WhatIfSuggestRelatedResponseSchema(Schema):
+    """Schema for suggest_related response."""
+
+    suggested_modifications = fields.List(
+        fields.Nested(SuggestedModificationSchema),
+        required=True,
+        metadata={"description": "List of AI-suggested column modifications"},
+    )
+    explanation = fields.String(
+        required=False,
+        load_default=None,
+        metadata={"description": "Overall explanation of the relationship analysis"},
+    )